├── .gitignore ├── README.md ├── demo.ipynb ├── demo ├── canvas_3d_barebones.png ├── canvas_3d_flight.gif ├── canvas_3d_various_colors.png ├── canvas_bev_barebones.png ├── canvas_bev_various_colors.png └── data │ ├── 000008_calib.txt │ ├── 000008_image_2.png │ ├── 000008_label_2.txt │ └── 000008_velodyne.bin └── simple_plot3d ├── __init__.py ├── canvas_3d.py └── canvas_bev.py /.gitignore: -------------------------------------------------------------------------------- 1 | .ipynb_checkpoints/ 2 | __pycache__/ -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # A Simple Visualizer for 3D point clouds & bounding boxes 2 | ![gif_example](demo/canvas_3d_flight.gif) 3 | 4 | Supports 3D & BEV visualizations with a few lines: 5 | ```Python 6 | canvas = Canvas_BEV() 7 | canvas_xy, valid_mask = canvas.get_canvas_coords(pts_xyz) # Get Canvas coords 8 | canvas.draw_canvas_points(canvas_xy[valid_mask]) # Only draw valid points 9 | canvas.draw_boxes(gt_lidar_bboxes_3d, texts=gt_names) # Draw boxes 10 | ``` 11 | 12 | ## 3D 13 | ![3d_example](demo/canvas_3d_various_colors.png) 14 | 15 | ## BEV 16 | ![bev_example](demo/canvas_bev_various_colors.png) 17 | -------------------------------------------------------------------------------- /demo/canvas_3d_barebones.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Divadi/simple_plot3d/f9383587d69b0000ee3ae02823667e6915baec3b/demo/canvas_3d_barebones.png -------------------------------------------------------------------------------- /demo/canvas_3d_flight.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Divadi/simple_plot3d/f9383587d69b0000ee3ae02823667e6915baec3b/demo/canvas_3d_flight.gif -------------------------------------------------------------------------------- /demo/canvas_3d_various_colors.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Divadi/simple_plot3d/f9383587d69b0000ee3ae02823667e6915baec3b/demo/canvas_3d_various_colors.png -------------------------------------------------------------------------------- /demo/canvas_bev_barebones.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Divadi/simple_plot3d/f9383587d69b0000ee3ae02823667e6915baec3b/demo/canvas_bev_barebones.png -------------------------------------------------------------------------------- /demo/canvas_bev_various_colors.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Divadi/simple_plot3d/f9383587d69b0000ee3ae02823667e6915baec3b/demo/canvas_bev_various_colors.png -------------------------------------------------------------------------------- /demo/data/000008_calib.txt: -------------------------------------------------------------------------------- 1 | P0: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 0.000000000000e+00 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00 2 | P1: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 -3.875744000000e+02 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 0.000000000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 0.000000000000e+00 3 | P2: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 4.485728000000e+01 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 2.163791000000e-01 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 2.745884000000e-03 4 | P3: 7.215377000000e+02 0.000000000000e+00 6.095593000000e+02 -3.395242000000e+02 0.000000000000e+00 7.215377000000e+02 1.728540000000e+02 2.199936000000e+00 0.000000000000e+00 0.000000000000e+00 1.000000000000e+00 2.729905000000e-03 5 | R0_rect: 9.999239000000e-01 9.837760000000e-03 -7.445048000000e-03 -9.869795000000e-03 9.999421000000e-01 -4.278459000000e-03 7.402527000000e-03 4.351614000000e-03 9.999631000000e-01 6 | Tr_velo_to_cam: 7.533745000000e-03 -9.999714000000e-01 -6.166020000000e-04 -4.069766000000e-03 1.480249000000e-02 7.280733000000e-04 -9.998902000000e-01 -7.631618000000e-02 9.998621000000e-01 7.523790000000e-03 1.480755000000e-02 -2.717806000000e-01 7 | Tr_imu_to_velo: 9.999976000000e-01 7.553071000000e-04 -2.035826000000e-03 -8.086759000000e-01 -7.854027000000e-04 9.998898000000e-01 -1.482298000000e-02 3.195559000000e-01 2.024406000000e-03 1.482454000000e-02 9.998881000000e-01 -7.997231000000e-01 8 | 9 | -------------------------------------------------------------------------------- /demo/data/000008_image_2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Divadi/simple_plot3d/f9383587d69b0000ee3ae02823667e6915baec3b/demo/data/000008_image_2.png -------------------------------------------------------------------------------- /demo/data/000008_label_2.txt: -------------------------------------------------------------------------------- 1 | Car 0.88 3 -0.69 0.00 192.37 402.31 374.00 1.60 1.57 3.23 -2.70 1.74 3.68 -1.29 2 | Car 0.00 1 2.04 334.85 178.94 624.50 372.04 1.57 1.50 3.68 -1.17 1.65 7.86 1.90 3 | Car 0.34 3 -1.84 937.29 197.39 1241.00 374.00 1.39 1.44 3.08 3.81 1.64 6.15 -1.31 4 | Car 0.00 1 -1.33 597.59 176.18 720.90 261.14 1.47 1.60 3.66 1.07 1.55 14.44 -1.25 5 | Car 0.00 0 1.74 741.18 168.83 792.25 208.43 1.70 1.63 4.08 7.24 1.55 33.20 1.95 6 | Car 0.00 0 -1.65 884.52 178.31 956.41 240.18 1.59 1.59 2.47 8.48 1.75 19.96 -1.25 7 | DontCare -1 -1 -10 800.38 163.67 825.45 184.07 -1 -1 -1 -1000 -1000 -1000 -10 8 | DontCare -1 -1 -10 859.58 172.34 886.26 194.51 -1 -1 -1 -1000 -1000 -1000 -10 9 | DontCare -1 -1 -10 801.81 163.96 825.20 183.59 -1 -1 -1 -1000 -1000 -1000 -10 10 | DontCare -1 -1 -10 826.87 162.28 845.84 178.86 -1 -1 -1 -1000 -1000 -1000 -10 11 | -------------------------------------------------------------------------------- /demo/data/000008_velodyne.bin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Divadi/simple_plot3d/f9383587d69b0000ee3ae02823667e6915baec3b/demo/data/000008_velodyne.bin -------------------------------------------------------------------------------- /simple_plot3d/__init__.py: -------------------------------------------------------------------------------- 1 | from .canvas_3d import Canvas_3D 2 | from .canvas_bev import Canvas_BEV 3 | -------------------------------------------------------------------------------- /simple_plot3d/canvas_3d.py: -------------------------------------------------------------------------------- 1 | """ 2 | Written by Jinhyung Park 3 | 4 | Simple 3D visualization for 3D points & boxes. Intended as a simple, hackable 5 | alternative to mayavi for certain point cloud tasks. 6 | """ 7 | 8 | import numpy as np 9 | import cv2 10 | import copy 11 | from functools import partial 12 | import matplotlib 13 | 14 | class Canvas_3D(object): 15 | def __init__(self, 16 | canvas_shape=(500, 1000), 17 | camera_center_coords=(-4, 0, 4), 18 | camera_focus_coords=(-4 + 0.9396926, 0, 4 - 0.34202014), 19 | focal_length=None, 20 | canvas_bg_color=(0, 0, 0)): 21 | """ 22 | Args: 23 | canvas_shape (Tuple[Int]): Canvas image size - height & width. 24 | camera_center_coords (Tuple[Float]): Location of camera center in 25 | 3D space. 26 | camera_focus_coords (Tuple[Float]): Intuitively, what point in 3D 27 | space is the camera pointed at? These are absolute coordinates, 28 | *not* relative to camera center. 29 | focal_length (None | Int): 30 | None: Half of the max of height & width of canvas_shape. This 31 | seems to be a decent default. 32 | Int: Specified directly. 33 | canvas_bg_color (Tuple[Int]): RGB (0 ~ 255) of canvas background 34 | color. 35 | """ 36 | 37 | self.canvas_shape = canvas_shape 38 | self.H, self.W = self.canvas_shape 39 | self.canvas_bg_color = canvas_bg_color 40 | 41 | self.camera_center_coords = camera_center_coords 42 | self.camera_focus_coords = camera_focus_coords 43 | 44 | if focal_length is None: 45 | self.focal_length = max(self.H, self.W) // 2 46 | else: 47 | self.focal_length = focal_length 48 | 49 | # Setup extrinsics and intrinsics of this virtual camera. 50 | self.ext_matrix = self.get_extrinsic_matrix( 51 | self.camera_center_coords, self.camera_focus_coords) 52 | self.int_matrix = np.array([ 53 | [self.focal_length, 0, self.W // 2, 0], 54 | [0, self.focal_length, self.H // 2, 0], 55 | [0, 0, 1, 0], 56 | ]) 57 | 58 | self.clear_canvas() 59 | 60 | def get_canvas(self): 61 | return self.canvas 62 | 63 | def clear_canvas(self): 64 | self.canvas = np.zeros((self.H, self.W, 3), dtype=np.uint8) 65 | self.canvas[..., :] = self.canvas_bg_color 66 | 67 | def get_canvas_coords(self, 68 | xyz, 69 | depth_min=0.1, 70 | return_depth=False): 71 | """ 72 | Projects XYZ points onto the canvas and returns the projected canvas 73 | coordinates. 74 | 75 | Args: 76 | xyz (ndarray): (N, 3+) array of coordinates. Additional columns 77 | beyond the first three are ignored. 78 | depth_min (Float): Only points with a projected depth larger 79 | than this value are "valid". 80 | return_depth (Boolean): Whether to additionally return depth of 81 | projected points. 82 | Returns: 83 | canvas_xy (ndarray): (N, 2) array of projected canvas coordinates. 84 | "x" is dim0, "y" is dim1 of canvas. 85 | valid_mask (ndarray): (N,) boolean mask indicating which of 86 | canvas_xy fits into canvas (are visible from virtual camera). 87 | depth (ndarray): Optionally returned (N,) array of depth values 88 | """ 89 | xyz_hom = np.concatenate( 90 | [xyz, np.ones((xyz.shape[0], 1), dtype=np.float32)], axis=1) 91 | img_pts = (self.int_matrix @ self.ext_matrix @ xyz_hom.T).T 92 | 93 | depth = img_pts[:, 2] 94 | xy = img_pts[:, :2] / depth[:, None] 95 | xy_int = xy.round().astype(np.int32) 96 | 97 | # Flip X and Y so "x" is dim0, "y" is dim1 of canvas 98 | xy_int = xy_int[:, ::-1] 99 | 100 | valid_mask = ((depth > depth_min) & 101 | (xy_int[:, 0] >= 0) & (xy_int[:, 0] < self.H) & 102 | (xy_int[:, 1] >= 0) & (xy_int[:, 1] < self.W)) 103 | 104 | if return_depth: 105 | return xy_int, valid_mask, depth 106 | else: 107 | return xy_int, valid_mask 108 | 109 | 110 | def draw_canvas_points(self, 111 | canvas_xy, 112 | radius=-1, 113 | colors=None, 114 | colors_operand=None): 115 | """ 116 | Draws canvas_xy onto self.canvas. 117 | 118 | Args: 119 | canvas_xy (ndarray): (N, 2) array of *valid* canvas coordinates. 120 | "x" is dim0, "y" is dim1 of canvas. 121 | radius (Int): 122 | -1: Each point is visualized as a single pixel. 123 | r: Each point is visualized as a circle with radius r. 124 | colors: 125 | None: colors all points white. 126 | Tuple: RGB (0 ~ 255), indicating a single color for all points. 127 | ndarray: (N, 3) array of RGB values for each point. 128 | String: Such as "Spectral", uses a matplotlib cmap, with the 129 | operand (the value cmap is called on for each point) being 130 | colors_operand. 131 | colors_operand (ndarray): (N,) array of values cooresponding to 132 | canvas_xy, to be used only if colors is a cmap. Unlike 133 | Canvas_BEV, cannot be None if colors is a String. 134 | """ 135 | if len(canvas_xy) == 0: 136 | return 137 | 138 | if colors is None: 139 | colors = np.full( 140 | (len(canvas_xy), 3), fill_value=255, dtype=np.uint8) 141 | elif isinstance(colors, tuple): 142 | assert len(colors) == 3 143 | colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8) 144 | colors_tmp[..., :len(colors)] = np.array(colors) 145 | colors = colors_tmp 146 | elif isinstance(colors, np.ndarray): 147 | assert len(colors) == len(canvas_xy) 148 | colors = colors.astype(np.uint8) 149 | elif isinstance(colors, str): 150 | assert colors_operand is not None 151 | colors = matplotlib.cm.get_cmap(colors) 152 | 153 | # Normalize 0 ~ 1 for cmap 154 | colors_operand = colors_operand - colors_operand.min() 155 | colors_operand = colors_operand / colors_operand.max() 156 | 157 | # Get cmap colors - note that cmap returns (*input_shape, 4), with 158 | # colors scaled 0 ~ 1 159 | colors = (colors(colors_operand)[:, :3] * 255).astype(np.uint8) 160 | else: 161 | raise Exception( 162 | "colors type {} was not an expected type".format(type(colors))) 163 | 164 | if radius == -1: 165 | self.canvas[canvas_xy[:, 0], canvas_xy[:, 1], :] = colors 166 | else: 167 | for color, (x, y) in zip(colors.tolist(), canvas_xy.tolist()): 168 | self.canvas = cv2.circle(self.canvas, (y, x), radius, color, 169 | -1, lineType=cv2.LINE_AA) 170 | 171 | def draw_lines(self, 172 | start_xyz, 173 | end_xyz, 174 | colors=(255, 255, 255), 175 | thickness=1): 176 | """ 177 | Draws lines between provided 3D points. 178 | 179 | Args: 180 | start_xyz (ndarray): Shape (N, 3) of 3D points to start from. 181 | end_xyz (ndarray): Shape (N, 3) of 3D points to end at. Same length 182 | as start_xyz. 183 | colors: 184 | None: colors all points white. 185 | Tuple: RGB (0 ~ 255), indicating a single color for all points. 186 | ndarray: (N, 3) array of RGB values for each point. 187 | thickness (Int): 188 | Thickness of drawn cv2 line. 189 | """ 190 | if colors is None: 191 | colors = np.full( 192 | (len(canvas_xy), 3), fill_value=255, dtype=np.uint8) 193 | elif isinstance(colors, tuple): 194 | assert len(colors) == 3 195 | colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8) 196 | colors_tmp[..., :len(colors)] = np.array(colors) 197 | colors = colors_tmp 198 | elif isinstance(colors, np.ndarray): 199 | assert len(colors) == len(canvas_xy) 200 | colors = colors.astype(np.uint8) 201 | else: 202 | raise Exception( 203 | "colors type {} was not an expected type".format(type(colors))) 204 | 205 | start_pts_xy, start_pts_valid_mask, start_pts_d = \ 206 | self.get_canvas_coords(start_xyz, True) 207 | end_pts_xy, end_pts_valid_mask, end_pts_d = \ 208 | self.get_canvas_coords(end_xyz, True) 209 | 210 | for idx, (color, start_pt_xy, end_pt_xy) in enumerate( 211 | zip(colors.tolist(), start_pts_xy.tolist(), 212 | end_pts_xy.tolist())): 213 | 214 | if start_pts_valid_mask[idx] and end_pts_valid_mask[idx]: 215 | self.canvas = cv2.line(self.canvas, 216 | tuple(start_pt_xy[::-1]), 217 | tuple(end_pt_xy[::-1]), 218 | color=color, 219 | thickness=thickness, 220 | lineType=cv2.LINE_AA) 221 | 222 | def draw_boxes(self, 223 | boxes, 224 | colors=None, 225 | texts=None, 226 | depth_min=0.1, 227 | draw_incomplete_boxes=True, 228 | box_line_thickness=2, 229 | box_text_size=0.5, 230 | text_corner=1): 231 | """ 232 | Draws 3D boxes. 233 | 234 | Args: 235 | boxes (ndarray): Shape (N, 7), each row representing a box of 236 | format (x, y, z, x_size, y_size, z_size, yaw). This function 237 | assumes *bottom center* - the xyz center of the provided box 238 | is the center of the bottom face of the 3D box, not the 239 | floating true center of the 3D box. 240 | colors: 241 | None: colors all points white. 242 | Tuple: RGB (0 ~ 255), indicating a single color for all points. 243 | ndarray: (N, 3) array of RGB values for each point. 244 | texts (List[String]): Length N; text to write next to boxes. 245 | depth_min (Float): Only box corners with a projected depth larger 246 | than this value are drawn if draw_incomplete_boxes is True. 247 | draw_incomplete_boxes (Boolean): If any boxes are incomplete, 248 | meaning it has a corner out of view based on depth_min, decide 249 | whether to draw them at all. 250 | thickness (Int): 251 | Thickness of drawn cv2 box lines. 252 | box_line_thickness (int): cv2 line/text thickness 253 | box_text_size (float): cv2 putText size 254 | text_corner (int): 0 ~ 7. Which corner of 3D box to write text at. 255 | """ 256 | # Setup colors 257 | if colors is None: 258 | colors = np.full( 259 | (len(boxes), 3), fill_value=255, dtype=np.uint8) 260 | elif isinstance(colors, tuple): 261 | assert len(colors) == 3 262 | colors_tmp = np.zeros((len(boxes), 3), dtype=np.uint8) 263 | colors_tmp[..., :len(colors)] = np.array(colors) 264 | colors = colors_tmp 265 | elif isinstance(colors, np.ndarray): 266 | assert len(colors) == len(boxes) 267 | colors = colors.astype(np.uint8) 268 | else: 269 | raise Exception( 270 | "colors type {} was not an expected type".format(type(colors))) 271 | 272 | # boxes is N x 7 273 | boxes = np.copy(boxes) # prevent in-place modifications 274 | assert len(boxes.shape) == 2 275 | 276 | dims = boxes[:, 3:6] 277 | corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1) 278 | 279 | corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] 280 | # use relative origin [0.5, 0.5, 0], assuming bottom center 281 | corners_norm = corners_norm - np.array([0.5, 0.5, 0]) 282 | corners = dims.reshape(-1, 1, 3) * corners_norm.reshape([1, 8, 3]) 283 | 284 | # rotate around z axis 285 | angles = boxes[:, 6] 286 | rot_sin = np.sin(angles) 287 | rot_cos = np.cos(angles) 288 | ones = np.ones_like(rot_cos) 289 | zeros = np.zeros_like(rot_cos) 290 | rot_mat_T = np.stack([ 291 | np.stack([rot_cos, -rot_sin, zeros]), 292 | np.stack([rot_sin, rot_cos, zeros]), 293 | np.stack([zeros, zeros, ones]) 294 | ]) 295 | corners = np.einsum('aij,jka->aik', corners, rot_mat_T) 296 | corners += boxes[:, :3].reshape(-1, 1, 3) # N x 8 x 3 297 | 298 | # Now we have corners. Need them on the canvas 2D space. 299 | corners_xy, valid_mask = self.get_canvas_coords( 300 | corners.reshape(-1, 3), depth_min=depth_min) 301 | corners_xy = corners_xy.reshape(-1, 8, 2) 302 | valid_mask = valid_mask.reshape(-1, 8) 303 | 304 | # Now draw them with lines in correct places 305 | for i, (color, curr_corners_xy, curr_valid_mask) in enumerate( 306 | zip(colors.tolist(), corners_xy.tolist(), valid_mask.tolist())): 307 | 308 | if not draw_incomplete_boxes and sum(curr_valid_mask) != 8: 309 | # Some corner is invalid, don't draw the box at all. 310 | continue 311 | 312 | for start, end in [(0, 1), (1, 2), (2, 3), (3, 0), 313 | (0, 4), (1, 5), (2, 6), (3, 7), 314 | (4, 5), (5, 6), (6, 7), (7, 4)]: 315 | if not (curr_valid_mask[start] and curr_valid_mask[end]): 316 | continue # start or end is not valid 317 | 318 | self.canvas = cv2.line( 319 | self.canvas, 320 | (curr_corners_xy[start][1], curr_corners_xy[start][0]), 321 | (curr_corners_xy[end][1], curr_corners_xy[end][0]), 322 | color=color, 323 | thickness=box_line_thickness, 324 | lineType=cv2.LINE_AA) 325 | 326 | # If even a single line was drawn, add text as well. 327 | if sum(curr_valid_mask) > 0: 328 | if texts is not None: 329 | self.canvas = cv2.putText(self.canvas, 330 | str(texts[i]), 331 | (curr_corners_xy[text_corner][1], 332 | curr_corners_xy[text_corner][0]), 333 | cv2.FONT_HERSHEY_SIMPLEX, 334 | box_text_size, 335 | color, 336 | thickness=box_line_thickness) 337 | 338 | @staticmethod 339 | def cart2sph(xyz): 340 | x, y, z = xyz[:, 0], xyz[:, 1], xyz[:, 2] 341 | 342 | depth = np.linalg.norm(xyz, 2, axis=1) 343 | az = -np.arctan2(y, x) 344 | el = np.arcsin(z / depth) 345 | return az, el, depth 346 | 347 | @staticmethod 348 | def get_extrinsic_matrix( 349 | camera_center_coords, 350 | camera_focus_coords, 351 | ): 352 | """ 353 | Args: 354 | camera_center_coords: (x, y, z) of where camera should be located 355 | in 3D space. 356 | camera_focus_coords: (x, y, z) of where camera should look at from 357 | camera_center_coords 358 | 359 | Thoughts: 360 | Remember that in camera coordiantes, pos x is right, pos y is up, 361 | pos z is forward. 362 | """ 363 | center_x, center_y, center_z = camera_center_coords 364 | focus_x, focus_y, focus_z = camera_focus_coords 365 | az, el, depth = Canvas_3D.cart2sph(np.array([ 366 | [focus_x - center_x, focus_y - center_y, focus_z - center_z] 367 | ])) 368 | az = float(az) 369 | el = float(el) 370 | depth = float(depth) 371 | 372 | ### First, construct extrinsics 373 | ## Rotation matrix 374 | 375 | z_rot = np.array([ 376 | [np.cos(az), -np.sin(az), 0], 377 | [np.sin(az), np.cos(az), 0], 378 | [0, 0, 1] 379 | ]) 380 | 381 | # el is rotation around y axis. 382 | y_rot = np.array([ 383 | [np.cos(-el), 0, -np.sin(-el)], 384 | [0, 1, 0], 385 | [np.sin(-el), 0, np.cos(-el)], 386 | ]) 387 | 388 | 389 | ## Now, how the z_rot and y_rot work (spherical coordiantes), is it 390 | ## computes rotations starting from the positive x axis and rotates 391 | ## positive x axis to the desired direction. The desired direction is 392 | ## the "looking direction" of the camera, which should actually be the 393 | ## z-axis. So should convert the points so that the x axis is the new z 394 | ## axis, and after the transformations. 395 | ## Why x -> z for points? If we think about rotating the camera, z 396 | ## should become x, so reverse when moving points. 397 | last_rot = np.array([ 398 | [0, -1, 0], 399 | [0, 0, -1], 400 | [1, 0, 0] # x -> z 401 | ]) 402 | 403 | # Put them together. Order matters. Make it hom. 404 | rot_matrix = np.eye(4, dtype=np.float32) 405 | rot_matrix[:3, :3] = last_rot @ y_rot @ z_rot 406 | 407 | ## Translation matrix 408 | trans_matrix = np.array([ 409 | [1, 0, 0, -center_x], 410 | [0, 1, 0, -center_y], 411 | [0, 0, 1, -center_z], 412 | [0, 0, 0, 1], 413 | ]) 414 | 415 | ## Finally, extrinsics matrix. Order matters - do trans then rot 416 | ext_matrix = rot_matrix @ trans_matrix 417 | 418 | return ext_matrix 419 | -------------------------------------------------------------------------------- /simple_plot3d/canvas_bev.py: -------------------------------------------------------------------------------- 1 | """ 2 | Written by Jinhyung Park 3 | 4 | Simple BEV visualization for 3D points & boxes. 5 | """ 6 | 7 | import numpy as np 8 | import cv2 9 | import copy 10 | from functools import partial 11 | import matplotlib 12 | 13 | class Canvas_BEV(object): 14 | def __init__(self, 15 | canvas_shape=(1000, 1000), 16 | canvas_x_range=(-50, 50), 17 | canvas_y_range=(-50, 50), 18 | canvas_bg_color=(0, 0, 0)): 19 | """ 20 | Args: 21 | canvas_shape (Tuple[int]): Shape of BEV Canvas image. First element 22 | corresponds to X range, the second element to Y range. 23 | canvas_x_range (Tuple[int]): Range of X-coords to visualize. X is 24 | vertical: negative ~ positive is top ~ down. 25 | canvas_y_range (Tuple[int]): Range of Y-coords to visualize. Y is 26 | horizontal: negative ~ positive is left ~ right. 27 | canvas_bg_color (Tuple[int]): RGB (0 ~ 255) of Canvas background 28 | color. 29 | """ 30 | 31 | # Sanity check ratios 32 | if ((canvas_shape[0] / canvas_shape[1]) != 33 | ((canvas_x_range[0] - canvas_x_range[1]) / 34 | (canvas_y_range[0] - canvas_y_range[1]))): 35 | 36 | print("Not an error, but the x & y ranges are not "\ 37 | "proportional to canvas height & width.") 38 | 39 | self.canvas_shape = canvas_shape 40 | self.canvas_x_range = canvas_x_range 41 | self.canvas_y_range = canvas_y_range 42 | self.canvas_bg_color = canvas_bg_color 43 | 44 | self.clear_canvas() 45 | 46 | def get_canvas(self): 47 | return self.canvas 48 | 49 | def clear_canvas(self): 50 | self.canvas = np.zeros((*self.canvas_shape, 3), dtype=np.uint8) 51 | self.canvas[..., :] = self.canvas_bg_color 52 | 53 | def get_canvas_coords(self, xy): 54 | """ 55 | Args: 56 | xy (ndarray): (N, 2+) array of coordinates. Additional columns 57 | beyond the first two are ignored. 58 | 59 | Returns: 60 | canvas_xy (ndarray): (N, 2) array of xy scaled into canvas 61 | coordinates. Invalid locations of canvas_xy are clipped into 62 | range. "x" is dim0, "y" is dim1 of canvas. 63 | valid_mask (ndarray): (N,) boolean mask indicating which of 64 | canvas_xy fits into canvas. 65 | """ 66 | xy = np.copy(xy) # prevent in-place modifications 67 | 68 | x = xy[:, 0] 69 | y = xy[:, 1] 70 | 71 | # Get valid mask 72 | valid_mask = ((x > self.canvas_x_range[0]) & 73 | (x < self.canvas_x_range[1]) & 74 | (y > self.canvas_y_range[0]) & 75 | (y < self.canvas_y_range[1])) 76 | 77 | # Rescale points 78 | x = ((x - self.canvas_x_range[0]) / 79 | (self.canvas_x_range[1] - self.canvas_x_range[0])) 80 | x = x * self.canvas_shape[0] 81 | x = np.clip(np.around(x), 0, 82 | self.canvas_shape[0] - 1).astype(np.int32) 83 | 84 | y = ((y - self.canvas_y_range[0]) / 85 | (self.canvas_y_range[1] - self.canvas_y_range[0])) 86 | y = y * self.canvas_shape[1] 87 | y = np.clip(np.around(y), 0, 88 | self.canvas_shape[1] - 1).astype(np.int32) 89 | 90 | # Return 91 | canvas_xy = np.stack([x, y], axis=1) 92 | 93 | return canvas_xy, valid_mask 94 | 95 | 96 | def draw_canvas_points(self, 97 | canvas_xy, 98 | radius=-1, 99 | colors=None, 100 | colors_operand=None): 101 | """ 102 | Draws canvas_xy onto self.canvas. 103 | 104 | Args: 105 | canvas_xy (ndarray): (N, 2) array of *valid* canvas coordinates. 106 | "x" is dim0, "y" is dim1 of canvas. 107 | radius (Int): 108 | -1: Each point is visualized as a single pixel. 109 | r: Each point is visualized as a circle with radius r. 110 | colors: 111 | None: colors all points white. 112 | Tuple: RGB (0 ~ 255), indicating a single color for all points. 113 | ndarray: (N, 3) array of RGB values for each point. 114 | String: Such as "Spectral", uses a matplotlib cmap, with the 115 | operand (the value cmap is called on for each point) being 116 | colors_operand. If colors_operand is None, uses normalized 117 | distance from (0, 0) of XY point coords. 118 | colors_operand (ndarray | None): (N,) array of values cooresponding 119 | to canvas_xy, to be used only if colors is a cmap. 120 | """ 121 | if len(canvas_xy) == 0: 122 | return 123 | 124 | if colors is None: 125 | colors = np.full( 126 | (len(canvas_xy), 3), fill_value=255, dtype=np.uint8) 127 | elif isinstance(colors, tuple): 128 | assert len(colors) == 3 129 | colors_tmp = np.zeros((len(canvas_xy), 3), dtype=np.uint8) 130 | colors_tmp[..., :] = np.array(colors) 131 | colors = colors_tmp 132 | elif isinstance(colors, np.ndarray): 133 | assert len(colors) == len(canvas_xy) 134 | colors = colors.astype(np.uint8) 135 | elif isinstance(colors, str): 136 | colors = matplotlib.cm.get_cmap(colors) 137 | if colors_operand is None: 138 | # Get distances from (0, 0) (albeit potentially clipped) 139 | origin_center = self.get_canvas_coords(np.zeros((1, 2)))[0][0] 140 | colors_operand = np.sqrt( 141 | ((canvas_xy - origin_center) ** 2).sum(axis=1)) 142 | 143 | # Normalize 0 ~ 1 for cmap 144 | colors_operand = colors_operand - colors_operand.min() 145 | colors_operand = colors_operand / colors_operand.max() 146 | 147 | # Get cmap colors - note that cmap returns (*input_shape, 4), with 148 | # colors scaled 0 ~ 1 149 | colors = (colors(colors_operand)[:, :3] * 255).astype(np.uint8) 150 | else: 151 | raise Exception( 152 | "colors type {} was not an expected type".format(type(colors))) 153 | 154 | if radius == -1: 155 | self.canvas[canvas_xy[:, 0], canvas_xy[:, 1], :] = colors 156 | else: 157 | for color, (x, y) in zip(colors.tolist(), canvas_xy.tolist()): 158 | self.canvas = cv2.circle(self.canvas, (y, x), radius, color, 159 | -1, lineType=cv2.LINE_AA) 160 | 161 | def draw_boxes(self, 162 | boxes, 163 | colors=None, 164 | texts=None, 165 | box_line_thickness=2, 166 | box_text_size=0.5, 167 | text_corner=0): 168 | """ 169 | Draws a set of boxes onto the canvas. 170 | Args: 171 | boxes (ndarray): Can either be of shape: 172 | (N, 7): Then, assumes (x, y, z, x_size, y_size, z_size, yaw) 173 | (N, 5): Then, assumes (x, y, x_size, y_size, yaw) 174 | Everything is in the same coordinate system as points 175 | (not canvas coordinates) 176 | colors: 177 | None: colors all points white. 178 | Tuple: RGB (0 ~ 255), indicating a single color for all points. 179 | ndarray: (N, 3) array of RGB values for each point. 180 | texts (List[String]): Length N; text to write next to boxes. 181 | box_line_thickness (int): cv2 line/text thickness 182 | box_text_size (float): cv2 putText size 183 | text_corner (int): 0 ~ 3. Which corner of 3D box to write text at. 184 | """ 185 | # Setup colors 186 | if colors is None: 187 | colors = np.full((len(boxes), 3), fill_value=255, dtype=np.uint8) 188 | elif isinstance(colors, tuple): 189 | assert len(colors) == 3 190 | colors_tmp = np.zeros((len(boxes), 3), dtype=np.uint8) 191 | colors_tmp[..., :len(colors)] = np.array(colors) 192 | colors = colors_tmp 193 | elif isinstance(colors, np.ndarray): 194 | assert len(colors) == len(boxes) 195 | colors = colors.astype(np.uint8) 196 | else: 197 | raise Exception( 198 | "colors type {} was not an expected type".format(type(colors))) 199 | 200 | boxes = np.copy(boxes) # prevent in-place modifications 201 | assert len(boxes.shape) == 2 202 | 203 | if boxes.shape[-1] == 7: 204 | boxes = boxes[:, [0, 1, 3, 4, 6]] 205 | else: 206 | assert boxes.shape[-1] == 5 207 | 208 | ## Get the BEV four corners 209 | # Get BEV 4 corners in box canonical coordinates 210 | bev_corners = np.array([[ 211 | [0.5, 0.5], 212 | [-0.5, 0.5], 213 | [-0.5, -0.5], 214 | [0.5, -0.5] 215 | ]]) * boxes[:, None, [2, 3]] # N x 4 x 2 216 | 217 | # Get rotation matrix from yaw 218 | rot_sin = np.sin(boxes[:, -1]) 219 | rot_cos = np.cos(boxes[:, -1]) 220 | rot_matrix = np.stack([ 221 | [rot_cos, -rot_sin], 222 | [rot_sin, rot_cos] 223 | ]) # 2 x 2 x N 224 | 225 | # Rotate BEV 4 corners. Result: N x 4 x 2 226 | bev_corners = np.einsum('aij,jka->aik', bev_corners, rot_matrix) 227 | 228 | # Translate BEV 4 corners 229 | bev_corners = bev_corners + boxes[:, None, [0, 1]] 230 | 231 | ## Transform BEV 4 corners to canvas coords 232 | bev_corners_canvas, valid_mask = \ 233 | self.get_canvas_coords(bev_corners.reshape(-1, 2)) 234 | bev_corners_canvas = bev_corners_canvas.reshape(*bev_corners.shape) 235 | valid_mask = valid_mask.reshape(*bev_corners.shape[:-1]) 236 | 237 | # At least 1 corner in canvas to draw. 238 | valid_mask = valid_mask.sum(axis=1) > 0 239 | bev_corners_canvas = bev_corners_canvas[valid_mask] 240 | if texts is not None: 241 | texts = np.array(texts)[valid_mask] 242 | 243 | ## Draw onto canvas 244 | # Draw the outer boundaries 245 | idx_draw_pairs = [(0, 1), (1, 2), (2, 3), (3, 0)] 246 | for i, (color, curr_box_corners) in enumerate( 247 | zip(colors.tolist(), bev_corners_canvas)): 248 | 249 | curr_box_corners = curr_box_corners.astype(np.int32) 250 | for start, end in idx_draw_pairs: 251 | self.canvas = cv2.line(self.canvas, 252 | tuple(curr_box_corners[start][::-1]\ 253 | .tolist()), 254 | tuple(curr_box_corners[end][::-1]\ 255 | .tolist()), 256 | color=color, 257 | thickness=box_line_thickness) 258 | if texts is not None: 259 | self.canvas = cv2.putText(self.canvas, 260 | str(texts[i]), 261 | tuple(curr_box_corners[text_corner]\ 262 | [::-1].tolist()), 263 | cv2.FONT_HERSHEY_SIMPLEX, 264 | box_text_size, 265 | color=color, 266 | thickness=box_line_thickness) 267 | --------------------------------------------------------------------------------