local_visualizer.py 27.8 KB
Newer Older
ZCMax's avatar
ZCMax committed
1
2
3
4
5
6
7
# 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
8
from mmengine import mkdir_or_exist
ZCMax's avatar
ZCMax committed
9
10
11
from mmengine.dist import master_only
from torch import Tensor

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

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

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

from mmdet3d.registry import VISUALIZERS
25
26
27
28
from mmdet3d.structures import (BaseInstance3DBoxes, CameraInstance3DBoxes,
                                Coord3DMode, DepthInstance3DBoxes,
                                Det3DDataSample, LiDARInstance3DBoxes,
                                PointData)
ZCMax's avatar
ZCMax committed
29
30
31
32
33
34
35
36
37
38
39
40
41
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
42
      - draw_seg_mask: draw segmentation mask via per-point colorization
ZCMax's avatar
ZCMax committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61

    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
62
63
64
        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
65
66
67
68
69
70
        alpha (int, float): The transparency of bboxes or mask.
                Defaults to 0.8.

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

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

    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
99
                 vis_cfg: dict = dict(size=1, origin=[0, 0, 0]),
ZCMax's avatar
ZCMax committed
100
101
102
103
104
105
106
107
108
109
110
111
                 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)
112
        self.seg_num = 0
ZCMax's avatar
ZCMax committed
113
114
115
116
117
118
119
120
121
122
123
124
125
126

    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
127
        mesh_frame = geometry.TriangleMesh.create_coordinate_frame(**vis_cfg)
ZCMax's avatar
ZCMax committed
128
129
130
131
132
133
134
        o3d_vis.add_geometry(mesh_frame)

        return o3d_vis

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

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)

164
        if hasattr(self, 'pcd') and vis_task != 'lidar_seg':
ZCMax's avatar
ZCMax committed
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
            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
188
        self.points_colors = points_colors
ZCMax's avatar
ZCMax committed
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
258

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

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

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

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

318
    def draw_seg_mask(self, seg_mask_colors: np.array):
ZCMax's avatar
ZCMax committed
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
        """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
337
338
        self.set_points(
            seg_points, vis_task='lidar_seg', pcd_mode=2, mode='xyzrgb')
ZCMax's avatar
ZCMax committed
339
340
341
342
343
344
345
346
347
348
349
350

    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:
351
                'lidar_det', 'multi-modality_det', 'mono_det'.
ZCMax's avatar
ZCMax committed
352
353

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

        bboxes_3d = instances.bboxes_3d  # BaseInstance3DBoxes

359
        data_3d = dict()
ZCMax's avatar
ZCMax committed
360

361
        if vis_task in ['lidar_det', 'multi-modality_det']:
ZCMax's avatar
ZCMax committed
362
363
364
365
366
367
368
369
370
371
            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()

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

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

378
        if vis_task in ['mono_det', 'multi-modality_det']:
ZCMax's avatar
ZCMax committed
379
            assert 'img' in data_input
380
            img = data_input['img']
zhangshilong's avatar
zhangshilong committed
381
            if isinstance(data_input['img'], Tensor):
382
                img = img.permute(1, 2, 0).numpy()
zhangshilong's avatar
zhangshilong committed
383
384
                img = img[..., [2, 1, 0]]  # bgr to rgb
            self.set_image(img)
ZCMax's avatar
ZCMax committed
385
            self.draw_proj_bboxes_3d(bboxes_3d, input_meta)
386
387
388
            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
389
            drawn_img = self.get_image()
390
            data_3d['img'] = drawn_img
ZCMax's avatar
ZCMax committed
391
392
393
394

        return data_3d

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

411
412
413
        Returns:
            dict: the drawn points with color.
        """
ZCMax's avatar
ZCMax committed
414
415
416
417
        check_type('points', points, (np.ndarray, Tensor))

        points = tensor2ndarray(points)
        pts_sem_seg = tensor2ndarray(pts_seg.pts_semantic_mask)
418
        palette = np.array(palette)
ZCMax's avatar
ZCMax committed
419
420
421
422
423
424
425
426

        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)

427
        self.set_points(points, pcd_mode=2, vis_task='lidar_seg')
428
        self.draw_seg_mask(seg_color)
ZCMax's avatar
ZCMax committed
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445

        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:
446
                'lidar_det', 'multi-modality_det', 'mono_det', 'lidar_seg'.
ZCMax's avatar
ZCMax committed
447
448
449
450
451
452
453
454
455
456
            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.
        """
457
        if vis_task in ['lidar_det', 'multi-modality_det', 'lidar_seg']:
ZCMax's avatar
ZCMax committed
458
459
            self.o3d_vis.run()
            if out_file is not None:
460
                self.o3d_vis.capture_screen_image(out_file + '.png')
ZCMax's avatar
ZCMax committed
461
462
            self.o3d_vis.destroy_window()

463
        if vis_task in ['mono_det', 'multi-modality_det']:
ZCMax's avatar
ZCMax committed
464
465
466
467
468
            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)

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

519
520
521
522
523
524
525
        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

526
527
528
529
530
531
        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
532
533
534
535
536
537
538
539
                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
540
541
542
543
544
545
546
                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'],
547
                                           data_sample.pred_pts_seg,
548
                                           palette, ignore_index)
ZCMax's avatar
ZCMax committed
549

550
551
552
        if draw_pred and data_sample is not None:
            if 'pred_instances_3d' in data_sample:
                pred_instances_3d = data_sample.pred_instances_3d
553
554
                # .cpu can not be used for BaseInstancesBoxes3D
                # so we need to use .to('cpu')
ZCMax's avatar
ZCMax committed
555
                pred_instances_3d = pred_instances_3d[
556
                    pred_instances_3d.scores_3d > pred_score_thr].to('cpu')
ZCMax's avatar
ZCMax committed
557
558
                pred_data_3d = self._draw_instances_3d(data_input,
                                                       pred_instances_3d,
559
                                                       data_sample.metainfo,
ZCMax's avatar
ZCMax committed
560
                                                       vis_task, palette)
561
562
563
            if 'pred_instances' in data_sample:
                if 'img' in data_input and len(data_sample.pred_instances) > 0:
                    pred_instances = data_sample.pred_instances
564
565
566
567
568
569
570
                    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)
571
            if 'pred_pts_seg' in data_sample and vis_task == 'lidar_seg':
ZCMax's avatar
ZCMax committed
572
573
574
575
576
577
578
                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'],
579
                                           data_sample.pred_pts_seg,
580
                                           palette, ignore_index)
ZCMax's avatar
ZCMax committed
581
582

        # monocular 3d object detection image
583
584
585
586
587
588
589
590
        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
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
        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:
614
            mkdir_or_exist(out_file)
ZCMax's avatar
ZCMax committed
615
            if drawn_img_3d is not None:
616
                mmcv.imwrite(drawn_img_3d[..., ::-1], out_file + '.jpg')
ZCMax's avatar
ZCMax committed
617
            if drawn_img is not None:
618
                mmcv.imwrite(drawn_img[..., ::-1], out_file + '.jpg')
ZCMax's avatar
ZCMax committed
619
620
621
622
623
624
            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:
625
626
627
628
629
                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
630
631
632
633
634
635
636
637
638
639
640
641
            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)