"""3D visualization plotting utilities. Self-contained module extracted from yolov5-3d/utils/plots.py. Contains only the functions required by visualize_single_frame.py and visualize_batch.py for ground-truth 3D visualization, with no cross-repo dependencies. """ import cv2 import numpy as np import torch from ultralytics.utils.plotting import Annotator # noqa: F401 (re-exported) # --------------------------------------------------------------------------- # Color helpers # --------------------------------------------------------------------------- class Colors: """Provides an RGB color palette derived from Ultralytics color scheme.""" def __init__(self): hexs = ( "FF3838", "FF9D97", "FF701F", "FFB21D", "CFD231", "48F90A", "92CC17", "3DDB86", "1A9334", "00D4BB", "2C99A8", "00C2FF", "344593", "6473FF", "0018EC", "8438FF", "520085", "CB38FF", "FF95C8", "FF37C7", ) self.palette = [self.hex2rgb(f"#{c}") for c in hexs] self.n = len(self.palette) def __call__(self, i, bgr=False): """Return colour for index ``i`` (BGR if ``bgr=True``, else RGB).""" c = self.palette[int(i) % self.n] return (c[2], c[1], c[0]) if bgr else c @staticmethod def hex2rgb(h): """Convert hex colour string to (R, G, B) tuple.""" return tuple(int(h[1 + i: 1 + i + 2], 16) for i in (0, 2, 4)) colors = Colors() # module-level singleton; callers: from plots_3d import colors # --------------------------------------------------------------------------- # 3D geometry helpers # --------------------------------------------------------------------------- def rotation_3d_in_axis(points, angles, axis=0): """Rotate *points* around a specified camera-frame axis. Args: points (np.ndarray): (N, 3) array of 3D points. angles (float): Rotation angle in radians. axis (int): 0=X, 1=Y, 2=Z. Returns: np.ndarray: Rotated points (N, 3). """ rot_sin = np.sin(angles) rot_cos = np.cos(angles) ones = np.ones_like(rot_cos) zeros = np.zeros_like(rot_cos) if axis == 1: rot_mat = np.stack([ np.stack([rot_cos, zeros, -rot_sin]), np.stack([zeros, ones, zeros]), np.stack([rot_sin, zeros, rot_cos]), ]) elif axis == 2: rot_mat = np.stack([ np.stack([rot_cos, rot_sin, zeros]), np.stack([-rot_sin, rot_cos, zeros]), np.stack([zeros, zeros, ones]), ]) elif axis == 0: rot_mat = np.stack([ np.stack([ones, zeros, zeros]), np.stack([zeros, rot_cos, rot_sin]), np.stack([zeros, -rot_sin, rot_cos]), ]) else: raise ValueError(f"axis must be in {{0, 1, 2}}, got {axis}") return np.dot(points, rot_mat) def compute_3d_box_corners_4face(center_3d, dimensions, rotation, face_type=0): """Compute the 8 corners of a 3D bounding box from a face-center point. Args: center_3d (array-like): (x, y, z) centre of the specified face in camera coords. dimensions (array-like): (length, height, width) of the box in metres. rotation (float): rot_y — rotation about the Y axis in radians. face_type (int): 0=front, 1=tail/rear, 2=left, 3=right, -1=box centre. Returns: np.ndarray: (8, 3) corner coordinates in camera frame. """ l, h, w = dimensions corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1).astype(np.float64) corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]] offsets = { 1: [0, 0.5, 0.5], # tail 0: [1, 0.5, 0.5], # front 3: [0.5, 0.5, 0], # right 2: [0.5, 0.5, 1], # left -1: [0.5, 0.5, 0.5], # whole centre } corners_norm = corners_norm - offsets.get(face_type, offsets[-1]) corners = np.array([l, h, w]).reshape(1, 3) * corners_norm.reshape(8, 3) corners = rotation_3d_in_axis(corners, rotation, axis=1) corners += np.array(center_3d).reshape(1, 3) return corners # --------------------------------------------------------------------------- # Fisheye (KB) distortion helpers # --------------------------------------------------------------------------- def apply_fisheye_distortion(x, y, distort_coeffs): """Apply Kannala-Brandt fisheye distortion to normalised camera coordinates. Args: x (float): Normalised x coordinate (x3d / z3d). y (float): Normalised y coordinate (y3d / z3d). distort_coeffs (list): [k1, k2, k3, k4] KB coefficients. Returns: tuple[float, float]: Distorted normalised coordinates (xd, yd). """ if not distort_coeffs or len(distort_coeffs) < 4: return x, y k1, k2, k3, k4 = distort_coeffs[:4] r = np.sqrt(x * x + y * y) if r < 1e-8: return x, y theta = np.arctan(r) theta_d = theta * (1 + k1 * theta**2 + k2 * theta**4 + k3 * theta**6 + k4 * theta**8) scale = theta_d / r return x * scale, y * scale def remove_fisheye_distortion(xd, yd, distort_coeffs, max_iter=20): """Remove KB fisheye distortion from normalised camera coordinates. Args: xd (float): Distorted normalised x coordinate. yd (float): Distorted normalised y coordinate. distort_coeffs (list): [k1, k2, k3, k4] KB coefficients. max_iter (int): Maximum Newton-Raphson iterations. Returns: tuple[float, float]: Undistorted normalised coordinates (xn, yn). """ if not distort_coeffs or len(distort_coeffs) < 4: return xd, yd k1, k2, k3, k4 = distort_coeffs[:4] r_d = np.sqrt(xd * xd + yd * yd) if r_d < 1e-8: return xd, yd theta_d = r_d theta = theta_d / (1 + k1 * theta_d * theta_d) for _ in range(max_iter): theta2 = theta * theta theta4 = theta2 * theta2 theta6 = theta4 * theta2 theta8 = theta4 * theta4 f = theta * (1 + k1 * theta2 + k2 * theta4 + k3 * theta6 + k4 * theta8) - theta_d f_prime = 1 + 3 * k1 * theta2 + 5 * k2 * theta4 + 7 * k3 * theta6 + 9 * k4 * theta8 theta_new = theta - f / f_prime if abs(theta_new - theta) < 1e-8: theta = theta_new break theta = theta_new r = np.tan(theta) scale = r / r_d return xd * scale, yd * scale # --------------------------------------------------------------------------- # 3D-to-2D projection helpers # --------------------------------------------------------------------------- def project_3d_to_2d_with_distortion(points_3d, calib): """Project 3D points to 2D using KB fisheye camera calibration. Args: points_3d (np.ndarray): (N, 3) points in camera coordinates. calib (dict): Camera parameters ``fx``, ``fy``, ``cx``, ``cy``, and optional ``distort_coeffs``. Returns: np.ndarray: (N, 2) image coordinates (NaN for behind-camera points). """ fx, fy = calib['fx'], calib['fy'] cx, cy = calib['cx'], calib['cy'] distort_coeffs = calib.get('distort_coeffs', []) points_2d = [] for x, y, z in points_3d: if z > 0.1: xn, yn = x / z, y / z xd, yd = apply_fisheye_distortion(xn, yn, distort_coeffs) points_2d.append([fx * xd + cx, fy * yd + cy]) else: points_2d.append([np.nan, np.nan]) return np.array(points_2d) def sample_3d_edge(p1, p2, num_samples=10): """Uniformly sample *num_samples* points along the 3D edge from *p1* to *p2*. Args: p1 (array-like): Start point (x, y, z). p2 (array-like): End point (x, y, z). num_samples (int): Number of sample points. Returns: np.ndarray: (num_samples, 3) sampled 3D points. """ t = np.linspace(0, 1, num_samples).reshape(-1, 1) return p1 + t * (p2 - p1) def project_3d_box_edges_with_distortion(corners_3d, calib, samples_per_edge=10): """Project 3D box edges to 2D by sampling, handling fisheye distortion. Args: corners_3d (np.ndarray): (8, 3) 3D corner coordinates. calib (dict): Camera calibration dict. samples_per_edge (int): Number of samples per edge. Returns: dict: Mapping edge_name → (N, 2) 2D projected points. """ edges = { 'back_0': (4, 5), 'back_1': (5, 6), 'back_2': (6, 7), 'back_3': (7, 4), 'connect_0': (0, 4), 'connect_1': (1, 5), 'connect_2': (2, 6), 'connect_3': (3, 7), 'front_0': (0, 1), 'front_1': (1, 2), 'front_2': (2, 3), 'front_3': (3, 0), 'front_x1': (0, 2), 'front_x2': (1, 3), } return { name: project_3d_to_2d_with_distortion( sample_3d_edge(corners_3d[i], corners_3d[j], samples_per_edge), calib ) for name, (i, j) in edges.items() } def plot_box3d_on_img_with_distortion(img, edge_points_2d, color_front=(255, 0, 0), color_back=(0, 0, 255), color_side=(0, 255, 255), thickness=1): """Draw a 3D box on *img* using pre-projected edge point lists (fisheye-aware). Args: img (np.ndarray): BGR image to draw on. edge_points_2d (dict): Output of :func:`project_3d_box_edges_with_distortion`. color_front (tuple): BGR colour for front-face edges. color_back (tuple): BGR colour for back-face edges. color_side (tuple): BGR colour for side connecting edges. thickness (int): Line thickness in pixels. Returns: np.ndarray: Modified image. """ front_edges = {'front_0', 'front_1', 'front_2', 'front_3', 'front_x1', 'front_x2'} back_edges = {'back_0', 'back_1', 'back_2', 'back_3'} for edge_name, points in edge_points_2d.items(): if np.any(np.isnan(points)): continue pts = points.astype(np.int32) if edge_name in front_edges: color = color_front elif edge_name in back_edges: color = color_back else: color = color_side cv2.polylines(img, [pts], isClosed=False, color=color, thickness=thickness, lineType=cv2.LINE_AA) return img def project_3d_to_2d_with_calib(points_3d, calib): """Project 3D points to 2D using pinhole calibration (no distortion). Args: points_3d (np.ndarray): (N, 3) points in camera coordinates. calib (dict): Camera parameters ``fx``, ``fy``, ``cx``, ``cy``. Returns: np.ndarray: (N, 2) image coordinates (NaN for behind-camera points). """ fx, fy = calib['fx'], calib['fy'] cx, cy = calib['cx'], calib['cy'] points_2d = [] for x, y, z in points_3d: if z > 0.1: points_2d.append([fx * x / z + cx, fy * y / z + cy]) else: points_2d.append([np.nan, np.nan]) return np.array(points_2d) def project_3d_to_2d_simple(points_3d, img_size): """Project 3D points to 2D using a simple estimated pinhole model. Args: points_3d (np.ndarray): (N, 3) points in camera coordinates. img_size (tuple[int, int]): ``(width, height)`` of the image. Returns: np.ndarray: (N, 2) image coordinates. """ w, h = img_size fx = fy = w * 1.2 cx, cy = w / 2, h / 2 points_2d = [] for x, y, z in points_3d: if z > 0.1: points_2d.append([fx * x / z + cx, fy * y / z + cy]) else: points_2d.append([np.nan, np.nan]) return np.array(points_2d) def plot_box3d_on_img(img, corners_2d, color_front=(255, 0, 0), color_back=(0, 0, 255), color_side=(0, 255, 255), thickness=1): """Draw a 3D bounding box on *img* from 2D projected corners. Args: img (np.ndarray): BGR image to draw on. corners_2d (np.ndarray): (8, 2) projected corner coordinates. color_front (tuple): BGR colour for front-face edges (indices 0-3). color_back (tuple): BGR colour for back-face edges (indices 4-7). color_side (tuple): BGR colour for connecting side edges. thickness (int): Line thickness in pixels. Returns: np.ndarray: Modified image. """ line_indices = ( (4, 5), (5, 6), (6, 7), (7, 4), # back face (0, 4), (1, 5), (2, 6), (3, 7), # side edges (0, 1), (1, 2), (2, 3), (3, 0), (0, 2), (1, 3), # front face + X mark ) front_edges = {(0, 1), (1, 2), (2, 3), (3, 0), (0, 2), (1, 3)} back_edges = {(4, 5), (5, 6), (6, 7), (7, 4)} corners = corners_2d.astype(np.int32) for start, end in line_indices: try: pt1 = (corners[start, 0], corners[start, 1]) pt2 = (corners[end, 0], corners[end, 1]) if (start, end) in front_edges: cv2.line(img, pt1, pt2, color_front, thickness, cv2.LINE_AA) elif (start, end) in back_edges: cv2.line(img, pt1, pt2, color_back, thickness, cv2.LINE_AA) else: cv2.line(img, pt1, pt2, color_side, thickness, cv2.LINE_AA) except Exception: pass return img # --------------------------------------------------------------------------- # 3D box reconstruction from target label format # --------------------------------------------------------------------------- def _reconstruct_3d_box_from_face(face_uv, face_z, dims, rot_y, face_type, calib): """Reconstruct 3D box corners from a visible face centre. Args: face_uv (tuple[float, float]): Pixel coordinates (u, v) of the face centre. face_z (float): Depth of the face centre in metres. dims (array-like): (length, height, width) in metres. rot_y (float): Yaw rotation in radians. face_type (int): 0=front, 1=rear, 2=left, 3=right. calib (dict): Camera calibration dict. Returns: tuple[np.ndarray, list] | None: ``(corners_3d, object_3d)`` or ``None`` on failure. """ if calib is None: return None fx, fy = calib['fx'], calib['fy'] cx, cy = calib['cx'], calib['cy'] distort_coeffs = calib.get('distort_coeffs', []) u_face, v_face = face_uv xd = (u_face - cx) / fx yd = (v_face - cy) / fy if distort_coeffs and len(distort_coeffs) >= 4: xn, yn = remove_fisheye_distortion(xd, yd, distort_coeffs) else: xn, yn = xd, yd l, h, w = dims if np.isnan(l) or np.isnan(h) or np.isnan(w) or np.isnan(rot_y): return None face_center_3d = np.array([xn * face_z, yn * face_z, face_z]) corners_3d = compute_3d_box_corners_4face(face_center_3d, dims, rot_y, face_type=face_type) object_3d = [face_center_3d[0], face_center_3d[1], face_center_3d[2], l, h, w, rot_y, face_type] return corners_3d, object_3d def decode_and_reconstruct_3d_box_from_target(target, calib, img_width, img_height, face_3d_classes=None, complete_3d_classes=None): """Decode a ground-truth target vector and reconstruct its 3D box. The target array follows the 48-column format used in ``YOLOGround3DDataset``: * col 0 — image index * col 1 — class id * cols 2-5 — normalised 2D bbox [x, y, w, h] * cols 6-8 — 3D centre [x3d, y3d, z3d] in camera coords * cols 9-11 — dimensions [l, h, w] * col 12 — rot_y * cols 13-14 — normalised UV projection of 3D centre * col 15 — alpha * cols 16-23 — front face [x3d, y3d, z3d, alpha, xc, yc, score, visible] * cols 24-31 — rear face (same layout) * cols 32-39 — left face * cols 40-47 — right face Args: target (np.ndarray): (48,) label vector. calib (dict): Camera calibration dict; must contain ``depth_scale``. img_width (int): Image width in pixels. img_height (int): Image height in pixels. face_3d_classes (list | None): Class IDs using face-based reconstruction. Defaults to ``[0, 13]`` (vehicles, tricycles). complete_3d_classes (list | None): Class IDs using complete-box reconstruction. Defaults to ``[1, 2, 3]`` (pedestrians, bicycles, riders). Returns: dict | None: Result dict with keys ``should_draw``, ``cls``, ``corners_3d``, ``face_center_2d``, ``face_color``, ``center_2d``, ``object_3d``; or ``None`` if the target is invalid / unsupported. """ if face_3d_classes is None: face_3d_classes = [0, 1, 2, 3, 4, 5, 6, 7, 8] # vehicles if complete_3d_classes is None: complete_3d_classes = [9, 10, 11, 12] # pedestrian / cyclists if len(target) < 15 or np.isnan(target[1]): return None cls = int(target[1]) if cls not in face_3d_classes and cls not in complete_3d_classes: return None depth_scale = calib['depth_scale'] result = { 'should_draw': True, 'cls': cls, 'corners_3d': None, 'face_center_2d': None, 'face_color': None, 'center_2d': None, } if cls in face_3d_classes: if len(target) < 48: return None face_offsets = [16, 24, 32, 40] # BGR: front=red, rear=blue, left=green, right=yellow face_colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (0, 255, 255)] best_face_type = -1 best_score = -1.0 best_face_data = None for face_type, face_offset in enumerate(face_offsets): if target.shape[0] < face_offset + 8: continue face_data = target[face_offset:face_offset + 8] _, _, z3d_face, _, xc_face, yc_face, score, is_visible = face_data if is_visible == -1 or np.isnan(is_visible) or is_visible != 1: continue if np.isnan(score) or score < 0.3: continue if np.isnan(xc_face) or np.isnan(yc_face) or np.isnan(z3d_face) or z3d_face <= 0: continue if score > best_score: best_score = score best_face_type = face_type best_face_data = face_data if best_face_type != -1 and best_face_data is not None: xc_face = best_face_data[4] yc_face = best_face_data[5] z3d_face = best_face_data[2] * depth_scale u_face = xc_face * img_width v_face = yc_face * img_height dims = target[9:12] rot_y = target[12] result_face = _reconstruct_3d_box_from_face( (u_face, v_face), z3d_face, dims, rot_y, best_face_type, calib ) if result_face is not None: corners_3d, object_3d = result_face result['corners_3d'] = corners_3d result['face_center_2d'] = (u_face, v_face) result['face_color'] = face_colors[best_face_type] result['object_3d'] = object_3d elif cls in complete_3d_classes: x3d, y3d, z3d = target[6:9] dimensions = target[9:12] rot_y = target[12] xc_norm, yc_norm = target[13:15] z3d = z3d * depth_scale if np.isnan(z3d) or z3d <= 0 or np.any(np.isnan(dimensions)): return None if np.isnan(x3d) or np.isnan(y3d): return None corners_3d = compute_3d_box_corners_4face( np.array([x3d, y3d, z3d]), dimensions, rot_y, face_type=-1 ) result['corners_3d'] = corners_3d result['center_2d'] = (xc_norm * img_width, yc_norm * img_height) result['object_3d'] = [ x3d, y3d, z3d, dimensions[0], dimensions[1], dimensions[2], rot_y, -1, ] return result # --------------------------------------------------------------------------- # Drawing helpers # --------------------------------------------------------------------------- def draw_3d_box_from_corners(im, corners_3d, calib, img_shape, face_center_2d=None, face_color=None, thickness=1): """Project and draw a 3D box given raw corner coordinates. Args: im (np.ndarray): BGR image array (H, W, 3). corners_3d (np.ndarray): (8, 3) corners in camera frame (output of :func:`compute_3d_box_corners_4face`). calib (dict | None): Camera calibration dict. img_shape (tuple[int, int]): ``(width, height)`` of *im*. face_center_2d (tuple[float, float] | None): Pixel coords to mark with a dot. face_color (tuple | None): BGR colour for the face-centre dot. thickness (int): Line thickness in pixels. Returns: np.ndarray: Modified image. """ w, h = img_shape # compute_3d_box_corners_4face places rear at indices 0-3 and front at 4-7. # plot_box3d_on_img expects front at indices 0-3 (drawn red) and rear at 4-7. corners_3d = corners_3d[[4, 5, 6, 7, 0, 1, 2, 3]] color_front = (0, 0, 255) # Red (BGR) color_back = (255, 0, 0) # Blue (BGR) color_side = (255, 255, 0) # Cyan (BGR) if calib and calib.get('distort_coeffs'): edge_pts = project_3d_box_edges_with_distortion(corners_3d, calib, samples_per_edge=15) im = plot_box3d_on_img_with_distortion( im, edge_pts, color_front, color_back, color_side, thickness=thickness ) else: corners_2d = (project_3d_to_2d_with_calib(corners_3d, calib) if calib is not None else project_3d_to_2d_simple(corners_3d, (w, h))) if not np.any(np.isnan(corners_2d)): im = plot_box3d_on_img(im, corners_2d, color_front, color_back, color_side, thickness=thickness) if face_center_2d is not None and face_color is not None: cv2.circle(im, (int(face_center_2d[0]), int(face_center_2d[1])), 2, face_color, -1, cv2.LINE_AA) return im def plot_3d_boxes_from_decoded_targets(im, decoded_results, paths, calib=None, names=None, label_text=None, scale_factor=2): """Render 3D boxes from pre-decoded ground-truth targets onto *im*. Args: im (torch.Tensor): Batch tensor (N, 3, H, W) normalised [0, 1] in BGR format (as returned by the 3D dataloader). decoded_results (list[list[dict]]): ``decoded_results[i]`` holds a list of result dicts for image *i*, each produced by :func:`decode_and_reconstruct_3d_box_from_target`. paths (list[str]): Image file paths (informational only). calib (dict | list[dict] | None): Camera calibration dict (or first element of a list) used for 3D→2D projection. names (dict | None): Mapping from class id to class name (unused here but kept for API symmetry). label_text (str | None): Text to overlay in the top-left corner, e.g. ``"3D GT"``. scale_factor (int): Upscale factor applied to the output image. Returns: np.ndarray | None: (H*scale_factor, W*scale_factor, 3) RGB image with 3D boxes drawn, or ``None`` if the input is empty. """ if im.ndim == 3: im = im.unsqueeze(0) im_np = im[0].cpu().numpy().transpose(1, 2, 0) im_np = np.ascontiguousarray(im_np * 255, dtype=np.uint8) h_orig, w_orig = im_np.shape[:2] h_new = h_orig * scale_factor w_new = w_orig * scale_factor im_bgr = cv2.resize(im_np, (w_new, h_new), interpolation=cv2.INTER_LINEAR) if isinstance(calib, list): calib = calib[0] if calib else None calib_scaled = None if calib is not None: calib_scaled = { 'fx': calib['fx'] * scale_factor, 'fy': calib['fy'] * scale_factor, 'cx': calib['cx'] * scale_factor, 'cy': calib['cy'] * scale_factor, 'distort_coeffs': calib.get('distort_coeffs', []), } for decoded in decoded_results[0]: if not decoded or not decoded['should_draw'] or decoded['corners_3d'] is None: continue cls = decoded['cls'] if cls in range(9): # vehicles (car/suv/pickup/...special_vehicle/unknown) — face-based face_center_2d = decoded.get('face_center_2d') if face_center_2d is not None: face_center_2d = ( face_center_2d[0] * scale_factor, face_center_2d[1] * scale_factor, ) im_bgr = draw_3d_box_from_corners( im_bgr, decoded['corners_3d'], calib_scaled, (w_new, h_new), face_center_2d=face_center_2d, face_color=decoded.get('face_color'), ) elif cls in (9, 10, 11, 12): # pedestrian / bicyclists / bicycles / tricycles — complete box im_bgr = draw_3d_box_from_corners( im_bgr, decoded['corners_3d'], calib_scaled, (w_new, h_new), thickness=2, ) if decoded.get('center_2d'): color = colors(cls) u, v = decoded['center_2d'] u_s, v_s = int(u * scale_factor), int(v * scale_factor) cv2.circle(im_bgr, (u_s, v_s), 4, color, -1, cv2.LINE_AA) cv2.circle(im_bgr, (u_s, v_s), 6, (255, 255, 255), 2, cv2.LINE_AA) if label_text: cv2.putText(im_bgr, label_text, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3, cv2.LINE_AA) return cv2.cvtColor(im_bgr, cv2.COLOR_BGR2RGB) # --------------------------------------------------------------------------- # Bird's-Eye View (BEV) visualization # --------------------------------------------------------------------------- def draw_bev_blank(max_range=200): """Create a blank BEV image with distance and lateral grid lines. Args: max_range (int): Maximum forward range in metres (default 200). Returns: np.ndarray: (H, W, 3) BGR image with grid and ego-vehicle box drawn. """ pixels_per_meter = 20 img_height = max_range * pixels_per_meter # e.g. 4000 px for 200 m img_width = 100 * pixels_per_meter # lateral range -50 m … +50 m bevimg = np.ones((img_height, img_width, 3), dtype=np.uint8) * 255 ego_center_x = img_width // 2 ego_center_y = img_height # ego position at bottom-centre ego_half_l = int(4.5 * pixels_per_meter) // 2 ego_half_w = int(1.8 * pixels_per_meter) // 2 ego_box = np.array([ [ego_center_x - ego_half_w, ego_center_y - ego_half_l], [ego_center_x + ego_half_w, ego_center_y - ego_half_l], [ego_center_x + ego_half_w, ego_center_y + ego_half_l], [ego_center_x - ego_half_w, ego_center_y + ego_half_l], ], dtype=np.int32) # Horizontal (range) grid grid_step_px = 20 * pixels_per_meter # every 20 m for i in range(max_range // 20 + 1): y_pos = ego_center_y - i * grid_step_px if y_pos >= 0: cv2.line(bevimg, (0, y_pos), (img_width, y_pos), (180, 180, 180), 3, cv2.LINE_AA) cv2.putText(bevimg, f"{i * 20}m", (ego_center_x + 15, y_pos - 15), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (80, 80, 80), 2, cv2.LINE_AA) # Vertical (lateral) grid lat_step_px = 10 * pixels_per_meter # every 10 m for i in range(-5, 6): x_pos = ego_center_x + i * lat_step_px if 0 <= x_pos < img_width: cv2.line(bevimg, (x_pos, 0), (x_pos, img_height), (180, 180, 180), 3, cv2.LINE_AA) if i != 0: cv2.putText(bevimg, f"{i * 10}m", (x_pos - 40, img_height - 20), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (80, 80, 80), 2, cv2.LINE_AA) cv2.drawContours(bevimg, [ego_box], 0, (255, 0, 0), -1) # blue-filled ego box return bevimg def cam_corners_front_rear(pred3d, facetype): """Compute 8 box corners from a face-centre 3D representation. Args: pred3d (array-like): ``[x, y, z, l, h, w, rot_y]`` (7 values). facetype (int | str): Face type identifier — ``0``/``'front'``, ``1``/``'tail'``, ``2``/``'left'``, ``3``/``'right'``, ``-1``/``'whole'``. Returns: np.ndarray: (8, 3) 3D corner coordinates in camera frame. """ dims = pred3d[3:6] corners_norm = np.stack(np.unravel_index(np.arange(8), [2] * 3), axis=1) corners_norm = corners_norm[[0, 1, 3, 2, 4, 5, 7, 6]].astype(float) if facetype in ('tail', 1): corners_norm -= [0, 0.5, 0.5] elif facetype in ('front', 0): corners_norm -= [1, 0.5, 0.5] elif facetype in ('right', 3): corners_norm -= [0.5, 0.5, 0] elif facetype in ('left', 2): corners_norm -= [0.5, 0.5, 1] elif facetype in ('whole', -1): corners_norm -= [0.5, 0.5, 0.5] else: raise AssertionError(f"Non-valid face type: {facetype}") corners = dims.reshape(1, 3) * corners_norm.reshape(8, 3) corners = rotation_3d_in_axis(corners, pred3d[6], axis=1) corners += pred3d[:3].reshape(1, 3) return corners def drawbev(bevimg, vehicle3d, is_pred=True): """Draw a single vehicle box on a BEV image. Args: bevimg (np.ndarray): BEV image produced by :func:`draw_bev_blank`. vehicle3d (list | np.ndarray): ``[x, y, z, l, h, w, ..., rot_y, face_type]``. The second-to-last element is ``rot_y`` and the last is ``face_type``. is_pred (bool): ``True`` → red box (prediction); ``False`` → green box (GT). Returns: np.ndarray: Modified BEV image. """ x, y, z = vehicle3d[0], vehicle3d[1], vehicle3d[2] l, h, w = vehicle3d[3], vehicle3d[4], vehicle3d[5] rotation_y = vehicle3d[-2] face_type = vehicle3d[-1] pixels_per_meter = 20 max_range = 200 lateral_range = 50 img_height = bevimg.shape[0] img_width = bevimg.shape[1] ego_center_x = img_width // 2 ego_center_y = img_height if x > lateral_range or x < -lateral_range or z > max_range or z < 0: return bevimg corners = cam_corners_front_rear(np.array([x, y, z, l, h, w, rotation_y]), face_type) xyz3d_front = np.mean(corners[4:8, :], axis=0) xyz3d_center = np.mean(corners[0:8, :], axis=0) center = ( int(ego_center_x + xyz3d_center[0] * pixels_per_meter), int(ego_center_y - xyz3d_center[2] * pixels_per_meter), ) front_point = ( int(ego_center_x + xyz3d_front[0] * pixels_per_meter), int(ego_center_y - xyz3d_front[2] * pixels_per_meter), ) rect = (center, (l * pixels_per_meter, w * pixels_per_meter), np.degrees(rotation_y)) box = np.intp(cv2.boxPoints(rect)) color = (0, 0, 255) if is_pred else (0, 255, 0) cv2.drawContours(bevimg, [box], 0, color, 3, cv2.LINE_AA) cv2.arrowedLine(bevimg, center, front_point, color, thickness=3, tipLength=0.3, line_type=cv2.LINE_AA) return bevimg