[bevformer渲染可视化] 2d框可视化 并可视化出小目标

可视化代码：

代码使用方法：
1.复制代码全文到任意python文件中

2.下载nuscenes v1.0-mini数据集,修改数据集路径，保证能读取到数据集

3.按照需求修改代码（本文是2dbox面积的面积和整个图片的面积比小于0.03视为小目标）

if area_ratio < 0.03:

   label = "small target"

3.运行命令

 在bevfomer环境下运行此python文件

import os
import cv2
import numpy as np
from nuscenes.utils.geometry_utils import view_points
from pyquaternion import Quaternion
from nuscenes import NuScenes

# 初始化nuScenes
nusc = NuScenes(version='v1.0-mini', dataroot='/home/lyh/Documents/dataset/v1.0-mini/', verbose=True)

# 定义类别颜色映射
category_colors = {
    'vehicle': (255, 255, 0),    # Yellow
    'pedestrian': (0, 0, 255),   # Blue
    'bicycle': (0, 255, 0),      # Green
    'motorcycle': (255, 0, 255), # Magenta
    'bus': (0, 255, 255),        # Cyan
    'trailer': (255, 165, 0),    # Orange
    'truck': (255, 20, 147)      # Deep Pink
}

def get_color_for_category(category_name):
    """
    获取类别的颜色
    """
    for category, color in category_colors.items():
        if category in category_name:
            return color
    return (128, 128, 128)  # Default color: Gray

def draw_3d_box_on_image(image, corners_3d, camera_intrinsic, color, label=None):
    """
    在图像上绘制3D框
    """
    # 从相机坐标系->像素坐标系
    view = np.eye(4)
    view[:3, :3] = camera_intrinsic
    in_front = corners_3d[2, :] > 0.1
    if all(in_front) is False:
        return image
    points = corners_3d
    points = np.concatenate((points, np.ones((1, points.shape[1]))), axis=0)
    points = np.dot(view, points)[:3, :]
    points /= points[2, :]

    box_img = points.astype(np.int32)

    for i in range(4):
        j = (i + 1) % 4
        # 下底面
        cv2.line(image, (box_img[0, i], box_img[1, i]), (box_img[0, j], box_img[1, j]), color, thickness=2)
        # 上底面
        cv2.line(image, (box_img[0, i + 4], box_img[1, i + 4]), (box_img[0, j + 4], box_img[1, j + 4]), color, thickness=2)
        # 侧边线
        cv2.line(image, (box_img[0, i], box_img[1, i]), (box_img[0, i + 4], box_img[1, i + 4]), color, thickness=2)

    if label:
        # 在图像上绘制文本标签
        x_min = min(box_img[0])
        y_min = min(box_img[1])
        cv2.putText(image, label, (x_min, y_min - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    
    return image

def calculate_2d_bbox(corners_3d, camera_intrinsic):
    """
    计算2D边界框
    """
    view = np.eye(4)
    view[:3, :3] = camera_intrinsic
    in_front = corners_3d[2, :] > 0.1
    if all(in_front) is False:
        return None
    points = corners_3d
    points = np.concatenate((points, np.ones((1, points.shape[1]))), axis=0)
    points = np.dot(view, points)[:3, :]
    points /= points[2, :]

    box_img = points.astype(np.int32)
    x_min, y_min = np.min(box_img[0]), np.min(box_img[1])
    x_max, y_max = np.max(box_img[0]), np.max(box_img[1])

    return x_min, y_min, x_max, y_max

def visualize_sample(nusc, sample_token, save_path):
    sample = nusc.get('sample', sample_token)
    sample_data_token = sample['data']['CAM_FRONT']

    sd_rec = nusc.get('sample_data', sample_data_token)
    filename = sd_rec['filename']
    img_path = os.path.join(nusc.dataroot, filename)
    img = cv2.imread(img_path)

    assert sd_rec['sensor_modality'] == 'camera', 'Error: This function only works for camera sample_data!'
    if not sd_rec['is_key_frame']:
        raise ValueError('The 2D re-projections are available only for keyframes.')

    cs_rec = nusc.get('calibrated_sensor', sd_rec['calibrated_sensor_token'])
    pose_rec = nusc.get('ego_pose', sd_rec['ego_pose_token'])
    camera_intrinsic = np.array(cs_rec['camera_intrinsic'])

    img_height, img_width = img.shape[:2]
    img_area = img_height * img_width

    ann_recs = [nusc.get('sample_annotation', token) for token in sample['anns']]
    for ann_rec in ann_recs:
        box = nusc.get_box(ann_rec['token'])

        # 从世界坐标系->车身坐标系
        box.translate(-np.array(pose_rec['translation']))
        box.rotate(Quaternion(pose_rec['rotation']).inverse)

        # 从车身坐标系->相机坐标系
        box.translate(-np.array(cs_rec['translation']))
        box.rotate(Quaternion(cs_rec['rotation']).inverse)

        corners_3d = box.corners()
        bbox_2d = calculate_2d_bbox(corners_3d, camera_intrinsic)

        if bbox_2d:
            x_min, y_min, x_max, y_max = bbox_2d
            box_area = (x_max - x_min) * (y_max - y_min)
            area_ratio = np.sqrt(box_area / img_area)

            label = None
            if area_ratio < 0.03:
                label = "small target"

            color = get_color_for_category(ann_rec['category_name'])
            img = draw_3d_box_on_image(img, corners_3d, camera_intrinsic, color, label)

    save_filename = os.path.join(save_path, f"{sample_token}.png")
    cv2.imwrite(save_filename, img)
    print(f"Saved {save_filename}")

def visualize_all_samples(nusc, save_path):
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    for sample in nusc.sample:
        visualize_sample(nusc, sample['token'], save_path)


# 设置保存路径并可视化所有样本
save_path = '/home/lyh/Documents/small_target_Output/'
visualize_all_samples(nusc, save_path)

可视化效果：