loading.py 28 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
3
from typing import Sequence

zhangwenwei's avatar
zhangwenwei committed
4
5
import mmcv
import numpy as np
6
from mmcv.transforms import LoadImageFromFile
7
from mmcv.transforms.base import BaseTransform
zhangwenwei's avatar
zhangwenwei committed
8

9
from mmdet3d.core.points import BasePoints, get_points_type
10
11
from mmdet3d.registry import TRANSFORMS
from mmdet.datasets.pipelines import LoadAnnotations
zhangwenwei's avatar
zhangwenwei committed
12
13


14
@TRANSFORMS.register_module()
zhangwenwei's avatar
zhangwenwei committed
15
class LoadMultiViewImageFromFiles(object):
zhangwenwei's avatar
zhangwenwei committed
16
    """Load multi channel images from a list of separate channel files.
zhangwenwei's avatar
zhangwenwei committed
17

liyinhao's avatar
liyinhao committed
18
19
20
    Expects results['img_filename'] to be a list of filenames.

    Args:
21
        to_float32 (bool, optional): Whether to convert the img to float32.
liyinhao's avatar
liyinhao committed
22
            Defaults to False.
23
24
        color_type (str, optional): Color type of the file.
            Defaults to 'unchanged'.
zhangwenwei's avatar
zhangwenwei committed
25
    """
zhangwenwei's avatar
zhangwenwei committed
26

zhangwenwei's avatar
zhangwenwei committed
27
28
29
    def __init__(self, to_float32=False, color_type='unchanged'):
        self.to_float32 = to_float32
        self.color_type = color_type
zhangwenwei's avatar
zhangwenwei committed
30
31

    def __call__(self, results):
32
33
34
35
36
37
        """Call function to load multi-view image from files.

        Args:
            results (dict): Result dict containing multi-view image filenames.

        Returns:
38
            dict: The result dict containing the multi-view image data.
39
40
41
42
43
44
45
46
47
48
                Added keys and values are described below.

                - filename (str): Multi-view image filenames.
                - img (np.ndarray): Multi-view image arrays.
                - img_shape (tuple[int]): Shape of multi-view image arrays.
                - ori_shape (tuple[int]): Shape of original image arrays.
                - pad_shape (tuple[int]): Shape of padded image arrays.
                - scale_factor (float): Scale factor.
                - img_norm_cfg (dict): Normalization configuration of images.
        """
zhangwenwei's avatar
zhangwenwei committed
49
        filename = results['img_filename']
50
        # img is of shape (h, w, c, num_views)
zhangwenwei's avatar
zhangwenwei committed
51
52
53
54
55
        img = np.stack(
            [mmcv.imread(name, self.color_type) for name in filename], axis=-1)
        if self.to_float32:
            img = img.astype(np.float32)
        results['filename'] = filename
56
        # unravel to list, see `DefaultFormatBundle` in formatting.py
57
58
        # which will transpose each image separately and then stack into array
        results['img'] = [img[..., i] for i in range(img.shape[-1])]
zhangwenwei's avatar
zhangwenwei committed
59
60
61
62
63
64
65
66
67
68
        results['img_shape'] = img.shape
        results['ori_shape'] = img.shape
        # Set initial values for default meta_keys
        results['pad_shape'] = img.shape
        results['scale_factor'] = 1.0
        num_channels = 1 if len(img.shape) < 3 else img.shape[2]
        results['img_norm_cfg'] = dict(
            mean=np.zeros(num_channels, dtype=np.float32),
            std=np.ones(num_channels, dtype=np.float32),
            to_rgb=False)
zhangwenwei's avatar
zhangwenwei committed
69
70
71
        return results

    def __repr__(self):
72
        """str: Return a string that describes the module."""
73
74
75
76
        repr_str = self.__class__.__name__
        repr_str += f'(to_float32={self.to_float32}, '
        repr_str += f"color_type='{self.color_type}')"
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
77
78


79
@TRANSFORMS.register_module()
80
81
82
83
84
class LoadImageFromFileMono3D(LoadImageFromFile):
    """Load an image from file in monocular 3D object detection. Compared to 2D
    detection, additional camera parameters need to be loaded.

    Args:
85
        kwargs (dict): Arguments are the same as those in
86
87
88
89
90
91
92
93
94
95
96
97
98
            :class:`LoadImageFromFile`.
    """

    def __call__(self, results):
        """Call functions to load image and get image meta information.

        Args:
            results (dict): Result dict from :obj:`mmdet.CustomDataset`.

        Returns:
            dict: The dict contains loaded image and meta information.
        """
        super().__call__(results)
99
        results['cam2img'] = results['img_info']['cam_intrinsic']
100
101
102
        return results


103
@TRANSFORMS.register_module()
VVsssssk's avatar
VVsssssk committed
104
class LoadPointsFromMultiSweeps(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
105
    """Load points from multiple sweeps.
zhangwenwei's avatar
zhangwenwei committed
106

zhangwenwei's avatar
zhangwenwei committed
107
108
109
    This is usually used for nuScenes dataset to utilize previous sweeps.

    Args:
110
111
112
113
114
115
116
        sweeps_num (int, optional): Number of sweeps. Defaults to 10.
        load_dim (int, optional): Dimension number of the loaded points.
            Defaults to 5.
        use_dim (list[int], optional): Which dimension to use.
            Defaults to [0, 1, 2, 4].
        file_client_args (dict, optional): Config dict of file clients,
            refer to
zhangwenwei's avatar
zhangwenwei committed
117
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
118
            for more details. Defaults to dict(backend='disk').
119
        pad_empty_sweeps (bool, optional): Whether to repeat keyframe when
120
            sweeps is empty. Defaults to False.
121
        remove_close (bool, optional): Whether to remove close points.
122
            Defaults to False.
123
        test_mode (bool, optional): If `test_mode=True`, it will not
124
125
            randomly sample sweeps but select the nearest N frames.
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
126
127
128
129
130
    """

    def __init__(self,
                 sweeps_num=10,
                 load_dim=5,
131
132
133
134
135
                 use_dim=[0, 1, 2, 4],
                 file_client_args=dict(backend='disk'),
                 pad_empty_sweeps=False,
                 remove_close=False,
                 test_mode=False):
zhangwenwei's avatar
zhangwenwei committed
136
        self.load_dim = load_dim
zhangwenwei's avatar
zhangwenwei committed
137
        self.sweeps_num = sweeps_num
138
        self.use_dim = use_dim
zhangwenwei's avatar
zhangwenwei committed
139
140
        self.file_client_args = file_client_args.copy()
        self.file_client = None
141
142
143
        self.pad_empty_sweeps = pad_empty_sweeps
        self.remove_close = remove_close
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
144
145

    def _load_points(self, pts_filename):
146
147
148
149
150
151
152
153
        """Private function to load point clouds data.

        Args:
            pts_filename (str): Filename of point clouds data.

        Returns:
            np.ndarray: An array containing point clouds data.
        """
zhangwenwei's avatar
zhangwenwei committed
154
155
156
157
158
159
160
161
162
163
164
165
        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)
        try:
            pts_bytes = self.file_client.get(pts_filename)
            points = np.frombuffer(pts_bytes, dtype=np.float32)
        except ConnectionError:
            mmcv.check_file_exist(pts_filename)
            if pts_filename.endswith('.npy'):
                points = np.load(pts_filename)
            else:
                points = np.fromfile(pts_filename, dtype=np.float32)
        return points
zhangwenwei's avatar
zhangwenwei committed
166

167
168
169
170
    def _remove_close(self, points, radius=1.0):
        """Removes point too close within a certain radius from origin.

        Args:
171
            points (np.ndarray | :obj:`BasePoints`): Sweep points.
172
            radius (float, optional): Radius below which points are removed.
173
174
175
176
177
                Defaults to 1.0.

        Returns:
            np.ndarray: Points after removing.
        """
178
179
180
181
182
183
184
185
        if isinstance(points, np.ndarray):
            points_numpy = points
        elif isinstance(points, BasePoints):
            points_numpy = points.tensor.numpy()
        else:
            raise NotImplementedError
        x_filt = np.abs(points_numpy[:, 0]) < radius
        y_filt = np.abs(points_numpy[:, 1]) < radius
186
        not_close = np.logical_not(np.logical_and(x_filt, y_filt))
187
        return points[not_close]
188

VVsssssk's avatar
VVsssssk committed
189
    def transform(self, results):
190
191
192
        """Call function to load multi-sweep point clouds from files.

        Args:
193
            results (dict): Result dict containing multi-sweep point cloud
194
195
196
                filenames.

        Returns:
197
            dict: The result dict containing the multi-sweep points data.
198
199
                Added key and value are described below.

200
                - points (np.ndarray | :obj:`BasePoints`): Multi-sweep point
201
                    cloud arrays.
202
        """
zhangwenwei's avatar
zhangwenwei committed
203
        points = results['points']
204
        points.tensor[:, 4] = 0
zhangwenwei's avatar
zhangwenwei committed
205
206
        sweep_points_list = [points]
        ts = results['timestamp']
VVsssssk's avatar
VVsssssk committed
207
208
209
210
211
212
213
        if 'lidar_sweeps' not in results:
            if self.pad_empty_sweeps:
                for i in range(self.sweeps_num):
                    if self.remove_close:
                        sweep_points_list.append(self._remove_close(points))
                    else:
                        sweep_points_list.append(points)
214
        else:
VVsssssk's avatar
VVsssssk committed
215
216
            if len(results['lidar_sweeps']) <= self.sweeps_num:
                choices = np.arange(len(results['lidar_sweeps']))
217
218
219
220
            elif self.test_mode:
                choices = np.arange(self.sweeps_num)
            else:
                choices = np.random.choice(
VVsssssk's avatar
VVsssssk committed
221
222
223
                    len(results['lidar_sweeps']),
                    self.sweeps_num,
                    replace=False)
224
            for idx in choices:
VVsssssk's avatar
VVsssssk committed
225
226
227
                sweep = results['lidar_sweeps'][idx]
                points_sweep = self._load_points(
                    sweep['lidar_points']['lidar_path'])
228
229
230
                points_sweep = np.copy(points_sweep).reshape(-1, self.load_dim)
                if self.remove_close:
                    points_sweep = self._remove_close(points_sweep)
VVsssssk's avatar
VVsssssk committed
231
232
233
234
235
                # bc-breaking: Timestamp has divided 1e6 in pkl infos.
                sweep_ts = sweep['timestamp']
                lidar2cam = np.array(sweep['lidar_points']['lidar2sensor'])
                points_sweep[:, :3] = points_sweep[:, :3] @ lidar2cam[:3, :3]
                points_sweep[:, :3] -= lidar2cam[:3, 3]
236
                points_sweep[:, 4] = ts - sweep_ts
237
                points_sweep = points.new_point(points_sweep)
238
239
                sweep_points_list.append(points_sweep)

240
241
        points = points.cat(sweep_points_list)
        points = points[:, self.use_dim]
zhangwenwei's avatar
zhangwenwei committed
242
243
244
245
        results['points'] = points
        return results

    def __repr__(self):
246
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
247
        return f'{self.__class__.__name__}(sweeps_num={self.sweeps_num})'
wuyuefeng's avatar
wuyuefeng committed
248
249


250
@TRANSFORMS.register_module()
251
class PointSegClassMapping(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
252
253
    """Map original semantic class to valid category ids.

254
255
256
257
258
259
260
261
262
263
    Required Keys:

    - lidar_points (dict)

        - lidar_path (str)

    Added Keys:

    - points (np.float32)

wuyuefeng's avatar
wuyuefeng committed
264
265
266
267
    Map valid classes as 0~len(valid_cat_ids)-1 and
    others as len(valid_cat_ids).

    Args:
268
        valid_cat_ids (tuple[int]): A tuple of valid category.
269
270
        max_cat_id (int, optional): The max possible cat_id in input
            segmentation mask. Defaults to 40.
wuyuefeng's avatar
wuyuefeng committed
271
272
    """

273
274
275
    def __init__(self,
                 valid_cat_ids: Sequence[int],
                 max_cat_id: int = 40) -> None:
276
277
278
        assert max_cat_id >= np.max(valid_cat_ids), \
            'max_cat_id should be greater than maximum id in valid_cat_ids'

wuyuefeng's avatar
wuyuefeng committed
279
        self.valid_cat_ids = valid_cat_ids
280
281
282
283
284
285
286
287
        self.max_cat_id = int(max_cat_id)

        # build cat_id to class index mapping
        neg_cls = len(valid_cat_ids)
        self.cat_id2class = np.ones(
            self.max_cat_id + 1, dtype=np.int) * neg_cls
        for cls_idx, cat_id in enumerate(valid_cat_ids):
            self.cat_id2class[cat_id] = cls_idx
wuyuefeng's avatar
wuyuefeng committed
288

289
    def transform(self, results: dict) -> None:
290
291
292
293
294
295
        """Call function to map original semantic class to valid category ids.

        Args:
            results (dict): Result dict containing point semantic masks.

        Returns:
296
            dict: The result dict containing the mapped category ids.
297
298
299
300
                Updated key and value are described below.

                - pts_semantic_mask (np.ndarray): Mapped semantic masks.
        """
wuyuefeng's avatar
wuyuefeng committed
301
302
303
        assert 'pts_semantic_mask' in results
        pts_semantic_mask = results['pts_semantic_mask']

304
        converted_pts_sem_mask = self.cat_id2class[pts_semantic_mask]
wuyuefeng's avatar
wuyuefeng committed
305

306
        results['pts_semantic_mask'] = converted_pts_sem_mask
wuyuefeng's avatar
wuyuefeng committed
307
308
309
        return results

    def __repr__(self):
310
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
311
        repr_str = self.__class__.__name__
312
313
        repr_str += f'(valid_cat_ids={self.valid_cat_ids}, '
        repr_str += f'max_cat_id={self.max_cat_id})'
wuyuefeng's avatar
wuyuefeng committed
314
315
316
        return repr_str


317
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
318
class NormalizePointsColor(object):
zhangwenwei's avatar
zhangwenwei committed
319
    """Normalize color of points.
wuyuefeng's avatar
wuyuefeng committed
320
321
322
323
324
325
326
327
328

    Args:
        color_mean (list[float]): Mean color of the point cloud.
    """

    def __init__(self, color_mean):
        self.color_mean = color_mean

    def __call__(self, results):
329
330
331
332
333
334
        """Call function to normalize color of points.

        Args:
            results (dict): Result dict containing point clouds data.

        Returns:
335
            dict: The result dict containing the normalized points.
336
337
                Updated key and value are described below.

338
                - points (:obj:`BasePoints`): Points after color normalization.
339
        """
wuyuefeng's avatar
wuyuefeng committed
340
        points = results['points']
341
        assert points.attribute_dims is not None and \
342
343
               'color' in points.attribute_dims.keys(), \
               'Expect points have color attribute'
344
345
        if self.color_mean is not None:
            points.color = points.color - \
346
                           points.color.new_tensor(self.color_mean)
347
        points.color = points.color / 255.0
wuyuefeng's avatar
wuyuefeng committed
348
349
350
351
        results['points'] = points
        return results

    def __repr__(self):
352
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
353
        repr_str = self.__class__.__name__
354
        repr_str += f'(color_mean={self.color_mean})'
wuyuefeng's avatar
wuyuefeng committed
355
356
357
        return repr_str


358
@TRANSFORMS.register_module()
jshilong's avatar
jshilong committed
359
class LoadPointsFromFile(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
360
361
    """Load Points From File.

jshilong's avatar
jshilong committed
362
363
364
365
366
367
368
369
370
    Required Keys:

    - lidar_points (dict)

        - lidar_path (str)

    Added Keys:

    - points (np.float32)
wuyuefeng's avatar
wuyuefeng committed
371
372

    Args:
373
374
375
376
377
        coord_type (str): The type of coordinates of points cloud.
            Available options includes:
            - 'LIDAR': Points in LiDAR coordinates.
            - 'DEPTH': Points in depth coordinates, usually for indoor dataset.
            - 'CAMERA': Points in camera coordinates.
378
        load_dim (int, optional): The dimension of the loaded points.
379
            Defaults to 6.
380
        use_dim (list[int], optional): Which dimensions of the points to use.
liyinhao's avatar
liyinhao committed
381
382
            Defaults to [0, 1, 2]. For KITTI dataset, set use_dim=4
            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
383
384
385
386
387
388
        shift_height (bool, optional): Whether to use shifted height.
            Defaults to False.
        use_color (bool, optional): Whether to use color features.
            Defaults to False.
        file_client_args (dict, optional): Config dict of file clients,
            refer to
wuyuefeng's avatar
wuyuefeng committed
389
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
390
            for more details. Defaults to dict(backend='disk').
wuyuefeng's avatar
wuyuefeng committed
391
392
    """

jshilong's avatar
jshilong committed
393
394
395
396
397
398
399
400
401
    def __init__(
        self,
        coord_type: str,
        load_dim: int = 6,
        use_dim: list = [0, 1, 2],
        shift_height: bool = False,
        use_color: bool = False,
        file_client_args: dict = dict(backend='disk')
    ) -> None:
wuyuefeng's avatar
wuyuefeng committed
402
        self.shift_height = shift_height
403
        self.use_color = use_color
wuyuefeng's avatar
wuyuefeng committed
404
405
406
407
        if isinstance(use_dim, int):
            use_dim = list(range(use_dim))
        assert max(use_dim) < load_dim, \
            f'Expect all used dimensions < {load_dim}, got {use_dim}'
408
        assert coord_type in ['CAMERA', 'LIDAR', 'DEPTH']
wuyuefeng's avatar
wuyuefeng committed
409

410
        self.coord_type = coord_type
wuyuefeng's avatar
wuyuefeng committed
411
412
413
414
415
        self.load_dim = load_dim
        self.use_dim = use_dim
        self.file_client_args = file_client_args.copy()
        self.file_client = None

jshilong's avatar
jshilong committed
416
    def _load_points(self, pts_filename: str) -> np.ndarray:
417
418
419
420
421
422
423
424
        """Private function to load point clouds data.

        Args:
            pts_filename (str): Filename of point clouds data.

        Returns:
            np.ndarray: An array containing point clouds data.
        """
wuyuefeng's avatar
wuyuefeng committed
425
426
427
428
429
430
431
432
433
434
435
        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)
        try:
            pts_bytes = self.file_client.get(pts_filename)
            points = np.frombuffer(pts_bytes, dtype=np.float32)
        except ConnectionError:
            mmcv.check_file_exist(pts_filename)
            if pts_filename.endswith('.npy'):
                points = np.load(pts_filename)
            else:
                points = np.fromfile(pts_filename, dtype=np.float32)
436

wuyuefeng's avatar
wuyuefeng committed
437
438
        return points

jshilong's avatar
jshilong committed
439
440
    def transform(self, results: dict) -> dict:
        """Method to load points data from file.
441
442
443
444
445

        Args:
            results (dict): Result dict containing point clouds data.

        Returns:
446
            dict: The result dict containing the point clouds data.
447
448
                Added key and value are described below.

449
                - points (:obj:`BasePoints`): Point clouds data.
450
        """
jshilong's avatar
jshilong committed
451
452
        pts_file_path = results['lidar_points']['lidar_path']
        points = self._load_points(pts_file_path)
wuyuefeng's avatar
wuyuefeng committed
453
454
        points = points.reshape(-1, self.load_dim)
        points = points[:, self.use_dim]
455
        attribute_dims = None
wuyuefeng's avatar
wuyuefeng committed
456
457
458
459

        if self.shift_height:
            floor_height = np.percentile(points[:, 2], 0.99)
            height = points[:, 2] - floor_height
460
461
462
            points = np.concatenate(
                [points[:, :3],
                 np.expand_dims(height, 1), points[:, 3:]], 1)
463
464
            attribute_dims = dict(height=3)

465
466
467
468
469
470
471
472
473
474
475
        if self.use_color:
            assert len(self.use_dim) >= 6
            if attribute_dims is None:
                attribute_dims = dict()
            attribute_dims.update(
                dict(color=[
                    points.shape[1] - 3,
                    points.shape[1] - 2,
                    points.shape[1] - 1,
                ]))

476
477
478
        points_class = get_points_type(self.coord_type)
        points = points_class(
            points, points_dim=points.shape[-1], attribute_dims=attribute_dims)
wuyuefeng's avatar
wuyuefeng committed
479
        results['points'] = points
480

wuyuefeng's avatar
wuyuefeng committed
481
482
483
        return results

    def __repr__(self):
484
        """str: Return a string that describes the module."""
liyinhao's avatar
liyinhao committed
485
        repr_str = self.__class__.__name__ + '('
486
487
488
489
490
        repr_str += f'shift_height={self.shift_height}, '
        repr_str += f'use_color={self.use_color}, '
        repr_str += f'file_client_args={self.file_client_args}, '
        repr_str += f'load_dim={self.load_dim}, '
        repr_str += f'use_dim={self.use_dim})'
wuyuefeng's avatar
wuyuefeng committed
491
492
493
        return repr_str


494
@TRANSFORMS.register_module()
495
496
497
498
499
500
501
502
class LoadPointsFromDict(LoadPointsFromFile):
    """Load Points From Dict."""

    def __call__(self, results):
        assert 'points' in results
        return results


503
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
504
505
506
507
508
509
class LoadAnnotations3D(LoadAnnotations):
    """Load Annotations3D.

    Load instance mask and semantic mask of points and
    encapsulate the items into related fields.

jshilong's avatar
jshilong committed
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
    Required Keys:

    - ann_info (dict)
        - gt_bboxes_3d (:obj:`LiDARInstance3DBoxes` |
          :obj:`DepthInstance3DBoxes` | :obj:`CameraInstance3DBoxes`):
          3D ground truth bboxes. Only when `with_bbox_3d` is True
        - gt_labels_3d (np.int64): Labels of ground truths.
          Only when `with_label_3d` is True.
        - gt_bboxes (np.float32): 2D ground truth bboxes.
          Only when `with_bbox` is True.
        - gt_labels (np.ndarray): Labels of ground truths.
          Only when `with_label` is True.
        - depths (np.ndarray): Only when
          `with_bbox_depth` is True.
        - centers_2d (np.ndarray): Only when
          `with_bbox_depth` is True.
        - attr_labels (np.ndarray): Attribute labels of instances.
          Only when `with_attr_label` is True.

    - pts_instance_mask_path (str): Path of instance mask file.
      Only when `with_mask_3d` is True.
    - pts_semantic_mask_path (str): Path of semantic mask file.
      Only when

    Added Keys:

    - gt_bboxes_3d (:obj:`LiDARInstance3DBoxes` |
      :obj:`DepthInstance3DBoxes` | :obj:`CameraInstance3DBoxes`):
      3D ground truth bboxes. Only when `with_bbox_3d` is True
    - gt_labels_3d (np.int64): Labels of ground truths.
      Only when `with_label_3d` is True.
    - gt_bboxes (np.float32): 2D ground truth bboxes.
      Only when `with_bbox` is True.
    - gt_labels (np.int64): Labels of ground truths.
      Only when `with_label` is True.
    - depths (np.float32): Only when
      `with_bbox_depth` is True.
    - centers_2d (np.ndarray): Only when
      `with_bbox_depth` is True.
    - attr_labels (np.int64): Attribute labels of instances.
      Only when `with_attr_label` is True.
    - pts_instance_mask (np.int64): Instance mask of each point.
      Only when `with_mask_3d` is True.
    - pts_semantic_mask (np.int64): Semantic mask of each point.
      Only when `with_seg_3d` is True.

wuyuefeng's avatar
wuyuefeng committed
556
557
558
559
560
    Args:
        with_bbox_3d (bool, optional): Whether to load 3D boxes.
            Defaults to True.
        with_label_3d (bool, optional): Whether to load 3D labels.
            Defaults to True.
561
562
        with_attr_label (bool, optional): Whether to load attribute label.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
563
564
565
566
567
568
569
570
571
572
573
574
        with_mask_3d (bool, optional): Whether to load 3D instance masks.
            for points. Defaults to False.
        with_seg_3d (bool, optional): Whether to load 3D semantic masks.
            for points. Defaults to False.
        with_bbox (bool, optional): Whether to load 2D boxes.
            Defaults to False.
        with_label (bool, optional): Whether to load 2D labels.
            Defaults to False.
        with_mask (bool, optional): Whether to load 2D instance masks.
            Defaults to False.
        with_seg (bool, optional): Whether to load 2D semantic masks.
            Defaults to False.
575
576
        with_bbox_depth (bool, optional): Whether to load 2.5D boxes.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
577
578
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
579
        seg_3d_dtype (dtype, optional): Dtype of 3D semantic masks.
jshilong's avatar
jshilong committed
580
            Defaults to int64.
wuyuefeng's avatar
wuyuefeng committed
581
582
583
584
585
        file_client_args (dict): Config dict of file clients, refer to
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
            for more details.
    """

jshilong's avatar
jshilong committed
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
    def __init__(
        self,
        with_bbox_3d: bool = True,
        with_label_3d: bool = True,
        with_attr_label: bool = False,
        with_mask_3d: bool = False,
        with_seg_3d: bool = False,
        with_bbox: bool = False,
        with_label: bool = False,
        with_mask: bool = False,
        with_seg: bool = False,
        with_bbox_depth: bool = False,
        poly2mask: bool = True,
        seg_3d_dtype: np.dtype = np.int64,
        file_client_args: dict = dict(backend='disk')
    ) -> None:
wuyuefeng's avatar
wuyuefeng committed
602
        super().__init__(
jshilong's avatar
jshilong committed
603
604
605
606
607
            with_bbox=with_bbox,
            with_label=with_label,
            with_mask=with_mask,
            with_seg=with_seg,
            poly2mask=poly2mask,
wuyuefeng's avatar
wuyuefeng committed
608
609
            file_client_args=file_client_args)
        self.with_bbox_3d = with_bbox_3d
610
        self.with_bbox_depth = with_bbox_depth
wuyuefeng's avatar
wuyuefeng committed
611
        self.with_label_3d = with_label_3d
612
        self.with_attr_label = with_attr_label
wuyuefeng's avatar
wuyuefeng committed
613
614
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d
615
        self.seg_3d_dtype = seg_3d_dtype
wuyuefeng's avatar
wuyuefeng committed
616

jshilong's avatar
jshilong committed
617
618
619
    def _load_bboxes_3d(self, results: dict) -> dict:
        """Private function to move the 3D bounding box annotation from
        `ann_info` field to the root of `results`.
620
621
622
623
624
625
626

        Args:
            results (dict): Result dict from :obj:`mmdet3d.CustomDataset`.

        Returns:
            dict: The dict containing loaded 3D bounding box annotations.
        """
jshilong's avatar
jshilong committed
627

wuyuefeng's avatar
wuyuefeng committed
628
629
630
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
        return results

jshilong's avatar
jshilong committed
631
    def _load_bboxes_depth(self, results: dict) -> dict:
632
633
634
635
636
637
638
639
        """Private function to load 2.5D bounding box annotations.

        Args:
            results (dict): Result dict from :obj:`mmdet3d.CustomDataset`.

        Returns:
            dict: The dict containing loaded 2.5D bounding box annotations.
        """
jshilong's avatar
jshilong committed
640

641
        results['depths'] = results['ann_info']['depths']
jshilong's avatar
jshilong committed
642
        results['centers_2d'] = results['ann_info']['centers_2d']
643
644
        return results

jshilong's avatar
jshilong committed
645
    def _load_labels_3d(self, results: dict) -> dict:
646
647
648
649
650
651
652
653
        """Private function to load label annotations.

        Args:
            results (dict): Result dict from :obj:`mmdet3d.CustomDataset`.

        Returns:
            dict: The dict containing loaded label annotations.
        """
jshilong's avatar
jshilong committed
654

wuyuefeng's avatar
wuyuefeng committed
655
656
657
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results

jshilong's avatar
jshilong committed
658
    def _load_attr_labels(self, results: dict) -> dict:
659
660
661
662
663
664
665
666
667
668
669
        """Private function to load label annotations.

        Args:
            results (dict): Result dict from :obj:`mmdet3d.CustomDataset`.

        Returns:
            dict: The dict containing loaded label annotations.
        """
        results['attr_labels'] = results['ann_info']['attr_labels']
        return results

jshilong's avatar
jshilong committed
670
    def _load_masks_3d(self, results: dict) -> dict:
671
672
673
674
675
676
677
678
        """Private function to load 3D mask annotations.

        Args:
            results (dict): Result dict from :obj:`mmdet3d.CustomDataset`.

        Returns:
            dict: The dict containing loaded 3D mask annotations.
        """
jshilong's avatar
jshilong committed
679
        pts_instance_mask_path = results['pts_instance_mask_path']
wuyuefeng's avatar
wuyuefeng committed
680
681
682
683
684

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)
        try:
            mask_bytes = self.file_client.get(pts_instance_mask_path)
685
            pts_instance_mask = np.frombuffer(mask_bytes, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
686
687
688
        except ConnectionError:
            mmcv.check_file_exist(pts_instance_mask_path)
            pts_instance_mask = np.fromfile(
WRH's avatar
WRH committed
689
                pts_instance_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
690
691

        results['pts_instance_mask'] = pts_instance_mask
jshilong's avatar
jshilong committed
692
693
694
        # 'eval_ann_info' will be passed to evaluator
        if 'eval_ann_info' in results:
            results['eval_ann_info']['pts_instance_mask'] = pts_instance_mask
wuyuefeng's avatar
wuyuefeng committed
695
696
        return results

jshilong's avatar
jshilong committed
697
    def _load_semantic_seg_3d(self, results: dict) -> dict:
698
699
700
701
702
703
704
705
        """Private function to load 3D semantic segmentation annotations.

        Args:
            results (dict): Result dict from :obj:`mmdet3d.CustomDataset`.

        Returns:
            dict: The dict containing the semantic segmentation annotations.
        """
jshilong's avatar
jshilong committed
706
        pts_semantic_mask_path = results['pts_semantic_mask_path']
wuyuefeng's avatar
wuyuefeng committed
707
708
709
710
711
712

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)
        try:
            mask_bytes = self.file_client.get(pts_semantic_mask_path)
            # add .copy() to fix read-only bug
713
714
            pts_semantic_mask = np.frombuffer(
                mask_bytes, dtype=self.seg_3d_dtype).copy()
wuyuefeng's avatar
wuyuefeng committed
715
716
717
        except ConnectionError:
            mmcv.check_file_exist(pts_semantic_mask_path)
            pts_semantic_mask = np.fromfile(
WRH's avatar
WRH committed
718
                pts_semantic_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
719
720

        results['pts_semantic_mask'] = pts_semantic_mask
jshilong's avatar
jshilong committed
721
722
723
        # 'eval_ann_info' will be passed to evaluator
        if 'eval_ann_info' in results:
            results['eval_ann_info']['pts_semantic_mask'] = pts_semantic_mask
wuyuefeng's avatar
wuyuefeng committed
724
725
        return results

jshilong's avatar
jshilong committed
726
727
    def transform(self, results: dict) -> dict:
        """Function to load multiple types annotations.
728
729
730
731
732
733

        Args:
            results (dict): Result dict from :obj:`mmdet3d.CustomDataset`.

        Returns:
            dict: The dict containing loaded 3D bounding box, label, mask and
jshilong's avatar
jshilong committed
734
            semantic segmentation annotations.
735
        """
jshilong's avatar
jshilong committed
736
        results = super().transform(results)
wuyuefeng's avatar
wuyuefeng committed
737
738
        if self.with_bbox_3d:
            results = self._load_bboxes_3d(results)
739
740
        if self.with_bbox_depth:
            results = self._load_bboxes_depth(results)
wuyuefeng's avatar
wuyuefeng committed
741
742
        if self.with_label_3d:
            results = self._load_labels_3d(results)
743
744
        if self.with_attr_label:
            results = self._load_attr_labels(results)
wuyuefeng's avatar
wuyuefeng committed
745
746
747
748
749
750
751
752
        if self.with_mask_3d:
            results = self._load_masks_3d(results)
        if self.with_seg_3d:
            results = self._load_semantic_seg_3d(results)

        return results

    def __repr__(self):
753
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
754
755
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
liyinhao's avatar
liyinhao committed
756
757
        repr_str += f'{indent_str}with_bbox_3d={self.with_bbox_3d}, '
        repr_str += f'{indent_str}with_label_3d={self.with_label_3d}, '
758
        repr_str += f'{indent_str}with_attr_label={self.with_attr_label}, '
liyinhao's avatar
liyinhao committed
759
760
761
762
763
764
        repr_str += f'{indent_str}with_mask_3d={self.with_mask_3d}, '
        repr_str += f'{indent_str}with_seg_3d={self.with_seg_3d}, '
        repr_str += f'{indent_str}with_bbox={self.with_bbox}, '
        repr_str += f'{indent_str}with_label={self.with_label}, '
        repr_str += f'{indent_str}with_mask={self.with_mask}, '
        repr_str += f'{indent_str}with_seg={self.with_seg}, '
765
        repr_str += f'{indent_str}with_bbox_depth={self.with_bbox_depth}, '
wuyuefeng's avatar
wuyuefeng committed
766
767
        repr_str += f'{indent_str}poly2mask={self.poly2mask})'
        return repr_str