点云检测框投影到图像上（mmdetection3d）

原模型检测时候只有点云的检测框，本文主要是将demo文件中的pcd_demo.py中的代码，将点云检测出的3d框投影到图像上面显示。 

# Copyright (c) OpenMMLab. All rights reserved.
from argparse import ArgumentParser
# import sys
# sys.path
# sys.path.append('D:\Aware_model\mmdetection3d\mmdet3d')
import os
import sys
dir_mytest = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, dir_mytest)

from mmdet3d.apis import inference_detector, init_model, show_result_meshlab

os.environ['CUDA_LAUNCH_B LOCKING'] = '1'  # 下面老是报错 shape 不一致

import numpy as np
import cv2


def project_velo2rgb(velo, calib):
    T = np.zeros([4, 4], dtype=np.float32)
    T[:3, :] = calib['Tr_velo2cam']
    T[3, 3] = 1
    R = np.zeros([4, 4], dtype=np.float32)
    R[:3, :3] = calib['R0']
    R[3, 3] = 1
    num = len(velo)
    projections = np.zeros((num, 8, 2), dtype=np.int32)
    for i in range(len(velo)):
        box3d = np.ones([8, 4], dtype=np.float32)
        box3d[:, :3] = velo[i]
        M = np.dot(calib['P2'], R)
        M = np.dot(M, T)
        box2d = np.dot(M, box3d.T)
        box2d = box2d[:2, :].T / box2d[2, :].reshape(8, 1)
        projections[i] = box2d
    return projections

def draw_rgb_projections(image, projections, color=(255, 255, 255), thickness=2, darker=1):

    img = image.copy() * darker
    num = len(projections)
    forward_color = (255, 255, 0)
    for n in range(num):
        qs = projections[n]
        for k in range(0, 4):
            i, j = k, (k + 1) % 4
            cv2.line(img, (qs[i, 0], qs[i, 1]), (qs[j, 0], qs[j, 1]), color, thickness, cv2.LINE_AA)
            i, j = k + 4, (k + 1) % 4 + 4
            cv2.line(img, (qs[i, 0], qs[i, 1]), (qs[j, 0], qs[j, 1]), color, thickness, cv2.LINE_AA)
            i, j = k, k + 4
            cv2.line(img, (qs[i, 0], qs[i, 1]), (qs[j, 0], qs[j, 1]), color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[3, 0], qs[3, 1]), (qs[7, 0], qs[7, 1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[7, 0], qs[7, 1]), (qs[6, 0], qs[6, 1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[6, 0], qs[6, 1]), (qs[2, 0], qs[2, 1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[2, 0], qs[2, 1]), (qs[3, 0], qs[3, 1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[3, 0], qs[3, 1]), (qs[6, 0], qs[6, 1]), forward_color, thickness, cv2.LINE_AA)
        cv2.line(img, (qs[2, 0], qs[2, 1]), (qs[7, 0], qs[7, 1]), forward_color, thickness, cv2.LINE_AA)
    return img


def load_kitti_calib(calib_file):
    """
    load projection matrix
    """
    with open(calib_file) as fi:
        lines = fi.readlines()
        assert (len(lines) == 8)
    obj = lines[0].strip().split(' ')[1:]
    P0 = np.array(obj, dtype=np.float32)
    obj = lines[1].strip().split(' ')[1:]
    P1 = np.array(obj, dtype=np.float32)
    obj = lines[2].strip().split(' ')[1:]
    P2 = np.array(obj, dtype=np.float32)
    obj = lines[3].strip().split(' ')[1:]
    P3 = np.array(obj, dtype=np.float32)
    obj = lines[4].strip().split(' ')[1:]
    R0 = np.array(obj, dtype=np.float32)
    obj = lines[5].strip().split(' ')[1:]
    Tr_velo_to_cam = np.array(obj, dtype=np.float32)
    obj = lines[6].strip().split(' ')[1:]
    Tr_imu_to_velo = np.array(obj, dtype=np.float32)
    return {'P2': P2.reshape(3, 4),
            'R0': R0.reshape(3, 3),
            'Tr_velo2cam': Tr_velo_to_cam.reshape(3, 4)}

def box3d_cam_to_velo(box3d,tr):

    tx, ty, tz, l, w, h, ry = [float(i) for i in box3d]     #都是激光坐标系下的参数
    cam = np.ones([3, 1])
    cam[0] = tx
    cam[1] = ty
    cam[2] = tz
    t_lidar = cam.reshape(1,3)#摄像头下转点云坐标系
    Box = np.array([[-l / 2, -l / 2, l / 2, l / 2, -l / 2, -l / 2, l / 2, l / 2],
                    [w / 2, -w / 2, -w / 2, w / 2, w / 2, -w / 2, -w / 2, w / 2],
                    [0, 0, 0, 0, h, h, h, h]])

    rotMat = np.array([
        [np.cos(ry), -np.sin(ry), 0.0],
        [np.sin(ry), np.cos(ry), 0.0],
        [0.0, 0.0, 1.0]])
    velo_box = np.dot(rotMat, Box)
    print(np.tile(t_lidar, (8, 1)).T)
    print(velo_box)
    cornerPosInVelo = velo_box + np.tile(t_lidar, (8, 1)).T
    box3d_corner = cornerPosInVelo.transpose()
    return box3d_corner.astype(np.float32)

def load_kitti_label(box, Tr):

    gt_boxes3d_corner=[]
    for j in range(len(box)):
        box3d_corner = box3d_cam_to_velo(box[j], Tr)
        gt_boxes3d_corner.append(box3d_corner)
    gt_boxes3d_corner = np.array(gt_boxes3d_corner).reshape(-1, 8, 3)
    return gt_boxes3d_corner


if __name__ == '__main__':
    parser = ArgumentParser()
    parser.add_argument('--pcd', default=r'D:\Aware_model\mmdetection3d\data\kitti\training\velodyne\000016.bin',
                        help='Point cloud file')
    parser.add_argument('--config',
                        default=r'D:\Aware_model\mmdetection3d\work_dirs\hv_pointpillars_secfpn_6x8_160e_kitti-3d-car\hv_pointpillars_secfpn_6x8_160e_kitti-3d-car.py',
                        help='Config file')
    parser.add_argument('--checkpoint',
                        default=r'D:\Aware_model\mmdetection3d\checkpoints\hv_pointpillars_secfpn_6x8_160e_kitti-3d-car_20220331_134606-d42d15ed.pth',
                        help='Checkpoint file')
    parser.add_argument('--img', default=r'D:\Aware_model\mmdetection3d\data\kitti\training\image_2\000016.png')
    parser.add_argument('--calib', default=r'D:\Aware_model\mmdetection3d\data\kitti\training\calib\000016.txt')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
        '--score-thr', type=float, default=0.0, help='bbox score threshold')
    parser.add_argument(
        '--out-dir', type=str, default='demo', help='dir to save results')
    parser.add_argument(
        '--show',
        action='store_true',
        help='show online visualization results')
    parser.add_argument(
        '--snapshot',
        action='store_true',
        help='whether to save online visualization results')
    args = parser.parse_args()
    # build the model from a config file and a checkpoint file
    model = init_model(args.config, args.checkpoint, device=args.device)
    # test a single image
    result, data = inference_detector(model, args.pcd)
    # show the results
    show_result_meshlab(
        data,
        result,
        args.out_dir,
        args.score_thr,
        # show=args.show,
        show=True,
        snapshot=args.snapshot,
        task='det')


    lidar = data['points'][0][0].cpu().numpy()
    image = cv2.imread(args.img)
    calib = load_kitti_calib(args.calib)
    gt_box3d = result[0]['boxes_3d'].tensor.numpy()

    # 在激光坐标系下预测框的八个顶点
    gt_box3d = load_kitti_label(gt_box3d, calib['Tr_velo2cam'])

    # view in point cloud，可视化
    gt_3dTo2D = project_velo2rgb(gt_box3d, calib)
    img_with_box = draw_rgb_projections(image, gt_3dTo2D, color=(0, 0, 255), thickness=1)
    cv2.imwrite('img_with_box.png', img_with_box)
    cv2.imshow('img_with_box.png', img_with_box)
    cv2.waitKey(0)

 

原作者：https://blog.csdn.net/suiyingy