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

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

10
from mmdet3d.registry import TRANSFORMS
zhangshilong's avatar
zhangshilong committed
11
12
from mmdet3d.structures.points import BasePoints, get_points_type
from mmdet.datasets.transforms import LoadAnnotations
zhangwenwei's avatar
zhangwenwei committed
13
14


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

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

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

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

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

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

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

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


80
@TRANSFORMS.register_module()
81
82
83
84
85
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:
86
        kwargs (dict): Arguments are the same as those in
87
88
89
            :class:`LoadImageFromFile`.
    """

ZCMax's avatar
ZCMax committed
90
    def transform(self, results: dict) -> dict:
91
92
93
94
95
96
97
98
        """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.
        """
ZCMax's avatar
ZCMax committed
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
        # TODO: load different camera image from data info,
        # for kitti dataset, we load 'CAM2' image.
        # for nuscenes dataset, we load 'CAM_FRONT' image.

        if 'CAM2' in results['images']:
            filename = results['images']['CAM2']['img_path']
            results['cam2img'] = results['images']['CAM2']['cam2img']
        elif len(list(results['images'].keys())) == 1:
            camera_type = list(results['images'].keys())[0]
            filename = results['images'][camera_type]['img_path']
            results['cam2img'] = results['images'][camera_type]['cam2img']
        else:
            raise NotImplementedError(
                'Currently we only support load image from kitti and'
                'nuscenes datasets')

        img_bytes = self.file_client.get(filename)
        img = mmcv.imfrombytes(
            img_bytes, flag=self.color_type, backend=self.imdecode_backend)
        if self.to_float32:
            img = img.astype(np.float32)

        results['img'] = img
        results['img_shape'] = img.shape[:2]
        results['ori_shape'] = img.shape[:2]

125
126
127
        return results


128
@TRANSFORMS.register_module()
VVsssssk's avatar
VVsssssk committed
129
class LoadPointsFromMultiSweeps(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
130
    """Load points from multiple sweeps.
zhangwenwei's avatar
zhangwenwei committed
131

zhangwenwei's avatar
zhangwenwei committed
132
133
134
    This is usually used for nuScenes dataset to utilize previous sweeps.

    Args:
135
136
137
138
139
140
141
        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
142
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
143
            for more details. Defaults to dict(backend='disk').
144
        pad_empty_sweeps (bool, optional): Whether to repeat keyframe when
145
            sweeps is empty. Defaults to False.
146
        remove_close (bool, optional): Whether to remove close points.
147
            Defaults to False.
148
        test_mode (bool, optional): If `test_mode=True`, it will not
149
150
            randomly sample sweeps but select the nearest N frames.
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
151
152
153
154
155
    """

    def __init__(self,
                 sweeps_num=10,
                 load_dim=5,
156
157
158
159
160
                 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
161
        self.load_dim = load_dim
zhangwenwei's avatar
zhangwenwei committed
162
        self.sweeps_num = sweeps_num
163
        self.use_dim = use_dim
zhangwenwei's avatar
zhangwenwei committed
164
165
        self.file_client_args = file_client_args.copy()
        self.file_client = None
166
167
168
        self.pad_empty_sweeps = pad_empty_sweeps
        self.remove_close = remove_close
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
169
170

    def _load_points(self, pts_filename):
171
172
173
174
175
176
177
178
        """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
179
        if self.file_client is None:
180
            self.file_client = mmengine.FileClient(**self.file_client_args)
zhangwenwei's avatar
zhangwenwei committed
181
182
183
184
        try:
            pts_bytes = self.file_client.get(pts_filename)
            points = np.frombuffer(pts_bytes, dtype=np.float32)
        except ConnectionError:
185
            mmengine.check_file_exist(pts_filename)
zhangwenwei's avatar
zhangwenwei committed
186
187
188
189
190
            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
191

192
193
194
195
    def _remove_close(self, points, radius=1.0):
        """Removes point too close within a certain radius from origin.

        Args:
196
            points (np.ndarray | :obj:`BasePoints`): Sweep points.
197
            radius (float, optional): Radius below which points are removed.
198
199
200
201
202
                Defaults to 1.0.

        Returns:
            np.ndarray: Points after removing.
        """
203
204
205
206
207
208
209
210
        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
211
        not_close = np.logical_not(np.logical_and(x_filt, y_filt))
212
        return points[not_close]
213

VVsssssk's avatar
VVsssssk committed
214
    def transform(self, results):
215
216
217
        """Call function to load multi-sweep point clouds from files.

        Args:
218
            results (dict): Result dict containing multi-sweep point cloud
219
220
221
                filenames.

        Returns:
222
            dict: The result dict containing the multi-sweep points data.
223
224
                Added key and value are described below.

225
                - points (np.ndarray | :obj:`BasePoints`): Multi-sweep point
226
                    cloud arrays.
227
        """
zhangwenwei's avatar
zhangwenwei committed
228
        points = results['points']
229
        points.tensor[:, 4] = 0
zhangwenwei's avatar
zhangwenwei committed
230
231
        sweep_points_list = [points]
        ts = results['timestamp']
VVsssssk's avatar
VVsssssk committed
232
233
234
235
236
237
238
        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)
239
        else:
VVsssssk's avatar
VVsssssk committed
240
241
            if len(results['lidar_sweeps']) <= self.sweeps_num:
                choices = np.arange(len(results['lidar_sweeps']))
242
243
244
245
            elif self.test_mode:
                choices = np.arange(self.sweeps_num)
            else:
                choices = np.random.choice(
VVsssssk's avatar
VVsssssk committed
246
247
248
                    len(results['lidar_sweeps']),
                    self.sweeps_num,
                    replace=False)
249
            for idx in choices:
VVsssssk's avatar
VVsssssk committed
250
251
252
                sweep = results['lidar_sweeps'][idx]
                points_sweep = self._load_points(
                    sweep['lidar_points']['lidar_path'])
253
254
255
                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
256
257
                # bc-breaking: Timestamp has divided 1e6 in pkl infos.
                sweep_ts = sweep['timestamp']
258
259
260
261
                lidar2sensor = np.array(sweep['lidar_points']['lidar2sensor'])
                points_sweep[:, :
                             3] = points_sweep[:, :3] @ lidar2sensor[:3, :3]
                points_sweep[:, :3] -= lidar2sensor[:3, 3]
262
                points_sweep[:, 4] = ts - sweep_ts
263
                points_sweep = points.new_point(points_sweep)
264
265
                sweep_points_list.append(points_sweep)

266
267
        points = points.cat(sweep_points_list)
        points = points[:, self.use_dim]
zhangwenwei's avatar
zhangwenwei committed
268
269
270
271
        results['points'] = points
        return results

    def __repr__(self):
272
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
273
        return f'{self.__class__.__name__}(sweeps_num={self.sweeps_num})'
wuyuefeng's avatar
wuyuefeng committed
274
275


276
@TRANSFORMS.register_module()
277
class PointSegClassMapping(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
278
279
    """Map original semantic class to valid category ids.

280
281
    Required Keys:

282
283
    - seg_label_mapping (np.ndarray)
    - pts_semantic_mask (np.ndarray)
284
285
286
287
288

    Added Keys:

    - points (np.float32)

wuyuefeng's avatar
wuyuefeng committed
289
290
291
292
    Map valid classes as 0~len(valid_cat_ids)-1 and
    others as len(valid_cat_ids).
    """

293
    def transform(self, results: dict) -> None:
294
295
296
297
298
299
        """Call function to map original semantic class to valid category ids.

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

        Returns:
300
            dict: The result dict containing the mapped category ids.
301
302
303
304
                Updated key and value are described below.

                - pts_semantic_mask (np.ndarray): Mapped semantic masks.
        """
wuyuefeng's avatar
wuyuefeng committed
305
306
307
        assert 'pts_semantic_mask' in results
        pts_semantic_mask = results['pts_semantic_mask']

308
309
310
        assert 'seg_label_mapping' in results
        label_mapping = results['seg_label_mapping']
        converted_pts_sem_mask = label_mapping[pts_semantic_mask]
wuyuefeng's avatar
wuyuefeng committed
311

312
        results['pts_semantic_mask'] = converted_pts_sem_mask
ZCMax's avatar
ZCMax committed
313
314
315
316
317
318
319

        # 'eval_ann_info' will be passed to evaluator
        if 'eval_ann_info' in results:
            assert 'pts_semantic_mask' in results['eval_ann_info']
            results['eval_ann_info']['pts_semantic_mask'] = \
                converted_pts_sem_mask

wuyuefeng's avatar
wuyuefeng committed
320
321
322
        return results

    def __repr__(self):
323
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
324
        repr_str = self.__class__.__name__
325
326
        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
327
328
329
        return repr_str


330
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
331
class NormalizePointsColor(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
332
    """Normalize color of points.
wuyuefeng's avatar
wuyuefeng committed
333
334
335
336
337

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

ZCMax's avatar
ZCMax committed
338
    def __init__(self, color_mean: List[float]) -> None:
wuyuefeng's avatar
wuyuefeng committed
339
340
        self.color_mean = color_mean

ZCMax's avatar
ZCMax committed
341
    def transform(self, input_dict: dict) -> dict:
342
343
344
345
346
347
        """Call function to normalize color of points.

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

        Returns:
348
            dict: The result dict containing the normalized points.
349
350
                Updated key and value are described below.

351
                - points (:obj:`BasePoints`): Points after color normalization.
352
        """
ZCMax's avatar
ZCMax committed
353
        points = input_dict['points']
354
        assert points.attribute_dims is not None and \
355
356
               'color' in points.attribute_dims.keys(), \
               'Expect points have color attribute'
357
358
        if self.color_mean is not None:
            points.color = points.color - \
359
                           points.color.new_tensor(self.color_mean)
360
        points.color = points.color / 255.0
ZCMax's avatar
ZCMax committed
361
362
        input_dict['points'] = points
        return input_dict
wuyuefeng's avatar
wuyuefeng committed
363
364

    def __repr__(self):
365
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
366
        repr_str = self.__class__.__name__
367
        repr_str += f'(color_mean={self.color_mean})'
wuyuefeng's avatar
wuyuefeng committed
368
369
370
        return repr_str


371
@TRANSFORMS.register_module()
jshilong's avatar
jshilong committed
372
class LoadPointsFromFile(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
373
374
    """Load Points From File.

jshilong's avatar
jshilong committed
375
376
377
378
379
380
381
382
383
    Required Keys:

    - lidar_points (dict)

        - lidar_path (str)

    Added Keys:

    - points (np.float32)
wuyuefeng's avatar
wuyuefeng committed
384
385

    Args:
386
387
388
389
390
        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.
391
        load_dim (int, optional): The dimension of the loaded points.
392
            Defaults to 6.
393
        use_dim (list[int], optional): Which dimensions of the points to use.
liyinhao's avatar
liyinhao committed
394
395
            Defaults to [0, 1, 2]. For KITTI dataset, set use_dim=4
            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
396
397
398
399
400
401
        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
402
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
403
            for more details. Defaults to dict(backend='disk').
wuyuefeng's avatar
wuyuefeng committed
404
405
    """

jshilong's avatar
jshilong committed
406
407
408
409
410
411
412
413
414
    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
415
        self.shift_height = shift_height
416
        self.use_color = use_color
wuyuefeng's avatar
wuyuefeng committed
417
418
419
420
        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}'
421
        assert coord_type in ['CAMERA', 'LIDAR', 'DEPTH']
wuyuefeng's avatar
wuyuefeng committed
422

423
        self.coord_type = coord_type
wuyuefeng's avatar
wuyuefeng committed
424
425
426
427
428
        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
429
    def _load_points(self, pts_filename: str) -> np.ndarray:
430
431
432
433
434
435
436
437
        """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
438
        if self.file_client is None:
439
            self.file_client = mmengine.FileClient(**self.file_client_args)
wuyuefeng's avatar
wuyuefeng committed
440
441
442
443
        try:
            pts_bytes = self.file_client.get(pts_filename)
            points = np.frombuffer(pts_bytes, dtype=np.float32)
        except ConnectionError:
444
            mmengine.check_file_exist(pts_filename)
wuyuefeng's avatar
wuyuefeng committed
445
446
447
448
            if pts_filename.endswith('.npy'):
                points = np.load(pts_filename)
            else:
                points = np.fromfile(pts_filename, dtype=np.float32)
449

wuyuefeng's avatar
wuyuefeng committed
450
451
        return points

jshilong's avatar
jshilong committed
452
453
    def transform(self, results: dict) -> dict:
        """Method to load points data from file.
454
455
456
457
458

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

        Returns:
459
            dict: The result dict containing the point clouds data.
460
461
                Added key and value are described below.

462
                - points (:obj:`BasePoints`): Point clouds data.
463
        """
jshilong's avatar
jshilong committed
464
465
        pts_file_path = results['lidar_points']['lidar_path']
        points = self._load_points(pts_file_path)
wuyuefeng's avatar
wuyuefeng committed
466
467
        points = points.reshape(-1, self.load_dim)
        points = points[:, self.use_dim]
468
        attribute_dims = None
wuyuefeng's avatar
wuyuefeng committed
469
470
471
472

        if self.shift_height:
            floor_height = np.percentile(points[:, 2], 0.99)
            height = points[:, 2] - floor_height
473
474
475
            points = np.concatenate(
                [points[:, :3],
                 np.expand_dims(height, 1), points[:, 3:]], 1)
476
477
            attribute_dims = dict(height=3)

478
479
480
481
482
483
484
485
486
487
488
        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,
                ]))

489
490
491
        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
492
        results['points'] = points
493

wuyuefeng's avatar
wuyuefeng committed
494
495
496
        return results

    def __repr__(self):
497
        """str: Return a string that describes the module."""
liyinhao's avatar
liyinhao committed
498
        repr_str = self.__class__.__name__ + '('
499
500
501
502
503
        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
504
505
506
        return repr_str


507
@TRANSFORMS.register_module()
508
509
510
class LoadPointsFromDict(LoadPointsFromFile):
    """Load Points From Dict."""

ChaimZhu's avatar
ChaimZhu committed
511
    def transform(self, results: dict) -> dict:
512
513
514
515
        assert 'points' in results
        return results


516
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
517
518
519
520
521
522
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
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
556
557
558
559
560
561
562
563
564
565
566
567
568
    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
569
570
571
572
573
    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.
574
575
        with_attr_label (bool, optional): Whether to load attribute label.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
576
577
578
579
580
581
582
583
584
585
586
587
        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.
588
589
        with_bbox_depth (bool, optional): Whether to load 2.5D boxes.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
590
591
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
592
        seg_3d_dtype (dtype, optional): Dtype of 3D semantic masks.
jshilong's avatar
jshilong committed
593
            Defaults to int64.
wuyuefeng's avatar
wuyuefeng committed
594
595
596
597
598
        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
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
    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
615
        super().__init__(
jshilong's avatar
jshilong committed
616
617
618
619
620
            with_bbox=with_bbox,
            with_label=with_label,
            with_mask=with_mask,
            with_seg=with_seg,
            poly2mask=poly2mask,
wuyuefeng's avatar
wuyuefeng committed
621
622
            file_client_args=file_client_args)
        self.with_bbox_3d = with_bbox_3d
623
        self.with_bbox_depth = with_bbox_depth
wuyuefeng's avatar
wuyuefeng committed
624
        self.with_label_3d = with_label_3d
625
        self.with_attr_label = with_attr_label
wuyuefeng's avatar
wuyuefeng committed
626
627
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d
628
        self.seg_3d_dtype = seg_3d_dtype
wuyuefeng's avatar
wuyuefeng committed
629

jshilong's avatar
jshilong committed
630
631
632
    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`.
633
634
635
636
637
638
639

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

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

wuyuefeng's avatar
wuyuefeng committed
641
642
643
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
        return results

jshilong's avatar
jshilong committed
644
    def _load_bboxes_depth(self, results: dict) -> dict:
645
646
647
648
649
650
651
652
        """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
653

654
        results['depths'] = results['ann_info']['depths']
jshilong's avatar
jshilong committed
655
        results['centers_2d'] = results['ann_info']['centers_2d']
656
657
        return results

jshilong's avatar
jshilong committed
658
    def _load_labels_3d(self, results: dict) -> dict:
659
660
661
662
663
664
665
666
        """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
667

wuyuefeng's avatar
wuyuefeng committed
668
669
670
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results

jshilong's avatar
jshilong committed
671
    def _load_attr_labels(self, results: dict) -> dict:
672
673
674
675
676
677
678
679
680
681
682
        """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
683
    def _load_masks_3d(self, results: dict) -> dict:
684
685
686
687
688
689
690
691
        """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
692
        pts_instance_mask_path = results['pts_instance_mask_path']
wuyuefeng's avatar
wuyuefeng committed
693
694

        if self.file_client is None:
695
            self.file_client = mmengine.FileClient(**self.file_client_args)
wuyuefeng's avatar
wuyuefeng committed
696
697
        try:
            mask_bytes = self.file_client.get(pts_instance_mask_path)
698
            pts_instance_mask = np.frombuffer(mask_bytes, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
699
        except ConnectionError:
700
            mmengine.check_file_exist(pts_instance_mask_path)
wuyuefeng's avatar
wuyuefeng committed
701
            pts_instance_mask = np.fromfile(
WRH's avatar
WRH committed
702
                pts_instance_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
703
704

        results['pts_instance_mask'] = pts_instance_mask
jshilong's avatar
jshilong committed
705
706
707
        # '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
708
709
        return results

jshilong's avatar
jshilong committed
710
    def _load_semantic_seg_3d(self, results: dict) -> dict:
711
712
713
714
715
716
717
718
        """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
719
        pts_semantic_mask_path = results['pts_semantic_mask_path']
wuyuefeng's avatar
wuyuefeng committed
720
721

        if self.file_client is None:
722
            self.file_client = mmengine.FileClient(**self.file_client_args)
wuyuefeng's avatar
wuyuefeng committed
723
724
725
        try:
            mask_bytes = self.file_client.get(pts_semantic_mask_path)
            # add .copy() to fix read-only bug
726
727
            pts_semantic_mask = np.frombuffer(
                mask_bytes, dtype=self.seg_3d_dtype).copy()
wuyuefeng's avatar
wuyuefeng committed
728
        except ConnectionError:
729
            mmengine.check_file_exist(pts_semantic_mask_path)
wuyuefeng's avatar
wuyuefeng committed
730
            pts_semantic_mask = np.fromfile(
WRH's avatar
WRH committed
731
                pts_semantic_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
732
733

        results['pts_semantic_mask'] = pts_semantic_mask
jshilong's avatar
jshilong committed
734
735
736
        # '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
737
738
        return results

zhangshilong's avatar
zhangshilong committed
739
740
741
742
743
744
745
746
747
748
749
750
    def _load_bboxes(self, results: dict) -> None:
        """Private function to load bounding box annotations.

        The only difference is it remove the proceess for
        `ignore_flag`

        Args:
            results (dict): Result dict from :obj:``mmcv.BaseDataset``.
        Returns:
            dict: The dict contains loaded bounding box annotations.
        """

751
        results['gt_bboxes'] = results['ann_info']['gt_bboxes']
zhangshilong's avatar
zhangshilong committed
752
753
754
755
756
757
758
759
760
761

    def _load_labels(self, results: dict) -> None:
        """Private function to load label annotations.

        Args:
            results (dict): Result dict from :obj :obj:``mmcv.BaseDataset``.

        Returns:
            dict: The dict contains loaded label annotations.
        """
762
        results['gt_bboxes_labels'] = results['ann_info']['gt_bboxes_labels']
zhangshilong's avatar
zhangshilong committed
763

jshilong's avatar
jshilong committed
764
765
    def transform(self, results: dict) -> dict:
        """Function to load multiple types annotations.
766
767
768
769
770
771

        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
772
            semantic segmentation annotations.
773
        """
jshilong's avatar
jshilong committed
774
        results = super().transform(results)
wuyuefeng's avatar
wuyuefeng committed
775
776
        if self.with_bbox_3d:
            results = self._load_bboxes_3d(results)
777
778
        if self.with_bbox_depth:
            results = self._load_bboxes_depth(results)
wuyuefeng's avatar
wuyuefeng committed
779
780
        if self.with_label_3d:
            results = self._load_labels_3d(results)
781
782
        if self.with_attr_label:
            results = self._load_attr_labels(results)
wuyuefeng's avatar
wuyuefeng committed
783
784
785
786
787
788
789
790
        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):
791
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
792
793
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
liyinhao's avatar
liyinhao committed
794
795
        repr_str += f'{indent_str}with_bbox_3d={self.with_bbox_3d}, '
        repr_str += f'{indent_str}with_label_3d={self.with_label_3d}, '
796
        repr_str += f'{indent_str}with_attr_label={self.with_attr_label}, '
liyinhao's avatar
liyinhao committed
797
798
799
800
801
802
        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}, '
803
        repr_str += f'{indent_str}with_bbox_depth={self.with_bbox_depth}, '
wuyuefeng's avatar
wuyuefeng committed
804
805
        repr_str += f'{indent_str}poly2mask={self.poly2mask})'
        return repr_str