local_visualizer.py 34.7 KB
Newer Older
ZCMax's avatar
ZCMax committed
1
2
3
4
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from typing import Dict, List, Optional, Tuple, Union

5
import matplotlib.pyplot as plt
ZCMax's avatar
ZCMax committed
6
7
import mmcv
import numpy as np
8
9
10
from matplotlib.collections import PatchCollection
from matplotlib.patches import PathPatch
from matplotlib.path import Path
ZCMax's avatar
ZCMax committed
11
12
13
from mmengine.dist import master_only
from torch import Tensor

14
from mmdet3d.structures.bbox_3d.box_3d_mode import Box3DMode
zhangshilong's avatar
zhangshilong committed
15
16
from mmdet.visualization import DetLocalVisualizer

ZCMax's avatar
ZCMax committed
17
18
19
20
21
22
23
try:
    import open3d as o3d
    from open3d import geometry
except ImportError:
    raise ImportError(
        'Please run "pip install open3d" to install open3d first.')

24
from mmengine.structures import InstanceData
ZCMax's avatar
ZCMax committed
25
26
27
from mmengine.visualization.utils import check_type, tensor2ndarray

from mmdet3d.registry import VISUALIZERS
28
29
30
from mmdet3d.structures import (BaseInstance3DBoxes, CameraInstance3DBoxes,
                                Coord3DMode, DepthInstance3DBoxes,
                                Det3DDataSample, LiDARInstance3DBoxes,
31
                                PointData, points_cam2img)
ZCMax's avatar
ZCMax committed
32
from .vis_utils import (proj_camera_bbox3d_to_img, proj_depth_bbox3d_to_img,
33
                        proj_lidar_bbox3d_to_img, to_depth_mode)
ZCMax's avatar
ZCMax committed
34
35
36
37
38
39
40
41
42
43


@VISUALIZERS.register_module()
class Det3DLocalVisualizer(DetLocalVisualizer):
    """MMDetection3D Local Visualizer.

    - 3D detection and segmentation drawing methods

      - draw_bboxes_3d: draw 3D bounding boxes on point clouds
      - draw_proj_bboxes_3d: draw projected 3D bounding boxes on image
zhangshilong's avatar
zhangshilong committed
44
      - draw_seg_mask: draw segmentation mask via per-point colorization
ZCMax's avatar
ZCMax committed
45
46
47

    Args:
        name (str): Name of the instance. Defaults to 'visualizer'.
48
        points (numpy.array, shape=[N, 3+C]): points to visualize.
ZCMax's avatar
ZCMax committed
49
50
        image (np.ndarray, optional): the origin image to draw. The format
            should be RGB. Defaults to None.
51
52
53
        pcd_mode (int): The point cloud mode (coordinates):
            0 represents LiDAR, 1 represents CAMERA, 2
            represents Depth. Defaults to 0.
ZCMax's avatar
ZCMax committed
54
55
56
57
58
59
60
61
62
63
64
65
66
67
        vis_backends (list, optional): Visual backend config list.
            Defaults to None.
        save_dir (str, optional): Save file dir for all storage backends.
            If it is None, the backend storage will not save any data.
        bbox_color (str, tuple(int), optional): Color of bbox lines.
            The tuple of color should be in BGR order. Defaults to None.
        text_color (str, tuple(int), optional): Color of texts.
            The tuple of color should be in BGR order.
            Defaults to (200, 200, 200).
        mask_color (str, tuple(int), optional): Color of masks.
            The tuple of color should be in BGR order.
            Defaults to None.
        line_width (int, float): The linewidth of lines.
            Defaults to 3.
68
        frame_cfg (dict): The coordinate frame config while Open3D
zhangshilong's avatar
zhangshilong committed
69
70
            visualization initialization.
            Defaults to dict(size=1, origin=[0, 0, 0]).
ZCMax's avatar
ZCMax committed
71
72
73
74
75
76
        alpha (int, float): The transparency of bboxes or mask.
                Defaults to 0.8.

    Examples:
        >>> import numpy as np
        >>> import torch
77
        >>> from mmengine.structures import InstanceData
zhangshilong's avatar
zhangshilong committed
78
79
        >>> from mmdet3d.structures import Det3DDataSample
        >>> from mmdet3d.visualization import Det3DLocalVisualizer
ZCMax's avatar
ZCMax committed
80
81
82
83

        >>> det3d_local_visualizer = Det3DLocalVisualizer()
        >>> image = np.random.randint(0, 256,
        ...                     size=(10, 12, 3)).astype('uint8')
zhangshilong's avatar
zhangshilong committed
84
        >>> points = np.random.rand((1000, ))
ZCMax's avatar
ZCMax committed
85
86
        >>> gt_instances_3d = InstanceData()
        >>> gt_instances_3d.bboxes_3d = BaseInstance3DBoxes(torch.rand((5, 7)))
zhangshilong's avatar
zhangshilong committed
87
        >>> gt_instances_3d.labels_3d = torch.randint(0, 2, (5,))
ZCMax's avatar
ZCMax committed
88
89
        >>> gt_det3d_data_sample = Det3DDataSample()
        >>> gt_det3d_data_sample.gt_instances_3d = gt_instances_3d
zhangshilong's avatar
zhangshilong committed
90
91
92
        >>> data_input = dict(img=image, points=points)
        >>> det3d_local_visualizer.add_datasample('3D Scene', data_input,
        ...                         gt_det3d_data_sample)
ZCMax's avatar
ZCMax committed
93
94
95
96
    """

    def __init__(self,
                 name: str = 'visualizer',
97
                 points: Optional[np.ndarray] = None,
ZCMax's avatar
ZCMax committed
98
                 image: Optional[np.ndarray] = None,
99
                 pcd_mode: int = 0,
ZCMax's avatar
ZCMax committed
100
101
102
103
104
105
106
                 vis_backends: Optional[Dict] = None,
                 save_dir: Optional[str] = None,
                 bbox_color: Optional[Union[str, Tuple[int]]] = None,
                 text_color: Optional[Union[str,
                                            Tuple[int]]] = (200, 200, 200),
                 mask_color: Optional[Union[str, Tuple[int]]] = None,
                 line_width: Union[int, float] = 3,
107
                 frame_cfg: dict = dict(size=1, origin=[0, 0, 0]),
ZCMax's avatar
ZCMax committed
108
109
110
111
112
113
114
115
116
117
118
                 alpha: float = 0.8):
        super().__init__(
            name=name,
            image=image,
            vis_backends=vis_backends,
            save_dir=save_dir,
            bbox_color=bbox_color,
            text_color=text_color,
            mask_color=mask_color,
            line_width=line_width,
            alpha=alpha)
119
120
121
        if points is not None:
            self.set_points(points, pcd_mode=pcd_mode, frame_cfg=frame_cfg)
        self.pts_seg_num = 0
ZCMax's avatar
ZCMax committed
122

123
124
    def _initialize_o3d_vis(self, frame_cfg) -> tuple:
        """Initialize open3d vis according to frame_cfg.
ZCMax's avatar
ZCMax committed
125
126

        Args:
127
128
            frame_cfg (dict): The config to create coordinate frame
                in open3d vis.
ZCMax's avatar
ZCMax committed
129
130

        Returns:
131
            :obj:`o3d.visualization.Visualizer`: Created open3d vis.
ZCMax's avatar
ZCMax committed
132
133
134
135
        """
        o3d_vis = o3d.visualization.Visualizer()
        o3d_vis.create_window()
        # create coordinate frame
136
        mesh_frame = geometry.TriangleMesh.create_coordinate_frame(**frame_cfg)
ZCMax's avatar
ZCMax committed
137
138
139
140
141
142
        o3d_vis.add_geometry(mesh_frame)
        return o3d_vis

    @master_only
    def set_points(self,
                   points: np.ndarray,
143
                   pcd_mode: int = 0,
144
                   frame_cfg: dict = dict(size=1, origin=[0, 0, 0]),
145
                   vis_task: str = 'lidar_det',
ZCMax's avatar
ZCMax committed
146
147
148
149
150
151
152
153
                   points_color: Tuple = (0.5, 0.5, 0.5),
                   points_size: int = 2,
                   mode: str = 'xyz') -> None:
        """Set the points to draw.

        Args:
            points (numpy.array, shape=[N, 3+C]):
                points to visualize.
154
155
            pcd_mode (int): The point cloud mode (coordinates):
                0 represents LiDAR, 1 represents CAMERA, 2
156
157
158
159
                represents Depth. Defaults to 0.
            frame_cfg (dict): The coordinate frame config while Open3D
                visualization initialization.
                Defaults to dict(size=1, origin=[0, 0, 0]).
ZCMax's avatar
ZCMax committed
160
            vis_task (str): Visualiztion task, it includes:
161
                'lidar_det', 'multi-modality_det', 'mono_det', 'lidar_seg'.
ZCMax's avatar
ZCMax committed
162
163
164
165
166
167
168
169
170
171
            point_color (tuple[float], optional): the color of points.
                Default: (0.5, 0.5, 0.5).
            points_size (int, optional): the size of points to show
                on visualizer. Default: 2.
            mode (str, optional):  indicate type of the input points,
                available mode ['xyz', 'xyzrgb']. Default: 'xyz'.
        """
        assert points is not None
        check_type('points', points, np.ndarray)

172
173
174
        if not hasattr(self, 'o3d_vis'):
            self.o3d_vis = self._initialize_o3d_vis(frame_cfg)

175
176
177
178
        # for now we convert points into depth mode for visualization
        if pcd_mode != Coord3DMode.DEPTH:
            points = Coord3DMode.convert(points, pcd_mode, Coord3DMode.DEPTH)

179
        if hasattr(self, 'pcd') and vis_task != 'lidar_seg':
ZCMax's avatar
ZCMax committed
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
            self.o3d_vis.remove_geometry(self.pcd)

        # set points size in Open3D
        self.o3d_vis.get_render_option().point_size = points_size

        points = points.copy()
        pcd = geometry.PointCloud()
        if mode == 'xyz':
            pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            points_colors = np.tile(
                np.array(points_color), (points.shape[0], 1))
        elif mode == 'xyzrgb':
            pcd.points = o3d.utility.Vector3dVector(points[:, :3])
            points_colors = points[:, 3:6]
            # normalize to [0, 1] for Open3D drawing
            if not ((points_colors >= 0.0) & (points_colors <= 1.0)).all():
                points_colors /= 255.0
        else:
            raise NotImplementedError

        pcd.colors = o3d.utility.Vector3dVector(points_colors)
        self.o3d_vis.add_geometry(pcd)
        self.pcd = pcd
203
        self.points_colors = points_colors
ZCMax's avatar
ZCMax committed
204
205
206
207
208

    # TODO: assign 3D Box color according to pred / GT labels
    # We draw GT / pred bboxes on the same point cloud scenes
    # for better detection performance comparison
    def draw_bboxes_3d(self,
209
                       bboxes_3d: BaseInstance3DBoxes,
ZCMax's avatar
ZCMax committed
210
211
212
213
214
215
216
217
218
                       bbox_color=(0, 1, 0),
                       points_in_box_color=(1, 0, 0),
                       rot_axis=2,
                       center_mode='lidar_bottom',
                       mode='xyz'):
        """Draw bbox on visualizer and change the color of points inside
        bbox3d.

        Args:
219
            bboxes_3d (:obj:`BaseInstance3DBoxes`, shape=[M, 7]):
ZCMax's avatar
ZCMax committed
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
                3d bbox (x, y, z, x_size, y_size, z_size, yaw) to visualize.
            bbox_color (tuple[float], optional): the color of 3D bboxes.
                Default: (0, 1, 0).
            points_in_box_color (tuple[float], optional):
                the color of points inside 3D bboxes. Default: (1, 0, 0).
            rot_axis (int, optional): rotation axis of 3D bboxes.
                Default: 2.
            center_mode (bool, optional): Indicates the center of bbox is
                bottom center or gravity center. available mode
                ['lidar_bottom', 'camera_bottom']. Default: 'lidar_bottom'.
            mode (str, optional):  Indicates type of input points,
                available mode ['xyz', 'xyzrgb']. Default: 'xyz'.
        """
        # Before visualizing the 3D Boxes in point cloud scene
        # we need to convert the boxes to Depth mode
235
236
237
238
        check_type('bboxes', bboxes_3d, BaseInstance3DBoxes)

        if not isinstance(bboxes_3d, DepthInstance3DBoxes):
            bboxes_3d = bboxes_3d.convert_to(Box3DMode.DEPTH)
ZCMax's avatar
ZCMax committed
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276

        # convert bboxes to numpy dtype
        bboxes_3d = tensor2ndarray(bboxes_3d.tensor)

        in_box_color = np.array(points_in_box_color)

        for i in range(len(bboxes_3d)):
            center = bboxes_3d[i, 0:3]
            dim = bboxes_3d[i, 3:6]
            yaw = np.zeros(3)
            yaw[rot_axis] = bboxes_3d[i, 6]
            rot_mat = geometry.get_rotation_matrix_from_xyz(yaw)

            if center_mode == 'lidar_bottom':
                # bottom center to gravity center
                center[rot_axis] += dim[rot_axis] / 2
            elif center_mode == 'camera_bottom':
                # bottom center to gravity center
                center[rot_axis] -= dim[rot_axis] / 2
            box3d = geometry.OrientedBoundingBox(center, rot_mat, dim)

            line_set = geometry.LineSet.create_from_oriented_bounding_box(
                box3d)
            line_set.paint_uniform_color(bbox_color)
            # draw bboxes on visualizer
            self.o3d_vis.add_geometry(line_set)

            # change the color of points which are in box
            if self.pcd is not None and mode == 'xyz':
                indices = box3d.get_point_indices_within_bounding_box(
                    self.pcd.points)
                self.points_colors[indices] = in_box_color

        # update points colors
        if self.pcd is not None:
            self.pcd.colors = o3d.utility.Vector3dVector(self.points_colors)
            self.o3d_vis.update_geometry(self.pcd)

277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
    def set_bev_image(self,
                      bev_image: Optional[np.ndarray] = None,
                      bev_shape: Optional[int] = 900) -> None:
        """Set the bev image to draw.

        Args:
            bev_image (np.ndarray, optional): The bev image to draw.
                Defaults to None.
            bev_shape (int): The bev image shape. Defaults to 900.
        """
        if bev_image is None:
            bev_image = np.zeros((bev_shape, bev_shape, 3), np.uint8)

        self._image = bev_image
        self.width, self.height = bev_image.shape[1], bev_image.shape[0]
        self._default_font_size = max(
            np.sqrt(self.height * self.width) // 90, 10)
        self.ax_save.cla()
        self.ax_save.axis(False)
        self.ax_save.imshow(bev_image, origin='lower')
        # plot camera view range
        x1 = np.linspace(0, self.width / 2)
        x2 = np.linspace(self.width / 2, self.width)
        self.ax_save.plot(
            x1,
            self.width / 2 - x1,
            ls='--',
            color='grey',
            linewidth=1,
            alpha=0.5)
        self.ax_save.plot(
            x2,
            x2 - self.width / 2,
            ls='--',
            color='grey',
            linewidth=1,
            alpha=0.5)
        self.ax_save.plot(
            self.width / 2,
            0,
            marker='+',
            markersize=16,
            markeredgecolor='red')

    # TODO: Support bev point cloud visualization
    @master_only
    def draw_bev_bboxes(self,
                        bboxes_3d: BaseInstance3DBoxes,
                        scale: int = 15,
                        edge_colors: Union[str, tuple, List[str],
                                           List[tuple]] = 'o',
                        line_styles: Union[str, List[str]] = '-',
                        line_widths: Union[Union[int, float],
                                           List[Union[int, float]]] = 1,
                        face_colors: Union[str, tuple, List[str],
                                           List[tuple]] = 'none',
                        alpha: Union[int, float] = 1):
        """Draw projected 3D boxes on the image.

        Args:
            bboxes_3d (:obj:`BaseInstance3DBoxes`, shape=[M, 7]):
                3d bbox (x, y, z, x_size, y_size, z_size, yaw) to visualize.
            scale (dict): Value to scale the bev bboxes for better
                visualization. Defaults to 15.
            edge_colors (Union[str, tuple, List[str], List[tuple]]): The
                colors of bboxes. ``colors`` can have the same length with
                lines or just single value. If ``colors`` is single value, all
                the lines will have the same colors. Refer to `matplotlib.
                colors` for full list of formats that are accepted.
                Defaults to 'o'.
            line_styles (Union[str, List[str]]): The linestyle
                of lines. ``line_styles`` can have the same length with
                texts or just single value. If ``line_styles`` is single
                value, all the lines will have the same linestyle.
                Reference to
                https://matplotlib.org/stable/api/collections_api.html?highlight=collection#matplotlib.collections.AsteriskPolygonCollection.set_linestyle
                for more details. Defaults to '-'.
            line_widths (Union[Union[int, float], List[Union[int, float]]]):
                The linewidth of lines. ``line_widths`` can have
                the same length with lines or just single value.
                If ``line_widths`` is single value, all the lines will
                have the same linewidth. Defaults to 2.
            face_colors (Union[str, tuple, List[str], List[tuple]]):
                The face colors. Default to 'none'.
            alpha (Union[int, float]): The transparency of bboxes.
                Defaults to 1.
        """

        check_type('bboxes', bboxes_3d, BaseInstance3DBoxes)
        bev_bboxes = tensor2ndarray(bboxes_3d.bev)
        # scale the bev bboxes for better visualization
        bev_bboxes[:, :4] *= scale
        ctr, w, h, theta = np.split(bev_bboxes, [2, 3, 4], axis=-1)
        cos_value, sin_value = np.cos(theta), np.sin(theta)
        vec1 = np.concatenate([w / 2 * cos_value, w / 2 * sin_value], axis=-1)
        vec2 = np.concatenate([-h / 2 * sin_value, h / 2 * cos_value], axis=-1)
        pt1 = ctr + vec1 + vec2
        pt2 = ctr + vec1 - vec2
        pt3 = ctr - vec1 - vec2
        pt4 = ctr - vec1 + vec2
        poly = np.stack([pt1, pt2, pt3, pt4], axis=-2)
        # move the object along x-axis
        poly[:, :, 0] += self.width / 2
        poly = [p for p in poly]
        return self.draw_polygons(
            poly,
            alpha=alpha,
            edge_colors=edge_colors,
            line_styles=line_styles,
            line_widths=line_widths,
            face_colors=face_colors)

    @master_only
    def draw_points_on_image(
            self,
            points: Union[np.ndarray, Tensor],
            pts2img: np.ndarray,
            sizes: Optional[Union[np.ndarray, Tensor, int]] = 10) -> None:
        """Draw projected points on the image.

        Args:
            positions (Union[np.ndarray, torch.Tensor]): Positions to draw.
            pts2imgs (np,ndarray): The transformatino matrix from the
                coordinate of point cloud to image plane.
            sizes (Optional[Union[np.ndarray, torch.Tensor, int]]): The
                marker size. Default to 10.
        """
        check_type('points', points, (np.ndarray, Tensor))
        points = tensor2ndarray(points)
        assert self._image is not None, 'Please set image using `set_image`'
        projected_points = points_cam2img(points, pts2img, with_depth=True)
        depths = projected_points[:, 2]
        colors = (depths % 20) / 20
        # use colormap to obtain the render color
        color_map = plt.get_cmap('jet')
        self.ax_save.scatter(
            projected_points[:, 0],
            projected_points[:, 1],
            c=colors,
            cmap=color_map,
            s=sizes,
            alpha=0.5,
            edgecolors='none')

421
    # TODO: set bbox color according to palette
422
    @master_only
ZCMax's avatar
ZCMax committed
423
424
425
    def draw_proj_bboxes_3d(self,
                            bboxes_3d: BaseInstance3DBoxes,
                            input_meta: dict,
426
427
                            edge_colors: Union[str, tuple, List[str],
                                               List[tuple]] = 'royalblue',
ZCMax's avatar
ZCMax committed
428
429
                            line_styles: Union[str, List[str]] = '-',
                            line_widths: Union[Union[int, float],
430
431
432
433
                                               List[Union[int, float]]] = 2,
                            face_colors: Union[str, tuple, List[str],
                                               List[tuple]] = 'royalblue',
                            alpha: Union[int, float] = 0.4):
ZCMax's avatar
ZCMax committed
434
435
436
437
438
439
        """Draw projected 3D boxes on the image.

        Args:
            bbox3d (:obj:`BaseInstance3DBoxes`, shape=[M, 7]):
                3d bbox (x, y, z, x_size, y_size, z_size, yaw) to visualize.
            input_meta (dict): Input meta information.
440
441
442
443
444
445
            edge_colors (Union[str, tuple, List[str], List[tuple]]): The
                colors of bboxes. ``colors`` can have the same length with
                lines or just single value. If ``colors`` is single value, all
                the lines will have the same colors. Refer to `matplotlib.
                colors` for full list of formats that are accepted.
                Defaults to 'royalblue'.
ZCMax's avatar
ZCMax committed
446
447
448
449
            line_styles (Union[str, List[str]]): The linestyle
                of lines. ``line_styles`` can have the same length with
                texts or just single value. If ``line_styles`` is single
                value, all the lines will have the same linestyle.
450
451
452
                Reference to
                https://matplotlib.org/stable/api/collections_api.html?highlight=collection#matplotlib.collections.AsteriskPolygonCollection.set_linestyle
                for more details. Defaults to '-'.
ZCMax's avatar
ZCMax committed
453
454
455
456
457
            line_widths (Union[Union[int, float], List[Union[int, float]]]):
                The linewidth of lines. ``line_widths`` can have
                the same length with lines or just single value.
                If ``line_widths`` is single value, all the lines will
                have the same linewidth. Defaults to 2.
458
459
460
461
            face_colors (Union[str, tuple, List[str], List[tuple]]):
                The face colors. Default to 'royalblue'.
            alpha (Union[int, float]): The transparency of bboxes.
                Defaults to 0.4.
ZCMax's avatar
ZCMax committed
462
463
464
465
        """

        check_type('bboxes', bboxes_3d, BaseInstance3DBoxes)

466
        if isinstance(bboxes_3d, DepthInstance3DBoxes):
ZCMax's avatar
ZCMax committed
467
            proj_bbox3d_to_img = proj_depth_bbox3d_to_img
468
        elif isinstance(bboxes_3d, LiDARInstance3DBoxes):
ZCMax's avatar
ZCMax committed
469
            proj_bbox3d_to_img = proj_lidar_bbox3d_to_img
470
        elif isinstance(bboxes_3d, CameraInstance3DBoxes):
ZCMax's avatar
ZCMax committed
471
472
            proj_bbox3d_to_img = proj_camera_bbox3d_to_img
        else:
473
            raise NotImplementedError('unsupported box type!')
ZCMax's avatar
ZCMax committed
474

475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
        corners_2d = proj_bbox3d_to_img(bboxes_3d, input_meta)

        lines_verts_idx = [0, 1, 2, 3, 7, 6, 5, 4, 0, 3, 7, 4, 5, 1, 2, 6]
        lines_verts = corners_2d[:, lines_verts_idx, :]
        front_polys = corners_2d[:, 4:, :]
        codes = [Path.LINETO] * lines_verts.shape[1]
        codes[0] = Path.MOVETO
        pathpatches = []
        for i in range(len(corners_2d)):
            verts = lines_verts[i]
            pth = Path(verts, codes)
            pathpatches.append(PathPatch(pth))

        p = PatchCollection(
            pathpatches,
            facecolors='none',
            edgecolors=edge_colors,
            linewidths=line_widths,
            linestyles=line_styles)

        self.ax_save.add_collection(p)

        # draw a mask on the front of project bboxes
        front_polys = [front_poly for front_poly in front_polys]
        return self.draw_polygons(
            front_polys,
            alpha=alpha,
            edge_colors=edge_colors,
            line_styles=line_styles,
            line_widths=line_widths,
            face_colors=face_colors)
ZCMax's avatar
ZCMax committed
506

507
    @master_only
508
    def draw_seg_mask(self, seg_mask_colors: np.array):
ZCMax's avatar
ZCMax committed
509
510
511
512
513
514
515
516
517
518
        """Add segmentation mask to visualizer via per-point colorization.

        Args:
            seg_mask_colors (numpy.array, shape=[N, 6]):
                The segmentation mask whose first 3 dims are point coordinates
                and last 3 dims are converted colors.
        """
        # we can't draw the colors on existing points
        # in case gt and pred mask would overlap
        # instead we set a large offset along x-axis for each seg mask
519
        self.pts_seg_num += 1
ZCMax's avatar
ZCMax committed
520
        offset = (np.array(self.pcd.points).max(0) -
521
                  np.array(self.pcd.points).min(0))[0] * 1.2 * self.pts_seg_num
ZCMax's avatar
ZCMax committed
522
523
524
525
526
        mesh_frame = geometry.TriangleMesh.create_coordinate_frame(
            size=1, origin=[offset, 0, 0])  # create coordinate frame for seg
        self.o3d_vis.add_geometry(mesh_frame)
        seg_points = copy.deepcopy(seg_mask_colors)
        seg_points[:, 0] += offset
527
528
        self.set_points(
            seg_points, vis_task='lidar_seg', pcd_mode=2, mode='xyzrgb')
ZCMax's avatar
ZCMax committed
529
530
531
532
533
534
535
536
537
538
539
540

    def _draw_instances_3d(self, data_input: dict, instances: InstanceData,
                           input_meta: dict, vis_task: str,
                           palette: Optional[List[tuple]]):
        """Draw 3D instances of GT or prediction.

        Args:
            data_input (dict): The input dict to draw.
            instances (:obj:`InstanceData`): Data structure for
                instance-level annotations or predictions.
            metainfo (dict): Meta information.
            vis_task (str): Visualiztion task, it includes:
541
                'lidar_det', 'multi-modality_det', 'mono_det'.
ZCMax's avatar
ZCMax committed
542
543

        Returns:
544
            dict: the drawn point cloud and image which channel is RGB.
ZCMax's avatar
ZCMax committed
545
546
547
548
        """

        bboxes_3d = instances.bboxes_3d  # BaseInstance3DBoxes

549
        data_3d = dict()
ZCMax's avatar
ZCMax committed
550

551
        if vis_task in ['lidar_det', 'multi-modality_det']:
ZCMax's avatar
ZCMax committed
552
553
554
555
556
557
558
559
560
561
            assert 'points' in data_input
            points = data_input['points']
            check_type('points', points, (np.ndarray, Tensor))
            points = tensor2ndarray(points)

            if not isinstance(bboxes_3d, DepthInstance3DBoxes):
                points, bboxes_3d_depth = to_depth_mode(points, bboxes_3d)
            else:
                bboxes_3d_depth = bboxes_3d.clone()

562
            self.set_points(points, pcd_mode=2, vis_task=vis_task)
ZCMax's avatar
ZCMax committed
563
564
            self.draw_bboxes_3d(bboxes_3d_depth)

565
566
            data_3d['bboxes_3d'] = tensor2ndarray(bboxes_3d_depth.tensor)
            data_3d['points'] = points
ZCMax's avatar
ZCMax committed
567

568
        if vis_task in ['mono_det', 'multi-modality_det']:
ZCMax's avatar
ZCMax committed
569
            assert 'img' in data_input
570
            img = data_input['img']
zhangshilong's avatar
zhangshilong committed
571
            if isinstance(data_input['img'], Tensor):
572
                img = img.permute(1, 2, 0).numpy()
zhangshilong's avatar
zhangshilong committed
573
574
                img = img[..., [2, 1, 0]]  # bgr to rgb
            self.set_image(img)
ZCMax's avatar
ZCMax committed
575
            self.draw_proj_bboxes_3d(bboxes_3d, input_meta)
576
577
578
            if vis_task == 'mono_det' and hasattr(instances, 'centers_2d'):
                centers_2d = instances.centers_2d
                self.draw_points(centers_2d)
ZCMax's avatar
ZCMax committed
579
            drawn_img = self.get_image()
580
            data_3d['img'] = drawn_img
ZCMax's avatar
ZCMax committed
581
582
583
584

        return data_3d

    def _draw_pts_sem_seg(self,
585
                          points: Union[Tensor, np.ndarray],
zhangshilong's avatar
zhangshilong committed
586
                          pts_seg: PointData,
ZCMax's avatar
ZCMax committed
587
588
                          palette: Optional[List[tuple]] = None,
                          ignore_index: Optional[int] = None):
589
590
591
592
593
594
595
596
597
598
599
        """Draw 3D semantic mask of GT or prediction.

        Args:
            points (Tensor | np.ndarray): The input point
                cloud to draw.
            pts_seg (:obj:`PointData`): Data structure for
                pixel-level annotations or predictions.
            palette (List[tuple], optional): Palette information
                corresponding to the category. Defaults to None.
            ignore_index (int, optional): Ignore category.
                Defaults to None.
ZCMax's avatar
ZCMax committed
600

601
602
603
        Returns:
            dict: the drawn points with color.
        """
ZCMax's avatar
ZCMax committed
604
605
606
607
        check_type('points', points, (np.ndarray, Tensor))

        points = tensor2ndarray(points)
        pts_sem_seg = tensor2ndarray(pts_seg.pts_semantic_mask)
608
        palette = np.array(palette)
ZCMax's avatar
ZCMax committed
609
610
611
612
613
614
615
616

        if ignore_index is not None:
            points = points[pts_sem_seg != ignore_index]
            pts_sem_seg = pts_sem_seg[pts_sem_seg != ignore_index]

        pts_color = palette[pts_sem_seg]
        seg_color = np.concatenate([points[:, :3], pts_color], axis=1)

617
        self.set_points(points, pcd_mode=2, vis_task='lidar_seg')
618
        self.draw_seg_mask(seg_color)
ZCMax's avatar
ZCMax committed
619
620
621

    @master_only
    def show(self,
622
             save_path: Optional[str] = None,
ZCMax's avatar
ZCMax committed
623
624
625
626
627
             drawn_img_3d: Optional[np.ndarray] = None,
             drawn_img: Optional[np.ndarray] = None,
             win_name: str = 'image',
             wait_time: int = 0,
             continue_key=' ') -> None:
628
        """Show the drawn point cloud/image.
ZCMax's avatar
ZCMax committed
629
630

        Args:
631
632
            save_path (str, optional): Path to save open3d visualized results.
                Default: None.
ZCMax's avatar
ZCMax committed
633
634
635
636
637
638
639
640
641
            drawn_img (np.ndarray, optional): The image to show. If drawn_img
                is None, it will show the image got by Visualizer. Defaults
                to None.
            win_name (str):  The image title. Defaults to 'image'.
            wait_time (int): Delay in milliseconds. 0 is the special
                value that means "forever". Defaults to 0.
            continue_key (str): The key for users to continue. Defaults to
                the space key.
        """
642
        if hasattr(self, 'o3d_vis'):
ZCMax's avatar
ZCMax committed
643
            self.o3d_vis.run()
644
645
            if save_path is not None:
                self.o3d_vis.capture_screen_image(save_path)
ZCMax's avatar
ZCMax committed
646
647
            self.o3d_vis.destroy_window()

648
649
650
        if hasattr(self, '_image'):
            if drawn_img_3d is None:
                super().show(drawn_img_3d, win_name, wait_time, continue_key)
ZCMax's avatar
ZCMax committed
651
652
            super().show(drawn_img, win_name, wait_time, continue_key)

653
654
    # TODO: Support Visualize the 3D results from image and point cloud
    # respectively
ZCMax's avatar
ZCMax committed
655
656
657
658
    @master_only
    def add_datasample(self,
                       name: str,
                       data_input: dict,
659
                       data_sample: Optional['Det3DDataSample'] = None,
ZCMax's avatar
ZCMax committed
660
661
662
663
664
                       draw_gt: bool = True,
                       draw_pred: bool = True,
                       show: bool = False,
                       wait_time: float = 0,
                       out_file: Optional[str] = None,
665
                       save_path: Optional[str] = None,
666
                       vis_task: str = 'mono_det',
ZCMax's avatar
ZCMax committed
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
                       pred_score_thr: float = 0.3,
                       step: int = 0) -> None:
        """Draw datasample and save to all backends.

        - If GT and prediction are plotted at the same time, they are
        displayed in a stitched image where the left image is the
        ground truth and the right image is the prediction.
        - If ``show`` is True, all storage backends are ignored, and
        the images will be displayed in a local window.
        - If ``out_file`` is specified, the drawn point cloud or
        image will be saved to ``out_file``. t is usually used when
        the display is not available.

        Args:
            name (str): The image identifier.
            data_input (dict): It should include the point clouds or image
                to draw.
684
            data_sample (:obj:`Det3DDataSample`, optional): Prediction
ZCMax's avatar
ZCMax committed
685
686
687
688
689
690
691
692
693
                Det3DDataSample. Defaults to None.
            draw_gt (bool): Whether to draw GT Det3DDataSample.
                Default to True.
            draw_pred (bool): Whether to draw Prediction Det3DDataSample.
                Defaults to True.
            show (bool): Whether to display the drawn point clouds and
                image. Default to False.
            wait_time (float): The interval of show (s). Defaults to 0.
            out_file (str): Path to output file. Defaults to None.
694
695
            save_path (str, optional): Path to save open3d visualized results.
                Default: None.
696
            vis-task (str): Visualization task. Defaults to 'mono_det'.
ZCMax's avatar
ZCMax committed
697
698
699
700
701
702
703
704
705
            pred_score_thr (float): The threshold to visualize the bboxes
                and masks. Defaults to 0.3.
            step (int): Global step value to record. Defaults to 0.
        """
        classes = self.dataset_meta.get('CLASSES', None)
        # For object detection datasets, no PALETTE is saved
        palette = self.dataset_meta.get('PALETTE', None)
        ignore_index = self.dataset_meta.get('ignore_index', None)

706
707
708
709
710
        gt_data_3d = None
        pred_data_3d = None
        gt_img_data = None
        pred_img_data = None

711
712
713
714
715
716
        if draw_gt and data_sample is not None:
            if 'gt_instances_3d' in data_sample:
                gt_data_3d = self._draw_instances_3d(
                    data_input, data_sample.gt_instances_3d,
                    data_sample.metainfo, vis_task, palette)
            if 'gt_instances' in data_sample:
ChaimZhu's avatar
ChaimZhu committed
717
718
719
720
721
722
723
724
                if len(data_sample.gt_instances) > 0:
                    assert 'img' in data_input
                    if isinstance(data_input['img'], Tensor):
                        img = data_input['img'].permute(1, 2, 0).numpy()
                        img = img[..., [2, 1, 0]]  # bgr to rgb
                    gt_img_data = self._draw_instances(
                        img, data_sample.gt_instances, classes, palette)
            if 'gt_pts_seg' in data_sample and vis_task == 'seg':
ZCMax's avatar
ZCMax committed
725
726
727
728
729
                assert classes is not None, 'class information is ' \
                                            'not provided when ' \
                                            'visualizing panoptic ' \
                                            'segmentation results.'
                assert 'points' in data_input
730
731
732
                self._draw_pts_sem_seg(data_input['points'],
                                       data_sample.pred_pts_seg, palette,
                                       ignore_index)
ZCMax's avatar
ZCMax committed
733

734
735
736
        if draw_pred and data_sample is not None:
            if 'pred_instances_3d' in data_sample:
                pred_instances_3d = data_sample.pred_instances_3d
737
738
                # .cpu can not be used for BaseInstancesBoxes3D
                # so we need to use .to('cpu')
ZCMax's avatar
ZCMax committed
739
                pred_instances_3d = pred_instances_3d[
740
                    pred_instances_3d.scores_3d > pred_score_thr].to('cpu')
ZCMax's avatar
ZCMax committed
741
742
                pred_data_3d = self._draw_instances_3d(data_input,
                                                       pred_instances_3d,
743
                                                       data_sample.metainfo,
ZCMax's avatar
ZCMax committed
744
                                                       vis_task, palette)
745
746
747
            if 'pred_instances' in data_sample:
                if 'img' in data_input and len(data_sample.pred_instances) > 0:
                    pred_instances = data_sample.pred_instances
748
749
750
751
752
753
754
                    pred_instances = pred_instances_3d[
                        pred_instances.scores > pred_score_thr].cpu()
                    if isinstance(data_input['img'], Tensor):
                        img = data_input['img'].permute(1, 2, 0).numpy()
                        img = img[..., [2, 1, 0]]  # bgr to rgb
                    pred_img_data = self._draw_instances(
                        img, pred_instances, classes, palette)
755
            if 'pred_pts_seg' in data_sample and vis_task == 'lidar_seg':
ZCMax's avatar
ZCMax committed
756
757
758
759
760
                assert classes is not None, 'class information is ' \
                                            'not provided when ' \
                                            'visualizing panoptic ' \
                                            'segmentation results.'
                assert 'points' in data_input
761
762
763
                self._draw_pts_sem_seg(data_input['points'],
                                       data_sample.pred_pts_seg, palette,
                                       ignore_index)
ZCMax's avatar
ZCMax committed
764
765

        # monocular 3d object detection image
766
        if vis_task in ['mono_det', 'multi-modality_det']:
767
768
769
770
771
772
773
            if gt_data_3d is not None and pred_data_3d is not None:
                drawn_img_3d = np.concatenate(
                    (gt_data_3d['img'], pred_data_3d['img']), axis=1)
            elif gt_data_3d is not None:
                drawn_img_3d = gt_data_3d['img']
            elif pred_data_3d is not None:
                drawn_img_3d = pred_data_3d['img']
ZCMax's avatar
ZCMax committed
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
        else:
            drawn_img_3d = None

        # 2d object detection image
        if gt_img_data is not None and pred_img_data is not None:
            drawn_img = np.concatenate((gt_img_data, pred_img_data), axis=1)
        elif gt_img_data is not None:
            drawn_img = gt_img_data
        elif pred_img_data is not None:
            drawn_img = pred_img_data
        else:
            drawn_img = None

        if show:
            self.show(
                vis_task,
790
                save_path,
ZCMax's avatar
ZCMax committed
791
792
793
794
795
796
797
                drawn_img_3d,
                drawn_img,
                win_name=name,
                wait_time=wait_time)

        if out_file is not None:
            if drawn_img_3d is not None:
798
                mmcv.imwrite(drawn_img_3d[..., ::-1], out_file)
ZCMax's avatar
ZCMax committed
799
            if drawn_img is not None:
800
                mmcv.imwrite(drawn_img[..., ::-1], out_file)
ZCMax's avatar
ZCMax committed
801
802
        else:
            self.add_image(name, drawn_img_3d, step)