2533 lines
97 KiB
Python
Executable File
2533 lines
97 KiB
Python
Executable File
# 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
|