loading.py 29.8 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
from typing import List, Union
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()
16
class LoadMultiViewImageFromFiles(BaseTransform):
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

28
29
30
    def __init__(self,
                 to_float32: bool = False,
                 color_type: str = 'unchanged') -> None:
zhangwenwei's avatar
zhangwenwei committed
31
32
        self.to_float32 = to_float32
        self.color_type = color_type
zhangwenwei's avatar
zhangwenwei committed
33

34
    def transform(self, results: dict) -> dict:
35
36
37
38
39
40
        """Call function to load multi-view image from files.

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

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

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


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

ZCMax's avatar
ZCMax committed
92
    def transform(self, results: dict) -> dict:
93
94
95
96
97
98
99
100
        """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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
        # 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]

127
128
129
        return results


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

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

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

155
156
157
158
159
160
161
162
    def __init__(self,
                 sweeps_num: int = 10,
                 load_dim: int = 5,
                 use_dim: List[int] = [0, 1, 2, 4],
                 file_client_args: dict = dict(backend='disk'),
                 pad_empty_sweeps: bool = False,
                 remove_close: bool = False,
                 test_mode: bool = False) -> None:
zhangwenwei's avatar
zhangwenwei committed
163
        self.load_dim = load_dim
zhangwenwei's avatar
zhangwenwei committed
164
        self.sweeps_num = sweeps_num
165
        self.use_dim = use_dim
zhangwenwei's avatar
zhangwenwei committed
166
167
        self.file_client_args = file_client_args.copy()
        self.file_client = None
168
169
170
        self.pad_empty_sweeps = pad_empty_sweeps
        self.remove_close = remove_close
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
171

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

194
195
196
    def _remove_close(self,
                      points: Union[np.ndarray, BasePoints],
                      radius: float = 1.0) -> Union[np.ndarray, BasePoints]:
197
198
199
        """Removes point too close within a certain radius from origin.

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

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

218
    def transform(self, results: dict) -> dict:
219
220
221
        """Call function to load multi-sweep point clouds from files.

        Args:
222
            results (dict): Result dict containing multi-sweep point cloud
223
224
225
                filenames.

        Returns:
226
            dict: The result dict containing the multi-sweep points data.
227
                Updated key and value are described below.
228

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

270
271
        points = points.cat(sweep_points_list)
        points = points[:, self.use_dim]
zhangwenwei's avatar
zhangwenwei committed
272
273
274
275
        results['points'] = points
        return results

    def __repr__(self):
276
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
277
        return f'{self.__class__.__name__}(sweeps_num={self.sweeps_num})'
wuyuefeng's avatar
wuyuefeng committed
278
279


280
@TRANSFORMS.register_module()
281
class PointSegClassMapping(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
282
283
    """Map original semantic class to valid category ids.

284
285
    Required Keys:

286
287
    - seg_label_mapping (np.ndarray)
    - pts_semantic_mask (np.ndarray)
288
289
290
291
292

    Added Keys:

    - points (np.float32)

wuyuefeng's avatar
wuyuefeng committed
293
294
295
296
    Map valid classes as 0~len(valid_cat_ids)-1 and
    others as len(valid_cat_ids).
    """

297
    def transform(self, results: dict) -> dict:
298
299
300
301
302
303
        """Call function to map original semantic class to valid category ids.

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

        Returns:
304
            dict: The result dict containing the mapped category ids.
305
306
307
308
                Updated key and value are described below.

                - pts_semantic_mask (np.ndarray): Mapped semantic masks.
        """
wuyuefeng's avatar
wuyuefeng committed
309
310
311
        assert 'pts_semantic_mask' in results
        pts_semantic_mask = results['pts_semantic_mask']

312
313
314
        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
315

316
        results['pts_semantic_mask'] = converted_pts_sem_mask
ZCMax's avatar
ZCMax committed
317
318
319
320
321
322
323

        # '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
324
325
326
        return results

    def __repr__(self):
327
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
328
329
330
331
        repr_str = self.__class__.__name__
        return repr_str


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

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

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

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

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

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

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

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


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

jshilong's avatar
jshilong committed
377
378
379
380
381
382
383
384
385
    Required Keys:

    - lidar_points (dict)

        - lidar_path (str)

    Added Keys:

    - points (np.float32)
wuyuefeng's avatar
wuyuefeng committed
386
387

    Args:
388
389
        coord_type (str): The type of coordinates of points cloud.
            Available options includes:
390

391
392
393
            - 'LIDAR': Points in LiDAR coordinates.
            - 'DEPTH': Points in depth coordinates, usually for indoor dataset.
            - 'CAMERA': Points in camera coordinates.
394
        load_dim (int, optional): The dimension of the loaded points.
395
            Defaults to 6.
396
397
        use_dim (list[int] | int, optional): Which dimensions of the points
            to use. Defaults to [0, 1, 2]. For KITTI dataset, set use_dim=4
liyinhao's avatar
liyinhao committed
398
            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
399
400
401
402
403
404
        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
405
            https://github.com/open-mmlab/mmengine/blob/main/mmengine/fileio/file_client.py
liyinhao's avatar
liyinhao committed
406
            for more details. Defaults to dict(backend='disk').
wuyuefeng's avatar
wuyuefeng committed
407
408
    """

jshilong's avatar
jshilong committed
409
410
411
412
    def __init__(
        self,
        coord_type: str,
        load_dim: int = 6,
413
        use_dim: Union[int, List[int]] = [0, 1, 2],
jshilong's avatar
jshilong committed
414
415
416
417
        shift_height: bool = False,
        use_color: bool = False,
        file_client_args: dict = dict(backend='disk')
    ) -> None:
wuyuefeng's avatar
wuyuefeng committed
418
        self.shift_height = shift_height
419
        self.use_color = use_color
wuyuefeng's avatar
wuyuefeng committed
420
421
422
423
        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}'
424
        assert coord_type in ['CAMERA', 'LIDAR', 'DEPTH']
wuyuefeng's avatar
wuyuefeng committed
425

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

wuyuefeng's avatar
wuyuefeng committed
453
454
        return points

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

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

        Returns:
462
            dict: The result dict containing the point clouds data.
463
464
                Added key and value are described below.

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

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

481
482
483
484
485
486
487
488
489
490
491
        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,
                ]))

492
493
494
        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
495
        results['points'] = points
496

wuyuefeng's avatar
wuyuefeng committed
497
498
499
        return results

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


510
@TRANSFORMS.register_module()
511
512
513
class LoadPointsFromDict(LoadPointsFromFile):
    """Load Points From Dict."""

ChaimZhu's avatar
ChaimZhu committed
514
    def transform(self, results: dict) -> dict:
515
516
517
518
        assert 'points' in results
        return results


519
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
520
521
522
523
524
525
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
526
527
528
    Required Keys:

    - ann_info (dict)
529

jshilong's avatar
jshilong committed
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
569
570
571
572
        - 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
573
574
575
576
577
    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.
578
579
        with_attr_label (bool, optional): Whether to load attribute label.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
580
581
582
583
584
585
586
587
588
589
590
591
        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.
592
593
        with_bbox_depth (bool, optional): Whether to load 2.5D boxes.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
594
595
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
596
        seg_3d_dtype (dtype, optional): Dtype of 3D semantic masks.
jshilong's avatar
jshilong committed
597
            Defaults to int64.
wuyuefeng's avatar
wuyuefeng committed
598
        file_client_args (dict): Config dict of file clients, refer to
599
            https://github.com/open-mmlab/mmengine/blob/main/mmengine/fileio/file_client.py
wuyuefeng's avatar
wuyuefeng committed
600
601
602
            for more details.
    """

jshilong's avatar
jshilong committed
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
    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
619
        super().__init__(
jshilong's avatar
jshilong committed
620
621
622
623
624
            with_bbox=with_bbox,
            with_label=with_label,
            with_mask=with_mask,
            with_seg=with_seg,
            poly2mask=poly2mask,
wuyuefeng's avatar
wuyuefeng committed
625
626
            file_client_args=file_client_args)
        self.with_bbox_3d = with_bbox_3d
627
        self.with_bbox_depth = with_bbox_depth
wuyuefeng's avatar
wuyuefeng committed
628
        self.with_label_3d = with_label_3d
629
        self.with_attr_label = with_attr_label
wuyuefeng's avatar
wuyuefeng committed
630
631
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d
632
        self.seg_3d_dtype = seg_3d_dtype
wuyuefeng's avatar
wuyuefeng committed
633

jshilong's avatar
jshilong committed
634
635
636
    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`.
637
638
639
640
641
642
643

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

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

wuyuefeng's avatar
wuyuefeng committed
645
646
647
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
        return results

jshilong's avatar
jshilong committed
648
    def _load_bboxes_depth(self, results: dict) -> dict:
649
650
651
652
653
654
655
656
        """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
657

658
        results['depths'] = results['ann_info']['depths']
jshilong's avatar
jshilong committed
659
        results['centers_2d'] = results['ann_info']['centers_2d']
660
661
        return results

jshilong's avatar
jshilong committed
662
    def _load_labels_3d(self, results: dict) -> dict:
663
664
665
666
667
668
669
670
        """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
671

wuyuefeng's avatar
wuyuefeng committed
672
673
674
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results

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

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

        results['pts_instance_mask'] = pts_instance_mask
jshilong's avatar
jshilong committed
709
710
711
        # '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
712
713
        return results

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

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

        results['pts_semantic_mask'] = pts_semantic_mask
jshilong's avatar
jshilong committed
738
739
740
        # '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
741
742
        return results

zhangshilong's avatar
zhangshilong committed
743
744
745
746
747
748
749
750
751
752
753
754
    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.
        """

755
        results['gt_bboxes'] = results['ann_info']['gt_bboxes']
zhangshilong's avatar
zhangshilong committed
756
757
758
759
760
761
762
763
764
765

    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.
        """
766
        results['gt_bboxes_labels'] = results['ann_info']['gt_bboxes_labels']
zhangshilong's avatar
zhangshilong committed
767

jshilong's avatar
jshilong committed
768
769
    def transform(self, results: dict) -> dict:
        """Function to load multiple types annotations.
770
771
772
773
774
775

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