loading.py 36.9 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
3
import copy
from typing import List, Optional, Union
4

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

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


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

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

    Args:
23
        to_float32 (bool, optional): Whether to convert the img to float32.
liyinhao's avatar
liyinhao committed
24
            Defaults to False.
25
26
        color_type (str, optional): Color type of the file.
            Defaults to 'unchanged'.
27
28
29
30
31
32
33
34
35
        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. Defaults to dict(backend='disk').
        num_views (int): num of view in a frame. Default to 5.
        num_ref_frames (int): num of frame in loading. Default to -1.
        test_mode (bool): Whether is test mode in loading. Default to False.
        set_default_scale (bool): Whether to set default scale. Default to
        True.
zhangwenwei's avatar
zhangwenwei committed
36
    """
zhangwenwei's avatar
zhangwenwei committed
37

38
39
    def __init__(self,
                 to_float32: bool = False,
40
41
42
43
44
45
                 color_type: str = 'unchanged',
                 file_client_args: dict = dict(backend='disk'),
                 num_views: int = 5,
                 num_ref_frames: int = -1,
                 test_mode: bool = False,
                 set_default_scale: bool = True) -> None:
zhangwenwei's avatar
zhangwenwei committed
46
47
        self.to_float32 = to_float32
        self.color_type = color_type
48
49
50
51
52
53
54
55
56
57
58
        self.file_client_args = file_client_args.copy()
        self.file_client = None
        self.num_views = num_views
        # num_ref_frames is used for multi-sweep loading
        self.num_ref_frames = num_ref_frames
        # when test_mode=False, we randomly select previous frames
        # otherwise, select the earliest one
        self.test_mode = test_mode
        self.set_default_scale = set_default_scale

    def transform(self, results: dict) -> Optional[dict]:
59
60
61
62
63
64
        """Call function to load multi-view image from files.

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

        Returns:
65
            dict: The result dict containing the multi-view image data.
66
67
68
69
70
71
72
73
74
75
                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.
        """
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
        # TODO: consider split the multi-sweep part out of this pipeline
        # Derive the mask and transform for loading of multi-sweep data
        if self.num_ref_frames > 0:
            # init choice with the current frame
            init_choice = np.array([0], dtype=np.int64)
            num_frames = len(results['img_filename']) // self.num_views - 1
            if num_frames == 0:  # no previous frame, then copy cur frames
                choices = np.random.choice(
                    1, self.num_ref_frames, replace=True)
            elif num_frames >= self.num_ref_frames:
                # NOTE: suppose the info is saved following the order
                # from latest to earlier frames
                if self.test_mode:
                    choices = np.arange(num_frames - self.num_ref_frames,
                                        num_frames) + 1
                # NOTE: +1 is for selecting previous frames
                else:
                    choices = np.random.choice(
                        num_frames, self.num_ref_frames, replace=False) + 1
            elif num_frames > 0 and num_frames < self.num_ref_frames:
                if self.test_mode:
                    base_choices = np.arange(num_frames) + 1
                    random_choices = np.random.choice(
                        num_frames,
                        self.num_ref_frames - num_frames,
                        replace=True) + 1
                    choices = np.concatenate([base_choices, random_choices])
                else:
                    choices = np.random.choice(
                        num_frames, self.num_ref_frames, replace=True) + 1
            else:
                raise NotImplementedError
            choices = np.concatenate([init_choice, choices])
            select_filename = []
            for choice in choices:
                select_filename += results['img_filename'][choice *
                                                           self.num_views:
                                                           (choice + 1) *
                                                           self.num_views]
            results['img_filename'] = select_filename
            for key in ['cam2img', 'lidar2cam']:
                if key in results:
                    select_results = []
                    for choice in choices:
                        select_results += results[key][choice *
                                                       self.num_views:(choice +
                                                                       1) *
                                                       self.num_views]
                    results[key] = select_results
            for key in ['ego2global']:
                if key in results:
                    select_results = []
                    for choice in choices:
                        select_results += [results[key][choice]]
                    results[key] = select_results
            # Transform lidar2cam to
            # [cur_lidar]2[prev_img] and [cur_lidar]2[prev_cam]
            for key in ['lidar2cam']:
                if key in results:
                    # only change matrices of previous frames
                    for choice_idx in range(1, len(choices)):
                        pad_prev_ego2global = np.eye(4)
                        prev_ego2global = results['ego2global'][choice_idx]
                        pad_prev_ego2global[:prev_ego2global.
                                            shape[0], :prev_ego2global.
                                            shape[1]] = prev_ego2global
                        pad_cur_ego2global = np.eye(4)
                        cur_ego2global = results['ego2global'][0]
                        pad_cur_ego2global[:cur_ego2global.
                                           shape[0], :cur_ego2global.
                                           shape[1]] = cur_ego2global
                        cur2prev = np.linalg.inv(pad_prev_ego2global).dot(
                            pad_cur_ego2global)
                        for result_idx in range(choice_idx * self.num_views,
                                                (choice_idx + 1) *
                                                self.num_views):
                            results[key][result_idx] = \
                                results[key][result_idx].dot(cur2prev)
        # Support multi-view images with different shapes
        # TODO: record the origin shape and padded shape
        filename, cam2img, lidar2cam = [], [], []
        for _, cam_item in results['images'].items():
            filename.append(cam_item['img_path'])
            cam2img.append(cam_item['cam2img'])
            lidar2cam.append(cam_item['lidar2cam'])
        results['filename'] = filename
        results['cam2img'] = cam2img
        results['lidar2cam'] = lidar2cam

        results['ori_cam2img'] = copy.deepcopy(results['cam2img'])

        if self.file_client is None:
            self.file_client = mmcv.FileClient(**self.file_client_args)

170
        # img is of shape (h, w, c, num_views)
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        # h and w can be different for different views
        img_bytes = [self.file_client.get(name) for name in filename]
        imgs = [
            mmcv.imfrombytes(img_byte, flag=self.color_type)
            for img_byte in img_bytes
        ]
        # handle the image with different shape
        img_shapes = np.stack([img.shape for img in imgs], axis=0)
        img_shape_max = np.max(img_shapes, axis=0)
        img_shape_min = np.min(img_shapes, axis=0)
        assert img_shape_min[-1] == img_shape_max[-1]
        if not np.all(img_shape_max == img_shape_min):
            pad_shape = img_shape_max[:2]
        else:
            pad_shape = None
        if pad_shape is not None:
            imgs = [
                mmcv.impad(img, shape=pad_shape, pad_val=0) for img in imgs
            ]
        img = np.stack(imgs, axis=-1)
zhangwenwei's avatar
zhangwenwei committed
191
192
        if self.to_float32:
            img = img.astype(np.float32)
193

zhangwenwei's avatar
zhangwenwei committed
194
        results['filename'] = filename
195
        # unravel to list, see `DefaultFormatBundle` in formating.py
196
197
        # 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
198
199
200
201
        results['img_shape'] = img.shape
        results['ori_shape'] = img.shape
        # Set initial values for default meta_keys
        results['pad_shape'] = img.shape
202
203
        if self.set_default_scale:
            results['scale_factor'] = 1.0
zhangwenwei's avatar
zhangwenwei committed
204
205
206
207
208
        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)
209
210
        results['num_views'] = self.num_views
        results['num_ref_frames'] = self.num_ref_frames
zhangwenwei's avatar
zhangwenwei committed
211
212
213
        return results

    def __repr__(self):
214
        """str: Return a string that describes the module."""
215
216
        repr_str = self.__class__.__name__
        repr_str += f'(to_float32={self.to_float32}, '
217
218
219
220
        repr_str += f"color_type='{self.color_type}', "
        repr_str += f'num_views={self.num_views}, '
        repr_str += f'num_ref_frames={self.num_ref_frames}, '
        repr_str += f'test_mode={self.test_mode})'
221
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
222
223


224
@TRANSFORMS.register_module()
225
226
227
228
229
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:
230
        kwargs (dict): Arguments are the same as those in
231
232
233
            :class:`LoadImageFromFile`.
    """

ZCMax's avatar
ZCMax committed
234
    def transform(self, results: dict) -> dict:
235
236
237
238
239
240
241
242
        """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
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
        # 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]

269
270
271
        return results


272
@TRANSFORMS.register_module()
VVsssssk's avatar
VVsssssk committed
273
class LoadPointsFromMultiSweeps(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
274
    """Load points from multiple sweeps.
zhangwenwei's avatar
zhangwenwei committed
275

zhangwenwei's avatar
zhangwenwei committed
276
277
278
    This is usually used for nuScenes dataset to utilize previous sweeps.

    Args:
279
280
281
282
283
284
285
        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
286
            https://github.com/open-mmlab/mmengine/blob/main/mmengine/fileio/file_client.py
liyinhao's avatar
liyinhao committed
287
            for more details. Defaults to dict(backend='disk').
288
        pad_empty_sweeps (bool, optional): Whether to repeat keyframe when
289
            sweeps is empty. Defaults to False.
290
        remove_close (bool, optional): Whether to remove close points.
291
            Defaults to False.
292
        test_mode (bool, optional): If `test_mode=True`, it will not
293
294
            randomly sample sweeps but select the nearest N frames.
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
295
296
    """

297
298
299
300
301
302
303
304
    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
305
        self.load_dim = load_dim
zhangwenwei's avatar
zhangwenwei committed
306
        self.sweeps_num = sweeps_num
307
        self.use_dim = use_dim
zhangwenwei's avatar
zhangwenwei committed
308
309
        self.file_client_args = file_client_args.copy()
        self.file_client = None
310
311
312
        self.pad_empty_sweeps = pad_empty_sweeps
        self.remove_close = remove_close
        self.test_mode = test_mode
zhangwenwei's avatar
zhangwenwei committed
313

314
    def _load_points(self, pts_filename: str) -> np.ndarray:
315
316
317
318
319
320
321
322
        """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
323
        if self.file_client is None:
324
            self.file_client = mmengine.FileClient(**self.file_client_args)
zhangwenwei's avatar
zhangwenwei committed
325
326
327
328
        try:
            pts_bytes = self.file_client.get(pts_filename)
            points = np.frombuffer(pts_bytes, dtype=np.float32)
        except ConnectionError:
329
            mmengine.check_file_exist(pts_filename)
zhangwenwei's avatar
zhangwenwei committed
330
331
332
333
334
            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
335

336
337
338
    def _remove_close(self,
                      points: Union[np.ndarray, BasePoints],
                      radius: float = 1.0) -> Union[np.ndarray, BasePoints]:
339
340
341
        """Removes point too close within a certain radius from origin.

        Args:
342
            points (np.ndarray | :obj:`BasePoints`): Sweep points.
343
            radius (float, optional): Radius below which points are removed.
344
345
346
                Defaults to 1.0.

        Returns:
347
            np.ndarray | :obj:`BasePoints`: Points after removing.
348
        """
349
350
351
352
353
354
355
356
        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
357
        not_close = np.logical_not(np.logical_and(x_filt, y_filt))
358
        return points[not_close]
359

360
    def transform(self, results: dict) -> dict:
361
362
363
        """Call function to load multi-sweep point clouds from files.

        Args:
364
            results (dict): Result dict containing multi-sweep point cloud
365
366
367
                filenames.

        Returns:
368
            dict: The result dict containing the multi-sweep points data.
369
                Updated key and value are described below.
370

371
                - points (np.ndarray | :obj:`BasePoints`): Multi-sweep point
372
                    cloud arrays.
373
        """
zhangwenwei's avatar
zhangwenwei committed
374
        points = results['points']
375
        points.tensor[:, 4] = 0
zhangwenwei's avatar
zhangwenwei committed
376
377
        sweep_points_list = [points]
        ts = results['timestamp']
VVsssssk's avatar
VVsssssk committed
378
379
380
381
382
383
384
        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)
385
        else:
VVsssssk's avatar
VVsssssk committed
386
387
            if len(results['lidar_sweeps']) <= self.sweeps_num:
                choices = np.arange(len(results['lidar_sweeps']))
388
389
390
391
            elif self.test_mode:
                choices = np.arange(self.sweeps_num)
            else:
                choices = np.random.choice(
VVsssssk's avatar
VVsssssk committed
392
393
394
                    len(results['lidar_sweeps']),
                    self.sweeps_num,
                    replace=False)
395
            for idx in choices:
VVsssssk's avatar
VVsssssk committed
396
397
398
                sweep = results['lidar_sweeps'][idx]
                points_sweep = self._load_points(
                    sweep['lidar_points']['lidar_path'])
399
400
401
                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
402
403
                # bc-breaking: Timestamp has divided 1e6 in pkl infos.
                sweep_ts = sweep['timestamp']
404
405
406
407
                lidar2sensor = np.array(sweep['lidar_points']['lidar2sensor'])
                points_sweep[:, :
                             3] = points_sweep[:, :3] @ lidar2sensor[:3, :3]
                points_sweep[:, :3] -= lidar2sensor[:3, 3]
408
                points_sweep[:, 4] = ts - sweep_ts
409
                points_sweep = points.new_point(points_sweep)
410
411
                sweep_points_list.append(points_sweep)

412
413
        points = points.cat(sweep_points_list)
        points = points[:, self.use_dim]
zhangwenwei's avatar
zhangwenwei committed
414
415
416
417
        results['points'] = points
        return results

    def __repr__(self):
418
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
419
        return f'{self.__class__.__name__}(sweeps_num={self.sweeps_num})'
wuyuefeng's avatar
wuyuefeng committed
420
421


422
@TRANSFORMS.register_module()
423
class PointSegClassMapping(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
424
425
    """Map original semantic class to valid category ids.

426
427
    Required Keys:

428
429
    - seg_label_mapping (np.ndarray)
    - pts_semantic_mask (np.ndarray)
430
431
432
433
434

    Added Keys:

    - points (np.float32)

wuyuefeng's avatar
wuyuefeng committed
435
436
437
438
    Map valid classes as 0~len(valid_cat_ids)-1 and
    others as len(valid_cat_ids).
    """

439
    def transform(self, results: dict) -> dict:
440
441
442
443
444
445
        """Call function to map original semantic class to valid category ids.

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

        Returns:
446
            dict: The result dict containing the mapped category ids.
447
448
449
450
                Updated key and value are described below.

                - pts_semantic_mask (np.ndarray): Mapped semantic masks.
        """
wuyuefeng's avatar
wuyuefeng committed
451
452
453
        assert 'pts_semantic_mask' in results
        pts_semantic_mask = results['pts_semantic_mask']

454
455
456
        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
457

458
        results['pts_semantic_mask'] = converted_pts_sem_mask
ZCMax's avatar
ZCMax committed
459
460
461
462
463
464
465

        # '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
466
467
468
        return results

    def __repr__(self):
469
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
470
471
472
473
        repr_str = self.__class__.__name__
        return repr_str


474
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
475
class NormalizePointsColor(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
476
    """Normalize color of points.
wuyuefeng's avatar
wuyuefeng committed
477
478
479
480
481

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

ZCMax's avatar
ZCMax committed
482
    def __init__(self, color_mean: List[float]) -> None:
wuyuefeng's avatar
wuyuefeng committed
483
484
        self.color_mean = color_mean

ZCMax's avatar
ZCMax committed
485
    def transform(self, input_dict: dict) -> dict:
486
487
488
489
490
491
        """Call function to normalize color of points.

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

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

495
                - points (:obj:`BasePoints`): Points after color normalization.
496
        """
ZCMax's avatar
ZCMax committed
497
        points = input_dict['points']
498
        assert points.attribute_dims is not None and \
499
500
               'color' in points.attribute_dims.keys(), \
               'Expect points have color attribute'
501
502
        if self.color_mean is not None:
            points.color = points.color - \
503
                           points.color.new_tensor(self.color_mean)
504
        points.color = points.color / 255.0
ZCMax's avatar
ZCMax committed
505
506
        input_dict['points'] = points
        return input_dict
wuyuefeng's avatar
wuyuefeng committed
507
508

    def __repr__(self):
509
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
510
        repr_str = self.__class__.__name__
511
        repr_str += f'(color_mean={self.color_mean})'
wuyuefeng's avatar
wuyuefeng committed
512
513
514
        return repr_str


515
@TRANSFORMS.register_module()
jshilong's avatar
jshilong committed
516
class LoadPointsFromFile(BaseTransform):
wuyuefeng's avatar
wuyuefeng committed
517
518
    """Load Points From File.

jshilong's avatar
jshilong committed
519
520
521
522
523
524
525
526
527
    Required Keys:

    - lidar_points (dict)

        - lidar_path (str)

    Added Keys:

    - points (np.float32)
wuyuefeng's avatar
wuyuefeng committed
528
529

    Args:
530
531
        coord_type (str): The type of coordinates of points cloud.
            Available options includes:
532

533
534
535
            - 'LIDAR': Points in LiDAR coordinates.
            - 'DEPTH': Points in depth coordinates, usually for indoor dataset.
            - 'CAMERA': Points in camera coordinates.
536
        load_dim (int, optional): The dimension of the loaded points.
537
            Defaults to 6.
538
539
        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
540
            or use_dim=[0, 1, 2, 3] to use the intensity dimension.
541
542
543
544
545
546
        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
547
            https://github.com/open-mmlab/mmengine/blob/main/mmengine/fileio/file_client.py
liyinhao's avatar
liyinhao committed
548
            for more details. Defaults to dict(backend='disk').
wuyuefeng's avatar
wuyuefeng committed
549
550
    """

jshilong's avatar
jshilong committed
551
552
553
554
    def __init__(
        self,
        coord_type: str,
        load_dim: int = 6,
555
        use_dim: Union[int, List[int]] = [0, 1, 2],
jshilong's avatar
jshilong committed
556
557
558
559
        shift_height: bool = False,
        use_color: bool = False,
        file_client_args: dict = dict(backend='disk')
    ) -> None:
wuyuefeng's avatar
wuyuefeng committed
560
        self.shift_height = shift_height
561
        self.use_color = use_color
wuyuefeng's avatar
wuyuefeng committed
562
563
564
565
        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}'
566
        assert coord_type in ['CAMERA', 'LIDAR', 'DEPTH']
wuyuefeng's avatar
wuyuefeng committed
567

568
        self.coord_type = coord_type
wuyuefeng's avatar
wuyuefeng committed
569
570
571
572
573
        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
574
    def _load_points(self, pts_filename: str) -> np.ndarray:
575
576
577
578
579
580
581
582
        """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
583
        if self.file_client is None:
584
            self.file_client = mmengine.FileClient(**self.file_client_args)
wuyuefeng's avatar
wuyuefeng committed
585
586
587
588
        try:
            pts_bytes = self.file_client.get(pts_filename)
            points = np.frombuffer(pts_bytes, dtype=np.float32)
        except ConnectionError:
589
            mmengine.check_file_exist(pts_filename)
wuyuefeng's avatar
wuyuefeng committed
590
591
592
593
            if pts_filename.endswith('.npy'):
                points = np.load(pts_filename)
            else:
                points = np.fromfile(pts_filename, dtype=np.float32)
594

wuyuefeng's avatar
wuyuefeng committed
595
596
        return points

jshilong's avatar
jshilong committed
597
598
    def transform(self, results: dict) -> dict:
        """Method to load points data from file.
599
600
601
602
603

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

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

607
                - points (:obj:`BasePoints`): Point clouds data.
608
        """
jshilong's avatar
jshilong committed
609
610
        pts_file_path = results['lidar_points']['lidar_path']
        points = self._load_points(pts_file_path)
wuyuefeng's avatar
wuyuefeng committed
611
612
        points = points.reshape(-1, self.load_dim)
        points = points[:, self.use_dim]
613
        attribute_dims = None
wuyuefeng's avatar
wuyuefeng committed
614
615
616
617

        if self.shift_height:
            floor_height = np.percentile(points[:, 2], 0.99)
            height = points[:, 2] - floor_height
618
619
620
            points = np.concatenate(
                [points[:, :3],
                 np.expand_dims(height, 1), points[:, 3:]], 1)
621
622
            attribute_dims = dict(height=3)

623
624
625
626
627
628
629
630
631
632
633
        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,
                ]))

634
635
636
        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
637
        results['points'] = points
638

wuyuefeng's avatar
wuyuefeng committed
639
640
641
        return results

    def __repr__(self):
642
        """str: Return a string that describes the module."""
liyinhao's avatar
liyinhao committed
643
        repr_str = self.__class__.__name__ + '('
644
645
646
647
648
        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
649
650
651
        return repr_str


652
@TRANSFORMS.register_module()
653
654
655
class LoadPointsFromDict(LoadPointsFromFile):
    """Load Points From Dict."""

ChaimZhu's avatar
ChaimZhu committed
656
    def transform(self, results: dict) -> dict:
657
658
659
660
        assert 'points' in results
        return results


661
@TRANSFORMS.register_module()
wuyuefeng's avatar
wuyuefeng committed
662
663
664
665
666
667
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
668
669
670
    Required Keys:

    - ann_info (dict)
671

jshilong's avatar
jshilong committed
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
        - 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
715
716
717
718
719
    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.
720
721
        with_attr_label (bool, optional): Whether to load attribute label.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
722
723
724
725
726
727
728
729
730
731
732
733
        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.
734
735
        with_bbox_depth (bool, optional): Whether to load 2.5D boxes.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
736
737
        poly2mask (bool, optional): Whether to convert polygon annotations
            to bitmasks. Defaults to True.
738
        seg_3d_dtype (dtype, optional): Dtype of 3D semantic masks.
jshilong's avatar
jshilong committed
739
            Defaults to int64.
wuyuefeng's avatar
wuyuefeng committed
740
        file_client_args (dict): Config dict of file clients, refer to
741
            https://github.com/open-mmlab/mmengine/blob/main/mmengine/fileio/file_client.py
wuyuefeng's avatar
wuyuefeng committed
742
743
744
            for more details.
    """

jshilong's avatar
jshilong committed
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
    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
761
        super().__init__(
jshilong's avatar
jshilong committed
762
763
764
765
766
            with_bbox=with_bbox,
            with_label=with_label,
            with_mask=with_mask,
            with_seg=with_seg,
            poly2mask=poly2mask,
wuyuefeng's avatar
wuyuefeng committed
767
768
            file_client_args=file_client_args)
        self.with_bbox_3d = with_bbox_3d
769
        self.with_bbox_depth = with_bbox_depth
wuyuefeng's avatar
wuyuefeng committed
770
        self.with_label_3d = with_label_3d
771
        self.with_attr_label = with_attr_label
wuyuefeng's avatar
wuyuefeng committed
772
773
        self.with_mask_3d = with_mask_3d
        self.with_seg_3d = with_seg_3d
774
        self.seg_3d_dtype = seg_3d_dtype
wuyuefeng's avatar
wuyuefeng committed
775

jshilong's avatar
jshilong committed
776
777
778
    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`.
779
780
781
782
783
784
785

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

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

wuyuefeng's avatar
wuyuefeng committed
787
788
789
        results['gt_bboxes_3d'] = results['ann_info']['gt_bboxes_3d']
        return results

jshilong's avatar
jshilong committed
790
    def _load_bboxes_depth(self, results: dict) -> dict:
791
792
793
794
795
796
797
798
        """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
799

800
        results['depths'] = results['ann_info']['depths']
jshilong's avatar
jshilong committed
801
        results['centers_2d'] = results['ann_info']['centers_2d']
802
803
        return results

jshilong's avatar
jshilong committed
804
    def _load_labels_3d(self, results: dict) -> dict:
805
806
807
808
809
810
811
812
        """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
813

wuyuefeng's avatar
wuyuefeng committed
814
815
816
        results['gt_labels_3d'] = results['ann_info']['gt_labels_3d']
        return results

jshilong's avatar
jshilong committed
817
    def _load_attr_labels(self, results: dict) -> dict:
818
819
820
821
822
823
824
825
826
827
828
        """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
829
    def _load_masks_3d(self, results: dict) -> dict:
830
831
832
833
834
835
836
837
        """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
838
        pts_instance_mask_path = results['pts_instance_mask_path']
wuyuefeng's avatar
wuyuefeng committed
839
840

        if self.file_client is None:
841
            self.file_client = mmengine.FileClient(**self.file_client_args)
wuyuefeng's avatar
wuyuefeng committed
842
843
        try:
            mask_bytes = self.file_client.get(pts_instance_mask_path)
844
            pts_instance_mask = np.frombuffer(mask_bytes, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
845
        except ConnectionError:
846
            mmengine.check_file_exist(pts_instance_mask_path)
wuyuefeng's avatar
wuyuefeng committed
847
            pts_instance_mask = np.fromfile(
WRH's avatar
WRH committed
848
                pts_instance_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
849
850

        results['pts_instance_mask'] = pts_instance_mask
jshilong's avatar
jshilong committed
851
852
853
        # '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
854
855
        return results

jshilong's avatar
jshilong committed
856
    def _load_semantic_seg_3d(self, results: dict) -> dict:
857
858
859
860
861
862
863
864
        """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
865
        pts_semantic_mask_path = results['pts_semantic_mask_path']
wuyuefeng's avatar
wuyuefeng committed
866
867

        if self.file_client is None:
868
            self.file_client = mmengine.FileClient(**self.file_client_args)
wuyuefeng's avatar
wuyuefeng committed
869
870
871
        try:
            mask_bytes = self.file_client.get(pts_semantic_mask_path)
            # add .copy() to fix read-only bug
872
873
            pts_semantic_mask = np.frombuffer(
                mask_bytes, dtype=self.seg_3d_dtype).copy()
wuyuefeng's avatar
wuyuefeng committed
874
        except ConnectionError:
875
            mmengine.check_file_exist(pts_semantic_mask_path)
wuyuefeng's avatar
wuyuefeng committed
876
            pts_semantic_mask = np.fromfile(
WRH's avatar
WRH committed
877
                pts_semantic_mask_path, dtype=np.int64)
wuyuefeng's avatar
wuyuefeng committed
878
879

        results['pts_semantic_mask'] = pts_semantic_mask
jshilong's avatar
jshilong committed
880
881
882
        # '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
883
884
        return results

zhangshilong's avatar
zhangshilong committed
885
886
887
888
889
890
891
892
893
894
895
896
    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.
        """

897
        results['gt_bboxes'] = results['ann_info']['gt_bboxes']
zhangshilong's avatar
zhangshilong committed
898
899
900
901
902
903
904
905
906
907

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

jshilong's avatar
jshilong committed
910
911
    def transform(self, results: dict) -> dict:
        """Function to load multiple types annotations.
912
913
914
915
916
917

        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
918
            semantic segmentation annotations.
919
        """
jshilong's avatar
jshilong committed
920
        results = super().transform(results)
wuyuefeng's avatar
wuyuefeng committed
921
922
        if self.with_bbox_3d:
            results = self._load_bboxes_3d(results)
923
924
        if self.with_bbox_depth:
            results = self._load_bboxes_depth(results)
wuyuefeng's avatar
wuyuefeng committed
925
926
        if self.with_label_3d:
            results = self._load_labels_3d(results)
927
928
        if self.with_attr_label:
            results = self._load_attr_labels(results)
wuyuefeng's avatar
wuyuefeng committed
929
930
931
932
933
934
935
936
        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):
937
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
938
939
        indent_str = '    '
        repr_str = self.__class__.__name__ + '(\n'
liyinhao's avatar
liyinhao committed
940
941
        repr_str += f'{indent_str}with_bbox_3d={self.with_bbox_3d}, '
        repr_str += f'{indent_str}with_label_3d={self.with_label_3d}, '
942
        repr_str += f'{indent_str}with_attr_label={self.with_attr_label}, '
liyinhao's avatar
liyinhao committed
943
944
945
946
947
948
        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}, '
949
        repr_str += f'{indent_str}with_bbox_depth={self.with_bbox_depth}, '
wuyuefeng's avatar
wuyuefeng committed
950
951
        repr_str += f'{indent_str}poly2mask={self.poly2mask})'
        return repr_str