可视化代码:
代码使用方法:
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)