remove the visualization dependency on CUDA OpenPCDet

tools/visual_utils/visualize_utils.py

 import numpy as np
 import mayavi.mlab as mlab
-from pcdet.utils import box_utils
+import torch
 
 box_colormap = [
     [1, 1, 1],
@@ -10,6 +10,65 @@ box_colormap = [
 ]
 
 
+def check_numpy_to_torch(x):
+    if isinstance(x, np.ndarray):
+        return torch.from_numpy(x).float(), True
+    return x, False
+
+
+def rotate_points_along_z(points, angle):
+    """
+    Args:
+        points: (B, N, 3 + C)
+        angle: (B), angle along z-axis, angle increases x ==> y
+    Returns:
+
+    """
+    points, is_numpy = check_numpy_to_torch(points)
+    angle, _ = check_numpy_to_torch(angle)
+
+    cosa = torch.cos(angle)
+    sina = torch.sin(angle)
+    zeros = angle.new_zeros(points.shape[0])
+    ones = angle.new_ones(points.shape[0])
+    rot_matrix = torch.stack((
+        cosa,  sina, zeros,
+        -sina, cosa, zeros,
+        zeros, zeros, ones
+    ), dim=1).view(-1, 3, 3).float()
+    points_rot = torch.matmul(points[:, :, 0:3], rot_matrix)
+    points_rot = torch.cat((points_rot, points[:, :, 3:]), dim=-1)
+    return points_rot.numpy() if is_numpy else points_rot
+
+
+def boxes_to_corners_3d(boxes3d):
+    """
+        7 -------- 4
+       /|         /|
+      6 -------- 5 .
+      | |        | |
+      . 3 -------- 0
+      |/         |/
+      2 -------- 1
+    Args:
+        boxes3d:  (N, 7) [x, y, z, dx, dy, dz, heading], (x, y, z) is the box center
+
+    Returns:
+    """
+    boxes3d, is_numpy = check_numpy_to_torch(boxes3d)
+
+    template = boxes3d.new_tensor((
+        [1, 1, -1], [1, -1, -1], [-1, -1, -1], [-1, 1, -1],
+        [1, 1, 1], [1, -1, 1], [-1, -1, 1], [-1, 1, 1],
+    )) / 2
+
+    corners3d = boxes3d[:, None, 3:6].repeat(1, 8, 1) * template[None, :, :]
+    corners3d = rotate_points_along_z(corners3d.view(-1, 8, 3), boxes3d[:, 6]).view(-1, 8, 3)
+    corners3d += boxes3d[:, None, 0:3]
+
+    return corners3d.numpy() if is_numpy else corners3d
+
+
 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):
@@ -32,6 +91,38 @@ def visualize_pts(pts, fig=None, bgcolor=(0, 0, 0), fgcolor=(1.0, 1.0, 1.0),
     return fig
 
 
+def draw_sphere_pts(pts, color=(0, 1, 0), fig=None, bgcolor=(0, 0, 0), scale_factor=0.2):
+    if not isinstance(pts, np.ndarray):
+        pts = pts.cpu().numpy()
+
+    if fig is None:
+        fig = mlab.figure(figure=None, bgcolor=bgcolor, fgcolor=None, engine=None, size=(600, 600))
+
+    if isinstance(color, np.ndarray) and color.shape[0] == 1:
+        color = color[0]
+        color = (color[0] / 255.0, color[1] / 255.0, color[2] / 255.0)
+
+    if isinstance(color, np.ndarray):
+        pts_color = np.zeros((pts.__len__(), 4), dtype=np.uint8)
+        pts_color[:, 0:3] = color
+        pts_color[:, 3] = 255
+        G = mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], np.arange(0, pts_color.__len__()), mode='sphere',
+                          scale_factor=scale_factor, figure=fig)
+        G.glyph.color_mode = 'color_by_scalar'
+        G.glyph.scale_mode = 'scale_by_vector'
+        G.module_manager.scalar_lut_manager.lut.table = pts_color
+    else:
+        mlab.points3d(pts[:, 0], pts[:, 1], pts[:, 2], mode='sphere', color=color,
+                      colormap='gnuplot', scale_factor=scale_factor, figure=fig)
+
+    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), line_width=3, tube_radius=None, figure=fig)
+    mlab.plot3d([0, 0], [0, 3], [0, 0], color=(0, 1, 0), line_width=3, tube_radius=None, figure=fig)
+    mlab.plot3d([0, 0], [0, 0], [0, 3], color=(1, 0, 0), line_width=3, tube_radius=None, figure=fig)
+
+    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)
@@ -63,11 +154,11 @@ def draw_scenes(points, gt_boxes=None, ref_boxes=None, ref_scores=None, ref_labe
     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)
+        corners3d = 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)
+        ref_corners3d = 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: