import pyqtgraph.opengl as gl
import pyqtgraph as pg
from PyQt5.QtWidgets import QGraphicsRectItem, QGraphicsLineItem
from PyQt5 import QtCore
from cosense3d.utils.box_utils import *
from cosense3d.dataset.toolkit.cosense import csColors
from cosense3d.dataset.toolkit.cosense import CoSenseDataConverter as cs
CSCOLORS = (np.array([csColors[k] for k in cs.OBJ_LIST]) / 255.).tolist()
BOX_COLORs = {
'inactive': CSCOLORS,
'highlight': (0., 1, 1, 1),
'active': (0.9, 0, 1, 1),
'local_gt': (1, 1, 0, 1),
'global_gt': (0, 1, 0, 1),
'gt': (0, 1, 0, 1),
'det': (1, 0, 0, 1),
'pred': (1, 0, 1, 1),
'successor': (0, 0.5, 1, 1),
'successor_gt': (0, 1, 1, 1)
}
pens = {
'yellow_dashed': pg.mkPen('y', width=1, style=QtCore.Qt.DashLine),
'yellow_solid': pg.mkPen('y', width=1, style=QtCore.Qt.SolidLine),
'virtual': pg.mkPen(color=(0, 0, 0, 0), width=1),
}
[docs]class MeshBoxItem(gl.GLMeshItem):
def __init__(self, size=(1, 1, 1), color=(0.0, 1.0, 0.0, 0.25)):
l, w, h = size
verts = [
[0, 0, 0],
[l, 0, 0],
[l, 0, h],
[0, 0, h],
[0, w, 0],
[l, w, 0],
[l, w, h],
[0, w, h]
]
verts = np.array(verts)
faces = [
[0, 1, 2],
[0, 2, 3],
[1, 5, 6],
[1, 6, 2],
[5, 4, 7],
[5, 7, 6],
[4, 0, 3],
[4, 3, 7],
[3, 2, 6],
[3, 6, 7],
[0, 4, 5],
[0, 5, 1]
]
faces = np.array(faces)
normals = np.array([
[0, -1, 0],
[0, 1, 0],
[1, 0, 0],
[-1, 0, 0],
[0, 0, -1],
[0, 0, 1]
])
colors = [color] * len(faces)
meshdata = gl.MeshData(vertexes=verts, faces=faces, faceColors=colors)
super().__init__(meshdata=meshdata, shader='balloon', glOptions='translucent')
[docs]class LineBoxItem(gl.GLLinePlotItem):
ids = set() # TODO: need to be initialized by labeled data in the current scenario
id_ptr = 0
def __init__(self,
box,
status='inactive',
show_direction=False,
last_pose=None,
line_width=1.):
"""
:param box: ([id], type_id, x, y, z, l, w, h, roll, pitch, yaw)
:param color:
4 -------- 5 ^ z
/| /| |
7 -------- 6 . |
| | | | | . x
. 0 -------- 1 |/
|/ |/ +-------> y
3 -------- 2
"""
id = None
box_score = None
if len(box) == 11:
id = int(box[0])
type_id = int(box[1])
box = box[2:]
elif len(box) == 10:
type_id = int(box[0])
box = box[1:]
elif len(box) == 12:
id = int(box[0])
type_id = int(box[1])
box_score = box[-1]
box = box[2:-1]
else:
raise NotImplementedError
vertices = np.zeros((12, 3))
vertices[:8] = boxes_to_corners_3d(np.array([box]))[0]
if show_direction:
# -----
# | |---- direction on top
# -----
top_center = np.mean(vertices[4:], axis=0)
top_front = np.mean(vertices[[4, 5]], axis=0)
top_ff = top_front * 2 - top_center
vertices[8] = top_front
vertices[9] = top_ff
if last_pose is not None:
# -----
# last pose on bottom of the boxe o----| |
# -----
assert len(last_pose) == 3
bottom_center = np.mean(vertices[:4], axis=0)
last_pose[2] = bottom_center[2] # set last pose z to ground
vertices[10] = np.array(last_pose)
vertices[11] = np.array(bottom_center)
self.vertices = vertices
# Define the edges of the box
edges = [
[0, 1], # front-bottom
[1, 5], # front-right
[5, 4], # front-top
[4, 0], # front-left
[0, 3], # left-bottom
[1, 2], # right-bottom
[5, 6], # right-top
[4, 7], # left-top
[3, 2], # back-bottom
[2, 6], # back-right
[6, 7], # back-top
[7, 3], # back-left
]
if show_direction:
edges.append([8, 9])
if last_pose is not None:
edges.append([10, 11])
self.edges = np.array(edges)
vertices_pairs = self.vertices[self.edges.flatten()]
while id is None:
if LineBoxItem.id_ptr not in LineBoxItem.ids:
id = LineBoxItem.id_ptr
else:
LineBoxItem.id_ptr += 1
self.id = id
self.typeid = type_id
LineBoxItem.ids.add(id)
super().__init__(pos=vertices_pairs,
color=self.color(status),
width=line_width,
mode='lines',
glOptions='opaque')
[docs] def to_center(self):
"""Convert box to center format"""
transform = self.transform().matrix()
corners = (transform[:3, :3] @ self.vertices[:8].T) + transform[:3, 3:]
box_center = corners_to_boxes_3d(corners.T[None, :])
return box_center[0]
[docs] def activate(self):
self.setData(color=BOX_COLORs['active'], width=2.0)
[docs] def deactivate(self):
self.setData(color=BOX_COLORs['inactive'][self.typeid] + [0.5])
[docs] def highlight(self):
self.setData(color=BOX_COLORs['highlight'], width=2.0)
@property
def isActive(self):
return self.color == BOX_COLORs['active']
[docs] def color(self, status):
if status in ['inactive']:
return BOX_COLORs[status][self.typeid] + [0.5]
else:
return BOX_COLORs[status]
[docs]class LineItem(QGraphicsLineItem):
def __init__(self, line, parent=None):
super().__init__(parent)
self.inactive_pen = pens['yellow_dashed']
self.active_pen = pens['yellow_solid']
self.setLine(*line)
self.setPen(self.inactive_pen)
self.setZValue(5)
self.active = False
[docs] def hoverEvent(self, event):
if event.isExit():
self.setPen(self.inactive_pen)
self.active = False
else:
self.setPen(self.active_pen)
self.active = True
[docs]class RectangleItem(QGraphicsRectItem):
def __init__(self, rect):
super().__init__(*rect)
self.setPen(pens['virtual'])
self.setZValue(0)
self.active = False
[docs] def hoverEvent(self, event):
if event.isExit():
self.setPen(pens['virtual'])
self.active = False
else:
pos = event.pos()
if abs(pos.x()) < 0.3 and abs(pos.y()) < 0.3:
self.setPen(pens['yellow_solid'])
self.active = True
if __name__ == "__main__":
from PyQt5 import QtWidgets
app = QtWidgets.QApplication([])
w = gl.GLViewWidget()
w.opts['distance'] = 20
w.show()
boxItem = LineBoxItem(
box=[-5, 8, -1, 4, 3, 2, 0, 0, 0],
show_direction=True
)
w.addItem(boxItem)
app.exec_()