vis_utils.py 6.13 KB
Newer Older
ZCMax's avatar
ZCMax committed
1
2
# Copyright (c) OpenMMLab. All rights reserved.
import copy
3
from typing import Tuple
ZCMax's avatar
ZCMax committed
4
5
6
7
8

import numpy as np
import torch
import trimesh

9
10
from mmdet3d.structures import (BaseInstance3DBoxes, Box3DMode,
                                CameraInstance3DBoxes, Coord3DMode,
zhangshilong's avatar
zhangshilong committed
11
                                DepthInstance3DBoxes, LiDARInstance3DBoxes)
ZCMax's avatar
ZCMax committed
12
13


14
def write_obj(points: np.ndarray, out_filename: str) -> None:
ZCMax's avatar
ZCMax committed
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
    """Write points into ``obj`` format for meshlab visualization.

    Args:
        points (np.ndarray): Points in shape (N, dim).
        out_filename (str): Filename to be saved.
    """
    N = points.shape[0]
    fout = open(out_filename, 'w')
    for i in range(N):
        if points.shape[1] == 6:
            c = points[i, 3:].astype(int)
            fout.write(
                'v %f %f %f %d %d %d\n' %
                (points[i, 0], points[i, 1], points[i, 2], c[0], c[1], c[2]))

        else:
            fout.write('v %f %f %f\n' %
                       (points[i, 0], points[i, 1], points[i, 2]))
    fout.close()


36
def write_oriented_bbox(scene_bbox: np.ndarray, out_filename: str) -> None:
ZCMax's avatar
ZCMax committed
37
38
39
    """Export oriented (around Z axis) scene bbox to meshes.

    Args:
40
41
42
        scene_bbox (np.ndarray): xyz pos of center and 3 lengths
            (x_size, y_size, z_size) and heading angle around Z axis.
            Y forward, X right, Z upward, heading angle of positive X is 0,
ZCMax's avatar
ZCMax committed
43
            heading angle of positive Y is 90 degrees.
44
        out_filename (str): Filename.
ZCMax's avatar
ZCMax committed
45
46
    """

47
    def heading2rotmat(heading_angle: float) -> np.ndarray:
ZCMax's avatar
ZCMax committed
48
49
50
51
52
53
54
        rotmat = np.zeros((3, 3))
        rotmat[2, 2] = 1
        cosval = np.cos(heading_angle)
        sinval = np.sin(heading_angle)
        rotmat[0:2, 0:2] = np.array([[cosval, -sinval], [sinval, cosval]])
        return rotmat

55
56
    def convert_oriented_box_to_trimesh_fmt(
            box: np.ndarray) -> trimesh.base.Trimesh:
ZCMax's avatar
ZCMax committed
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
        ctr = box[:3]
        lengths = box[3:6]
        trns = np.eye(4)
        trns[0:3, 3] = ctr
        trns[3, 3] = 1.0
        trns[0:3, 0:3] = heading2rotmat(box[6])
        box_trimesh_fmt = trimesh.creation.box(lengths, trns)
        return box_trimesh_fmt

    if len(scene_bbox) == 0:
        scene_bbox = np.zeros((1, 7))
    scene = trimesh.scene.Scene()
    for box in scene_bbox:
        scene.add_geometry(convert_oriented_box_to_trimesh_fmt(box))

    mesh_list = trimesh.util.concatenate(scene.dump())
    # save to obj file
    trimesh.io.export.export_mesh(mesh_list, out_filename, file_type='obj')


77
78
79
def to_depth_mode(
        points: np.ndarray,
        bboxes: BaseInstance3DBoxes) -> Tuple[np.ndarray, BaseInstance3DBoxes]:
ZCMax's avatar
ZCMax committed
80
81
82
83
84
85
86
87
88
89
90
91
    """Convert points and bboxes to Depth Coord and Depth Box mode."""
    if points is not None:
        points = Coord3DMode.convert_point(points.copy(), Coord3DMode.LIDAR,
                                           Coord3DMode.DEPTH)
    if bboxes is not None:
        bboxes = Box3DMode.convert(bboxes.clone(), Box3DMode.LIDAR,
                                   Box3DMode.DEPTH)
    return points, bboxes


# TODO: refactor lidar2img to img_meta
def proj_lidar_bbox3d_to_img(bboxes_3d: LiDARInstance3DBoxes,
92
                             input_meta: dict) -> np.ndarray:
ZCMax's avatar
ZCMax committed
93
94
95
    """Project the 3D bbox on 2D plane.

    Args:
96
97
98
        bboxes_3d (:obj:`LiDARInstance3DBoxes`): 3D bbox in lidar coordinate
            system to visualize.
        input_meta (dict): Meta information.
ZCMax's avatar
ZCMax committed
99
    """
100
    corners_3d = bboxes_3d.corners.cpu().numpy()
ZCMax's avatar
ZCMax committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
    num_bbox = corners_3d.shape[0]
    pts_4d = np.concatenate(
        [corners_3d.reshape(-1, 3),
         np.ones((num_bbox * 8, 1))], axis=-1)
    lidar2img = copy.deepcopy(input_meta['lidar2img']).reshape(4, 4)
    if isinstance(lidar2img, torch.Tensor):
        lidar2img = lidar2img.cpu().numpy()
    pts_2d = pts_4d @ lidar2img.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 imgfov_pts_2d


# TODO: remove third parameter in all functions here in favour of img_metas
def proj_depth_bbox3d_to_img(bboxes_3d: DepthInstance3DBoxes,
120
                             input_meta: dict) -> np.ndarray:
ZCMax's avatar
ZCMax committed
121
122
123
    """Project the 3D bbox on 2D plane and draw on input image.

    Args:
124
125
126
        bboxes_3d (:obj:`DepthInstance3DBoxes`): 3D bbox in depth coordinate
            system to visualize.
        input_meta (dict): Meta information.
ZCMax's avatar
ZCMax committed
127
128
    """
    from mmdet3d.models import apply_3d_transformation
zhangshilong's avatar
zhangshilong committed
129
    from mmdet3d.structures import points_cam2img
ZCMax's avatar
ZCMax committed
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150

    input_meta = copy.deepcopy(input_meta)
    corners_3d = bboxes_3d.corners
    num_bbox = corners_3d.shape[0]
    points_3d = corners_3d.reshape(-1, 3)

    # first reverse the data transformations
    xyz_depth = apply_3d_transformation(
        points_3d, 'DEPTH', input_meta, reverse=True)

    # project to 2d to get image coords (uv)
    uv_origin = points_cam2img(xyz_depth,
                               xyz_depth.new_tensor(input_meta['depth2img']))
    uv_origin = (uv_origin - 1).round()
    imgfov_pts_2d = uv_origin[..., :2].reshape(num_bbox, 8, 2).numpy()

    return imgfov_pts_2d


# project the camera bboxes 3d to image
def proj_camera_bbox3d_to_img(bboxes_3d: CameraInstance3DBoxes,
151
                              input_meta: dict) -> np.ndarray:
ZCMax's avatar
ZCMax committed
152
153
154
    """Project the 3D bbox on 2D plane and draw on input image.

    Args:
155
156
157
        bboxes_3d (:obj:`CameraInstance3DBoxes`): 3D bbox in camera coordinate
            system to visualize.
        input_meta (dict): Meta information.
ZCMax's avatar
ZCMax committed
158
    """
zhangshilong's avatar
zhangshilong committed
159
    from mmdet3d.structures import points_cam2img
ZCMax's avatar
ZCMax committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177

    cam2img = copy.deepcopy(input_meta['cam2img'])
    corners_3d = bboxes_3d.corners
    num_bbox = corners_3d.shape[0]
    points_3d = corners_3d.reshape(-1, 3)
    if not isinstance(cam2img, torch.Tensor):
        cam2img = torch.from_numpy(np.array(cam2img))

    assert (cam2img.shape == torch.Size([3, 3])
            or cam2img.shape == torch.Size([4, 4]))
    cam2img = cam2img.float().cpu()

    # project to 2d to get image coords (uv)
    uv_origin = points_cam2img(points_3d, cam2img)
    uv_origin = (uv_origin - 1).round()
    imgfov_pts_2d = uv_origin[..., :2].reshape(num_bbox, 8, 2).numpy()

    return imgfov_pts_2d