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()