"vscode:/vscode.git/clone" did not exist on "499f5e620c243b6a9980b63f7aa54d096a9a3ddd"
loading.py 30.8 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
5
import mmcv
import numpy as np
6
from mmcv.transforms import LoadImageFromFile
7
from mmcv.transforms.base import BaseTransform
zhangwenwei's avatar
zhangwenwei committed
8

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


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

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

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

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

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

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

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

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

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


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

    Args:
85
        kwargs (dict): Arguments are the same as those in
86
87
88
            :class:`LoadImageFromFile`.
    """

ZCMax's avatar
ZCMax committed
89
    def transform(self, results: dict) -> dict:
90
91
92
93
94
95
96
97
        """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
98
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
        # 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]

124
125
126
        return results


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

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

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

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

    def _load_points(self, pts_filename):
170
171
172
173
174
175
176
177
        """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
178
179
180
181
182
183
184
185
186
187
188
189
        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
190

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

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

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

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

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

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

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

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

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


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

278
279
    Required Keys:

280
281
    - seg_label_mapping (np.ndarray)
    - pts_semantic_mask (np.ndarray)
282
283
284
285
286

    Added Keys:

    - points (np.float32)

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

291
292
293
294
295
296
297
298
299
300
301
302
303
304
    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'

        self.valid_cat_ids = valid_cat_ids
        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

305
    def transform(self, results: dict) -> None:
306
307
308
309
310
311
        """Call function to map original semantic class to valid category ids.

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

        Returns:
312
            dict: The result dict containing the mapped category ids.
313
314
315
316
                Updated key and value are described below.

                - pts_semantic_mask (np.ndarray): Mapped semantic masks.
        """
wuyuefeng's avatar
wuyuefeng committed
317
318
319
        assert 'pts_semantic_mask' in results
        pts_semantic_mask = results['pts_semantic_mask']

320
321
322
        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
323

324
        results['pts_semantic_mask'] = converted_pts_sem_mask
ZCMax's avatar
ZCMax committed
325
326
327
328
329
330
331

        # '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
332
333
334
        return results

    def __repr__(self):
335
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
336
        repr_str = self.__class__.__name__
337
338
        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
339
340
341
        return repr_str


342
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
343
class NormalizePointsColor(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
344
    """Normalize color of points.
wuyuefeng's avatar
wuyuefeng committed
345
346
347
348
349

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

ZCMax's avatar
ZCMax committed
350
    def __init__(self, color_mean: List[float]) -> None:
wuyuefeng's avatar
wuyuefeng committed
351
352
        self.color_mean = color_mean

ZCMax's avatar
ZCMax committed
353
    def transform(self, input_dict: dict) -> dict:
354
355
356
357
358
359
        """Call function to normalize color of points.

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

        Returns:
360
            dict: The result dict containing the normalized points.
361
362
                Updated key and value are described below.

363
                - points (:obj:`BasePoints`): Points after color normalization.
364
        """
ZCMax's avatar
ZCMax committed
365
        points = input_dict['points']
366
        assert points.attribute_dims is not None and \
367
368
               'color' in points.attribute_dims.keys(), \
               'Expect points have color attribute'
369
370
        if self.color_mean is not None:
            points.color = points.color - \
371
                           points.color.new_tensor(self.color_mean)
372
        points.color = points.color / 255.0
ZCMax's avatar
ZCMax committed
373
374
        input_dict['points'] = points
        return input_dict
wuyuefeng's avatar
wuyuefeng committed
375
376

    def __repr__(self):
377
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
378
        repr_str = self.__class__.__name__
379
        repr_str += f'(color_mean={self.color_mean})'
wuyuefeng's avatar
wuyuefeng committed
380
381
382
        return repr_str


383
@TRANSFORMS.register_module()
jshilong's avatar
jshilong committed
384
class LoadPointsFromFile(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
385
386
    """Load Points From File.

jshilong's avatar
jshilong committed
387
388
389
390
391
392
393
394
395
    Required Keys:

    - lidar_points (dict)

        - lidar_path (str)

    Added Keys:

    - points (np.float32)
wuyuefeng's avatar
wuyuefeng committed
396
397

    Args:
398
399
400
401
402
        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.
403
        load_dim (int, optional): The dimension of the loaded points.
404
            Defaults to 6.
405
        use_dim (list[int], optional): Which dimensions of the points to use.
liyinhao's avatar
liyinhao committed
406
407
            Defaults to [0, 1, 2]. For KITTI dataset, set use_dim=4
            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
408
409
410
411
412
413
        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
414
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
415
            for more details. Defaults to dict(backend='disk').
wuyuefeng's avatar
wuyuefeng committed
416
417
    """

jshilong's avatar
jshilong committed
418
419
420
421
422
423
424
425
426
    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
427
        self.shift_height = shift_height
428
        self.use_color = use_color
wuyuefeng's avatar
wuyuefeng committed
429
430
431
432
        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}'
433
        assert coord_type in ['CAMERA', 'LIDAR', 'DEPTH']
wuyuefeng's avatar
wuyuefeng committed
434

435
        self.coord_type = coord_type
wuyuefeng's avatar
wuyuefeng committed
436
437
438
439
440
        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
441
    def _load_points(self, pts_filename: str) -> np.ndarray:
442
443
444
445
446
447
448
449
        """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
450
451
452
453
454
455
456
457
458
459
460
        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)
461

wuyuefeng's avatar
wuyuefeng committed
462
463
        return points

jshilong's avatar
jshilong committed
464
465
    def transform(self, results: dict) -> dict:
        """Method to load points data from file.
466
467
468
469
470

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

        Returns:
471
            dict: The result dict containing the point clouds data.
472
473
                Added key and value are described below.

474
                - points (:obj:`BasePoints`): Point clouds data.
475
        """
jshilong's avatar
jshilong committed
476
477
        pts_file_path = results['lidar_points']['lidar_path']
        points = self._load_points(pts_file_path)
wuyuefeng's avatar
wuyuefeng committed
478
479
        points = points.reshape(-1, self.load_dim)
        points = points[:, self.use_dim]
480
        attribute_dims = None
wuyuefeng's avatar
wuyuefeng committed
481
482
483
484

        if self.shift_height:
            floor_height = np.percentile(points[:, 2], 0.99)
            height = points[:, 2] - floor_height
485
486
487
            points = np.concatenate(
                [points[:, :3],
                 np.expand_dims(height, 1), points[:, 3:]], 1)
488
489
            attribute_dims = dict(height=3)

490
491
492
493
494
495
496
497
498
499
500
        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,
                ]))

501
502
503
        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
504
        results['points'] = points
505

wuyuefeng's avatar
wuyuefeng committed
506
507
508
        return results

    def __repr__(self):
509
        """str: Return a string that describes the module."""
liyinhao's avatar
liyinhao committed
510
        repr_str = self.__class__.__name__ + '('
511
512
513
514
515
        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
516
517
518
        return repr_str


519
@TRANSFORMS.register_module()
520
521
522
class LoadPointsFromDict(LoadPointsFromFile):
    """Load Points From Dict."""

ChaimZhu's avatar
ChaimZhu committed
523
    def transform(self, results: dict) -> dict:
524
525
526
527
        assert 'points' in results
        return results


528
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
529
530
531
532
533
534
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
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
573
574
575
576
577
578
579
580
    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
581
582
583
584
585
    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.
586
587
        with_attr_label (bool, optional): Whether to load attribute label.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
588
589
590
591
592
593
594
595
596
597
598
599
        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.
600
601
        with_bbox_depth (bool, optional): Whether to load 2.5D boxes.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
602
603
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
604
        seg_3d_dtype (dtype, optional): Dtype of 3D semantic masks.
jshilong's avatar
jshilong committed
605
            Defaults to int64.
wuyuefeng's avatar
wuyuefeng committed
606
607
608
609
610
        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
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
    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
627
        super().__init__(
jshilong's avatar
jshilong committed
628
629
630
631
632
            with_bbox=with_bbox,
            with_label=with_label,
            with_mask=with_mask,
            with_seg=with_seg,
            poly2mask=poly2mask,
wuyuefeng's avatar
wuyuefeng committed
633
634
            file_client_args=file_client_args)
        self.with_bbox_3d = with_bbox_3d
635
        self.with_bbox_depth = with_bbox_depth
wuyuefeng's avatar
wuyuefeng committed
636
        self.with_label_3d = with_label_3d
637
        self.with_attr_label = with_attr_label
wuyuefeng's avatar
wuyuefeng committed
638
639
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d
640
        self.seg_3d_dtype = seg_3d_dtype
wuyuefeng's avatar
wuyuefeng committed
641

jshilong's avatar
jshilong committed
642
643
644
    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`.
645
646
647
648
649
650
651

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

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

wuyuefeng's avatar
wuyuefeng committed
653
654
655
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
        return results

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

666
        results['depths'] = results['ann_info']['depths']
jshilong's avatar
jshilong committed
667
        results['centers_2d'] = results['ann_info']['centers_2d']
668
669
        return results

jshilong's avatar
jshilong committed
670
    def _load_labels_3d(self, results: dict) -> dict:
671
672
673
674
675
676
677
678
        """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
679

wuyuefeng's avatar
wuyuefeng committed
680
681
682
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results

jshilong's avatar
jshilong committed
683
    def _load_attr_labels(self, results: dict) -> dict:
684
685
686
687
688
689
690
691
692
693
694
        """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
695
    def _load_masks_3d(self, results: dict) -> dict:
696
697
698
699
700
701
702
703
        """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
704
        pts_instance_mask_path = results['pts_instance_mask_path']
wuyuefeng's avatar
wuyuefeng committed
705
706
707
708
709

        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)
710
            pts_instance_mask = np.frombuffer(mask_bytes, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
711
712
713
        except ConnectionError:
            mmcv.check_file_exist(pts_instance_mask_path)
            pts_instance_mask = np.fromfile(
WRH's avatar
WRH committed
714
                pts_instance_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
715
716

        results['pts_instance_mask'] = pts_instance_mask
jshilong's avatar
jshilong committed
717
718
719
        # '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
720
721
        return results

jshilong's avatar
jshilong committed
722
    def _load_semantic_seg_3d(self, results: dict) -> dict:
723
724
725
726
727
728
729
730
        """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
731
        pts_semantic_mask_path = results['pts_semantic_mask_path']
wuyuefeng's avatar
wuyuefeng committed
732
733
734
735
736
737

        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
738
739
            pts_semantic_mask = np.frombuffer(
                mask_bytes, dtype=self.seg_3d_dtype).copy()
wuyuefeng's avatar
wuyuefeng committed
740
741
742
        except ConnectionError:
            mmcv.check_file_exist(pts_semantic_mask_path)
            pts_semantic_mask = np.fromfile(
WRH's avatar
WRH committed
743
                pts_semantic_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
744
745

        results['pts_semantic_mask'] = pts_semantic_mask
jshilong's avatar
jshilong committed
746
747
748
        # '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
749
750
        return results

zhangshilong's avatar
zhangshilong committed
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
    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.
        """
        gt_bboxes = []
        for instance in results['instances']:
            gt_bboxes.append(instance['bbox'])
        if len(gt_bboxes) == 0:
            results['gt_bboxes'] = np.zeros((0, 4), dtype=np.float32)
        else:
            results['gt_bboxes'] = np.array(
                gt_bboxes, dtype=np.float32).reshape((-1, 4))

        if 'eval_ann_info' in results:
            results['eval_ann_info']['gt_bboxes'] = results['gt_bboxes']

    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.
        """
        gt_bboxes_labels = []
        for instance in results['instances']:
            gt_bboxes_labels.append(instance['bbox_label'])
        if len(gt_bboxes_labels) == 0:
            results['gt_bboxes_labels'] = np.zeros((0, ), dtype=np.int64)
        else:
            results['gt_bboxes_labels'] = np.array(
                gt_bboxes_labels, dtype=np.int64)
        if 'eval_ann_info' in results:
            results['eval_ann_info']['gt_bboxes_labels'] = results[
                'gt_bboxes_labels']

jshilong's avatar
jshilong committed
795
796
    def transform(self, results: dict) -> dict:
        """Function to load multiple types annotations.
797
798
799
800
801
802

        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
803
            semantic segmentation annotations.
804
        """
jshilong's avatar
jshilong committed
805
        results = super().transform(results)
wuyuefeng's avatar
wuyuefeng committed
806
807
        if self.with_bbox_3d:
            results = self._load_bboxes_3d(results)
808
809
        if self.with_bbox_depth:
            results = self._load_bboxes_depth(results)
wuyuefeng's avatar
wuyuefeng committed
810
811
        if self.with_label_3d:
            results = self._load_labels_3d(results)
812
813
        if self.with_attr_label:
            results = self._load_attr_labels(results)
wuyuefeng's avatar
wuyuefeng committed
814
815
816
817
818
819
820
821
        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):
822
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
823
824
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
liyinhao's avatar
liyinhao committed
825
826
        repr_str += f'{indent_str}with_bbox_3d={self.with_bbox_3d}, '
        repr_str += f'{indent_str}with_label_3d={self.with_label_3d}, '
827
        repr_str += f'{indent_str}with_attr_label={self.with_attr_label}, '
liyinhao's avatar
liyinhao committed
828
829
830
831
832
833
        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}, '
834
        repr_str += f'{indent_str}with_bbox_depth={self.with_bbox_depth}, '
wuyuefeng's avatar
wuyuefeng committed
835
836
        repr_str += f'{indent_str}poly2mask={self.poly2mask})'
        return repr_str