在mmdet3d中使用petr

from mmdet3d.apis import init_model
from copy import deepcopy
from mmengine.dataset import Compose, pseudo_collate
import mmengine
from mmdet3d.structures import Box3DMode, Det3DDataSample, get_box_type
from os import path as osp
import numpy as np
import torch
import mmcv
from mmdet3d.registry import VISUALIZERS
from demo.inference_img import create_bbox_lines, draw_bboxes_with_color
import open3d as o3d
import cv2


config_file = '/home/robotics/mmdetection3d/projects/PETR/configs/petr_vovnet_gridmask_p4_800x320.py'
checkpoint_file = '/home/robotics/Downloads/mmdet3d_weights/petr_vovnet_gridmask_p4_800x320-e2191752.pth'

model = init_model(config_file, checkpoint_file, device='cuda:0')

ann_file = '/home/robotics/mmdetection3d/demo/data/nuscenes/n015-2018-07-24-11-22-45+0800.pkl'
imgs = '/home/robotics/mmdetection3d/demo/data/nuscenes'
pcds = '/home/robotics/mmdetection3d/demo/data/nuscenes/n015-2018-07-24-11-22-45+0800__LIDAR_TOP__1532402927647951.pcd.bin'
pcds = [pcds]
imgs = [imgs]

cfg = model.cfg

# build the data pipeline
test_pipeline = deepcopy(cfg.test_dataloader.dataset.pipeline)
test_pipeline = Compose(test_pipeline)
box_type_3d, box_mode_3d = \
	get_box_type(cfg.test_dataloader.dataset.box_type_3d)

data_list = mmengine.load(ann_file)['data_list']
# change my own intrincs
data_list[0]['images']['CAM_FRONT']['cam2img'][0][0] = 1216.27
data_list[0]['images']['CAM_FRONT']['cam2img'][0][2] = 933.174
data_list[0]['images']['CAM_FRONT']['cam2img'][1][1] = 1220.32
data_list[0]['images']['CAM_FRONT']['cam2img'][1][2] = 639.048

data = []

for index, pcd in enumerate(pcds):
	# get data info containing calib
	data_info = data_list[index]
	img = imgs[index]

	for _, img_info in data_info['images'].items():
		img_info['img_path'] = osp.join(img, img_info['img_path'])
		assert osp.isfile(img_info['img_path']
						  ), f'{img_info["img_path"]} does not exist.'
	data_ = dict(
		lidar_points=dict(lidar_path=pcd),
		images=data_info['images'],
		box_type_3d=box_type_3d,
		box_mode_3d=box_mode_3d)

	data_ = test_pipeline(data_)
	data.append(data_)

collate_data = pseudo_collate(data)
with torch.no_grad():
	result = model.test_step(collate_data)[0]


boxes_2d = result.pred_instances['bboxes'].cpu().numpy()
scores_2d = result.pred_instances['scores'].cpu().numpy()
boxes_2d = boxes_2d[np.where((scores_2d > 0.1))[0]]
img = cv2.imread('/home/robotics/mmdetection3d/demo/data/nuscenes/n015-2018-07-24-11-22-45+0800__CAM_FRONT__1532402927612460.jpg')
for i in range(boxes_2d.shape[0]):
	topleft = (int(boxes_2d[i, 0]), int(boxes_2d[i, 1]))
	bottomright = (int(boxes_2d[i, 2]), int(boxes_2d[i, 3]))
	cv2.rectangle(img, topleft, bottomright, (255, 0, 0), 2)
cv2.imwrite('/home/robotics/mmdetection3d/projects/PETR/test.jpg', img)

boxes = result.pred_instances_3d['bboxes_3d'].tensor.cpu().numpy()      # LiDARInstance3DBoxes
labels = result.pred_instances_3d['labels_3d'].cpu().numpy()     # /home/koheikubo/mmdetection3d/mmdet3d/datasets/nuscenes_dataset.py line 60
classes = ['car', 'truck', 'trailer', 'bus', 'construction_vehicle',
'bicycle', 'motorcycle', 'pedestrian', 'traffic_cone', 'barrier']
label_idx_dict = dict(zip(range(len(classes)), classes))
colors = [
			[1, 0, 0],  # Red car
			[0, 1, 0],  # Green truck
			[0, 0, 1],  # Blue trailer
			[1, 1, 0],  # Yellow bus
			[1, 0, 1],  # Magenta construction_vehicle
			[0, 1, 1],  # Cyan bicycle
			[0.5, 0.3, 0],  # Brown motorcycle
			[1, 0.5, 0],  # Orange pedestrian
			[0.5, 0, 0.5],  # Purple traffic_cone
			[0, 0.8, 0.8],  # Teal barrier
		]
scores = result.pred_instances_3d['scores_3d'].cpu().numpy()     # select score > 0.3

threshold = 0.1
boxes = boxes[np.where((scores > threshold))[0]]
rotate_x = 90 / 180*np.pi
Rx = np.array([
	[1, 0, 0],
	[0, np.cos(rotate_x), -np.sin(rotate_x)],
	[0, np.sin(rotate_x), np.cos(rotate_x)]
])
rotate_y = 90 / 180*np.pi
Ry = np.array([
    [np.cos(rotate_y), 0, np.sin(rotate_y)],
    [0, 1, 0],
    [-np.sin(rotate_y), 0, np.cos(rotate_y)]
])
# boxes[:, :3] = np.dot(boxes[:, :3], Ry)
boxes[:, :3] = np.dot(boxes[:, :3], Rx)

labels = labels[np.where((scores > threshold))[0]]
detected_labels = [label_idx_dict[i] for i in labels]
# print(labels)
# print(detected_labels)

objects = []
for i in range(boxes.shape[0]):
	objects.append(draw_bboxes_with_color(boxes[i, :7], color=colors[labels[i]]))

points = np.fromfile(pcds[0], dtype=np.float32).reshape(-1, 5)
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points[:, :3])  # Assuming the first three columns are x, y, z
colors = np.zeros((points.shape[0], 3))
pcd.colors = o3d.utility.Vector3dVector(colors)

vis = o3d.visualization.Visualizer()
vis.create_window()

vis.add_geometry(pcd)
for obj in objects:
	vis.add_geometry(obj)

vis.get_render_option().point_size = 2.0  # 改变点的大小

# 创建一个TriangleMesh对象来表示坐标轴
coord_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=0.6, origin=[0, 0, 0])

# 将坐标轴添加到可视化窗口中
vis.add_geometry(coord_frame)

vis.run()
vis.destroy_window()