loading.py 27.2 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangwenwei's avatar
zhangwenwei committed
2
3
import mmcv
import numpy as np
jshilong's avatar
jshilong committed
4
from mmcv import BaseTransform
5
from mmcv.transforms import LoadImageFromFile
zhangwenwei's avatar
zhangwenwei committed
6

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


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

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

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

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

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

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

        Returns:
36
            dict: The result dict containing the multi-view image data.
37
38
39
40
41
42
43
44
45
46
                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
47
        filename = results['img_filename']
48
        # img is of shape (h, w, c, num_views)
zhangwenwei's avatar
zhangwenwei committed
49
50
51
52
53
        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
54
        # unravel to list, see `DefaultFormatBundle` in formatting.py
55
56
        # 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
57
58
59
60
61
62
63
64
65
66
        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
67
68
69
        return results

    def __repr__(self):
70
        """str: Return a string that describes the module."""
71
72
73
74
        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
75
76


77
@TRANSFORMS.register_module()
78
79
80
81
82
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:
83
        kwargs (dict): Arguments are the same as those in
84
85
86
87
88
89
90
91
92
93
94
95
96
            :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)
97
        results['cam2img'] = results['img_info']['cam_intrinsic']
98
99
100
        return results


101
@TRANSFORMS.register_module()
zhangwenwei's avatar
zhangwenwei committed
102
class LoadPointsFromMultiSweeps(object):
zhangwenwei's avatar
zhangwenwei committed
103
    """Load points from multiple sweeps.
zhangwenwei's avatar
zhangwenwei committed
104

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

    Args:
108
109
110
111
112
113
114
        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
115
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
116
            for more details. Defaults to dict(backend='disk').
117
        pad_empty_sweeps (bool, optional): Whether to repeat keyframe when
118
            sweeps is empty. Defaults to False.
119
        remove_close (bool, optional): Whether to remove close points.
120
            Defaults to False.
121
        test_mode (bool, optional): If `test_mode=True`, it will not
122
123
            randomly sample sweeps but select the nearest N frames.
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
124
125
126
127
128
    """

    def __init__(self,
                 sweeps_num=10,
                 load_dim=5,
129
130
131
132
133
                 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
134
        self.load_dim = load_dim
zhangwenwei's avatar
zhangwenwei committed
135
        self.sweeps_num = sweeps_num
136
        self.use_dim = use_dim
zhangwenwei's avatar
zhangwenwei committed
137
138
        self.file_client_args = file_client_args.copy()
        self.file_client = None
139
140
141
        self.pad_empty_sweeps = pad_empty_sweeps
        self.remove_close = remove_close
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
142
143

    def _load_points(self, pts_filename):
144
145
146
147
148
149
150
151
        """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
152
153
154
155
156
157
158
159
160
161
162
163
        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
164

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

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

        Returns:
            np.ndarray: Points after removing.
        """
176
177
178
179
180
181
182
183
        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
184
        not_close = np.logical_not(np.logical_and(x_filt, y_filt))
185
        return points[not_close]
186

zhangwenwei's avatar
zhangwenwei committed
187
    def __call__(self, results):
188
189
190
        """Call function to load multi-sweep point clouds from files.

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

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

198
                - points (np.ndarray | :obj:`BasePoints`): Multi-sweep point
199
                    cloud arrays.
200
        """
zhangwenwei's avatar
zhangwenwei committed
201
        points = results['points']
202
        points.tensor[:, 4] = 0
zhangwenwei's avatar
zhangwenwei committed
203
204
        sweep_points_list = [points]
        ts = results['timestamp']
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
        if self.pad_empty_sweeps and len(results['sweeps']) == 0:
            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)
        else:
            if len(results['sweeps']) <= self.sweeps_num:
                choices = np.arange(len(results['sweeps']))
            elif self.test_mode:
                choices = np.arange(self.sweeps_num)
            else:
                choices = np.random.choice(
                    len(results['sweeps']), self.sweeps_num, replace=False)
            for idx in choices:
                sweep = results['sweeps'][idx]
                points_sweep = self._load_points(sweep['data_path'])
                points_sweep = np.copy(points_sweep).reshape(-1, self.load_dim)
                if self.remove_close:
                    points_sweep = self._remove_close(points_sweep)
                sweep_ts = sweep['timestamp'] / 1e6
                points_sweep[:, :3] = points_sweep[:, :3] @ sweep[
                    'sensor2lidar_rotation'].T
                points_sweep[:, :3] += sweep['sensor2lidar_translation']
                points_sweep[:, 4] = ts - sweep_ts
230
                points_sweep = points.new_point(points_sweep)
231
232
                sweep_points_list.append(points_sweep)

233
234
        points = points.cat(sweep_points_list)
        points = points[:, self.use_dim]
zhangwenwei's avatar
zhangwenwei committed
235
236
237
238
        results['points'] = points
        return results

    def __repr__(self):
239
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
240
        return f'{self.__class__.__name__}(sweeps_num={self.sweeps_num})'
wuyuefeng's avatar
wuyuefeng committed
241
242


243
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
244
245
246
247
248
249
250
class PointSegClassMapping(object):
    """Map original semantic class to valid category ids.

    Map valid classes as 0~len(valid_cat_ids)-1 and
    others as len(valid_cat_ids).

    Args:
251
        valid_cat_ids (tuple[int]): A tuple of valid category.
252
253
        max_cat_id (int, optional): The max possible cat_id in input
            segmentation mask. Defaults to 40.
wuyuefeng's avatar
wuyuefeng committed
254
255
    """

256
257
258
259
    def __init__(self, valid_cat_ids, max_cat_id=40):
        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
260
        self.valid_cat_ids = valid_cat_ids
261
262
263
264
265
266
267
268
        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
269
270

    def __call__(self, results):
271
272
273
274
275
276
        """Call function to map original semantic class to valid category ids.

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

        Returns:
277
            dict: The result dict containing the mapped category ids.
278
279
280
281
                Updated key and value are described below.

                - pts_semantic_mask (np.ndarray): Mapped semantic masks.
        """
wuyuefeng's avatar
wuyuefeng committed
282
283
284
        assert 'pts_semantic_mask' in results
        pts_semantic_mask = results['pts_semantic_mask']

285
        converted_pts_sem_mask = self.cat_id2class[pts_semantic_mask]
wuyuefeng's avatar
wuyuefeng committed
286

287
        results['pts_semantic_mask'] = converted_pts_sem_mask
wuyuefeng's avatar
wuyuefeng committed
288
289
290
        return results

    def __repr__(self):
291
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
292
        repr_str = self.__class__.__name__
293
294
        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
295
296
297
        return repr_str


298
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
299
class NormalizePointsColor(object):
zhangwenwei's avatar
zhangwenwei committed
300
    """Normalize color of points.
wuyuefeng's avatar
wuyuefeng committed
301
302
303
304
305
306
307
308
309

    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):
310
311
312
313
314
315
        """Call function to normalize color of points.

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

        Returns:
316
            dict: The result dict containing the normalized points.
317
318
                Updated key and value are described below.

319
                - points (:obj:`BasePoints`): Points after color normalization.
320
        """
wuyuefeng's avatar
wuyuefeng committed
321
        points = results['points']
322
323
324
325
326
327
328
        assert points.attribute_dims is not None and \
            'color' in points.attribute_dims.keys(), \
            'Expect points have color attribute'
        if self.color_mean is not None:
            points.color = points.color - \
                points.color.new_tensor(self.color_mean)
        points.color = points.color / 255.0
wuyuefeng's avatar
wuyuefeng committed
329
330
331
332
        results['points'] = points
        return results

    def __repr__(self):
333
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
334
        repr_str = self.__class__.__name__
335
        repr_str += f'(color_mean={self.color_mean})'
wuyuefeng's avatar
wuyuefeng committed
336
337
338
        return repr_str


339
@TRANSFORMS.register_module()
jshilong's avatar
jshilong committed
340
class LoadPointsFromFile(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
341
342
    """Load Points From File.

jshilong's avatar
jshilong committed
343
344
345
346
347
348
349
350
351
    Required Keys:

    - lidar_points (dict)

        - lidar_path (str)

    Added Keys:

    - points (np.float32)
wuyuefeng's avatar
wuyuefeng committed
352
353

    Args:
354
355
356
357
358
        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.
359
        load_dim (int, optional): The dimension of the loaded points.
360
            Defaults to 6.
361
        use_dim (list[int], optional): Which dimensions of the points to use.
liyinhao's avatar
liyinhao committed
362
363
            Defaults to [0, 1, 2]. For KITTI dataset, set use_dim=4
            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
364
365
366
367
368
369
        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
370
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
371
            for more details. Defaults to dict(backend='disk').
wuyuefeng's avatar
wuyuefeng committed
372
373
    """

jshilong's avatar
jshilong committed
374
375
376
377
378
379
380
381
382
    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
383
        self.shift_height = shift_height
384
        self.use_color = use_color
wuyuefeng's avatar
wuyuefeng committed
385
386
387
388
        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}'
389
        assert coord_type in ['CAMERA', 'LIDAR', 'DEPTH']
wuyuefeng's avatar
wuyuefeng committed
390

391
        self.coord_type = coord_type
wuyuefeng's avatar
wuyuefeng committed
392
393
394
395
396
        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
397
    def _load_points(self, pts_filename: str) -> np.ndarray:
398
399
400
401
402
403
404
405
        """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
406
407
408
409
410
411
412
413
414
415
416
        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)
417

wuyuefeng's avatar
wuyuefeng committed
418
419
        return points

jshilong's avatar
jshilong committed
420
421
    def transform(self, results: dict) -> dict:
        """Method to load points data from file.
422
423
424
425
426

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

        Returns:
427
            dict: The result dict containing the point clouds data.
428
429
                Added key and value are described below.

430
                - points (:obj:`BasePoints`): Point clouds data.
431
        """
jshilong's avatar
jshilong committed
432
433
        pts_file_path = results['lidar_points']['lidar_path']
        points = self._load_points(pts_file_path)
wuyuefeng's avatar
wuyuefeng committed
434
435
        points = points.reshape(-1, self.load_dim)
        points = points[:, self.use_dim]
436
        attribute_dims = None
wuyuefeng's avatar
wuyuefeng committed
437
438
439
440

        if self.shift_height:
            floor_height = np.percentile(points[:, 2], 0.99)
            height = points[:, 2] - floor_height
441
442
443
            points = np.concatenate(
                [points[:, :3],
                 np.expand_dims(height, 1), points[:, 3:]], 1)
444
445
            attribute_dims = dict(height=3)

446
447
448
449
450
451
452
453
454
455
456
        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,
                ]))

457
458
459
        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
460
        results['points'] = points
461

wuyuefeng's avatar
wuyuefeng committed
462
463
464
        return results

    def __repr__(self):
465
        """str: Return a string that describes the module."""
liyinhao's avatar
liyinhao committed
466
        repr_str = self.__class__.__name__ + '('
467
468
469
470
471
        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
472
473
474
        return repr_str


475
@TRANSFORMS.register_module()
476
477
478
479
480
481
482
483
class LoadPointsFromDict(LoadPointsFromFile):
    """Load Points From Dict."""

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


484
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
485
486
487
488
489
490
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
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
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
    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
537
538
539
540
541
    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.
542
543
        with_attr_label (bool, optional): Whether to load attribute label.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
544
545
546
547
548
549
550
551
552
553
554
555
        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.
556
557
        with_bbox_depth (bool, optional): Whether to load 2.5D boxes.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
558
559
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
560
        seg_3d_dtype (dtype, optional): Dtype of 3D semantic masks.
jshilong's avatar
jshilong committed
561
            Defaults to int64.
wuyuefeng's avatar
wuyuefeng committed
562
563
564
565
566
        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
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
    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
583
        super().__init__(
jshilong's avatar
jshilong committed
584
585
586
587
588
            with_bbox=with_bbox,
            with_label=with_label,
            with_mask=with_mask,
            with_seg=with_seg,
            poly2mask=poly2mask,
wuyuefeng's avatar
wuyuefeng committed
589
590
            file_client_args=file_client_args)
        self.with_bbox_3d = with_bbox_3d
591
        self.with_bbox_depth = with_bbox_depth
wuyuefeng's avatar
wuyuefeng committed
592
        self.with_label_3d = with_label_3d
593
        self.with_attr_label = with_attr_label
wuyuefeng's avatar
wuyuefeng committed
594
595
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d
596
        self.seg_3d_dtype = seg_3d_dtype
wuyuefeng's avatar
wuyuefeng committed
597

jshilong's avatar
jshilong committed
598
599
600
    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`.
601
602
603
604
605
606
607

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

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

wuyuefeng's avatar
wuyuefeng committed
609
610
611
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
        return results

jshilong's avatar
jshilong committed
612
    def _load_bboxes_depth(self, results: dict) -> dict:
613
614
615
616
617
618
619
620
        """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
621

622
        results['depths'] = results['ann_info']['depths']
jshilong's avatar
jshilong committed
623
        results['centers_2d'] = results['ann_info']['centers_2d']
624
625
        return results

jshilong's avatar
jshilong committed
626
    def _load_labels_3d(self, results: dict) -> dict:
627
628
629
630
631
632
633
634
        """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
635

wuyuefeng's avatar
wuyuefeng committed
636
637
638
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results

jshilong's avatar
jshilong committed
639
    def _load_attr_labels(self, results: dict) -> dict:
640
641
642
643
644
645
646
647
648
649
650
        """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
651
    def _load_masks_3d(self, results: dict) -> dict:
652
653
654
655
656
657
658
659
        """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
660
        pts_instance_mask_path = results['pts_instance_mask_path']
wuyuefeng's avatar
wuyuefeng committed
661
662
663
664
665

        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)
666
            pts_instance_mask = np.frombuffer(mask_bytes, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
667
668
669
        except ConnectionError:
            mmcv.check_file_exist(pts_instance_mask_path)
            pts_instance_mask = np.fromfile(
WRH's avatar
WRH committed
670
                pts_instance_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
671
672
673
674

        results['pts_instance_mask'] = pts_instance_mask
        return results

jshilong's avatar
jshilong committed
675
    def _load_semantic_seg_3d(self, results: dict) -> dict:
676
677
678
679
680
681
682
683
        """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
684
        pts_semantic_mask_path = results['pts_semantic_mask_path']
wuyuefeng's avatar
wuyuefeng committed
685
686
687
688
689
690

        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
691
692
            pts_semantic_mask = np.frombuffer(
                mask_bytes, dtype=self.seg_3d_dtype).copy()
wuyuefeng's avatar
wuyuefeng committed
693
694
695
        except ConnectionError:
            mmcv.check_file_exist(pts_semantic_mask_path)
            pts_semantic_mask = np.fromfile(
WRH's avatar
WRH committed
696
                pts_semantic_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
697
698
699
700

        results['pts_semantic_mask'] = pts_semantic_mask
        return results

jshilong's avatar
jshilong committed
701
702
    def transform(self, results: dict) -> dict:
        """Function to load multiple types annotations.
703
704
705
706
707
708

        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
709
            semantic segmentation annotations.
710
        """
jshilong's avatar
jshilong committed
711
        results = super().transform(results)
wuyuefeng's avatar
wuyuefeng committed
712
713
        if self.with_bbox_3d:
            results = self._load_bboxes_3d(results)
714
715
        if self.with_bbox_depth:
            results = self._load_bboxes_depth(results)
wuyuefeng's avatar
wuyuefeng committed
716
717
        if self.with_label_3d:
            results = self._load_labels_3d(results)
718
719
        if self.with_attr_label:
            results = self._load_attr_labels(results)
wuyuefeng's avatar
wuyuefeng committed
720
721
722
723
724
725
726
727
        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):
728
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
729
730
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
liyinhao's avatar
liyinhao committed
731
732
        repr_str += f'{indent_str}with_bbox_3d={self.with_bbox_3d}, '
        repr_str += f'{indent_str}with_label_3d={self.with_label_3d}, '
733
        repr_str += f'{indent_str}with_attr_label={self.with_attr_label}, '
liyinhao's avatar
liyinhao committed
734
735
736
737
738
739
        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}, '
740
        repr_str += f'{indent_str}with_bbox_depth={self.with_bbox_depth}, '
wuyuefeng's avatar
wuyuefeng committed
741
742
        repr_str += f'{indent_str}poly2mask={self.poly2mask})'
        return repr_str