add simple visualization codes

tools/vis_utils/visualize_utils.py

+import numpy as np
+import mayavi.mlab as mlab
+from pcdet.utils import box_utils
+
+box_colormap = [
+    [1, 1, 1],
+    [0, 1, 0],
+    [0, 1, 1],
+    [1, 1, 0],
+]
+
+
+def visualize_pts(pts, fig=None, bgcolor=(0, 0, 0), fgcolor=(1.0, 1.0, 1.0),
+                  show_intensity=False, size=(600, 600), draw_origin=True):
+    if not isinstance(pts, np.ndarray):
+        pts = pts.cpu().numpy()
+    if fig is None:
+        fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=fgcolor, engine=None, size=size)
+
+    if show_intensity:
+        G = mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], pts[:, 3], mode='point',
+                          colormap='gnuplot', scale_factor=1, figure=fig)
+    else:
+        G = mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], mode='point',
+                          colormap='gnuplot', scale_factor=1, figure=fig)
+    if draw_origin:
+        mlab.points3d(0, 0, 0, color=(1, 1, 1), mode='cube', scale_factor=0.2)
+        mlab.plot3d([0, 3], [0, 0], [0, 0], color=(0, 0, 1), tube_radius=0.1)
+        mlab.plot3d([0, 0], [0, 3], [0, 0], color=(0, 1, 0), tube_radius=0.1)
+        mlab.plot3d([0, 0], [0, 0], [0, 3], color=(1, 0, 0), tube_radius=0.1)
+
+    return fig
+
+
+def draw_grid(x1, y1, x2, y2, fig, tube_radius=None, color=(0.5, 0.5, 0.5)):
+    mlab.plot3d([x1, x1], [y1, y2], [0, 0], color=color, tube_radius=tube_radius, line_width=1, figure=fig)
+    mlab.plot3d([x2, x2], [y1, y2], [0, 0], color=color, tube_radius=tube_radius, line_width=1, figure=fig)
+    mlab.plot3d([x1, x2], [y1, y1], [0, 0], color=color, tube_radius=tube_radius, line_width=1, figure=fig)
+    mlab.plot3d([x1, x2], [y2, y2], [0, 0], color=color, tube_radius=tube_radius, line_width=1, figure=fig)
+    return fig
+
+
+def draw_multi_grid_range(fig, grid_size=20, bv_range=(-60, -60, 60, 60)):
+    for x in range(bv_range[0], bv_range[2], grid_size):
+        for y in range(bv_range[1], bv_range[3], grid_size):
+            fig = draw_grid(x, y, x + grid_size, y + grid_size, fig)
+
+    return fig
+
+
+def draw_scenes(points, gt_boxes=None, ref_boxes=None, ref_scores=None, ref_labels=None):
+    if not isinstance(points, np.ndarray):
+        points = points.cpu().numpy()
+    if ref_boxes is not None and not isinstance(ref_boxes, np.ndarray):
+        ref_boxes = ref_boxes.cpu().numpy()
+    if gt_boxes is not None and not isinstance(gt_boxes, np.ndarray):
+        gt_boxes = gt_boxes.cpu().numpy()
+    if ref_scores is not None and not isinstance(ref_scores, np.ndarray):
+        ref_scores = ref_scores.cpu().numpy()
+    if ref_labels is not None and not isinstance(ref_labels, np.ndarray):
+        ref_labels = ref_labels.cpu().numpy()
+
+    fig = visualize_pts(points)
+    fig = draw_multi_grid_range(fig, bv_range=(0, -40, 80, 40))
+    if gt_boxes is not None:
+        corners3d = box_utils.boxes_to_corners_3d(gt_boxes)
+        fig = draw_corners3d(corners3d, fig=fig, color=(0, 0, 1), max_num=100)
+
+    if ref_boxes is not None:
+        ref_corners3d = box_utils.boxes_to_corners_3d(ref_boxes)
+        if ref_labels is None:
+            fig = draw_corners3d(ref_corners3d, fig=fig, color=(0, 1, 0), cls=ref_scores, max_num=100)
+        else:
+            for k in range(ref_labels.min(), ref_labels.max() + 1):
+                cur_color = tuple(box_colormap[k % len(box_colormap)])
+                mask = (ref_labels == k)
+                fig = draw_corners3d(ref_corners3d[mask], fig=fig, color=cur_color, cls=ref_scores[mask], max_num=100)
+    mlab.view(azimuth=-179, elevation=54.0, distance=104.0, roll=90.0)
+    return fig
+
+
+def draw_corners3d(corners3d, fig, color=(1, 1, 1), line_width=2, cls=None, tag='', max_num=500, tube_radius=None):
+    """
+    :param corners3d: (N, 8, 3)
+    :param fig:
+    :param color:
+    :param line_width:
+    :param cls:
+    :param tag:
+    :param max_num:
+    :return:
+    """
+    import mayavi.mlab as mlab
+    num = min(max_num, len(corners3d))
+    for n in range(num):
+        b = corners3d[n]  # (8, 3)
+
+        if cls is not None:
+            if isinstance(cls, np.ndarray):
+                mlab.text3d(b[6, 0], b[6, 1], b[6, 2], '%.2f' % cls[n], scale=(0.3, 0.3, 0.3), color=color, figure=fig)
+            else:
+                mlab.text3d(b[6, 0], b[6, 1], b[6, 2], '%s' % cls[n], scale=(0.3, 0.3, 0.3), color=color, figure=fig)
+
+        for k in range(0, 4):
+            i, j = k, (k + 1) % 4
+            mlab.plot3d([b[i, 0], b[j, 0]], [b[i, 1], b[j, 1]], [b[i, 2], b[j, 2]], color=color, tube_radius=tube_radius,
+                        line_width=line_width, figure=fig)
+
+            i, j = k + 4, (k + 1) % 4 + 4
+            mlab.plot3d([b[i, 0], b[j, 0]], [b[i, 1], b[j, 1]], [b[i, 2], b[j, 2]], color=color, tube_radius=tube_radius,
+                        line_width=line_width, figure=fig)
+
+            i, j = k, k + 4
+            mlab.plot3d([b[i, 0], b[j, 0]], [b[i, 1], b[j, 1]], [b[i, 2], b[j, 2]], color=color, tube_radius=tube_radius,
+                        line_width=line_width, figure=fig)
+
+        i, j = 0, 5
+        mlab.plot3d([b[i, 0], b[j, 0]], [b[i, 1], b[j, 1]], [b[i, 2], b[j, 2]], color=color, tube_radius=tube_radius,
+                    line_width=line_width, figure=fig)
+        i, j = 1, 4
+        mlab.plot3d([b[i, 0], b[j, 0]], [b[i, 1], b[j, 1]], [b[i, 2], b[j, 2]], color=color, tube_radius=tube_radius,
+                    line_width=line_width, figure=fig)
+
+    return fig