# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license """Plotting utils.""" import contextlib import math import os from copy import copy from pathlib import Path import cv2 import matplotlib import matplotlib.pyplot as plt import numpy as np import pandas as pd import seaborn as sn import torch from PIL import Image, ImageDraw from scipy.ndimage.filters import gaussian_filter1d from ultralytics.utils.plotting import Annotator from utils import TryExcept, threaded from utils.general import LOGGER, clip_boxes, increment_path, xywh2xyxy, xyxy2xywh, denormalize_3d_predictions from utils.metrics import fitness # Settings RANK = int(os.getenv("RANK", -1)) matplotlib.rc("font", **{"size": 11}) matplotlib.use("Agg") # for writing to files only # Global color map for 2D bounding boxes (RGB format) CLASS_COLOR_MAP = { 0: (255, 100, 100), # vehicle - light red 1: (100, 255, 100), # pedestrian - light green 2: (100, 100, 255), # bicycle - light blue 3: (255, 255, 100), # rider - yellow 4: (255, 100, 255), # motorcycle - magenta 5: (100, 255, 255), # tricycle - cyan 6: (255, 150, 50), # bus - orange 7: (150, 50, 255), # truck - purple 8: (50, 255, 150), # van - teal 9: (255, 200, 100), # car - gold 10: (200, 100, 255), # barrier - lavender 11: (100, 200, 255), # cone - sky blue 12: (255, 100, 200), # sign - pink 13: (200, 255, 100), # tricycle - lime } def get_class_color(cls_id): """Get consistent color for a given class ID. Args: cls_id: Class ID (integer) Returns: RGB color tuple """ return CLASS_COLOR_MAP.get(cls_id, (200, 200, 200)) # Default gray for unknown classes class Colors: """Provides an RGB color palette derived from Ultralytics color scheme for visualization tasks.""" def __init__(self): """Initializes the Colors class with a palette derived from Ultralytics color scheme, converting hex codes to RGB. Colors derived from `hex = matplotlib.colors.TABLEAU_COLORS.values()`. """ 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): """Returns color from palette by index `i`, in BGR format if `bgr=True`, else RGB; `i` is an integer index.""" c = self.palette[int(i) % self.n] return (c[2], c[1], c[0]) if bgr else c @staticmethod def hex2rgb(h): """Converts hexadecimal color `h` to an RGB tuple (PIL-compatible) with order (R, G, B).""" return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4)) colors = Colors() # create instance for 'from utils.plots import colors' def feature_visualization(x, module_type, stage, n=32, save_dir=Path("runs/detect/exp")): """ Args: x: Features to be visualized module_type: Module type stage: Module stage within model n: Maximum number of feature maps to plot save_dir: Directory to save results. """ if ("Detect" not in module_type) and ( "Segment" not in module_type ): # 'Detect' for Object Detect task,'Segment' for Segment task _batch, channels, height, width = x.shape # batch, channels, height, width if height > 1 and width > 1: f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename blocks = torch.chunk(x[0].cpu(), channels, dim=0) # select batch index 0, block by channels n = min(n, channels) # number of plots _fig, ax = plt.subplots(math.ceil(n / 8), 8, tight_layout=True) ax = ax.ravel() plt.subplots_adjust(wspace=0.05, hspace=0.05) for i in range(n): ax[i].imshow(blocks[i].squeeze()) # cmap='gray' ax[i].axis("off") LOGGER.info(f"Saving {f}... ({n}/{channels})") plt.savefig(f, dpi=300, bbox_inches="tight") plt.close() np.save(str(f.with_suffix(".npy")), x[0].cpu().numpy()) # npy save def hist2d(x, y, n=100): """Generates a logarithmic 2D histogram, useful for visualizing label or evolution distributions. Used in used in labels.png and evolve.png. """ xedges, yedges = np.linspace(x.min(), x.max(), n), np.linspace(y.min(), y.max(), n) hist, xedges, yedges = np.histogram2d(x, y, (xedges, yedges)) xidx = np.clip(np.digitize(x, xedges) - 1, 0, hist.shape[0] - 1) yidx = np.clip(np.digitize(y, yedges) - 1, 0, hist.shape[1] - 1) return np.log(hist[xidx, yidx]) def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5): """Applies a low-pass Butterworth filter to `data` with specified `cutoff`, `fs`, and `order`.""" from scipy.signal import butter, filtfilt # https://stackoverflow.com/questions/28536191/how-to-filter-smooth-with-scipy-numpy def butter_lowpass(cutoff, fs, order): """Applies a low-pass Butterworth filter to a signal with specified cutoff frequency, sample rate, and filter order. """ nyq = 0.5 * fs normal_cutoff = cutoff / nyq return butter(order, normal_cutoff, btype="low", analog=False) b, a = butter_lowpass(cutoff, fs, order=order) return filtfilt(b, a, data) # forward-backward filter def output_to_target(output, max_det=300): """Converts YOLOv5 model output to [batch_id, class_id, x, y, w, h, conf] format for plotting, limiting detections to `max_det`. """ targets = [] for i, o in enumerate(output): box, conf, cls = o[:max_det, :6].cpu().split((4, 1, 1), 1) j = torch.full((conf.shape[0], 1), i) targets.append(torch.cat((j, cls, xyxy2xywh(box), conf), 1)) return torch.cat(targets, 0).numpy() @threaded def plot_images(images, targets, paths=None, fname="images.jpg", names=None, is_labels=True): """Plots an image grid with labels from YOLOv5 predictions or targets, saving to `fname`. Args: images: Batch of images targets: Target labels or predictions paths: Image paths fname: Output filename names: Class names is_labels: If True, treats targets as ground truth labels (no conf filtering). If False, treats as predictions (applies conf threshold > 0.25). """ if isinstance(images, torch.Tensor): images = images.cpu().float().numpy().copy() # Make a copy to avoid modifying original if isinstance(targets, torch.Tensor): targets = targets.cpu().numpy() max_size = 1920 # max image size max_subplots = 16 # max image subplots, i.e. 4x4 bs, _, h, w = images.shape # batch size, _, height, width bs = min(bs, max_subplots) # limit plot images ns = np.ceil(bs**0.5) # number of subplots (square) if np.max(images[0]) <= 1: images = images * 255 # de-normalise (create new array, don't modify in place) # Build Image mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8) # init for i, im in enumerate(images): if i == max_subplots: # if last batch has fewer images than we expect break x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin im = im.transpose(1, 2, 0) mosaic[y : y + h, x : x + w, :] = im # Resize (optional) scale = max_size / ns / max(h, w) if scale < 1: h = math.ceil(scale * h) w = math.ceil(scale * w) # Use INTER_AREA for downscaling to avoid artifacts mosaic = cv2.resize(mosaic, tuple(int(x * ns) for x in (w, h)), interpolation=cv2.INTER_AREA) # Annotate fs = int((h + w) * ns * 0.01) # font size # Convert BGR to RGB for Annotator which expects RGB input (pil=True) mosaic = mosaic[..., ::-1] # BGR to RGB annotator = Annotator(mosaic, line_width=round(fs / 10), font_size=fs, pil=True, example=names) for i in range(bs): x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin annotator.rectangle([x, y, x + w, y + h], None, (255, 255, 255), width=2) # borders if paths: annotator.text([x + 5, y + 5], text=Path(paths[i]).name[:40], txt_color=(220, 220, 220)) # filenames if len(targets) > 0: ti = targets[targets[:, 0] == i] # image targets boxes = xywh2xyxy(ti[:, 2:6]).T classes = ti[:, 1].astype("int") conf = None if is_labels else ti[:, 6] # confidence only for predictions if boxes.shape[1]: if boxes.max() <= 1.01: # if normalized with tolerance 0.01 boxes[[0, 2]] *= w # scale to pixels boxes[[1, 3]] *= h elif scale < 1: # absolute coords need scale if image scales boxes *= scale boxes[[0, 2]] += x boxes[[1, 3]] += y for j, box in enumerate(boxes.T.tolist()): cls = classes[j] color = colors(cls) from utils.general import get_class_name cls = get_class_name(names, cls) if names else cls if is_labels or conf[j] > 0.25: # 0.25 conf thresh for predictions label = f"{cls}" if is_labels else f"{cls} {conf[j]:.1f}" annotator.box_label(box, label, color=color) annotator.im.save(fname) # save def rotation_3d_in_axis(points, angles, axis=0): """Rotate points around a specified axis. Args: points: (N, 3) array of 3D points angles: rotation angle in radians axis: 0=X, 1=Y, 2=Z Returns: 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: # Y axis (right-handed system: X=right, Y=down, Z=forward) # Standard right-handed Y-axis rotation matrix (transposed for post-multiplication) # R_y(θ)^T = [cos(θ) 0 -sin(θ)] # [ 0 1 0 ] # [sin(θ) 0 cos(θ)] # Using points @ rot_mat (post-multiplication), so we use the transposed form 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: # Z axis 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: # X axis 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 should in range [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 8 corners of a 3D bounding box from a face center point. The center_3d is the center of a specific face, not the box center. Args: center_3d: (x, y, z) center of the face in camera coordinates dimensions: (length, height, width) of the box rotation: rot_y (rotation around y-axis in radians) face_type: 0=front, 1=tail, 2=left, 3=right Returns: corners: (8, 3) array of corner coordinates Corner ordering: [0,1,3,2,4,5,7,6] pattern from unravel_index """ l, h, w = dimensions # Create 8 corners using unravel_index pattern (same as reference) 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]] # Offset based on face type (face center becomes origin) # face_type: 0=front, 1=tail, 2=left, 3=right # IMPORTANT: In object local frame, X=0 is REAR, X=length is FRONT # After rotation, this aligns correctly with camera frame if face_type == 1: # tail (back/rear face) - at X=0 (REAR in object frame) corners_norm = corners_norm - [0, 0.5, 0.5] elif face_type == 0: # front face - at X=length (FRONT in object frame) corners_norm = corners_norm - [1, 0.5, 0.5] elif face_type == 3: # right face corners_norm = corners_norm - [0.5, 0.5, 0] elif face_type == 2: # left face corners_norm = corners_norm - [0.5, 0.5, 1] else: # whole object center corners_norm = corners_norm - [0.5, 0.5, 0.5] # Scale by dimensions [l, h, w] dims = np.array([l, h, w]) corners = dims.reshape([1, 3]) * corners_norm.reshape([8, 3]) # Rotate around Y axis corners = rotation_3d_in_axis(corners, rotation, axis=1) # Translate to face center position corners += np.array(center_3d).reshape(1, 3) return corners def read_calib_from_path(img_path): """Read calibration parameters from the dataset directory structure. Args: img_path: Path to the image file Returns: dict with calibration parameters or None if not found """ import json base_path = os.path.dirname(img_path) calib_path = base_path.replace('images', 'calib') + '/L2_calib/camera4.json' if not os.path.exists(calib_path): return None try: with open(calib_path, 'r') as f: calib_params = json.load(f) return { 'fx': calib_params['focal_u'], 'fy': calib_params['focal_v'], 'cx': calib_params['cu'], 'cy': calib_params['cv'], 'yaw': calib_params.get('yaw', 0), 'pitch': calib_params.get('pitch', 0), } except Exception: return None def apply_fisheye_distortion(x, y, distort_coeffs): """Apply fisheye distortion to normalized camera coordinates. Args: x, y: normalized camera coordinates (x/z, y/z) distort_coeffs: [k1, k2, k3, k4] fisheye distortion coefficients Returns: xd, yd: distorted normalized coordinates """ 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) # Avoid division by zero if r < 1e-8: return x, y theta = np.arctan(r) # KB (Kannala-Brandt) fisheye distortion model: # theta_d = theta * (1 + k1*theta^2 + k2*theta^4 + k3*theta^6 + k4*theta^8) # For KB model: r_d = theta_d (NOT tan(theta_d)) 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 fisheye distortion from normalized camera coordinates. Args: xd, yd: distorted normalized camera coordinates distort_coeffs: [k1, k2, k3, k4] fisheye distortion coefficients max_iter: maximum iterations for Newton-Raphson solver Returns: xn, yn: undistorted normalized coordinates """ if not distort_coeffs or len(distort_coeffs) < 4: return xd, yd k1, k2, k3, k4 = distort_coeffs[:4] # Compute distorted radius r_d = np.sqrt(xd * xd + yd * yd) # Handle center point if r_d < 1e-8: return xd, yd # For KB model: r_d = theta_d, so theta_d = r_d directly theta_d = r_d # Solve for theta using Newton-Raphson iteration # We need to solve: theta_d = theta * (1 + k1*theta^2 + k2*theta^4 + k3*theta^6 + k4*theta^8) # Better initial guess: use theta_d / (1 + k1*theta_d^2) as starting point theta_d2 = theta_d * theta_d theta = theta_d / (1 + k1*theta_d2) # Better initial guess for _ in range(max_iter): # f(theta) = theta * (1 + k1*theta^2 + k2*theta^4 + k3*theta^6 + k4*theta^8) - theta_d 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'(theta) = 1 + 3*k1*theta^2 + 5*k2*theta^4 + 7*k3*theta^6 + 9*k4*theta^8 f_prime = 1 + 3*k1*theta2 + 5*k2*theta4 + 7*k3*theta6 + 9*k4*theta8 # Newton-Raphson update theta_new = theta - f / f_prime # Check convergence if abs(theta_new - theta) < 1e-8: theta = theta_new break theta = theta_new # For KB model: r = tan(theta) r = np.tan(theta) # Scale back to get undistorted coordinates scale = r / r_d return xd * scale, yd * scale def project_3d_to_2d_with_distortion(points_3d, calib): """Project 3D points to 2D using pre-adjusted camera calibration with fisheye distortion. Args: points_3d: (N, 3) array of 3D points in camera coordinates calib: dict with fx, fy, cx, cy, distort_coeffs (already adjusted for ROI-resized image) Returns: points_2d: (N, 2) array of 2D points in image coordinates """ fx, fy = calib['fx'], calib['fy'] cx, cy = calib['cx'], calib['cy'] distort_coeffs = calib.get('distort_coeffs', []) points_2d = [] for pt in points_3d: x, y, z = pt if z > 0.1: # Only project points in front of camera # Normalize coordinates xn = x / z yn = y / z # Apply fisheye distortion xd, yd = apply_fisheye_distortion(xn, yn, distort_coeffs) # Project to image coordinates (calib already adjusted for ROI-resized image) u = fx * xd + cx v = fy * yd + cy points_2d.append([u, v]) else: points_2d.append([np.nan, np.nan]) return np.array(points_2d) def sample_3d_edge(p1, p2, num_samples=10): """Sample points along a 3D edge between two points. Args: p1, p2: 3D points (x, y, z) num_samples: number of samples along the edge Returns: (num_samples, 3) array of 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 with distortion by sampling along edges. Args: corners_3d: (8, 3) array of 3D corner coordinates calib: dict with fx, fy, cx, cy, distort_coeffs (already adjusted for ROI-resized image) samples_per_edge: number of samples per edge for curved projection Returns: dict with edge_name -> (N, 2) array of 2D points for each edge """ # Define edges: (start_idx, end_idx) # Back face edges (indices 4-7) # Front face edges (indices 0-3) # Connecting edges 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), # X marks on front face } edge_points_2d = {} for edge_name, (i, j) in edges.items(): # Sample 3D points along the edge p1, p2 = corners_3d[i], corners_3d[j] sampled_3d = sample_3d_edge(p1, p2, samples_per_edge) # Project with distortion sampled_2d = project_3d_to_2d_with_distortion(sampled_3d, calib) edge_points_2d[edge_name] = sampled_2d return edge_points_2d 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 3D bounding box on image using edge points (handles distortion). Args: img: image to draw on edge_points_2d: dict with edge_name -> (N, 2) array of 2D points color_front: color for front face edges (default: red) color_back: color for back face edges (default: blue) color_side: color for connecting side edges (default: cyan) thickness: line thickness Returns: img with 3D box drawn """ front_edges = ['front_0', 'front_1', 'front_2', 'front_3', 'front_x1', 'front_x2'] back_edges = ['back_0', 'back_1', 'back_2', 'back_3', 'back_x1', 'back_x2'] for edge_name, points in edge_points_2d.items(): # Skip if any points are invalid if np.any(np.isnan(points)): continue pts = points.astype(np.int32) # Determine color based on edge type if edge_name in front_edges: color = color_front line_thickness = thickness elif edge_name in back_edges: color = color_back line_thickness = thickness else: # Connecting edges (side edges) color = color_side line_thickness = thickness # Draw polyline through sampled points cv2.polylines(img, [pts], isClosed=False, color=color, thickness=line_thickness, lineType=cv2.LINE_AA) return img def project_3d_to_2d_with_calib(points_3d, calib): """Project 3D points to 2D using pre-adjusted camera calibration. Args: points_3d: (N, 3) array of 3D points in camera coordinates calib: dict with fx, fy, cx, cy (already adjusted for ROI-resized image) Returns: points_2d: (N, 2) array of 2D points in image coordinates """ fx, fy = calib['fx'], calib['fy'] cx, cy = calib['cx'], calib['cy'] points_2d = [] for pt in points_3d: x, y, z = pt if z > 0.1: # Only project points in front of camera u = fx * x / z + cx v = fy * y / z + cy points_2d.append([u, v]) 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 simple pinhole camera model. Args: points_3d: (N, 3) array of 3D points in camera coordinates img_size: (width, height) of image Returns: points_2d: (N, 2) array of 2D points in pixel coordinates """ w, h = img_size # Estimate focal length from image size (typical for automotive cameras) # Using a reasonable approximation based on common camera setups fx = fy = w * 1.2 cx, cy = w / 2, h / 2 points_2d = [] for pt in points_3d: x, y, z = pt if z > 0.1: # Only project points in front of camera u = fx * x / z + cx v = fy * y / z + cy points_2d.append([u, v]) 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 3D bounding box on image. Args: img: image to draw on corners_2d: (8, 2) array of 2D corner coordinates Corner order matches compute_3d_box_corners_4face output color_front: color for front face edges (default: red) color_back: color for back face edges (default: blue) color_side: color for connecting side edges (default: cyan) thickness: line thickness Returns: img with 3D box drawn """ # Line indices for drawing the box # Back face (indices 4-7), connecting edges, front face (indices 0-3) line_indices = ( (4, 5), (5, 6), (6, 7), (7, 4), # back face (0, 4), (1, 5), (2, 6), (3, 7), # connecting edges (0, 1), (1, 2), (2, 3), (3, 0), (0, 2), (1, 3) # front face with X ) 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]) # Determine color based on edge type 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: # Connecting side edges cv2.line(img, pt1, pt2, color_side, thickness, cv2.LINE_AA) except Exception: pass return img def convert_predictions_to_target_format(pred_3d, calib, image_size): """Convert 41-channel post-processed predictions to 47-dimension target format. This utility function transforms post-processed model predictions (with UV in absolute pixels and rot_y decoded) to match the target label format, enabling consistent plotting. Args: pred_3d: Post-processed 3D predictions array with shape (N, 41) containing: - 0-5: Front face (z, u_abs, v_abs, l, h, w) - 6-11: Rear face (z, u_abs, v_abs, l, h, w) - 12-17: Left face (z, u_abs, v_abs, l, h, w) - 18-23: Right face (z, u_abs, v_abs, l, h, w) - 24-40: Whole object (z, u_abs, v_abs, l, h, w, yaw_cls×4, rot_y, ..., cut_cls×3) Note: UV coordinates should already be in absolute pixels (after post-processing) Note: rot_y should already be decoded (channel 34 contains rot_y, not yaw_reg) calib: Calibration dict with 'fx', 'fy', 'cx', 'cy' keys image_size: Tuple of (width, height) in pixels for UV normalization Returns: Array with shape (N, 47) matching target label format (without batch index): - 0: class (filled with 0, should be set externally from 2D predictions) - 1-4: [x, y, w, h] 2D bbox (filled with 0, should be set externally) - 5-7: [x3d, y3d, z3d] whole object 3D center - 8-10: [l, h, w] whole object size - 11: rot_y - 12-13: [xc, yc] whole object UV (normalized [0,1]) - 14: alpha (observation angle) - 15-22: front face [x3d, y3d, z3d, alpha, xc, yc, score, visible] - 23-30: rear face - 31-38: left face - 39-46: right face Note: This function expects post-processed predictions (UV in absolute pixels, rot_y decoded). It is used for plotting inference results where predictions have already been post-processed. """ if isinstance(pred_3d, torch.Tensor): pred_3d = pred_3d.cpu().numpy() n = pred_3d.shape[0] out = np.zeros((n, 47), dtype=pred_3d.dtype) # Get calibration parameters fx = calib.get('fx', image_size[0] / 2) fy = calib.get('fy', image_size[1] / 2) cx = calib.get('cx', image_size[0] / 2) cy = calib.get('cy', image_size[1] / 2) # Get image dimensions for UV normalization img_w, img_h = image_size # Helper function: inverse projection (u, v, z) -> (x3d, y3d, z3d) def inverse_project(u, v, z): x3d = (u - cx) * z / fx y3d = (v - cy) * z / fy z3d = z return x3d, y3d, z3d # Helper function: compute alpha from rot_y and position def compute_alpha(rot_y, x3d, z3d): # alpha = rot_y - atan2(x3d, z3d) return rot_y - np.arctan2(x3d, z3d) # --- Whole object (channels 24-40 -> dims 5-14) --- whole_z = pred_3d[:, 24] whole_u = pred_3d[:, 25] # Already in absolute pixels whole_v = pred_3d[:, 26] # Already in absolute pixels whole_l = pred_3d[:, 27] whole_h = pred_3d[:, 28] whole_w = pred_3d[:, 29] rot_y = pred_3d[:, 34] # Already decoded rot_y # Inverse projection for whole object whole_x3d, whole_y3d, whole_z3d = inverse_project(whole_u, whole_v, whole_z) whole_alpha = compute_alpha(rot_y, whole_x3d, whole_z3d) out[:, 5] = whole_x3d out[:, 6] = whole_y3d out[:, 7] = whole_z3d out[:, 8] = whole_l out[:, 9] = whole_h out[:, 10] = whole_w out[:, 11] = rot_y out[:, 12] = whole_u / img_w # Normalize to [0, 1] out[:, 13] = whole_v / img_h # Normalize to [0, 1] out[:, 14] = whole_alpha # --- Front face (channels 0-5 -> dims 15-22) --- front_z = pred_3d[:, 0] front_u = pred_3d[:, 1] # Already in absolute pixels front_v = pred_3d[:, 2] # Already in absolute pixels front_x3d, front_y3d, front_z3d = inverse_project(front_u, front_v, front_z) front_alpha = compute_alpha(rot_y, front_x3d, front_z3d) out[:, 15] = front_x3d out[:, 16] = front_y3d out[:, 17] = front_z3d out[:, 18] = front_alpha out[:, 19] = front_u / img_w # Normalize to [0, 1] out[:, 20] = front_v / img_h # Normalize to [0, 1] out[:, 21] = 1.0 # score (placeholder) out[:, 22] = 1.0 # visible (placeholder) # --- Rear face (channels 6-11 -> dims 23-30) --- rear_z = pred_3d[:, 6] rear_u = pred_3d[:, 7] # Already in absolute pixels rear_v = pred_3d[:, 8] # Already in absolute pixels rear_x3d, rear_y3d, rear_z3d = inverse_project(rear_u, rear_v, rear_z) rear_alpha = compute_alpha(rot_y, rear_x3d, rear_z3d) out[:, 23] = rear_x3d out[:, 24] = rear_y3d out[:, 25] = rear_z3d out[:, 26] = rear_alpha out[:, 27] = rear_u / img_w # Normalize to [0, 1] out[:, 28] = rear_v / img_h # Normalize to [0, 1] out[:, 29] = 1.0 # score out[:, 30] = 1.0 # visible # --- Left face (channels 12-17 -> dims 31-38) --- left_z = pred_3d[:, 12] left_u = pred_3d[:, 13] # Already in absolute pixels left_v = pred_3d[:, 14] # Already in absolute pixels left_x3d, left_y3d, left_z3d = inverse_project(left_u, left_v, left_z) left_alpha = compute_alpha(rot_y, left_x3d, left_z3d) out[:, 31] = left_x3d out[:, 32] = left_y3d out[:, 33] = left_z3d out[:, 34] = left_alpha out[:, 35] = left_u / img_w # Normalize to [0, 1] out[:, 36] = left_v / img_h # Normalize to [0, 1] out[:, 37] = 1.0 # score out[:, 38] = 1.0 # visible # --- Right face (channels 18-23 -> dims 39-46) --- right_z = pred_3d[:, 18] right_u = pred_3d[:, 19] # Already in absolute pixels right_v = pred_3d[:, 20] # Already in absolute pixels right_x3d, right_y3d, right_z3d = inverse_project(right_u, right_v, right_z) right_alpha = compute_alpha(rot_y, right_x3d, right_z3d) out[:, 39] = right_x3d out[:, 40] = right_y3d out[:, 41] = right_z3d out[:, 42] = right_alpha out[:, 43] = right_u / img_w # Normalize to [0, 1] out[:, 44] = right_v / img_h # Normalize to [0, 1] out[:, 45] = 1.0 # score out[:, 46] = 1.0 # visible return out def _reconstruct_3d_box_from_face(face_uv, face_z, dims, rot_y, face_type, calib): """Reconstruct 3D box corners from visible face center. Args: face_uv: (u, v) 2D pixel coordinates of face center face_z: Depth of face center in meters dims: (l, h, w) object dimensions in meters rot_y: Rotation yaw angle in radians face_type: 0=front, 1=rear, 2=left, 3=right calib: Calibration dict with fx, fy, cx, cy, distort_coeffs Returns: corners_3d: (8, 3) array of 3D corner coordinates in camera frame or None if invalid """ 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 # Back-project face center to 3D # Normalize pixel coordinates xd = (u_face - cx) / fx yd = (v_face - cy) / fy # Remove fisheye distortion if present if distort_coeffs and len(distort_coeffs) >= 4: xn, yn = remove_fisheye_distortion(xd, yd, distort_coeffs) else: xn, yn = xd, yd # Back-project to 3D x3d_face = xn * face_z y3d_face = yn * face_z z3d_face = face_z # Extract dimensions l, h, w = dims if np.isnan(l) or np.isnan(h) or np.isnan(w) or np.isnan(rot_y): return None # Use existing helper function that takes face center as input face_center_3d = np.array([x3d_face, y3d_face, z3d_face]) corners_3d = compute_3d_box_corners_4face( face_center_3d, dims, rot_y, face_type=face_type ) object_3d = [x3d_face, y3d_face, z3d_face, l, h, w, rot_y, face_type] # Return corners directly without reordering 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 and reconstruct 3D box from ground truth target format. Args: target: (48,) array in target format: - col 0: image index - col 1: class - cols 2-5: [x, y, w, h] normalized 2D bbox - cols 6-8: [x3d, y3d, z3d] 3D center in camera coords - cols 9-11: [length, height, width] 3D dimensions - col 12: rot_y (yaw rotation) - cols 13-14: [xc, yc] 2D center of 3D box projection (normalized) - col 15: alpha - cols 16-23: front face [x3d, y3d, z3d, alpha, xc, yc, score, visible] - cols 24-31: rear face - cols 32-39: left face - cols 40-47: right face calib: Calibration dict img_width: Image width in pixels img_height: Image height in pixels face_3d_classes: List/set of class IDs that use face-based reconstruction (default: [0, 1, 2, 3, 4, 5]) complete_3d_classes: List/set of class IDs that use complete 3D box (default: [6, 7, 8, 9]) Returns: dict or None: { 'should_draw': bool, 'cls': int, 'corners_3d': np.ndarray (8, 3) or None, 'face_center_2d': tuple or None, # For face-based classes 'face_color': tuple or None, # For face-based classes 'center_2d': tuple or None, # For complete box classes } """ # Default class lists: face-based reconstruction for vehicles (0) and tricycles (13), # complete 3D box for pedestrian (1), bicycle (2), rider (3) if face_3d_classes is None: face_3d_classes = [0, 13] if complete_3d_classes is None: complete_3d_classes = [1, 2, 3] if len(target) < 15 or np.isnan(target[1]): return None cls = int(target[1]) # Skip unsupported classes (not in either list) if cls not in face_3d_classes and cls not in complete_3d_classes: return None # Get depth_scale for converting normalized depth back to original calibration space # In ROI mode with target_fx, xyz3d are scaled by target_fx/fx_orig during training # depth_scale = fx_orig/target_fx converts them back depth_scale = calib['depth_scale'] # Required field, no default 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: # Classes with face annotations - use best visible face if len(target) < 48: return None # Find best face face_offsets = [16, 24, 32, 40] # Face colors in BGR format for OpenCV: 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 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] x3d_face, y3d_face, z3d_face, alpha, xc_face, yc_face, score, is_visible = face_data # Skip invalid faces 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 # Track best face by score if score > best_score: best_score = score best_face_type = face_type best_face_data = face_data # Reconstruct from best face 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] # Scale z back to original calibration space (from normalized depth) z3d_face = z3d_face * depth_scale # Convert normalized UV to absolute pixels u_face = xc_face * img_width v_face = yc_face * img_height # Get dimensions and rotation dims = target[9:12] # [l, h, w] rot_y = target[12] corners_3d, object_3d = _reconstruct_3d_box_from_face( (u_face, v_face), z3d_face, dims, rot_y, best_face_type, calib ) if corners_3d is not None: # No reordering here - will be done in visualization 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: # Classes with complete 3D box but no face annotations # Extract 3D box parameters x3d, y3d, z3d = target[6:9] dimensions = target[9:12] # [l, h, w] rot_y = target[12] xc_norm, yc_norm = target[13:15] # Scale only z3d back to original calibration space (from normalized depth) # X and Y don't need scaling - only Z is normalized for training z3d = z3d * depth_scale # Validate 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 # Compute 3D box corners center_3d = np.array([x3d, y3d, z3d]) corners_3d = compute_3d_box_corners_4face( center_3d, dimensions, rot_y, face_type=-1 ) # NOTE: GT does NOT need corner reordering (unlike predictions) # This suggests GT rot_y convention is different from predicted rot_y result['corners_3d'] = corners_3d result['center_2d'] = (xc_norm * img_width, yc_norm * img_height) # Create object_3d for BEV visualization: [x, y, z, l, h, w, rot_y, face_type] result['object_3d'] = [x3d, y3d, z3d, dimensions[0], dimensions[1], dimensions[2], rot_y, -1] return result def decode_and_reconstruct_3d_box(pred_2d_box, pred_3d_box, calib, conf_thres, face_3d_classes=None, complete_3d_classes=None): """Decode and reconstruct 3D box from compact 34-channel predictions. Args: pred_2d_box: (6,) array [x1, y1, x2, y2, conf, cls] pred_3d_box: (34,) array in compact format calib: Calibration dict conf_thres: Confidence threshold face_3d_classes: List/set of class IDs that use face-based reconstruction (default: [0, 1, 2, 3, 4, 5]) complete_3d_classes: List/set of class IDs that use complete 3D box (default: [6, 7, 8, 9]) Returns: dict or None: { 'should_draw': bool, 'cls': int, 'corners_3d': np.ndarray (8, 3) or None, 'face_center_2d': tuple or None, 'face_color': tuple or None, } """ # Default class lists: face-based reconstruction for vehicles (0) and tricycles (13), # complete 3D box for pedestrian (1), bicycle (2), rider (3) if face_3d_classes is None: face_3d_classes = [0, 13] if complete_3d_classes is None: complete_3d_classes = [1, 2, 3] conf = pred_2d_box[4] if conf < conf_thres: return None cls = int(pred_2d_box[5]) # Skip unsupported classes (not in either list) if cls not in face_3d_classes and cls not in complete_3d_classes: return None # NOTE: depth_scale is NOT applied here for predictions because it's already # applied in the model forward pass (models/yolo.py:436) # Only targets need depth_scale applied (in decode_and_reconstruct_3d_box_from_target) result = { 'should_draw': True, 'cls': cls, 'corners_3d': None, 'face_center_2d': None, 'face_color': None, 'cut_cls': None, } if cls in face_3d_classes: # Classes with face annotations - face-based reconstruction face_offsets = [0, 6, 12, 18] # Face colors in BGR format for OpenCV: front=red, rear=blue, left=green, right=yellow face_colors = [(0, 0, 255), (255, 0, 0), (0, 255, 0), (0, 255, 255)] # Parse cut_cls # NOTE: Decoding method MUST match training loss type: # - BCE loss training -> use sigmoid for independent class probabilities # - CrossEntropy loss training -> use softmax for normalized probabilities # Current default: softmax (matches CrossEntropy) cut_cls_logits = pred_3d_box[31:34] # Method 1: Softmax (for CrossEntropy trained models) - CURRENT DEFAULT # cut_cls_probs = np.exp(cut_cls_logits) / np.sum(np.exp(cut_cls_logits)) # Method 2: Sigmoid (for BCE trained models) - UNCOMMENT if using BCE loss cut_cls_probs = 1 / (1 + np.exp(-cut_cls_logits)) cut_cls_pred = np.argmax(cut_cls_probs) # Determine candidate faces if cut_cls_pred == 1: # cut_in candidate_faces = [0] elif cut_cls_pred == 2: # cut_out candidate_faces = [1] else: # normal candidate_faces = [0, 1, 2, 3] # Select best face best_face_type = -1 best_visible_score = -1.0 for face_type in candidate_faces: face_offset = face_offsets[face_type] z_face = pred_3d_box[face_offset] u_face = pred_3d_box[face_offset + 1] v_face = pred_3d_box[face_offset + 2] visible_score = pred_3d_box[face_offset + 5] if np.isnan(z_face) or z_face <= 0 or np.isnan(u_face) or np.isnan(v_face): continue if np.isnan(visible_score): continue if visible_score > best_visible_score: best_visible_score = visible_score best_face_type = face_type if best_face_type == -1: return None # Reconstruct 3D box face_offset = face_offsets[best_face_type] z_face = pred_3d_box[face_offset] u_face = pred_3d_box[face_offset + 1] v_face = pred_3d_box[face_offset + 2] dims = pred_3d_box[27:30] # [l, h, w] rot_y = pred_3d_box[30] # z is already scaled in model forward pass, use directly corners_3d, object_3d = _reconstruct_3d_box_from_face( (u_face, v_face), z_face, dims, rot_y, best_face_type, calib ) if corners_3d is not None: # No reordering here - will be done in visualization 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 result['cut_cls'] = cut_cls_pred else: return None elif cls in complete_3d_classes: # Classes with complete 3D box but no face annotations # Extract whole box info z_whole = pred_3d_box[24] u_whole = pred_3d_box[25] v_whole = pred_3d_box[26] l, h, w = pred_3d_box[27:30] rot_y = pred_3d_box[30] # z is already scaled in model forward pass, use directly if calib is None: return None fx, fy = calib['fx'], calib['fy'] cx, cy = calib['cx'], calib['cy'] distort_coeffs = calib.get('distort_coeffs', []) # Back-project to 3D xd = (u_whole - cx) / fx yd = (v_whole - cy) / fy if distort_coeffs and len(distort_coeffs) >= 4: xn, yn = remove_fisheye_distortion(xd, yd, distort_coeffs) else: xn, yn = xd, yd x_center = xn * z_whole y_center = yn * z_whole # Reconstruct 8 corners corners_3d = compute_3d_box_corners_4face( (x_center, y_center, z_whole), (l, h, w), rot_y, face_type=-1 ) result['corners_3d'] = corners_3d result['object_3d'] = [x_center, y_center, z_whole, l, h, w, rot_y, -1] return result def draw_3d_box_from_corners(im, corners_3d, calib, img_shape, face_center_2d=None, face_color=None, thickness=1): """Project and draw 3D box from corner coordinates. Args: im: Image array (H, W, 3) BGR corners_3d: (8, 3) 3D corner coordinates in camera frame calib: Calibration dict img_shape: (w, h) face_center_2d: (u, v) or None, face center marker for face-based face_color: tuple or None, face center marker color thickness: line thickness Returns: im: Modified image with 3D box drawn """ w, h = img_shape # Reorder corners: compute_3d_box_corners_4face outputs indices 0-3 at rear, 4-7 at front # plot_box3d_on_img expects indices 0-3 = front (red), indices 4-7 = rear (blue) # So we swap: [4,5,6,7,0,1,2,3] -> indices 0-3 = front, indices 4-7 = rear corners_3d = corners_3d[[4, 5, 6, 7, 0, 1, 2, 3]] # Colors in BGR format for OpenCV: # plot_box3d_on_img draws indices 0-3 with color_front and 4-7 with color_back color_front = (0, 0, 255) # Red in BGR (B=0, G=0, R=255) for front/head (indices 0-3) color_back = (255, 0, 0) # Blue in BGR (B=255, G=0, R=0) for rear/back (indices 4-7) color_side = (255, 255, 0) # Cyan in BGR (B=255, G=255, R=0) for side edges # Project to 2D if calib and calib.get('distort_coeffs'): edge_points_2d = project_3d_box_edges_with_distortion( corners_3d, calib, samples_per_edge=15 ) im = plot_box3d_on_img_with_distortion( im, edge_points_2d, color_front, color_back, color_side, thickness=thickness ) else: if calib is not None: corners_2d = project_3d_to_2d_with_calib(corners_3d, calib) else: corners_2d = 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 ) # Draw face center marker if provided if face_center_2d is not None and face_color is not None: u, v = face_center_2d cv2.circle(im, (int(u), int(v)), 2, face_color, -1, cv2.LINE_AA) return im def plot_3d_boxes_from_decoded(im, decoded_results, paths, calib=None, names=None, label_text=None, scale_factor=2): """Plot 3D boxes from pre-decoded results (inference script layer). This function only handles projection and drawing, not decoding/reconstruction. The 3D decoding should be done at the inference script level before calling this. Args: im: Batch tensor (N, 3, H, W) normalized [0, 1] in BGR format (from dataloader3d) decoded_results: List of lists, where decoded_results[i] contains decoded results for image i. Each decoded result is a dict from decode_and_reconstruct_3d_box: { 'should_draw': bool, 'cls': int, 'corners_3d': np.ndarray (8, 3) or None, 'face_center_2d': tuple or None, 'face_color': tuple or None, } paths: List of image paths calib: Calibration dict or list of calib dicts names: Class names dict label_text: Optional text to display on top-left (e.g., "3D GT" or "3D Pred") scale_factor: Scale factor for output image resolution (default 2 for clearer images) Returns: np.ndarray: (H*scale_factor, W*scale_factor, 3) RGB image with 3D boxes drawn, or None """ import cv2 if im.ndim == 3: im = im.unsqueeze(0) im_np = im[0].cpu().numpy().transpose(1, 2, 0) im_np = np.ascontiguousarray(im_np * 256, dtype=np.uint8) # Input is BGR from dataloader3d, keep as BGR for cv2 drawing h_orig, w_orig = im_np.shape[:2] # Upscale image for clearer visualization 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) # Handle calibration - scale calibration parameters for upscaled image if isinstance(calib, list): calib = calib[0] if len(calib) > 0 else None # Create scaled calibration for drawing 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', []), } else: calib_scaled = None # Draw each decoded 3D box for decoded in decoded_results[0]: # First image in batch if not decoded or not decoded['should_draw']: continue if decoded['corners_3d'] is not None: corners_3d = decoded['corners_3d'] # Scale face_center_2d if present 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) # No reordering here - draw_3d_box_from_corners handles it for all classes im_bgr = draw_3d_box_from_corners( im_bgr, corners_3d, calib_scaled, (w_new, h_new), face_center_2d=face_center_2d, face_color=decoded.get('face_color') ) # Draw label text on top-left (after drawing boxes so it's on top) if label_text: cv2.putText(im_bgr, label_text, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3, cv2.LINE_AA) # Convert back to RGB im_rgb = cv2.cvtColor(im_bgr, cv2.COLOR_BGR2RGB) return im_rgb def plot_3d_boxes_from_decoded_targets(im, decoded_results, paths, calib=None, names=None, label_text=None, scale_factor=2): """Plot 3D boxes from pre-decoded ground truth targets. This function only handles projection and drawing, not decoding/reconstruction. The 3D decoding should be done at the script level before calling this. Args: im: Batch tensor (N, 3, H, W) normalized [0, 1] in BGR format (from dataloader3d) decoded_results: List of lists, where decoded_results[i] contains decoded results for image i. Each decoded result is a dict from decode_and_reconstruct_3d_box_from_target: { 'should_draw': bool, 'cls': int, 'corners_3d': np.ndarray (8, 3) or None, 'face_center_2d': tuple or None, # For vehicles/tricycles 'face_color': tuple or None, # For vehicles/tricycles 'center_2d': tuple or None, # For pedestrians/bicycles/riders } paths: List of image paths calib: Calibration dict or list of calib dicts names: Class names dict label_text: Optional text to display on top-left (e.g., "3D GT" or "3D Pred") scale_factor: Scale factor for output image resolution (default 2 for clearer images) Returns: np.ndarray: (H*scale_factor, W*scale_factor, 3) RGB image with 3D boxes drawn, or None """ import cv2 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) # Input is BGR from dataloader3d, keep as BGR for cv2 drawing h_orig, w_orig = im_np.shape[:2] # Upscale image for clearer visualization 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) # Handle calibration - scale calibration parameters for upscaled image if isinstance(calib, list): calib = calib[0] if len(calib) > 0 else None # Create scaled calibration for drawing 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', []), } else: calib_scaled = None # Draw each decoded 3D box for decoded in decoded_results[0]: # First image in batch if not decoded or not decoded['should_draw']: continue if decoded['corners_3d'] is not None: cls = decoded['cls'] if cls in [0, 13]: # vehicles/tricycles - face-based # Scale face_center_2d if present 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) # Draw 3D box from corners 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 [1, 2, 3]: # pedestrians/bicycles/riders - complete box # Use centralized draw function (handles reordering and colors) im_bgr = draw_3d_box_from_corners( im_bgr, decoded['corners_3d'], calib_scaled, (w_new, h_new), thickness=2 ) # Draw center point (scaled) if decoded.get('center_2d'): color = colors(cls) u, v = decoded['center_2d'] u_scaled, v_scaled = int(u * scale_factor), int(v * scale_factor) cv2.circle(im_bgr, (u_scaled, v_scaled), 4, color, -1, cv2.LINE_AA) cv2.circle(im_bgr, (u_scaled, v_scaled), 6, (255, 255, 255), 2, cv2.LINE_AA) # Draw label text on top-left (after drawing boxes so it's on top) if label_text: cv2.putText(im_bgr, label_text, (10, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.5, (0, 255, 255), 3, cv2.LINE_AA) # Convert back to RGB im_rgb = cv2.cvtColor(im_bgr, cv2.COLOR_BGR2RGB) return im_rgb def _plot_images_3d_core(images, targets, paths=None, fname="images.jpg", names=None, calibs=None, return_array=False, face_3d_classes=None, complete_3d_classes=None): """Core implementation for plotting 2x2 grid: target 2D, target 3D, prediction 2D, prediction 3D. Layout for TensorBoard logging (single sample): (1,1) Target 2D | (1,2) Target 3D (2,1) Prediction 2D | (2,2) Prediction 3D Args: images: Batch of images (B, C, H, W) in BGR format (from dataloader3d) targets: Target tensor with 48 columns per target paths: List of image paths fname: Output filename names: Class names dict calibs: Tuple of calibration dicts from dataloader (pre-adjusted for ROI-resized image) return_array: If True, return numpy array instead of saving to file predictions_2d: List of 2D predictions per image (optional, for combined view) predictions_3d: 3D predictions (optional, for combined view) conf_thres: Confidence threshold for predictions face_3d_classes: List/set of class IDs that use face-based reconstruction complete_3d_classes: List/set of class IDs that use complete 3D box """ if isinstance(images, torch.Tensor): images = images.cpu().float().numpy() if isinstance(targets, torch.Tensor): targets = targets.cpu().numpy() # Only plot first sample for 2x2 grid _, _, h, w = images.shape if np.max(images[0]) <= 1: images = images * 255 # de-normalise # Convert BGR to RGB (dataloader3d returns BGR format) images = images[:, ::-1, :, :] # Reverse channel order: BGR -> RGB # Get first image im = images[0].transpose(1, 2, 0).astype(np.uint8) im_target_2d = im.copy() im_target_3d = im.copy() # Get targets for first image ti = targets[targets[:, 0] == 0] if len(targets) > 0 else np.array([]) # Get calibration for first image calib = None if calibs is not None and len(calibs) > 0: calib = calibs[0] elif paths: calib = read_calib_from_path(paths[0]) # --- Draw Target 2D boxes --- for t in ti: if len(t) < 6 or np.isnan(t[1]): continue cls = int(t[1]) color = get_class_color(cls) x_center, y_center, bw, bh = t[2:6] if np.any(np.isnan([x_center, y_center, bw, bh])): continue x1 = int((x_center - bw / 2) * w) y1 = int((y_center - bh / 2) * h) x2 = int((x_center + bw / 2) * w) y2 = int((y_center + bh / 2) * h) cv2.rectangle(im_target_2d, (x1, y1), (x2, y2), color, 1) # --- Draw Target 3D boxes --- # Resolve defaults here (same values used inside decode_and_reconstruct_3d_box_from_target) _face_3d_classes = face_3d_classes if face_3d_classes is not None else [0, 13] _complete_3d_classes = complete_3d_classes if complete_3d_classes is not None else [1, 2, 3] for t in ti: if len(t) < 15 or np.isnan(t[1]): continue cls = int(t[1]) decoded = decode_and_reconstruct_3d_box_from_target(t, calib, w, h, _face_3d_classes, _complete_3d_classes) if decoded and decoded['should_draw'] and decoded['corners_3d'] is not None: if cls in _face_3d_classes: # Face-based classes (vehicles/tricycles) im_target_3d = draw_3d_box_from_corners( im_target_3d, decoded['corners_3d'], calib, (w, h), face_center_2d=decoded.get('face_center_2d'), face_color=decoded.get('face_color') ) elif cls in _complete_3d_classes: # Complete box classes (pedestrians/bicycles/riders) corners_3d = decoded['corners_3d'] color = colors(cls) if calib and calib.get('distort_coeffs'): corners_3d_reordered = corners_3d[[4, 5, 6, 7, 0, 1, 2, 3]] edge_points_2d = project_3d_box_edges_with_distortion( corners_3d_reordered, calib, samples_per_edge=15 ) im_target_3d = plot_box3d_on_img_with_distortion( im_target_3d, edge_points_2d, (255, 0, 0), (0, 0, 255), (0, 255, 255), thickness=1 ) else: if calib is not None: corners_2d = project_3d_to_2d_with_calib(corners_3d, calib) else: corners_2d = project_3d_to_2d_simple(corners_3d, (w, h)) if not np.any(np.isnan(corners_2d)): im_target_3d = plot_box3d_on_img( im_target_3d, corners_2d, (255, 0, 0), (0, 0, 255), (0, 255, 255), thickness=1 ) if decoded.get('center_2d'): u, v = decoded['center_2d'] cv2.circle(im_target_3d, (int(u), int(v)), 3, color, -1, cv2.LINE_AA) # Create 2x2 figure: target 2D, target 3D (top row) fig, axes = plt.subplots(2, 2, figsize=(12, 8)) # Add labels cv2.putText(im_target_2d, 'Target 2D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2) cv2.putText(im_target_3d, 'Target 3D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2) axes[0, 0].imshow(im_target_2d) axes[0, 0].set_title('Target 2D', fontsize=12) axes[0, 0].axis('off') axes[0, 1].imshow(im_target_3d) axes[0, 1].set_title('Target 3D', fontsize=12) axes[0, 1].axis('off') # Bottom row: placeholder for predictions (empty images with text) im_pred_2d = im.copy() im_pred_3d = im.copy() cv2.putText(im_pred_2d, 'Prediction 2D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2) cv2.putText(im_pred_3d, 'Prediction 3D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2) axes[1, 0].imshow(im_pred_2d) axes[1, 0].set_title('Prediction 2D', fontsize=12) axes[1, 0].axis('off') axes[1, 1].imshow(im_pred_3d) axes[1, 1].set_title('Prediction 3D', fontsize=12) axes[1, 1].axis('off') if paths: fig.suptitle(Path(paths[0]).name[:50], fontsize=10) plt.tight_layout() if return_array: dpi = 150 fig.set_dpi(dpi) fig.canvas.draw() buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8) fig_w, fig_h = fig.canvas.get_width_height() img_array = buf.reshape(fig_h, fig_w, 4) img_array = img_array[:, :, :3] plt.close(fig) return img_array else: # Use fig.savefig (not plt.savefig) to avoid matplotlib global-state race condition. # plt.savefig() saves the global "current figure", which can be overwritten by another # thread calling plt.subplots() concurrently (ni=0,1,2 all trigger threads in quick # succession). fig.savefig() operates on the local figure object and is thread-safe. fig.savefig(fname, dpi=150, bbox_inches='tight') plt.close(fig) @threaded def _plot_images_3d_threaded(images, targets, paths=None, fname="images.jpg", names=None, calibs=None, face_3d_classes=None, complete_3d_classes=None): """Threaded wrapper for plot_images_3d when saving to file.""" _plot_images_3d_core(images, targets, paths, fname, names, calibs, return_array=False, face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes) def plot_images_3d(images, targets, paths=None, fname="images.jpg", names=None, calibs=None, return_array=False, face_3d_classes=None, complete_3d_classes=None): """Plots 2D and 3D labels in separate subplots for visualization. When return_array=False (default): runs in a separate thread like plot_images, returns Thread. When return_array=True: runs synchronously and returns numpy array for TensorBoard. Args: face_3d_classes: List/set of class IDs that use face-based reconstruction complete_3d_classes: List/set of class IDs that use complete 3D box """ if return_array: # Synchronous call - return numpy array directly return _plot_images_3d_core(images, targets, paths, fname, names, calibs, return_array=True, face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes) else: # Threaded call - same behavior as plot_images return _plot_images_3d_threaded(images, targets, paths, fname, names, calibs, face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes) def plot_predictions_3d(images, predictions_2d, predictions_3d=None, paths=None, fname="predictions.jpg", names=None, calibs=None, conf_thres=0.25, return_array=False, face_3d_classes=None, complete_3d_classes=None): """Plots 3D predictions on images. Args: images: Batch of images (B, C, H, W) predictions_2d: 2D detection predictions (after NMS) - list of tensors per image, each (N, 6) [x1, y1, x2, y2, conf, cls] predictions_3d: 3D predictions tensor with shape (B, N, 34) or list of (N, 34) arrays per image paths: Image paths fname: Output filename names: Class names calibs: Calibration parameters for each image conf_thres: Confidence threshold for filtering predictions return_array: If True, return numpy array instead of saving to file face_3d_classes: List/set of class IDs that use face-based reconstruction complete_3d_classes: List/set of class IDs that use complete 3D box Returns: Thread object (when running in background) or numpy array if return_array=True """ if return_array: # Synchronous call - return numpy array directly return _plot_predictions_3d_core(images, predictions_2d, predictions_3d, paths, fname, names, calibs, conf_thres, return_array=True, face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes) else: # Threaded call return _plot_predictions_3d_threaded(images, predictions_2d, predictions_3d, paths, fname, names, calibs, conf_thres, face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes) def _plot_predictions_3d_core(images, predictions_2d, predictions_3d=None, paths=None, fname="predictions.jpg", names=None, calibs=None, conf_thres=0.25, return_array=False, face_3d_classes=None, complete_3d_classes=None): """Core implementation for plotting 2x2 grid with predictions: target 2D, target 3D, prediction 2D, prediction 3D. Layout for TensorBoard logging (single sample): (1,1) Target 2D | (1,2) Target 3D (2,1) Prediction 2D | (2,2) Prediction 3D Note: This function only draws predictions. For targets, use plot_val_comparison_3d. Args: images: Batch of images (B, C, H, W) in BGR format (from dataloader3d) predictions_2d: List of 2D predictions per image, each (N, 6) [x1, y1, x2, y2, conf, cls] predictions_3d: 3D predictions - list of (N, 34) arrays per image paths: List of image paths fname: Output filename names: Class names dict calibs: Tuple of calibration dicts from dataloader conf_thres: Confidence threshold return_array: If True, return numpy array instead of saving to file """ if isinstance(images, torch.Tensor): images = images.cpu().float().numpy() # Only plot first sample for 2x2 grid _, _, h, w = images.shape if np.max(images[0]) <= 1: images = images * 255 # de-normalise # Convert BGR to RGB (dataloader3d returns BGR format) images = images[:, ::-1, :, :] # Reverse channel order: BGR -> RGB # Get first image im = images[0].transpose(1, 2, 0).astype(np.uint8) im_pred_2d = im.copy() im_pred_3d = im.copy() # Get calibration for first image calib = None if calibs is not None and len(calibs) > 0: calib = calibs[0] elif paths: calib = read_calib_from_path(paths[0]) # Get predictions for first image pred_2d = None if predictions_2d is not None and len(predictions_2d) > 0: pred_2d = predictions_2d[0] if isinstance(pred_2d, torch.Tensor): pred_2d = pred_2d.cpu().numpy() pred_3d = None if predictions_3d is not None: if isinstance(predictions_3d, list) and len(predictions_3d) > 0: pred_3d = predictions_3d[0] if isinstance(pred_3d, torch.Tensor): pred_3d = pred_3d.cpu().numpy() elif isinstance(predictions_3d, torch.Tensor): pred_3d = predictions_3d[0].cpu().numpy() elif isinstance(predictions_3d, np.ndarray) and len(predictions_3d) > 0: pred_3d = predictions_3d[0] # --- Draw Prediction 2D boxes --- if pred_2d is not None: for pred in pred_2d: if len(pred) < 6: continue x1, y1, x2, y2, conf, cls = pred[:6] cls = int(cls) color = get_class_color(cls) cv2.rectangle(im_pred_2d, (int(x1), int(y1)), (int(x2), int(y2)), color, 1) # --- Draw Prediction 3D boxes --- if pred_2d is not None and pred_3d is not None and len(pred_3d) > 0: for j, pred_2d_box in enumerate(pred_2d): if len(pred_2d_box) < 6 or j >= len(pred_3d): continue box_info = decode_and_reconstruct_3d_box( pred_2d_box, pred_3d[j], calib, conf_thres, face_3d_classes, complete_3d_classes ) if box_info is None or not box_info['should_draw']: continue if box_info['corners_3d'] is None: continue im_pred_3d = draw_3d_box_from_corners( im_pred_3d, box_info['corners_3d'], calib, (w, h), face_center_2d=box_info['face_center_2d'], face_color=box_info['face_color'] ) # Create 2x2 figure with placeholders for targets (top row) fig, axes = plt.subplots(2, 2, figsize=(12, 8)) # Top row: placeholder for targets im_target_2d = im.copy() im_target_3d = im.copy() cv2.putText(im_target_2d, 'Target 2D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2) cv2.putText(im_target_3d, 'Target 3D (N/A)', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (128, 128, 128), 2) axes[0, 0].imshow(im_target_2d) axes[0, 0].set_title('Target 2D', fontsize=12) axes[0, 0].axis('off') axes[0, 1].imshow(im_target_3d) axes[0, 1].set_title('Target 3D', fontsize=12) axes[0, 1].axis('off') # Bottom row: predictions cv2.putText(im_pred_2d, 'Prediction 2D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2) cv2.putText(im_pred_3d, 'Prediction 3D', (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (255, 255, 0), 2) axes[1, 0].imshow(im_pred_2d) axes[1, 0].set_title('Prediction 2D', fontsize=12) axes[1, 0].axis('off') axes[1, 1].imshow(im_pred_3d) axes[1, 1].set_title('Prediction 3D', fontsize=12) axes[1, 1].axis('off') if paths: fig.suptitle(Path(paths[0]).name[:50], fontsize=10) plt.tight_layout() if return_array: dpi = 150 fig.set_dpi(dpi) fig.canvas.draw() buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8) fig_w, fig_h = fig.canvas.get_width_height() img_array = buf.reshape(fig_h, fig_w, 4) img_array = img_array[:, :, :3] plt.close(fig) return img_array else: plt.savefig(fname, dpi=150, bbox_inches='tight') plt.close(fig) @threaded def _plot_predictions_3d_threaded(images, predictions_2d, predictions_3d=None, paths=None, fname="predictions.jpg", names=None, calibs=None, conf_thres=0.25, face_3d_classes=None, complete_3d_classes=None): """Threaded wrapper for plotting 3D predictions.""" _plot_predictions_3d_core(images, predictions_2d, predictions_3d, paths, fname, names, calibs, conf_thres, return_array=False, face_3d_classes=face_3d_classes, complete_3d_classes=complete_3d_classes) def plot_val_comparison_3d(images, targets, predictions_2d, predictions_3d=None, paths=None, fname="val_comparison.jpg", names=None, calibs=None, return_array=False, face_3d_classes=None, complete_3d_classes=None): """Plot 2x2 comparison grid: target 2D, target 3D, prediction 2D, prediction 3D. This is the main function for TensorBoard logging during validation. Layout: (1,1) Target 2D | (1,2) Target 3D (2,1) Prediction 2D | (2,2) Prediction 3D Args: images: Batch tensor (B, 3, H, W) normalized [0, 1] in BGR format targets: Target tensor with 48 columns per target predictions_2d: List of 2D predictions per image, each (N, 6) predictions_3d: 3D predictions - list of (N, 34) arrays per image paths: List of image paths fname: Output filename names: Class names dict calibs: Tuple of calibration dicts return_array: If True, return numpy array instead of saving to file face_3d_classes: List/set of class IDs that use face-based reconstruction complete_3d_classes: List/set of class IDs that use complete 3D box Returns: numpy array (H, W, 3) RGB if return_array=True, else None """ if isinstance(images, torch.Tensor): images = images.cpu().float().numpy() if isinstance(targets, torch.Tensor): targets = targets.cpu().numpy() _, _, h, w = images.shape if np.max(images[0]) <= 1: images = images * 255 # Keep BGR format for cv2 drawing (will convert to RGB at the end for matplotlib) # images = images[:, ::-1, :, :] # Don't convert here # Get first image (BGR format) im = images[0].transpose(1, 2, 0).astype(np.uint8) im_target_2d = im.copy() im_target_3d = im.copy() im_pred_2d = im.copy() im_pred_3d = im.copy() # Get targets for first image ti = targets[targets[:, 0] == 0] if len(targets) > 0 else np.array([]) # Get calibration calib = None if calibs is not None and len(calibs) > 0: calib = calibs[0] elif paths: calib = read_calib_from_path(paths[0]) # Apply project-correct defaults for class categorization (face-based vs complete 3D box) # face_3d_classes: only vehicles (0) and tricycles (13) have GT face annotations # complete_3d_classes: pedestrian (1), bicycle (2), rider (3) use full 3D box without face data _face_3d_classes = face_3d_classes if face_3d_classes is not None else [0, 13] _complete_3d_classes = complete_3d_classes if complete_3d_classes is not None else [1, 2, 3] # --- Draw Target 2D boxes --- for t in ti: if len(t) < 6 or np.isnan(t[1]): continue cls = int(t[1]) color = get_class_color(cls) x_center, y_center, bw, bh = t[2:6] if np.any(np.isnan([x_center, y_center, bw, bh])): continue x1 = int((x_center - bw / 2) * w) y1 = int((y_center - bh / 2) * h) x2 = int((x_center + bw / 2) * w) y2 = int((y_center + bh / 2) * h) cv2.rectangle(im_target_2d, (x1, y1), (x2, y2), color, 1) # --- Draw Target 3D boxes --- for t in ti: if len(t) < 15 or np.isnan(t[1]): continue cls = int(t[1]) decoded = decode_and_reconstruct_3d_box_from_target(t, calib, w, h, _face_3d_classes, _complete_3d_classes) if decoded and decoded['should_draw'] and decoded['corners_3d'] is not None: if cls in _face_3d_classes: im_target_3d = draw_3d_box_from_corners( im_target_3d, decoded['corners_3d'], calib, (w, h), face_center_2d=decoded.get('face_center_2d'), face_color=decoded.get('face_color') ) elif cls in _complete_3d_classes: # Use draw_3d_box_from_corners which handles corner reordering internally # This ensures consistent front/rear face colors (red=front, blue=rear) # and fixes the head/rear reversal bug in TensorBoard logs im_target_3d = draw_3d_box_from_corners( im_target_3d, decoded['corners_3d'], calib, (w, h), face_center_2d=decoded.get('center_2d'), face_color=colors(cls) ) # --- Draw Prediction 2D boxes --- pred_2d = None if predictions_2d is not None and len(predictions_2d) > 0: pred_2d = predictions_2d[0] if isinstance(pred_2d, torch.Tensor): pred_2d = pred_2d.cpu().numpy() for pred in pred_2d: if len(pred) < 6: continue x1, y1, x2, y2, conf, cls = pred[:6] cls = int(cls) color = get_class_color(cls) cv2.rectangle(im_pred_2d, (int(x1), int(y1)), (int(x2), int(y2)), color, 1) # --- Draw Prediction 3D boxes --- pred_3d = None if predictions_3d is not None: if isinstance(predictions_3d, list) and len(predictions_3d) > 0: pred_3d = predictions_3d[0] if isinstance(pred_3d, torch.Tensor): pred_3d = pred_3d.cpu().numpy() elif isinstance(predictions_3d, torch.Tensor): pred_3d = predictions_3d[0].cpu().numpy() elif isinstance(predictions_3d, np.ndarray) and len(predictions_3d) > 0: pred_3d = predictions_3d[0] if pred_2d is not None and pred_3d is not None and len(pred_3d) > 0: for j, pred_2d_box in enumerate(pred_2d): if len(pred_2d_box) < 6 or j >= len(pred_3d): continue box_info = decode_and_reconstruct_3d_box( pred_2d_box, pred_3d[j], calib, 0.0, face_3d_classes, complete_3d_classes # Predictions already filtered by conf_thres_plot ) if box_info is None or not box_info['should_draw']: continue if box_info['corners_3d'] is None: continue im_pred_3d = draw_3d_box_from_corners( im_pred_3d, box_info['corners_3d'], calib, (w, h), face_center_2d=box_info['face_center_2d'], face_color=box_info['face_color'] ) # Create 2x2 figure fig, axes = plt.subplots(2, 2, figsize=(12, 8)) # Convert BGR to RGB for matplotlib display axes[0, 0].imshow(im_target_2d[:, :, ::-1]) axes[0, 0].set_title('Target 2D', fontsize=14, fontweight='bold', color='darkblue') axes[0, 0].axis('off') axes[0, 1].imshow(im_target_3d[:, :, ::-1]) axes[0, 1].set_title('Target 3D', fontsize=14, fontweight='bold', color='darkblue') axes[0, 1].axis('off') axes[1, 0].imshow(im_pred_2d[:, :, ::-1]) axes[1, 0].set_title('Prediction 2D', fontsize=14, fontweight='bold', color='darkgreen') axes[1, 0].axis('off') axes[1, 1].imshow(im_pred_3d[:, :, ::-1]) axes[1, 1].set_title('Prediction 3D', fontsize=14, fontweight='bold', color='darkgreen') axes[1, 1].axis('off') if paths: fig.suptitle(Path(paths[0]).name[:50], fontsize=12, fontweight='bold') plt.tight_layout() if return_array: dpi = 300 fig.set_dpi(dpi) fig.canvas.draw() buf = np.frombuffer(fig.canvas.buffer_rgba(), dtype=np.uint8) fig_w, fig_h = fig.canvas.get_width_height() img_array = buf.reshape(fig_h, fig_w, 4) img_array = img_array[:, :, :3] plt.close(fig) return img_array else: plt.savefig(fname, dpi=300, bbox_inches='tight') plt.close(fig) return None def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=""): """Plots learning rate schedule for given optimizer and scheduler, saving plot to `save_dir`.""" optimizer, scheduler = copy(optimizer), copy(scheduler) # do not modify originals y = [] for _ in range(epochs): scheduler.step() y.append(optimizer.param_groups[0]["lr"]) plt.plot(y, ".-", label="LR") plt.xlabel("epoch") plt.ylabel("LR") plt.grid() plt.xlim(0, epochs) plt.ylim(0) plt.savefig(Path(save_dir) / "LR.png", dpi=200) plt.close() def plot_val_txt(): """Plots 2D and 1D histograms of bounding box centers from 'val.txt' using matplotlib, saving as 'hist2d.png' and 'hist1d.png'. Example: from utils.plots import *; plot_val() """ x = np.loadtxt("val.txt", dtype=np.float32) box = xyxy2xywh(x[:, :4]) cx, cy = box[:, 0], box[:, 1] fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True) ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0) ax.set_aspect("equal") plt.savefig("hist2d.png", dpi=300) _fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True) ax[0].hist(cx, bins=600) ax[1].hist(cy, bins=600) plt.savefig("hist1d.png", dpi=200) def plot_targets_txt(): """Plots histograms of object detection targets from 'targets.txt', saving the figure as 'targets.jpg'. Example: from utils.plots import *; plot_targets_txt() """ x = np.loadtxt("targets.txt", dtype=np.float32).T s = ["x targets", "y targets", "width targets", "height targets"] _fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True) ax = ax.ravel() for i in range(4): ax[i].hist(x[i], bins=100, label=f"{x[i].mean():.3g} +/- {x[i].std():.3g}") ax[i].legend() ax[i].set_title(s[i]) plt.savefig("targets.jpg", dpi=200) def plot_val_study(file="", dir="", x=None): """Plots validation study results from 'study*.txt' files in a directory or a specific file, comparing model performance and speed. Example: from utils.plots import *; plot_val_study() """ save_dir = Path(file).parent if file else Path(dir) plot2 = False # plot additional results if plot2: ax = plt.subplots(2, 4, figsize=(10, 6), tight_layout=True)[1].ravel() _fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True) # for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov5n6', 'yolov5s6', 'yolov5m6', 'yolov5l6', 'yolov5x6']]: for f in sorted(save_dir.glob("study*.txt")): y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T x = np.arange(y.shape[1]) if x is None else np.array(x) if plot2: s = ["P", "R", "mAP@.5", "mAP@.5:.95", "t_preprocess (ms/img)", "t_inference (ms/img)", "t_NMS (ms/img)"] for i in range(7): ax[i].plot(x, y[i], ".-", linewidth=2, markersize=8) ax[i].set_title(s[i]) j = y[3].argmax() + 1 ax2.plot( y[5, 1:j], y[3, 1:j] * 1e2, ".-", linewidth=2, markersize=8, label=f.stem.replace("study_coco_", "").replace("yolo", "YOLO"), ) ax2.plot( 1e3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5], "k.-", linewidth=2, markersize=8, alpha=0.25, label="EfficientDet", ) ax2.grid(alpha=0.2) ax2.set_yticks(np.arange(20, 60, 5)) ax2.set_xlim(0, 57) ax2.set_ylim(25, 55) ax2.set_xlabel("GPU Speed (ms/img)") ax2.set_ylabel("COCO AP val") ax2.legend(loc="lower right") f = save_dir / "study.png" print(f"Saving {f}...") plt.savefig(f, dpi=300) @TryExcept() # known issue https://github.com/ultralytics/yolov5/issues/5395 def plot_labels(labels, names=(), save_dir=Path("")): """Plots dataset labels, saving correlogram and label images, handles classes, and visualizes bounding boxes.""" LOGGER.info(f"Plotting labels to {save_dir / 'labels.jpg'}... ") c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes nc = int(c.max() + 1) # number of classes x = pd.DataFrame(b.transpose(), columns=["x", "y", "width", "height"]) labels_plot = labels # seaborn correlogram (can be slow with large datasets) sn.pairplot(x, corner=True, diag_kind="auto", kind="hist", diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9)) plt.savefig(save_dir / "labels_correlogram.jpg", dpi=200) plt.close() # matplotlib labels matplotlib.use("svg") # faster ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel() y = ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8) with contextlib.suppress(Exception): # color histogram bars by class [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)] # known issue #3195 ax[0].set_ylabel("instances") if 0 < len(names) < 30: ax[0].set_xticks(range(len(names))) ax[0].set_xticklabels(list(names.values()), rotation=90, fontsize=10) else: ax[0].set_xlabel("classes") sn.histplot(x, x="x", y="y", ax=ax[2], bins=50, pmax=0.9) sn.histplot(x, x="width", y="height", ax=ax[3], bins=50, pmax=0.9) # rectangles - use subsampled labels labels_rect = labels_plot.copy() labels_rect[:, 1:3] = 0.5 # center labels_rect[:, 1:] = xywh2xyxy(labels_rect[:, 1:]) * 2000 img = Image.fromarray(np.ones((2000, 2000, 3), dtype=np.uint8) * 255) for cls, *box in labels_rect[:1000]: ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls)) # plot ax[1].imshow(img) ax[1].axis("off") for a in [0, 1, 2, 3]: for s in ["top", "right", "left", "bottom"]: ax[a].spines[s].set_visible(False) plt.savefig(save_dir / "labels.jpg", dpi=200) matplotlib.use("Agg") plt.close() def imshow_cls(im, labels=None, pred=None, names=None, nmax=25, verbose=False, f=Path("images.jpg")): """Displays a grid of images with optional labels and predictions, saving to a file.""" from utils.augmentations import denormalize names = names or [f"class{i}" for i in range(1000)] blocks = torch.chunk( denormalize(im.clone()).cpu().float(), len(im), dim=0 ) # select batch index 0, block by channels n = min(len(blocks), nmax) # number of plots m = min(8, round(n**0.5)) # 8 x 8 default _fig, ax = plt.subplots(math.ceil(n / m), m) # 8 rows x n/8 cols ax = ax.ravel() if m > 1 else [ax] # plt.subplots_adjust(wspace=0.05, hspace=0.05) for i in range(n): ax[i].imshow(blocks[i].squeeze().permute((1, 2, 0)).numpy().clip(0.0, 1.0)) ax[i].axis("off") if labels is not None: from utils.general import get_class_name s = get_class_name(names, labels[i]) + (f"—{get_class_name(names, pred[i])}" if pred is not None else "") ax[i].set_title(s, fontsize=8, verticalalignment="top") plt.savefig(f, dpi=300, bbox_inches="tight") plt.close() if verbose: LOGGER.info(f"Saving {f}") if labels is not None: LOGGER.info("True: " + " ".join(f"{get_class_name(names, i):3s}" for i in labels[:nmax])) if pred is not None: LOGGER.info("Predicted:" + " ".join(f"{get_class_name(names, i):3s}" for i in pred[:nmax])) return f def plot_evolve(evolve_csv="path/to/evolve.csv"): """Plots hyperparameter evolution results from a given CSV, saving the plot and displaying best results. Example: from utils.plots import *; plot_evolve() """ evolve_csv = Path(evolve_csv) data = pd.read_csv(evolve_csv) keys = [x.strip() for x in data.columns] x = data.values f = fitness(x) j = np.argmax(f) # max fitness index plt.figure(figsize=(10, 12), tight_layout=True) matplotlib.rc("font", **{"size": 8}) print(f"Best results from row {j} of {evolve_csv}:") for i, k in enumerate(keys[7:]): v = x[:, 7 + i] mu = v[j] # best single result plt.subplot(6, 5, i + 1) plt.scatter(v, f, c=hist2d(v, f, 20), cmap="viridis", alpha=0.8, edgecolors="none") plt.plot(mu, f.max(), "k+", markersize=15) plt.title(f"{k} = {mu:.3g}", fontdict={"size": 9}) # limit to 40 characters if i % 5 != 0: plt.yticks([]) print(f"{k:>15}: {mu:.3g}") f = evolve_csv.with_suffix(".png") # filename plt.savefig(f, dpi=200) plt.close() print(f"Saved {f}") def plot_results(file="path/to/results.csv", dir=""): """Plots training results from a 'results.csv' file; accepts file path and directory as arguments. Example: from utils.plots import *; plot_results('path/to/results.csv') """ save_dir = Path(file).parent if file else Path(dir) fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True) ax = ax.ravel() files = list(save_dir.glob("results*.csv")) assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot." for f in files: try: data = pd.read_csv(f) s = [x.strip() for x in data.columns] x = data.values[:, 0] for i, j in enumerate([1, 2, 3, 4, 5, 8, 9, 10, 6, 7]): y = data.values[:, j].astype("float") # y[y == 0] = np.nan # don't show zero values ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8) # actual results ax[i].plot(x, gaussian_filter1d(y, sigma=3), ":", label="smooth", linewidth=2) # smoothing line ax[i].set_title(s[j], fontsize=12) # if j in [8, 9, 10]: # share train and val loss y axes # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) except Exception as e: LOGGER.info(f"Warning: Plotting error for {f}: {e}") ax[1].legend() fig.savefig(save_dir / "results.png", dpi=200) plt.close() def profile_idetection(start=0, stop=0, labels=(), save_dir=""): """Plots per-image iDetection logs, comparing metrics like storage and performance over time. Example: from utils.plots import *; profile_idetection() """ ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel() s = ["Images", "Free Storage (GB)", "RAM Usage (GB)", "Battery", "dt_raw (ms)", "dt_smooth (ms)", "real-world FPS"] files = list(Path(save_dir).glob("frames*.txt")) for fi, f in enumerate(files): try: results = np.loadtxt(f, ndmin=2).T[:, 90:-30] # clip first and last rows n = results.shape[1] # number of rows x = np.arange(start, min(stop, n) if stop else n) results = results[:, x] t = results[0] - results[0].min() # set t0=0s results[0] = x for i, a in enumerate(ax): if i < len(results): label = labels[fi] if len(labels) else f.stem.replace("frames_", "") a.plot(t, results[i], marker=".", label=label, linewidth=1, markersize=5) a.set_title(s[i]) a.set_xlabel("time (s)") # if fi == len(files) - 1: # a.set_ylim(bottom=0) for side in ["top", "right"]: a.spines[side].set_visible(False) else: a.remove() except Exception as e: print(f"Warning: Plotting error for {f}; {e}") ax[1].legend() plt.savefig(Path(save_dir) / "idetection_profile.png", dpi=200) def save_one_box(xyxy, im, file=Path("im.jpg"), gain=1.02, pad=10, square=False, BGR=False, save=True): """Crops and saves an image from bounding box `xyxy`, applied with `gain` and `pad`, optionally squares and adjusts for BGR. """ xyxy = torch.tensor(xyxy).view(-1, 4) b = xyxy2xywh(xyxy) # boxes if square: b[:, 2:] = b[:, 2:].max(1)[0].unsqueeze(1) # attempt rectangle to square b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad xyxy = xywh2xyxy(b).long() clip_boxes(xyxy, im.shape) crop = im[int(xyxy[0, 1]) : int(xyxy[0, 3]), int(xyxy[0, 0]) : int(xyxy[0, 2]), :: (1 if BGR else -1)] if save: file.parent.mkdir(parents=True, exist_ok=True) # make directory f = str(increment_path(file).with_suffix(".jpg")) # cv2.imwrite(f, crop) # save BGR, https://github.com/ultralytics/yolov5/issues/7007 chroma subsampling issue Image.fromarray(crop[..., ::-1]).save(f, quality=95, subsampling=0) # save RGB return crop # bev视角画板 def draw_bev_blank(max_range=200): """Create a blank BEV image for visualization. Args: max_range: Maximum forward range in meters (default 200m) Returns: BEV image with grid lines """ # Scale: 20 pixels per meter for high resolution (2x for clearer saved images) pixels_per_meter = 20 img_height = max_range * pixels_per_meter # 4000 for 200m img_width = 100 * pixels_per_meter # 2000 for 100m lateral range (-50m to +50m) bevimg = np.ones((img_height, img_width, 3), dtype=np.uint8) * 255 ego_center_x = img_width // 2 ego_center_y = img_height # Ego car position at bottom center # Dimensions of ego car in pixels (e.g., 4.5m long, 1.8m wide) ego_l_px = int(4.5 * pixels_per_meter) ego_w_px = int(1.8 * pixels_per_meter) ego_half_l, ego_half_w = ego_l_px // 2, ego_w_px // 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) # Plot horizontal grid by distance (every 20m for 200m range) grid_step_m = 20 # meters per grid line grid_step_px = grid_step_m * pixels_per_meter num_grids = max_range // grid_step_m for i in range(num_grids + 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 * grid_step_m}m", (ego_center_x + 15, y_pos - 15), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (80, 80, 80), 2, cv2.LINE_AA) # Plot vertical grid (every 10m, range from -50m to +50m) lateral_range = 50 # meters each side lateral_step_m = 10 lateral_step_px = lateral_step_m * pixels_per_meter for i in range(-lateral_range // lateral_step_m, lateral_range // lateral_step_m + 1): x_pos = ego_center_x + i * lateral_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) # Only draw label for non-zero positions to avoid clutter if i != 0: cv2.putText(bevimg, f"{i * lateral_step_m}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 rectangle for ego return bevimg def cam_corners_front_rear(pred3d, facetype): 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]] # use relative origin [0.5, 1, 0.5] if facetype in ['tail', 1]: # front corners_norm = corners_norm - [0, 0.5, 0.5] elif facetype in ['front', 0]: # tail corners_norm = corners_norm - [1, 0.5, 0.5] elif facetype in ['right', 3]: # left corners_norm = corners_norm - [0.5, 0.5, 0] elif facetype in ['left', 2]: # right corners_norm = corners_norm - [0.5, 0.5, 1] elif facetype in ['whole', -1]: # center corners_norm = corners_norm - [0.5, 0.5, 0.5] else: raise AssertionError('Non valid face type') corners = dims.reshape([1, 3]) * corners_norm.reshape([8, 3]) # rotate around y axis 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 vehicle on BEV image. Args: bevimg: BEV image to draw on vehicle3d: [x, y, z, l, h, w, ..., rotation_y, face_type] is_pred: True for predictions (red), False for GT (green) """ x, y, z, l, h, w = vehicle3d[0], vehicle3d[1], vehicle3d[2], vehicle3d[3], vehicle3d[4], vehicle3d[5] rotation_y = vehicle3d[-2] face_type = vehicle3d[-1] # 'front','rear','left','right','whole' # BEV image parameters - must match draw_bev_blank() # Scale: 20 pixels per meter for high resolution (2x for clearer saved images) pixels_per_meter = 20 max_range = 200 # meters forward lateral_range = 50 # meters each side img_height = bevimg.shape[0] img_width = bevimg.shape[1] ego_center_x = img_width // 2 ego_center_y = img_height # Filter out objects outside range 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_0 = np.mean(corners[4:8, :], axis=0) # Front_point xyz3d_center = np.mean(corners[0:8, :], axis=0) # center # Convert 3D coordinates to BEV pixel coordinates # x (lateral) -> horizontal in BEV, z (forward) -> vertical in BEV (inverted) 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_0[0] * pixels_per_meter), int(ego_center_y - xyz3d_0[2] * pixels_per_meter)) H, W = l * pixels_per_meter, w * pixels_per_meter size = (H, W) angle = np.degrees(rotation_y) rect = (center, size, angle) box = cv2.boxPoints(rect) box = np.intp(box) if is_pred: color = (0, 0, 255) # Red for predictions else: color = (0, 255, 0) # Green for ground truth 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