[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 colorreturn (128, 128, 128)  # Default color: Graydef draw_3d_box_on_image(image, corners_3d, camera_intrinsic, color, label=None):"""在图像上绘制3D框"""# 从相机坐标系->像素坐标系view = np.eye(4)view[:3, :3] = camera_intrinsicin_front = corners_3d[2, :] > 0.1if all(in_front) is False:return imagepoints = corners_3dpoints = 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 imagedef calculate_2d_bbox(corners_3d, camera_intrinsic):"""计算2D边界框"""view = np.eye(4)view[:3, :3] = camera_intrinsicin_front = corners_3d[2, :] > 0.1if all(in_front) is False:return Nonepoints = corners_3dpoints = 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_maxdef 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_widthann_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_2dbox_area = (x_max - x_min) * (y_max - y_min)area_ratio = np.sqrt(box_area / img_area)label = Noneif 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)