loading.py 24.8 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangwenwei's avatar
zhangwenwei committed
2
3
4
import mmcv
import numpy as np

5
from mmdet3d.core.points import BasePoints, get_points_type
6
from mmdet.datasets.builder import PIPELINES
7
from mmdet.datasets.pipelines import LoadAnnotations, LoadImageFromFile
zhangwenwei's avatar
zhangwenwei committed
8
9


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

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

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

zhangwenwei's avatar
zhangwenwei committed
22
23
24
    def __init__(self, to_float32=False, color_type='unchanged'):
        self.to_float32 = to_float32
        self.color_type = color_type
zhangwenwei's avatar
zhangwenwei committed
25
26

    def __call__(self, results):
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
        """Call function to load multi-view image from files.

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

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

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


74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
@PIPELINES.register_module()
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:
        kwargs (dict): Arguments are the same as those in \
            :class:`LoadImageFromFile`.
    """

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

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

        Returns:
            dict: The dict contains loaded image and meta information.
        """
        super().__call__(results)
94
        results['cam2img'] = results['img_info']['cam_intrinsic']
95
96
97
        return results


zhangwenwei's avatar
zhangwenwei committed
98
99
@PIPELINES.register_module()
class LoadPointsFromMultiSweeps(object):
zhangwenwei's avatar
zhangwenwei committed
100
    """Load points from multiple sweeps.
zhangwenwei's avatar
zhangwenwei committed
101

zhangwenwei's avatar
zhangwenwei committed
102
103
104
    This is usually used for nuScenes dataset to utilize previous sweeps.

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

    def __init__(self,
                 sweeps_num=10,
                 load_dim=5,
123
124
125
126
127
                 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
128
        self.load_dim = load_dim
zhangwenwei's avatar
zhangwenwei committed
129
        self.sweeps_num = sweeps_num
130
        self.use_dim = use_dim
zhangwenwei's avatar
zhangwenwei committed
131
132
        self.file_client_args = file_client_args.copy()
        self.file_client = None
133
134
135
        self.pad_empty_sweeps = pad_empty_sweeps
        self.remove_close = remove_close
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
136
137

    def _load_points(self, pts_filename):
138
139
140
141
142
143
144
145
        """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
146
147
148
149
150
151
152
153
154
155
156
157
        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
158

159
160
161
162
    def _remove_close(self, points, radius=1.0):
        """Removes point too close within a certain radius from origin.

        Args:
163
            points (np.ndarray | :obj:`BasePoints`): Sweep points.
164
165
166
167
168
169
            radius (float): Radius below which points are removed.
                Defaults to 1.0.

        Returns:
            np.ndarray: Points after removing.
        """
170
171
172
173
174
175
176
177
        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
178
        not_close = np.logical_not(np.logical_and(x_filt, y_filt))
179
        return points[not_close]
180

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

        Args:
            results (dict): Result dict containing multi-sweep point cloud \
                filenames.

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

192
193
                - points (np.ndarray | :obj:`BasePoints`): Multi-sweep point \
                    cloud arrays.
194
        """
zhangwenwei's avatar
zhangwenwei committed
195
        points = results['points']
196
        points.tensor[:, 4] = 0
zhangwenwei's avatar
zhangwenwei committed
197
198
        sweep_points_list = [points]
        ts = results['timestamp']
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
        if self.pad_empty_sweeps and len(results['sweeps']) == 0:
            for i in range(self.sweeps_num):
                if self.remove_close:
                    sweep_points_list.append(self._remove_close(points))
                else:
                    sweep_points_list.append(points)
        else:
            if len(results['sweeps']) <= self.sweeps_num:
                choices = np.arange(len(results['sweeps']))
            elif self.test_mode:
                choices = np.arange(self.sweeps_num)
            else:
                choices = np.random.choice(
                    len(results['sweeps']), self.sweeps_num, replace=False)
            for idx in choices:
                sweep = results['sweeps'][idx]
                points_sweep = self._load_points(sweep['data_path'])
                points_sweep = np.copy(points_sweep).reshape(-1, self.load_dim)
                if self.remove_close:
                    points_sweep = self._remove_close(points_sweep)
                sweep_ts = sweep['timestamp'] / 1e6
                points_sweep[:, :3] = points_sweep[:, :3] @ sweep[
                    'sensor2lidar_rotation'].T
                points_sweep[:, :3] += sweep['sensor2lidar_translation']
                points_sweep[:, 4] = ts - sweep_ts
224
                points_sweep = points.new_point(points_sweep)
225
226
                sweep_points_list.append(points_sweep)

227
228
        points = points.cat(sweep_points_list)
        points = points[:, self.use_dim]
zhangwenwei's avatar
zhangwenwei committed
229
230
231
232
        results['points'] = points
        return results

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


@PIPELINES.register_module()
class PointSegClassMapping(object):
    """Map original semantic class to valid category ids.

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

    Args:
245
        valid_cat_ids (tuple[int]): A tuple of valid category.
246
247
        max_cat_id (int): The max possible cat_id in input segmentation mask.
            Defaults to 40.
wuyuefeng's avatar
wuyuefeng committed
248
249
    """

250
251
252
253
    def __init__(self, valid_cat_ids, max_cat_id=40):
        assert max_cat_id >= np.max(valid_cat_ids), \
            'max_cat_id should be greater than maximum id in valid_cat_ids'

wuyuefeng's avatar
wuyuefeng committed
254
        self.valid_cat_ids = valid_cat_ids
255
256
257
258
259
260
261
262
        self.max_cat_id = int(max_cat_id)

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

    def __call__(self, results):
265
266
267
268
269
270
271
272
273
274
275
        """Call function to map original semantic class to valid category ids.

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

        Returns:
            dict: The result dict containing the mapped category ids. \
                Updated key and value are described below.

                - pts_semantic_mask (np.ndarray): Mapped semantic masks.
        """
wuyuefeng's avatar
wuyuefeng committed
276
277
278
        assert 'pts_semantic_mask' in results
        pts_semantic_mask = results['pts_semantic_mask']

279
        converted_pts_sem_mask = self.cat_id2class[pts_semantic_mask]
wuyuefeng's avatar
wuyuefeng committed
280

281
        results['pts_semantic_mask'] = converted_pts_sem_mask
wuyuefeng's avatar
wuyuefeng committed
282
283
284
        return results

    def __repr__(self):
285
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
286
        repr_str = self.__class__.__name__
287
288
        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
289
290
291
292
293
        return repr_str


@PIPELINES.register_module()
class NormalizePointsColor(object):
zhangwenwei's avatar
zhangwenwei committed
294
    """Normalize color of points.
wuyuefeng's avatar
wuyuefeng committed
295
296
297
298
299
300
301
302
303

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

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

    def __call__(self, results):
304
305
306
307
308
309
310
311
312
        """Call function to normalize color of points.

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

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

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

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


@PIPELINES.register_module()
class LoadPointsFromFile(object):
    """Load Points From File.

    Load sunrgbd and scannet points from file.

    Args:
340
341
342
343
344
        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.
345
346
        load_dim (int): The dimension of the loaded points.
            Defaults to 6.
wuyuefeng's avatar
wuyuefeng committed
347
        use_dim (list[int]): Which dimensions of the points to be used.
liyinhao's avatar
liyinhao committed
348
349
350
            Defaults to [0, 1, 2]. For KITTI dataset, set use_dim=4
            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
        shift_height (bool): Whether to use shifted height. Defaults to False.
351
        use_color (bool): Whether to use color features. Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
352
353
        file_client_args (dict): Config dict of file clients, refer to
            https://github.com/open-mmlab/mmcv/blob/master/mmcv/fileio/file_client.py
liyinhao's avatar
liyinhao committed
354
            for more details. Defaults to dict(backend='disk').
wuyuefeng's avatar
wuyuefeng committed
355
356
357
    """

    def __init__(self,
358
                 coord_type,
wuyuefeng's avatar
wuyuefeng committed
359
360
361
                 load_dim=6,
                 use_dim=[0, 1, 2],
                 shift_height=False,
362
                 use_color=False,
wuyuefeng's avatar
wuyuefeng committed
363
364
                 file_client_args=dict(backend='disk')):
        self.shift_height = shift_height
365
        self.use_color = use_color
wuyuefeng's avatar
wuyuefeng committed
366
367
368
369
        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}'
370
        assert coord_type in ['CAMERA', 'LIDAR', 'DEPTH']
wuyuefeng's avatar
wuyuefeng committed
371

372
        self.coord_type = coord_type
wuyuefeng's avatar
wuyuefeng committed
373
374
375
376
377
378
        self.load_dim = load_dim
        self.use_dim = use_dim
        self.file_client_args = file_client_args.copy()
        self.file_client = None

    def _load_points(self, pts_filename):
379
380
381
382
383
384
385
386
        """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
387
388
389
390
391
392
393
394
395
396
397
        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)
398

wuyuefeng's avatar
wuyuefeng committed
399
400
401
        return points

    def __call__(self, results):
402
403
404
405
406
407
408
409
410
        """Call function to load points data from file.

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

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

411
                - points (:obj:`BasePoints`): Point clouds data.
412
        """
wuyuefeng's avatar
wuyuefeng committed
413
414
415
416
        pts_filename = results['pts_filename']
        points = self._load_points(pts_filename)
        points = points.reshape(-1, self.load_dim)
        points = points[:, self.use_dim]
417
        attribute_dims = None
wuyuefeng's avatar
wuyuefeng committed
418
419
420
421

        if self.shift_height:
            floor_height = np.percentile(points[:, 2], 0.99)
            height = points[:, 2] - floor_height
422
423
424
            points = np.concatenate(
                [points[:, :3],
                 np.expand_dims(height, 1), points[:, 3:]], 1)
425
426
            attribute_dims = dict(height=3)

427
428
429
430
431
432
433
434
435
436
437
        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,
                ]))

438
439
440
        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
441
        results['points'] = points
442

wuyuefeng's avatar
wuyuefeng committed
443
444
445
        return results

    def __repr__(self):
446
        """str: Return a string that describes the module."""
liyinhao's avatar
liyinhao committed
447
        repr_str = self.__class__.__name__ + '('
448
449
450
451
452
        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
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
        return repr_str


@PIPELINES.register_module()
class LoadAnnotations3D(LoadAnnotations):
    """Load Annotations3D.

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

    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.
468
469
        with_attr_label (bool, optional): Whether to load attribute label.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
470
471
472
473
474
475
476
477
478
479
480
481
        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.
482
483
        with_bbox_depth (bool, optional): Whether to load 2.5D boxes.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
484
485
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
486
487
        seg_3d_dtype (dtype, optional): Dtype of 3D semantic masks.
            Defaults to int64
wuyuefeng's avatar
wuyuefeng committed
488
489
490
491
492
493
494
495
        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.
    """

    def __init__(self,
                 with_bbox_3d=True,
                 with_label_3d=True,
496
                 with_attr_label=False,
wuyuefeng's avatar
wuyuefeng committed
497
498
499
500
501
502
                 with_mask_3d=False,
                 with_seg_3d=False,
                 with_bbox=False,
                 with_label=False,
                 with_mask=False,
                 with_seg=False,
503
                 with_bbox_depth=False,
wuyuefeng's avatar
wuyuefeng committed
504
                 poly2mask=True,
505
                 seg_3d_dtype='int',
wuyuefeng's avatar
wuyuefeng committed
506
507
508
509
510
511
512
513
514
                 file_client_args=dict(backend='disk')):
        super().__init__(
            with_bbox,
            with_label,
            with_mask,
            with_seg,
            poly2mask,
            file_client_args=file_client_args)
        self.with_bbox_3d = with_bbox_3d
515
        self.with_bbox_depth = with_bbox_depth
wuyuefeng's avatar
wuyuefeng committed
516
        self.with_label_3d = with_label_3d
517
        self.with_attr_label = with_attr_label
wuyuefeng's avatar
wuyuefeng committed
518
519
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d
520
        self.seg_3d_dtype = seg_3d_dtype
wuyuefeng's avatar
wuyuefeng committed
521
522

    def _load_bboxes_3d(self, results):
523
524
525
526
527
528
529
530
        """Private function to load 3D bounding box annotations.

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

        Returns:
            dict: The dict containing loaded 3D bounding box annotations.
        """
wuyuefeng's avatar
wuyuefeng committed
531
532
533
534
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
        results['bbox3d_fields'].append('gt_bboxes_3d')
        return results

535
536
537
538
539
540
541
542
543
544
545
546
547
    def _load_bboxes_depth(self, results):
        """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.
        """
        results['centers2d'] = results['ann_info']['centers2d']
        results['depths'] = results['ann_info']['depths']
        return results

wuyuefeng's avatar
wuyuefeng committed
548
    def _load_labels_3d(self, results):
549
550
551
552
553
554
555
556
        """Private function to load label annotations.

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

        Returns:
            dict: The dict containing loaded label annotations.
        """
wuyuefeng's avatar
wuyuefeng committed
557
558
559
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results

560
561
562
563
564
565
566
567
568
569
570
571
    def _load_attr_labels(self, results):
        """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

wuyuefeng's avatar
wuyuefeng committed
572
    def _load_masks_3d(self, results):
573
574
575
576
577
578
579
580
        """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.
        """
wuyuefeng's avatar
wuyuefeng committed
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
        pts_instance_mask_path = results['ann_info']['pts_instance_mask_path']

        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)
            pts_instance_mask = np.frombuffer(mask_bytes, dtype=np.int)
        except ConnectionError:
            mmcv.check_file_exist(pts_instance_mask_path)
            pts_instance_mask = np.fromfile(
                pts_instance_mask_path, dtype=np.long)

        results['pts_instance_mask'] = pts_instance_mask
        results['pts_mask_fields'].append('pts_instance_mask')
        return results

    def _load_semantic_seg_3d(self, results):
598
599
600
601
602
603
604
605
        """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.
        """
wuyuefeng's avatar
wuyuefeng committed
606
607
608
609
610
611
612
        pts_semantic_mask_path = results['ann_info']['pts_semantic_mask_path']

        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
613
614
            pts_semantic_mask = np.frombuffer(
                mask_bytes, dtype=self.seg_3d_dtype).copy()
wuyuefeng's avatar
wuyuefeng committed
615
616
617
618
619
620
621
622
623
624
        except ConnectionError:
            mmcv.check_file_exist(pts_semantic_mask_path)
            pts_semantic_mask = np.fromfile(
                pts_semantic_mask_path, dtype=np.long)

        results['pts_semantic_mask'] = pts_semantic_mask
        results['pts_seg_fields'].append('pts_semantic_mask')
        return results

    def __call__(self, results):
625
626
627
628
629
630
631
632
633
        """Call function to load multiple types annotations.

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

        Returns:
            dict: The dict containing loaded 3D bounding box, label, mask and
                semantic segmentation annotations.
        """
wuyuefeng's avatar
wuyuefeng committed
634
635
636
637
638
        results = super().__call__(results)
        if self.with_bbox_3d:
            results = self._load_bboxes_3d(results)
            if results is None:
                return None
639
640
641
642
        if self.with_bbox_depth:
            results = self._load_bboxes_depth(results)
            if results is None:
                return None
wuyuefeng's avatar
wuyuefeng committed
643
644
        if self.with_label_3d:
            results = self._load_labels_3d(results)
645
646
        if self.with_attr_label:
            results = self._load_attr_labels(results)
wuyuefeng's avatar
wuyuefeng committed
647
648
649
650
651
652
653
654
        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):
655
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
656
657
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
liyinhao's avatar
liyinhao committed
658
659
        repr_str += f'{indent_str}with_bbox_3d={self.with_bbox_3d}, '
        repr_str += f'{indent_str}with_label_3d={self.with_label_3d}, '
660
        repr_str += f'{indent_str}with_attr_label={self.with_attr_label}, '
liyinhao's avatar
liyinhao committed
661
662
663
664
665
666
        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}, '
667
        repr_str += f'{indent_str}with_bbox_depth={self.with_bbox_depth}, '
wuyuefeng's avatar
wuyuefeng committed
668
669
        repr_str += f'{indent_str}poly2mask={self.poly2mask})'
        return repr_str