local_visualizer.py 27.5 KB
Newer Older
ZCMax's avatar
ZCMax committed
1
2
3
4
5
6
7
8
9
10
# Copyright (c) OpenMMLab. All rights reserved.
import copy
from os import path as osp
from typing import Dict, List, Optional, Tuple, Union

import mmcv
import numpy as np
from mmengine.dist import master_only
from torch import Tensor

zhangshilong's avatar
zhangshilong committed
11
12
from mmdet.visualization import DetLocalVisualizer

ZCMax's avatar
ZCMax committed
13
14
15
16
17
18
19
try:
    import open3d as o3d
    from open3d import geometry
except ImportError:
    raise ImportError(
        'Please run "pip install open3d" to install open3d first.')

20
from mmengine.structures import InstanceData
ZCMax's avatar
ZCMax committed
21
22
23
from mmengine.visualization.utils import check_type, tensor2ndarray

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


@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
41
      - draw_seg_mask: draw segmentation mask via per-point colorization
ZCMax's avatar
ZCMax committed
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

    Args:
        name (str): Name of the instance. Defaults to 'visualizer'.
        image (np.ndarray, optional): the origin image to draw. The format
            should be RGB. Defaults to None.
        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.
zhangshilong's avatar
zhangshilong committed
61
62
63
        vis_cfg (dict): The coordinate frame config while Open3D
            visualization initialization.
            Defaults to dict(size=1, origin=[0, 0, 0]).
ZCMax's avatar
ZCMax committed
64
65
66
67
68
69
        alpha (int, float): The transparency of bboxes or mask.
                Defaults to 0.8.

    Examples:
        >>> import numpy as np
        >>> import torch
70
        >>> from mmengine.structures import InstanceData
zhangshilong's avatar
zhangshilong committed
71
72
        >>> from mmdet3d.structures import Det3DDataSample
        >>> from mmdet3d.visualization import Det3DLocalVisualizer
ZCMax's avatar
ZCMax committed
73
74
75
76

        >>> det3d_local_visualizer = Det3DLocalVisualizer()
        >>> image = np.random.randint(0, 256,
        ...                     size=(10, 12, 3)).astype('uint8')
zhangshilong's avatar
zhangshilong committed
77
        >>> points = np.random.rand((1000, ))
ZCMax's avatar
ZCMax committed
78
79
        >>> gt_instances_3d = InstanceData()
        >>> gt_instances_3d.bboxes_3d = BaseInstance3DBoxes(torch.rand((5, 7)))
zhangshilong's avatar
zhangshilong committed
80
        >>> gt_instances_3d.labels_3d = torch.randint(0, 2, (5,))
ZCMax's avatar
ZCMax committed
81
82
        >>> gt_det3d_data_sample = Det3DDataSample()
        >>> gt_det3d_data_sample.gt_instances_3d = gt_instances_3d
zhangshilong's avatar
zhangshilong committed
83
84
85
        >>> 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
86
87
88
89
90
91
92
93
94
95
96
97
    """

    def __init__(self,
                 name: str = 'visualizer',
                 image: Optional[np.ndarray] = None,
                 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,
zhangshilong's avatar
zhangshilong committed
98
                 vis_cfg: dict = dict(size=1, origin=[0, 0, 0]),
ZCMax's avatar
ZCMax committed
99
100
101
102
103
104
105
106
107
108
109
110
                 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)
        self.o3d_vis = self._initialize_o3d_vis(vis_cfg)
111
        self.seg_num = 0
ZCMax's avatar
ZCMax committed
112
113
114
115
116
117
118
119
120
121
122
123
124
125

    def _initialize_o3d_vis(self, vis_cfg) -> tuple:
        """Build open3d vis according to vis_cfg.

        Args:
            vis_cfg (dict): The config to build open3d vis.

        Returns:
             tuple: build open3d vis.
        """
        # init open3d visualizer
        o3d_vis = o3d.visualization.Visualizer()
        o3d_vis.create_window()
        # create coordinate frame
zhangshilong's avatar
zhangshilong committed
126
        mesh_frame = geometry.TriangleMesh.create_coordinate_frame(**vis_cfg)
ZCMax's avatar
ZCMax committed
127
128
129
130
131
132
133
        o3d_vis.add_geometry(mesh_frame)

        return o3d_vis

    @master_only
    def set_points(self,
                   points: np.ndarray,
134
135
                   pcd_mode: int = 0,
                   vis_task: str = 'det',
ZCMax's avatar
ZCMax committed
136
137
138
139
140
141
142
143
                   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.
144
145
146
            pcd_mode (int): The point cloud mode (coordinates):
                0 represents LiDAR, 1 represents CAMERA, 2
                represents Depth.
ZCMax's avatar
ZCMax committed
147
148
149
150
151
152
153
154
155
156
157
158
            vis_task (str): Visualiztion task, it includes:
                'det', 'multi_modality-det', 'mono-det', 'seg'.
            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)

159
160
161
162
163
        # for now we convert points into depth mode for visualization
        if pcd_mode != Coord3DMode.DEPTH:
            points = Coord3DMode.convert(points, pcd_mode, Coord3DMode.DEPTH)

        if hasattr(self, 'pcd') and vis_task != 'seg':
ZCMax's avatar
ZCMax committed
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
            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
187
        self.points_colors = points_colors
ZCMax's avatar
ZCMax committed
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257

    # 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,
                       bboxes_3d: DepthInstance3DBoxes,
                       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:
            bboxes_3d (:obj:`DepthInstance3DBoxes`, shape=[M, 7]):
                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
        check_type('bboxes', bboxes_3d, (DepthInstance3DBoxes))

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

258
    # TODO: set bbox color according to palette
ZCMax's avatar
ZCMax committed
259
260
261
    def draw_proj_bboxes_3d(self,
                            bboxes_3d: BaseInstance3DBoxes,
                            input_meta: dict,
262
                            bbox_color: Tuple[float] = 'b',
ZCMax's avatar
ZCMax committed
263
264
                            line_styles: Union[str, List[str]] = '-',
                            line_widths: Union[Union[int, float],
265
                                               List[Union[int, float]]] = 1):
ZCMax's avatar
ZCMax committed
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
        """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.
            bbox_color (tuple[float], optional): the color of bbox.
                Default: (0, 1, 0).
            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.
            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.
        """

        check_type('bboxes', bboxes_3d, BaseInstance3DBoxes)

287
        if isinstance(bboxes_3d, DepthInstance3DBoxes):
ZCMax's avatar
ZCMax committed
288
            proj_bbox3d_to_img = proj_depth_bbox3d_to_img
289
        elif isinstance(bboxes_3d, LiDARInstance3DBoxes):
ZCMax's avatar
ZCMax committed
290
            proj_bbox3d_to_img = proj_lidar_bbox3d_to_img
291
        elif isinstance(bboxes_3d, CameraInstance3DBoxes):
ZCMax's avatar
ZCMax committed
292
293
            proj_bbox3d_to_img = proj_camera_bbox3d_to_img
        else:
294
            raise NotImplementedError('unsupported box type!')
ZCMax's avatar
ZCMax committed
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316

        # (num_bboxes_3d, 8, 2)
        proj_bboxes_3d = proj_bbox3d_to_img(bboxes_3d, input_meta)
        num_bboxes_3d = proj_bboxes_3d.shape[0]

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

        # TODO: assign each projected 3d bboxes color
        # according to pred / gt class.
        for i in range(num_bboxes_3d):
            x_datas = []
            y_datas = []
            corners = proj_bboxes_3d[i].astype(np.int)  # (8, 2)
            for start, end in line_indices:
                x_datas.append([corners[start][0], corners[end][0]])
                y_datas.append([corners[start][1], corners[end][1]])
            x_datas = np.array(x_datas)
            y_datas = np.array(y_datas)
            self.draw_lines(x_datas, y_datas, bbox_color, line_styles,
                            line_widths)

317
    def draw_seg_mask(self, seg_mask_colors: np.array):
ZCMax's avatar
ZCMax committed
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
        """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
        self.seg_num += 1
        offset = (np.array(self.pcd.points).max(0) -
                  np.array(self.pcd.points).min(0))[0] * 1.2 * self.seg_num
        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
336
        self.set_points(seg_points, vis_task='seg', pcd_mode=2, mode='xyzrgb')
ZCMax's avatar
ZCMax committed
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351

    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:
                'det', 'multi_modality-det', 'mono-det'.

        Returns:
352
            dict: the drawn point cloud and image which channel is RGB.
ZCMax's avatar
ZCMax committed
353
354
355
356
        """

        bboxes_3d = instances.bboxes_3d  # BaseInstance3DBoxes

357
        data_3d = dict()
ZCMax's avatar
ZCMax committed
358
359
360
361
362
363
364
365
366
367
368
369

        if vis_task in ['det', 'multi_modality-det']:
            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()

370
            self.set_points(points, pcd_mode=2, vis_task=vis_task)
ZCMax's avatar
ZCMax committed
371
372
            self.draw_bboxes_3d(bboxes_3d_depth)

373
374
            data_3d['bboxes_3d'] = tensor2ndarray(bboxes_3d_depth.tensor)
            data_3d['points'] = points
ZCMax's avatar
ZCMax committed
375
376
377

        if vis_task in ['mono-det', 'multi_modality-det']:
            assert 'img' in data_input
378
            img = data_input['img']
zhangshilong's avatar
zhangshilong committed
379
            if isinstance(data_input['img'], Tensor):
380
                img = img.permute(1, 2, 0).numpy()
zhangshilong's avatar
zhangshilong committed
381
382
                img = img[..., [2, 1, 0]]  # bgr to rgb
            self.set_image(img)
ZCMax's avatar
ZCMax committed
383
384
            self.draw_proj_bboxes_3d(bboxes_3d, input_meta)
            drawn_img = self.get_image()
385
            data_3d['img'] = drawn_img
ZCMax's avatar
ZCMax committed
386
387
388
389

        return data_3d

    def _draw_pts_sem_seg(self,
390
                          points: Union[Tensor, np.ndarray],
zhangshilong's avatar
zhangshilong committed
391
                          pts_seg: PointData,
ZCMax's avatar
ZCMax committed
392
393
                          palette: Optional[List[tuple]] = None,
                          ignore_index: Optional[int] = None):
394
395
396
397
398
399
400
401
402
403
404
        """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
405

406
407
408
        Returns:
            dict: the drawn points with color.
        """
ZCMax's avatar
ZCMax committed
409
410
411
412
        check_type('points', points, (np.ndarray, Tensor))

        points = tensor2ndarray(points)
        pts_sem_seg = tensor2ndarray(pts_seg.pts_semantic_mask)
413
        palette = np.array(palette)
ZCMax's avatar
ZCMax committed
414
415
416
417
418
419
420
421

        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)

422
423
        self.set_points(points, pcd_mode=2, vis_task='seg')
        self.draw_seg_mask(seg_color)
ZCMax's avatar
ZCMax committed
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440

        seg_data_3d = dict(points=points, seg_color=seg_color)
        return seg_data_3d

    @master_only
    def show(self,
             vis_task: str = None,
             out_file: str = None,
             drawn_img_3d: Optional[np.ndarray] = None,
             drawn_img: Optional[np.ndarray] = None,
             win_name: str = 'image',
             wait_time: int = 0,
             continue_key=' ') -> None:
        """Show the drawn image.

        Args:
            vis_task (str): Visualiztion task, it includes:
441
                'det', 'multi_modality-det', 'mono-det', 'seg'.
ZCMax's avatar
ZCMax committed
442
443
444
445
446
447
448
449
450
451
            out_file (str): Output file path.
            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.
        """
452
        if vis_task in ['det', 'multi_modality-det', 'seg']:
ZCMax's avatar
ZCMax committed
453
454
            self.o3d_vis.run()
            if out_file is not None:
455
                self.o3d_vis.capture_screen_image(out_file + '.png')
ZCMax's avatar
ZCMax committed
456
457
458
459
460
461
462
463
            self.o3d_vis.destroy_window()

        if vis_task in ['mono-det', 'multi_modality-det']:
            super().show(drawn_img_3d, win_name, wait_time, continue_key)

        if drawn_img is not None:
            super().show(drawn_img, win_name, wait_time, continue_key)

464
465
    # TODO: Support Visualize the 3D results from image and point cloud
    # respectively
ZCMax's avatar
ZCMax committed
466
467
468
469
    @master_only
    def add_datasample(self,
                       name: str,
                       data_input: dict,
470
                       data_sample: Optional['Det3DDataSample'] = None,
ZCMax's avatar
ZCMax committed
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
                       draw_gt: bool = True,
                       draw_pred: bool = True,
                       show: bool = False,
                       wait_time: float = 0,
                       out_file: Optional[str] = None,
                       vis_task: str = 'mono-det',
                       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.
494
            data_sample (:obj:`Det3DDataSample`, optional): Prediction
ZCMax's avatar
ZCMax committed
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
                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.
            vis-task (str): Visualization task. Defaults to 'mono-det'.
            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)

514
515
516
517
518
519
520
        gt_data_3d = None
        pred_data_3d = None
        gt_seg_data_3d = None
        pred_seg_data_3d = None
        gt_img_data = None
        pred_img_data = None

521
522
523
524
525
526
        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:
ZCMax's avatar
ZCMax committed
527
                assert 'img' in data_input
zhangshilong's avatar
zhangshilong committed
528
529
530
                if isinstance(data_input['img'], Tensor):
                    img = data_input['img'].permute(1, 2, 0).numpy()
                    img = img[..., [2, 1, 0]]  # bgr to rgb
531
532
                gt_img_data = self._draw_instances(img,
                                                   data_sample.gt_instances,
ZCMax's avatar
ZCMax committed
533
                                                   classes, palette)
534
            if 'gt_pts_seg' in data_sample:
ZCMax's avatar
ZCMax committed
535
536
537
538
539
540
541
                assert classes is not None, 'class information is ' \
                                            'not provided when ' \
                                            'visualizing panoptic ' \
                                            'segmentation results.'
                assert 'points' in data_input
                gt_seg_data_3d = \
                    self._draw_pts_sem_seg(data_input['points'],
542
                                           data_sample.pred_pts_seg,
543
                                           palette, ignore_index)
ZCMax's avatar
ZCMax committed
544

545
546
547
        if draw_pred and data_sample is not None:
            if 'pred_instances_3d' in data_sample:
                pred_instances_3d = data_sample.pred_instances_3d
548
549
                # .cpu can not be used for BaseInstancesBoxes3D
                # so we need to use .to('cpu')
ZCMax's avatar
ZCMax committed
550
                pred_instances_3d = pred_instances_3d[
551
                    pred_instances_3d.scores_3d > pred_score_thr].to('cpu')
ZCMax's avatar
ZCMax committed
552
553
                pred_data_3d = self._draw_instances_3d(data_input,
                                                       pred_instances_3d,
554
                                                       data_sample.metainfo,
ZCMax's avatar
ZCMax committed
555
                                                       vis_task, palette)
556
557
558
            if 'pred_instances' in data_sample:
                if 'img' in data_input and len(data_sample.pred_instances) > 0:
                    pred_instances = data_sample.pred_instances
559
560
561
562
563
564
565
                    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)
566
            if 'pred_pts_seg' in data_sample:
ZCMax's avatar
ZCMax committed
567
568
569
570
571
572
573
                assert classes is not None, 'class information is ' \
                                            'not provided when ' \
                                            'visualizing panoptic ' \
                                            'segmentation results.'
                assert 'points' in data_input
                pred_seg_data_3d = \
                    self._draw_pts_sem_seg(data_input['points'],
574
                                           data_sample.pred_pts_seg,
575
                                           palette, ignore_index)
ZCMax's avatar
ZCMax committed
576
577

        # monocular 3d object detection image
578
579
580
581
582
583
584
585
        if vis_task in ['mono-det', 'multi_modality-det']:
            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
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
        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,
                out_file,
                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:
610
                mmcv.imwrite(drawn_img_3d[..., ::-1], out_file + '.jpg')
ZCMax's avatar
ZCMax committed
611
            if drawn_img is not None:
612
                mmcv.imwrite(drawn_img[..., ::-1], out_file + '.jpg')
ZCMax's avatar
ZCMax committed
613
614
615
616
617
618
            if gt_data_3d is not None:
                write_obj(gt_data_3d['points'],
                          osp.join(out_file, 'points.obj'))
                write_oriented_bbox(gt_data_3d['bboxes_3d'],
                                    osp.join(out_file, 'gt_bbox.obj'))
            if pred_data_3d is not None:
619
620
621
622
623
                if 'points' in pred_data_3d:
                    write_obj(pred_data_3d['points'],
                              osp.join(out_file, 'points.obj'))
                    write_oriented_bbox(pred_data_3d['bboxes_3d'],
                                        osp.join(out_file, 'pred_bbox.obj'))
ZCMax's avatar
ZCMax committed
624
625
626
627
628
629
630
631
632
633
634
635
            if gt_seg_data_3d is not None:
                write_obj(gt_seg_data_3d['points'],
                          osp.join(out_file, 'points.obj'))
                write_obj(gt_seg_data_3d['seg_color'],
                          osp.join(out_file, 'gt_seg.obj'))
            if pred_seg_data_3d is not None:
                write_obj(pred_seg_data_3d['points'],
                          osp.join(out_file, 'points.obj'))
                write_obj(pred_seg_data_3d['seg_color'],
                          osp.join(out_file, 'pred_seg.obj'))
        else:
            self.add_image(name, drawn_img_3d, step)