MMDet3D——可视化库：3D Lidar Boxes坐标投影到Image上并可视化

本代码适用于mmdet3d默认生成的pkl获取到的数据，并使用类似于DETR3D pipline的代码。
代码参考：from mmdet3d.core.visualizer.image_vis import plot_rect3d_on_img, draw_lidar_bbox3d_on_img

可视化主函数

• 需要img_metas、gt_bboxes_3d即可完成画图功能；
• 可以插在任意部分的代码中，只要有以上两个变量即可。

# -------------------------- vis -----------------------------
import cv2
img = cv2.imread(kwargs['img_metas'][0]['filename'][0]) # front 
r_h, r_w,_ = kwargs['img_metas'][0]['img_shape'][0]
img = cv2.resize(img, (r_w, r_h))
box = kwargs['gt_bboxes_3d'][0]
lidar2img_rt = kwargs['img_metas'][0]['lidar2img'][0]
plot_img = draw_lidar_bbox3d_on_img(box,
                                    raw_img=img,
                                    lidar2img_rt=lidar2img_rt,
                                    img_metas=None,
                                    color=(0, 255, 0),
                                    thickness=1)
cv2.imwrite("plot_img.jpg", plot_img)
# ------------------------------------------------------------

Lidar Box可视化功能函数

import cv2
def plot_rect3d_on_img(img,
                       num_rects,
                       rect_corners,
                       color=(0, 255, 0),
                       thickness=1):
    """Plot the boundary lines of 3D rectangular on 2D images.

    Args:
        img (numpy.array): The numpy array of image.
        num_rects (int): Number of 3D rectangulars.
        rect_corners (numpy.array): Coordinates of the corners of 3D
            rectangulars. Should be in the shape of [num_rect, 8, 2].
        color (tuple[int]): The color to draw bboxes. Default: (0, 255, 0).
        thickness (int, optional): The thickness of bboxes. Default: 1.
    """
    line_indices = ((0, 1), (0, 3), (0, 4), (1, 2), (1, 5), (3, 2), (3, 7),
                    (4, 5), (4, 7), (2, 6), (5, 6), (6, 7))
    # filter boxes out of range
    h,w,c = img.shape
    for i in range(num_rects):
        corners = rect_corners[i].astype(np.int)
        for idx, corner in enumerate(corners):
            corners[idx][0] = w if corner[0] > w else corner[0]
            corners[idx][0] = 0 if corner[0] < 0 else corner[0]
            corners[idx][1] = w if corner[1] > h else corner[1]
            corners[idx][1] = 0 if corner[1] < 0 else corner[1]
        # draw
        for start, end in line_indices:
            cv2.line(img, (corners[start, 0], corners[start, 1]),
                     (corners[end, 0], corners[end, 1]), color, thickness,
                     cv2.LINE_AA)

    return img.astype(np.uint8)

def draw_lidar_bbox3d_on_img(bboxes3d,
                             raw_img,
                             lidar2img_rt,
                             img_metas,
                             color=(0, 255, 0),
                             thickness=1):
    """Project the 3D bbox on 2D plane and draw on input image.

    Args:
        bboxes3d (:obj:`LiDARInstance3DBoxes`):
            3d bbox in lidar coordinate system to visualize.
        raw_img (numpy.array): The numpy array of image.
        lidar2img_rt (numpy.array, shape=[4, 4]): The projection matrix
            according to the camera intrinsic parameters.
        img_metas (dict): Useless here.
        color (tuple[int]): The color to draw bboxes. Default: (0, 255, 0).
        thickness (int, optional): The thickness of bboxes. Default: 1.
    """
    img = raw_img.copy()
    corners_3d = bboxes3d.corners
    num_bbox = corners_3d.shape[0]
    pts_4d = np.concatenate(
        [corners_3d.reshape(-1, 3),
         np.ones((num_bbox * 8, 1))], axis=-1)
    lidar2img_rt = copy.deepcopy(lidar2img_rt).reshape(4, 4)
    if isinstance(lidar2img_rt, torch.Tensor):
        lidar2img_rt = lidar2img_rt.cpu().numpy()
    pts_2d = pts_4d @ lidar2img_rt.T

    pts_2d[:, 2] = np.clip(pts_2d[:, 2], a_min=1e-5, a_max=1e5)
    pts_2d[:, 0] /= pts_2d[:, 2]
    pts_2d[:, 1] /= pts_2d[:, 2]
    imgfov_pts_2d = pts_2d[..., :2].reshape(num_bbox, 8, 2)

    return plot_rect3d_on_img(img, num_bbox, imgfov_pts_2d, color, thickness)

Image Pts可视化

def draw_pts_on_img(img, imgfov_pts_2d, thickness=3, color=[0,255,0]):
	'''
	imgfov_pts_2d: [n, 2] 
	'''
    if isinstance(imgfov_pts_2d, torch.Tensor):
        imgfov_pts_2d = imgfov_pts_2d.detach().cpu().numpy()
    
    h,w,c = img.shape
    for i in range(imgfov_pts_2d.shape[0]):
        if imgfov_pts_2d[i, 0] > w or imgfov_pts_2d[i, 0] < 0:
            continue
        if imgfov_pts_2d[i, 1] > h or imgfov_pts_2d[i, 1] < 0:
            continue
        cv2.circle(
            img,
            center=(int(np.round(imgfov_pts_2d[i, 0])),
                    int(np.round(imgfov_pts_2d[i, 1]))),
            radius=thickness,
            color=color,
            thickness=thickness,
        )
    return img