"tests/models/autoencoders/test_models_vq.py" did not exist on "f4781a0b27ffb3ea61ecd25b0b87305e0960304e"
transforms_3d.py 88 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import random
3
import warnings
4
from typing import List, Optional, Tuple, Union
5
6

import cv2
7
import mmcv
8
import numpy as np
9
from mmcv.transforms import BaseTransform, RandomResize, Resize
10
from mmengine import is_tuple_of
zhangwenwei's avatar
zhangwenwei committed
11

zhangshilong's avatar
zhangshilong committed
12
from mmdet3d.models.task_modules import VoxelGenerator
13
from mmdet3d.registry import TRANSFORMS
zhangshilong's avatar
zhangshilong committed
14
15
16
17
from mmdet3d.structures import (CameraInstance3DBoxes, DepthInstance3DBoxes,
                                LiDARInstance3DBoxes)
from mmdet3d.structures.ops import box_np_ops
from mmdet3d.structures.points import BasePoints
18
19
20
from mmdet.datasets.transforms import (PhotoMetricDistortion, RandomCrop,
                                       RandomFlip)
from .compose import Compose
zhangwenwei's avatar
zhangwenwei committed
21
22
23
from .data_augment_utils import noise_per_object_v3_


24
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
25
class RandomDropPointsColor(BaseTransform):
26
27
28
29
30
31
32
    r"""Randomly set the color of points to all zeros.

    Once this transform is executed, all the points' color will be dropped.
    Refer to `PAConv <https://github.com/CVMI-Lab/PAConv/blob/main/scene_seg/
    util/transform.py#L223>`_ for more details.

    Args:
33
        drop_ratio (float, optional): The probability of dropping point colors.
34
35
36
            Defaults to 0.2.
    """

ZCMax's avatar
ZCMax committed
37
    def __init__(self, drop_ratio: float = 0.2) -> None:
38
39
40
41
        assert isinstance(drop_ratio, (int, float)) and 0 <= drop_ratio <= 1, \
            f'invalid drop_ratio value {drop_ratio}'
        self.drop_ratio = drop_ratio

ZCMax's avatar
ZCMax committed
42
    def transform(self, input_dict: dict) -> dict:
43
44
45
46
47
48
        """Call function to drop point colors.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
49
            dict: Results after color dropping,
50
51
52
53
54
55
56
                'points' key is updated in the result dict.
        """
        points = input_dict['points']
        assert points.attribute_dims is not None and \
            'color' in points.attribute_dims, \
            'Expect points have color attribute'

57
58
59
60
61
62
63
        # this if-expression is a bit strange
        # `RandomDropPointsColor` is used in training 3D segmentor PAConv
        # we discovered in our experiments that, using
        # `if np.random.rand() > 1.0 - self.drop_ratio` consistently leads to
        # better results than using `if np.random.rand() < self.drop_ratio`
        # so we keep this hack in our codebase
        if np.random.rand() > 1.0 - self.drop_ratio:
64
65
66
67
68
69
70
71
72
73
            points.color = points.color * 0.0
        return input_dict

    def __repr__(self):
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
        repr_str += f'(drop_ratio={self.drop_ratio})'
        return repr_str


74
@TRANSFORMS.register_module()
zhangwenwei's avatar
zhangwenwei committed
75
76
77
78
79
80
81
class RandomFlip3D(RandomFlip):
    """Flip the points & bbox.

    If the input dict contains the key "flip", then the flag will be used,
    otherwise it will be randomly decided by a ratio specified in the init
    method.

jshilong's avatar
jshilong committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    Required Keys:

    - points (np.float32)
    - gt_bboxes_3d (np.float32)

    Modified Keys:

    - points (np.float32)
    - gt_bboxes_3d (np.float32)

    Added Keys:

    - points (np.float32)
    - pcd_trans (np.float32)
    - pcd_rotation (np.float32)
    - pcd_rotation_angle (np.float32)
    - pcd_scale_factor (np.float32)

zhangwenwei's avatar
zhangwenwei committed
100
    Args:
101
        sync_2d (bool): Whether to apply flip according to the 2D
zhangwenwei's avatar
zhangwenwei committed
102
103
            images. If True, it will apply the same flip as that to 2D images.
            If False, it will decide whether to flip randomly and independently
liyinhao's avatar
liyinhao committed
104
            to that of 2D images. Defaults to True.
105
        flip_ratio_bev_horizontal (float): The flipping probability
liyinhao's avatar
liyinhao committed
106
            in horizontal direction. Defaults to 0.0.
107
        flip_ratio_bev_vertical (float): The flipping probability
liyinhao's avatar
liyinhao committed
108
            in vertical direction. Defaults to 0.0.
109
110
111
112
        flip_box3d (bool): Whether to flip bounding box. In most of the case,
            the box should be fliped. In cam-based bev detection, this is set
            to false, since the flip of 2D images does not influence the 3D
            box. Default to True.
zhangwenwei's avatar
zhangwenwei committed
113
114
    """

wuyuefeng's avatar
wuyuefeng committed
115
    def __init__(self,
jshilong's avatar
jshilong committed
116
117
118
                 sync_2d: bool = True,
                 flip_ratio_bev_horizontal: float = 0.0,
                 flip_ratio_bev_vertical: float = 0.0,
119
                 flip_box3d: bool = True,
jshilong's avatar
jshilong committed
120
121
122
123
                 **kwargs) -> None:
        # `flip_ratio_bev_horizontal` is equal to
        # for flip prob of 2d image when
        # `sync_2d` is True
wuyuefeng's avatar
wuyuefeng committed
124
        super(RandomFlip3D, self).__init__(
jshilong's avatar
jshilong committed
125
            prob=flip_ratio_bev_horizontal, direction='horizontal', **kwargs)
zhangwenwei's avatar
zhangwenwei committed
126
        self.sync_2d = sync_2d
jshilong's avatar
jshilong committed
127
        self.flip_ratio_bev_horizontal = flip_ratio_bev_horizontal
wuyuefeng's avatar
wuyuefeng committed
128
        self.flip_ratio_bev_vertical = flip_ratio_bev_vertical
129
        self.flip_box3d = flip_box3d
wuyuefeng's avatar
wuyuefeng committed
130
131
132
133
134
135
136
137
138
        if flip_ratio_bev_horizontal is not None:
            assert isinstance(
                flip_ratio_bev_horizontal,
                (int, float)) and 0 <= flip_ratio_bev_horizontal <= 1
        if flip_ratio_bev_vertical is not None:
            assert isinstance(
                flip_ratio_bev_vertical,
                (int, float)) and 0 <= flip_ratio_bev_vertical <= 1

jshilong's avatar
jshilong committed
139
140
141
    def random_flip_data_3d(self,
                            input_dict: dict,
                            direction: str = 'horizontal') -> None:
142
143
        """Flip 3D data randomly.

jshilong's avatar
jshilong committed
144
145
146
147
148
149
150
        `random_flip_data_3d` should take these situations into consideration:

        - 1. LIDAR-based 3d detection
        - 2. LIDAR-based 3d segmentation
        - 3. vision-only detection
        - 4. multi-modality 3d detection.

151
152
        Args:
            input_dict (dict): Result dict from loading pipeline.
153
154
            direction (str, optional): Flip direction.
                Default: 'horizontal'.
155
156

        Returns:
157
            dict: Flipped results, 'points', 'bbox3d_fields' keys are
158
159
                updated in the result dict.
        """
wuyuefeng's avatar
wuyuefeng committed
160
        assert direction in ['horizontal', 'vertical']
161
162
163
164
165
166
167
168
        if self.flip_box3d:
            if 'gt_bboxes_3d' in input_dict:
                if 'points' in input_dict:
                    input_dict['points'] = input_dict['gt_bboxes_3d'].flip(
                        direction, points=input_dict['points'])
                else:
                    # vision-only detection
                    input_dict['gt_bboxes_3d'].flip(direction)
169
            else:
170
                input_dict['points'].flip(direction)
jshilong's avatar
jshilong committed
171
172

        if 'centers_2d' in input_dict:
173
174
            assert self.sync_2d is True and direction == 'horizontal', \
                'Only support sync_2d=True and horizontal flip with images'
175
            w = input_dict['img_shape'][1]
jshilong's avatar
jshilong committed
176
177
            input_dict['centers_2d'][..., 0] = \
                w - input_dict['centers_2d'][..., 0]
178
179
            # need to modify the horizontal position of camera center
            # along u-axis in the image (flip like centers2d)
180
            # ['cam2img'][0][2] = c_u
181
182
            # see more details and examples at
            # https://github.com/open-mmlab/mmdetection3d/pull/744
183
            input_dict['cam2img'][0][2] = w - input_dict['cam2img'][0][2]
zhangwenwei's avatar
zhangwenwei committed
184

185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    def _flip_on_direction(self, results: dict) -> None:
        """Function to flip images, bounding boxes, semantic segmentation map
        and keypoints.

        Add the override feature that if 'flip' is already in results, use it
        to do the augmentation.
        """
        if 'flip' not in results:
            cur_dir = self._choose_direction()
        else:
            cur_dir = results['flip_direction']
        if cur_dir is None:
            results['flip'] = False
            results['flip_direction'] = None
        else:
            results['flip'] = True
            results['flip_direction'] = cur_dir
            self._flip(results)

jshilong's avatar
jshilong committed
204
    def transform(self, input_dict: dict) -> dict:
205
        """Call function to flip points, values in the ``bbox3d_fields`` and
206
207
208
209
210
211
        also flip 2D image and its annotations.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
212
213
            dict: Flipped results, 'flip', 'flip_direction',
                'pcd_horizontal_flip' and 'pcd_vertical_flip' keys are added
214
215
                into result dict.
        """
216
        # flip 2D image and its annotations
jshilong's avatar
jshilong committed
217
218
        if 'img' in input_dict:
            super(RandomFlip3D, self).transform(input_dict)
zhangwenwei's avatar
zhangwenwei committed
219

jshilong's avatar
jshilong committed
220
        if self.sync_2d and 'img' in input_dict:
wuyuefeng's avatar
wuyuefeng committed
221
222
            input_dict['pcd_horizontal_flip'] = input_dict['flip']
            input_dict['pcd_vertical_flip'] = False
zhangwenwei's avatar
zhangwenwei committed
223
        else:
wuyuefeng's avatar
wuyuefeng committed
224
225
            if 'pcd_horizontal_flip' not in input_dict:
                flip_horizontal = True if np.random.rand(
jshilong's avatar
jshilong committed
226
                ) < self.flip_ratio_bev_horizontal else False
wuyuefeng's avatar
wuyuefeng committed
227
228
229
230
231
232
                input_dict['pcd_horizontal_flip'] = flip_horizontal
            if 'pcd_vertical_flip' not in input_dict:
                flip_vertical = True if np.random.rand(
                ) < self.flip_ratio_bev_vertical else False
                input_dict['pcd_vertical_flip'] = flip_vertical

233
234
235
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

wuyuefeng's avatar
wuyuefeng committed
236
237
        if input_dict['pcd_horizontal_flip']:
            self.random_flip_data_3d(input_dict, 'horizontal')
238
            input_dict['transformation_3d_flow'].extend(['HF'])
wuyuefeng's avatar
wuyuefeng committed
239
240
        if input_dict['pcd_vertical_flip']:
            self.random_flip_data_3d(input_dict, 'vertical')
241
            input_dict['transformation_3d_flow'].extend(['VF'])
zhangwenwei's avatar
zhangwenwei committed
242
243
        return input_dict

zhangwenwei's avatar
zhangwenwei committed
244
    def __repr__(self):
245
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
246
        repr_str = self.__class__.__name__
247
        repr_str += f'(sync_2d={self.sync_2d},'
248
        repr_str += f' flip_ratio_bev_vertical={self.flip_ratio_bev_vertical})'
wuyuefeng's avatar
wuyuefeng committed
249
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
250

zhangwenwei's avatar
zhangwenwei committed
251

252
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
253
class RandomJitterPoints(BaseTransform):
254
255
    """Randomly jitter point coordinates.

256
    Different from the global translation in ``GlobalRotScaleTrans``, here we
257
258
259
260
        apply different noises to each point in a scene.

    Args:
        jitter_std (list[float]): The standard deviation of jittering noise.
261
262
            This applies random noise to all points in a 3D scene, which is
            sampled from a gaussian distribution whose standard deviation is
263
            set by ``jitter_std``. Defaults to [0.01, 0.01, 0.01]
264
        clip_range (list[float]): Clip the randomly generated jitter
265
266
267
268
            noise into this range. If None is given, don't perform clipping.
            Defaults to [-0.05, 0.05]

    Note:
269
        This transform should only be used in point cloud segmentation tasks
270
271
272
273
274
            because we don't transform ground-truth bboxes accordingly.
        For similar transform in detection task, please refer to `ObjectNoise`.
    """

    def __init__(self,
ZCMax's avatar
ZCMax committed
275
276
                 jitter_std: List[float] = [0.01, 0.01, 0.01],
                 clip_range: List[float] = [-0.05, 0.05]) -> None:
277
278
279
280
281
282
283
284
285
286
287
288
289
290
        seq_types = (list, tuple, np.ndarray)
        if not isinstance(jitter_std, seq_types):
            assert isinstance(jitter_std, (int, float)), \
                f'unsupported jitter_std type {type(jitter_std)}'
            jitter_std = [jitter_std, jitter_std, jitter_std]
        self.jitter_std = jitter_std

        if clip_range is not None:
            if not isinstance(clip_range, seq_types):
                assert isinstance(clip_range, (int, float)), \
                    f'unsupported clip_range type {type(clip_range)}'
                clip_range = [-clip_range, clip_range]
        self.clip_range = clip_range

ZCMax's avatar
ZCMax committed
291
    def transform(self, input_dict: dict) -> dict:
292
293
294
295
296
297
        """Call function to jitter all the points in the scene.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
298
            dict: Results after adding noise to each point,
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
                'points' key is updated in the result dict.
        """
        points = input_dict['points']
        jitter_std = np.array(self.jitter_std, dtype=np.float32)
        jitter_noise = \
            np.random.randn(points.shape[0], 3) * jitter_std[None, :]
        if self.clip_range is not None:
            jitter_noise = np.clip(jitter_noise, self.clip_range[0],
                                   self.clip_range[1])

        points.translate(jitter_noise)
        return input_dict

    def __repr__(self):
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
        repr_str += f'(jitter_std={self.jitter_std},'
        repr_str += f' clip_range={self.clip_range})'
        return repr_str


320
321
@TRANSFORMS.register_module()
class ObjectSample(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
322
    """Sample GT objects to the data.
zhangwenwei's avatar
zhangwenwei committed
323

324
325
326
327
328
329
330
331
332
333
    Required Keys:

    - points
    - ann_info
    - gt_bboxes_3d
    - gt_labels_3d
    - img (optional)
    - gt_bboxes (optional)

    Modified Keys:
334

335
336
337
338
339
340
341
342
343
344
    - points
    - gt_bboxes_3d
    - gt_labels_3d
    - img (optional)
    - gt_bboxes (optional)

    Added Keys:

    - plane (optional)

zhangwenwei's avatar
zhangwenwei committed
345
346
347
348
    Args:
        db_sampler (dict): Config dict of the database sampler.
        sample_2d (bool): Whether to also paste 2D image patch to the images
            This should be true when applying multi-modality cut-and-paste.
liyinhao's avatar
liyinhao committed
349
            Defaults to False.
350
        use_ground_plane (bool): Whether to use ground plane to adjust the
351
            3D labels.
zhangwenwei's avatar
zhangwenwei committed
352
    """
zhangwenwei's avatar
zhangwenwei committed
353

354
355
356
    def __init__(self,
                 db_sampler: dict,
                 sample_2d: bool = False,
357
                 use_ground_plane: bool = False) -> None:
zhangwenwei's avatar
zhangwenwei committed
358
359
360
361
        self.sampler_cfg = db_sampler
        self.sample_2d = sample_2d
        if 'type' not in db_sampler.keys():
            db_sampler['type'] = 'DataBaseSampler'
362
        self.db_sampler = TRANSFORMS.build(db_sampler)
363
        self.use_ground_plane = use_ground_plane
zhangwenwei's avatar
zhangwenwei committed
364
365

    @staticmethod
366
367
    def remove_points_in_boxes(points: BasePoints,
                               boxes: np.ndarray) -> np.ndarray:
368
369
370
        """Remove the points in the sampled bounding boxes.

        Args:
371
            points (:obj:`BasePoints`): Input point cloud array.
372
373
374
375
376
            boxes (np.ndarray): Sampled ground truth boxes.

        Returns:
            np.ndarray: Points with those in the boxes removed.
        """
377
        masks = box_np_ops.points_in_rbbox(points.coord.numpy(), boxes)
zhangwenwei's avatar
zhangwenwei committed
378
379
380
        points = points[np.logical_not(masks.any(-1))]
        return points

381
382
    def transform(self, input_dict: dict) -> dict:
        """Transform function to sample ground truth objects to the data.
383
384
385
386
387

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
388
389
            dict: Results after object sampling augmentation,
                'points', 'gt_bboxes_3d', 'gt_labels_3d' keys are updated
390
391
                in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
392
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
393
394
        gt_labels_3d = input_dict['gt_labels_3d']

ChaimZhu's avatar
ChaimZhu committed
395
396
        if self.use_ground_plane:
            ground_plane = input_dict.get('plane', None)
397
398
            assert ground_plane is not None, '`use_ground_plane` is True ' \
                                             'but find plane is None'
399
400
        else:
            ground_plane = None
zhangwenwei's avatar
zhangwenwei committed
401
402
403
        # change to float for blending operation
        points = input_dict['points']
        if self.sample_2d:
wuyuefeng's avatar
wuyuefeng committed
404
            img = input_dict['img']
zhangwenwei's avatar
zhangwenwei committed
405
406
407
            gt_bboxes_2d = input_dict['gt_bboxes']
            # Assume for now 3D & 2D bboxes are the same
            sampled_dict = self.db_sampler.sample_all(
408
409
410
411
                gt_bboxes_3d.tensor.numpy(),
                gt_labels_3d,
                gt_bboxes_2d=gt_bboxes_2d,
                img=img)
zhangwenwei's avatar
zhangwenwei committed
412
413
        else:
            sampled_dict = self.db_sampler.sample_all(
414
415
416
417
                gt_bboxes_3d.tensor.numpy(),
                gt_labels_3d,
                img=None,
                ground_plane=ground_plane)
zhangwenwei's avatar
zhangwenwei committed
418
419
420
421

        if sampled_dict is not None:
            sampled_gt_bboxes_3d = sampled_dict['gt_bboxes_3d']
            sampled_points = sampled_dict['points']
zhangwenwei's avatar
zhangwenwei committed
422
            sampled_gt_labels = sampled_dict['gt_labels_3d']
zhangwenwei's avatar
zhangwenwei committed
423

zhangwenwei's avatar
zhangwenwei committed
424
425
            gt_labels_3d = np.concatenate([gt_labels_3d, sampled_gt_labels],
                                          axis=0)
426
427
428
            gt_bboxes_3d = gt_bboxes_3d.new_box(
                np.concatenate(
                    [gt_bboxes_3d.tensor.numpy(), sampled_gt_bboxes_3d]))
zhangwenwei's avatar
zhangwenwei committed
429

zhangwenwei's avatar
zhangwenwei committed
430
431
            points = self.remove_points_in_boxes(points, sampled_gt_bboxes_3d)
            # check the points dimension
432
            points = points.cat([sampled_points, points])
zhangwenwei's avatar
zhangwenwei committed
433
434
435
436
437

            if self.sample_2d:
                sampled_gt_bboxes_2d = sampled_dict['gt_bboxes_2d']
                gt_bboxes_2d = np.concatenate(
                    [gt_bboxes_2d, sampled_gt_bboxes_2d]).astype(np.float32)
zhangwenwei's avatar
zhangwenwei committed
438

zhangwenwei's avatar
zhangwenwei committed
439
                input_dict['gt_bboxes'] = gt_bboxes_2d
wuyuefeng's avatar
wuyuefeng committed
440
                input_dict['img'] = sampled_dict['img']
zhangwenwei's avatar
zhangwenwei committed
441
442

        input_dict['gt_bboxes_3d'] = gt_bboxes_3d
WRH's avatar
WRH committed
443
        input_dict['gt_labels_3d'] = gt_labels_3d.astype(np.int64)
zhangwenwei's avatar
zhangwenwei committed
444
        input_dict['points'] = points
zhangwenwei's avatar
zhangwenwei committed
445

zhangwenwei's avatar
zhangwenwei committed
446
447
448
        return input_dict

    def __repr__(self):
449
        """str: Return a string that describes the module."""
450
        repr_str = self.__class__.__name__
451
        repr_str += f'db_sampler={self.db_sampler},'
452
        repr_str += f' sample_2d={self.sample_2d},'
453
        repr_str += f' use_ground_plane={self.use_ground_plane}'
454
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
455
456


457
458
@TRANSFORMS.register_module()
class ObjectNoise(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
459
    """Apply noise to each GT objects in the scene.
zhangwenwei's avatar
zhangwenwei committed
460

461
462
463
464
465
466
467
468
469
470
    Required Keys:

    - points
    - gt_bboxes_3d

    Modified Keys:

    - points
    - gt_bboxes_3d

zhangwenwei's avatar
zhangwenwei committed
471
    Args:
472
        translation_std (list[float], optional): Standard deviation of the
zhangwenwei's avatar
zhangwenwei committed
473
474
            distribution where translation noise are sampled from.
            Defaults to [0.25, 0.25, 0.25].
475
        global_rot_range (list[float], optional): Global rotation to the scene.
zhangwenwei's avatar
zhangwenwei committed
476
            Defaults to [0.0, 0.0].
477
        rot_range (list[float], optional): Object rotation range.
zhangwenwei's avatar
zhangwenwei committed
478
479
480
481
            Defaults to [-0.15707963267, 0.15707963267].
        num_try (int, optional): Number of times to try if the noise applied is
            invalid. Defaults to 100.
    """
zhangwenwei's avatar
zhangwenwei committed
482
483

    def __init__(self,
484
485
486
487
                 translation_std: List[float] = [0.25, 0.25, 0.25],
                 global_rot_range: List[float] = [0.0, 0.0],
                 rot_range: List[float] = [-0.15707963267, 0.15707963267],
                 num_try: int = 100) -> None:
zhangwenwei's avatar
zhangwenwei committed
488
        self.translation_std = translation_std
zhangwenwei's avatar
zhangwenwei committed
489
        self.global_rot_range = global_rot_range
zhangwenwei's avatar
zhangwenwei committed
490
        self.rot_range = rot_range
zhangwenwei's avatar
zhangwenwei committed
491
492
        self.num_try = num_try

493
494
    def transform(self, input_dict: dict) -> dict:
        """Transform function to apply noise to each ground truth in the scene.
495
496
497
498
499

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
500
            dict: Results after adding noise to each object,
501
502
                'points', 'gt_bboxes_3d' keys are updated in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
503
504
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
        points = input_dict['points']
zhangwenwei's avatar
zhangwenwei committed
505

506
        # TODO: this is inplace operation
507
        numpy_box = gt_bboxes_3d.tensor.numpy()
508
509
        numpy_points = points.tensor.numpy()

zhangwenwei's avatar
zhangwenwei committed
510
        noise_per_object_v3_(
511
            numpy_box,
512
            numpy_points,
zhangwenwei's avatar
zhangwenwei committed
513
514
            rotation_perturb=self.rot_range,
            center_noise_std=self.translation_std,
zhangwenwei's avatar
zhangwenwei committed
515
516
            global_random_rot_range=self.global_rot_range,
            num_try=self.num_try)
517
518

        input_dict['gt_bboxes_3d'] = gt_bboxes_3d.new_box(numpy_box)
519
        input_dict['points'] = points.new_point(numpy_points)
zhangwenwei's avatar
zhangwenwei committed
520
521
522
        return input_dict

    def __repr__(self):
523
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
524
        repr_str = self.__class__.__name__
525
526
527
528
        repr_str += f'(num_try={self.num_try},'
        repr_str += f' translation_std={self.translation_std},'
        repr_str += f' global_rot_range={self.global_rot_range},'
        repr_str += f' rot_range={self.rot_range})'
zhangwenwei's avatar
zhangwenwei committed
529
530
531
        return repr_str


532
@TRANSFORMS.register_module()
533
class GlobalAlignment(BaseTransform):
534
535
536
537
538
539
    """Apply global alignment to 3D scene points by rotation and translation.

    Args:
        rotation_axis (int): Rotation axis for points and bboxes rotation.

    Note:
540
541
        We do not record the applied rotation and translation as in
            GlobalRotScaleTrans. Because usually, we do not need to reverse
542
            the alignment step.
543
        For example, ScanNet 3D detection task uses aligned ground-truth
544
545
546
            bounding boxes for evaluation.
    """

547
    def __init__(self, rotation_axis: int) -> None:
548
549
        self.rotation_axis = rotation_axis

550
    def _trans_points(self, results: dict, trans_factor: np.ndarray) -> None:
551
552
553
554
555
556
557
558
559
        """Private function to translate points.

        Args:
            input_dict (dict): Result dict from loading pipeline.
            trans_factor (np.ndarray): Translation vector to be applied.

        Returns:
            dict: Results after translation, 'points' is updated in the dict.
        """
560
        results['points'].translate(trans_factor)
561

562
    def _rot_points(self, results: dict, rot_mat: np.ndarray) -> None:
563
564
565
566
567
568
569
570
571
572
        """Private function to rotate bounding boxes and points.

        Args:
            input_dict (dict): Result dict from loading pipeline.
            rot_mat (np.ndarray): Rotation matrix to be applied.

        Returns:
            dict: Results after rotation, 'points' is updated in the dict.
        """
        # input should be rot_mat_T so I transpose it here
573
        results['points'].rotate(rot_mat.T)
574

575
    def _check_rot_mat(self, rot_mat: np.ndarray) -> None:
576
577
578
579
580
581
582
583
584
585
586
587
        """Check if rotation matrix is valid for self.rotation_axis.

        Args:
            rot_mat (np.ndarray): Rotation matrix to be applied.
        """
        is_valid = np.allclose(np.linalg.det(rot_mat), 1.0)
        valid_array = np.zeros(3)
        valid_array[self.rotation_axis] = 1.0
        is_valid &= (rot_mat[self.rotation_axis, :] == valid_array).all()
        is_valid &= (rot_mat[:, self.rotation_axis] == valid_array).all()
        assert is_valid, f'invalid rotation matrix {rot_mat}'

588
    def transform(self, results: dict) -> dict:
589
590
591
592
593
594
        """Call function to shuffle points.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
595
            dict: Results after global alignment, 'points' and keys in
596
597
                input_dict['bbox3d_fields'] are updated in the result dict.
        """
598
        assert 'axis_align_matrix' in results, \
599
600
            'axis_align_matrix is not provided in GlobalAlignment'

601
        axis_align_matrix = results['axis_align_matrix']
602
603
604
605
606
607
        assert axis_align_matrix.shape == (4, 4), \
            f'invalid shape {axis_align_matrix.shape} for axis_align_matrix'
        rot_mat = axis_align_matrix[:3, :3]
        trans_vec = axis_align_matrix[:3, -1]

        self._check_rot_mat(rot_mat)
608
609
        self._rot_points(results, rot_mat)
        self._trans_points(results, trans_vec)
610

611
        return results
612
613

    def __repr__(self):
614
        """str: Return a string that describes the module."""
615
616
617
618
619
        repr_str = self.__class__.__name__
        repr_str += f'(rotation_axis={self.rotation_axis})'
        return repr_str


620
@TRANSFORMS.register_module()
jshilong's avatar
jshilong committed
621
class GlobalRotScaleTrans(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
622
    """Apply global rotation, scaling and translation to a 3D scene.
zhangwenwei's avatar
zhangwenwei committed
623

jshilong's avatar
jshilong committed
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
    Required Keys:

    - points (np.float32)
    - gt_bboxes_3d (np.float32)

    Modified Keys:

    - points (np.float32)
    - gt_bboxes_3d (np.float32)

    Added Keys:

    - points (np.float32)
    - pcd_trans (np.float32)
    - pcd_rotation (np.float32)
    - pcd_rotation_angle (np.float32)
    - pcd_scale_factor (np.float32)

zhangwenwei's avatar
zhangwenwei committed
642
    Args:
643
        rot_range (list[float], optional): Range of rotation angle.
liyinhao's avatar
liyinhao committed
644
            Defaults to [-0.78539816, 0.78539816] (close to [-pi/4, pi/4]).
645
        scale_ratio_range (list[float], optional): Range of scale ratio.
liyinhao's avatar
liyinhao committed
646
            Defaults to [0.95, 1.05].
647
648
        translation_std (list[float], optional): The standard deviation of
            translation noise applied to a scene, which
zhangwenwei's avatar
zhangwenwei committed
649
            is sampled from a gaussian distribution whose standard deviation
liyinhao's avatar
liyinhao committed
650
            is set by ``translation_std``. Defaults to [0, 0, 0]
651
        shift_height (bool, optional): Whether to shift height.
wuyuefeng's avatar
wuyuefeng committed
652
            (the fourth dimension of indoor points) when scaling.
liyinhao's avatar
liyinhao committed
653
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
654
    """
zhangwenwei's avatar
zhangwenwei committed
655
656

    def __init__(self,
jshilong's avatar
jshilong committed
657
658
659
660
                 rot_range: List[float] = [-0.78539816, 0.78539816],
                 scale_ratio_range: List[float] = [0.95, 1.05],
                 translation_std: List[int] = [0, 0, 0],
                 shift_height: bool = False) -> None:
661
662
663
664
665
        seq_types = (list, tuple, np.ndarray)
        if not isinstance(rot_range, seq_types):
            assert isinstance(rot_range, (int, float)), \
                f'unsupported rot_range type {type(rot_range)}'
            rot_range = [-rot_range, rot_range]
zhangwenwei's avatar
zhangwenwei committed
666
        self.rot_range = rot_range
667
668
669

        assert isinstance(scale_ratio_range, seq_types), \
            f'unsupported scale_ratio_range type {type(scale_ratio_range)}'
jshilong's avatar
jshilong committed
670

zhangwenwei's avatar
zhangwenwei committed
671
        self.scale_ratio_range = scale_ratio_range
672
673
674
675
676
677
678

        if not isinstance(translation_std, seq_types):
            assert isinstance(translation_std, (int, float)), \
                f'unsupported translation_std type {type(translation_std)}'
            translation_std = [
                translation_std, translation_std, translation_std
            ]
679
680
        assert all([std >= 0 for std in translation_std]), \
            'translation_std should be positive'
zhangwenwei's avatar
zhangwenwei committed
681
        self.translation_std = translation_std
wuyuefeng's avatar
wuyuefeng committed
682
        self.shift_height = shift_height
zhangwenwei's avatar
zhangwenwei committed
683

jshilong's avatar
jshilong committed
684
    def _trans_bbox_points(self, input_dict: dict) -> None:
685
686
687
688
689
690
        """Private function to translate bounding boxes and points.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
691
            dict: Results after translation, 'points', 'pcd_trans'
jshilong's avatar
jshilong committed
692
693
            and `gt_bboxes_3d` is updated
            in the result dict.
694
        """
695
        translation_std = np.array(self.translation_std, dtype=np.float32)
zhangwenwei's avatar
zhangwenwei committed
696
697
        trans_factor = np.random.normal(scale=translation_std, size=3).T

698
        input_dict['points'].translate(trans_factor)
zhangwenwei's avatar
zhangwenwei committed
699
        input_dict['pcd_trans'] = trans_factor
jshilong's avatar
jshilong committed
700
701
        if 'gt_bboxes_3d' in input_dict:
            input_dict['gt_bboxes_3d'].translate(trans_factor)
zhangwenwei's avatar
zhangwenwei committed
702

jshilong's avatar
jshilong committed
703
    def _rot_bbox_points(self, input_dict: dict) -> None:
704
705
706
707
708
709
        """Private function to rotate bounding boxes and points.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
710
            dict: Results after rotation, 'points', 'pcd_rotation'
jshilong's avatar
jshilong committed
711
712
            and `gt_bboxes_3d` is updated
            in the result dict.
713
        """
zhangwenwei's avatar
zhangwenwei committed
714
        rotation = self.rot_range
zhangwenwei's avatar
zhangwenwei committed
715
        noise_rotation = np.random.uniform(rotation[0], rotation[1])
zhangwenwei's avatar
zhangwenwei committed
716

jshilong's avatar
jshilong committed
717
718
719
720
721
722
723
724
        if 'gt_bboxes_3d' in input_dict and \
                len(input_dict['gt_bboxes_3d'].tensor) != 0:
            # rotate points with bboxes
            points, rot_mat_T = input_dict['gt_bboxes_3d'].rotate(
                noise_rotation, input_dict['points'])
            input_dict['points'] = points
        else:
            # if no bbox in input_dict, only rotate points
725
            rot_mat_T = input_dict['points'].rotate(noise_rotation)
jshilong's avatar
jshilong committed
726
727
728
729
730

        input_dict['pcd_rotation'] = rot_mat_T
        input_dict['pcd_rotation_angle'] = noise_rotation

    def _scale_bbox_points(self, input_dict: dict) -> None:
731
732
733
734
735
736
        """Private function to scale bounding boxes and points.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
jshilong's avatar
jshilong committed
737
738
739
            dict: Results after scaling, 'points' and
            `gt_bboxes_3d` is updated
            in the result dict.
740
        """
zhangwenwei's avatar
zhangwenwei committed
741
        scale = input_dict['pcd_scale_factor']
742
743
        points = input_dict['points']
        points.scale(scale)
wuyuefeng's avatar
wuyuefeng committed
744
        if self.shift_height:
745
746
            assert 'height' in points.attribute_dims.keys(), \
                'setting shift_height=True but points have no height attribute'
747
748
            points.tensor[:, points.attribute_dims['height']] *= scale
        input_dict['points'] = points
wuyuefeng's avatar
wuyuefeng committed
749

jshilong's avatar
jshilong committed
750
751
752
        if 'gt_bboxes_3d' in input_dict and \
                len(input_dict['gt_bboxes_3d'].tensor) != 0:
            input_dict['gt_bboxes_3d'].scale(scale)
zhangwenwei's avatar
zhangwenwei committed
753

jshilong's avatar
jshilong committed
754
    def _random_scale(self, input_dict: dict) -> None:
755
756
757
758
759
760
        """Private function to randomly set the scale factor.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
jshilong's avatar
jshilong committed
761
762
            dict: Results after scaling, 'pcd_scale_factor'
            are updated in the result dict.
763
        """
zhangwenwei's avatar
zhangwenwei committed
764
765
766
        scale_factor = np.random.uniform(self.scale_ratio_range[0],
                                         self.scale_ratio_range[1])
        input_dict['pcd_scale_factor'] = scale_factor
zhangwenwei's avatar
zhangwenwei committed
767

jshilong's avatar
jshilong committed
768
    def transform(self, input_dict: dict) -> dict:
769
        """Private function to rotate, scale and translate bounding boxes and
770
771
772
773
774
775
776
        points.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
            dict: Results after scaling, 'points', 'pcd_rotation',
jshilong's avatar
jshilong committed
777
778
            'pcd_scale_factor', 'pcd_trans' and `gt_bboxes_3d` is updated
            in the result dict.
779
        """
780
781
782
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

zhangwenwei's avatar
zhangwenwei committed
783
        self._rot_bbox_points(input_dict)
zhangwenwei's avatar
zhangwenwei committed
784

zhangwenwei's avatar
zhangwenwei committed
785
786
787
        if 'pcd_scale_factor' not in input_dict:
            self._random_scale(input_dict)
        self._scale_bbox_points(input_dict)
zhangwenwei's avatar
zhangwenwei committed
788

zhangwenwei's avatar
zhangwenwei committed
789
        self._trans_bbox_points(input_dict)
790
791

        input_dict['transformation_3d_flow'].extend(['R', 'S', 'T'])
zhangwenwei's avatar
zhangwenwei committed
792
793
794
        return input_dict

    def __repr__(self):
795
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
796
        repr_str = self.__class__.__name__
797
798
799
800
        repr_str += f'(rot_range={self.rot_range},'
        repr_str += f' scale_ratio_range={self.scale_ratio_range},'
        repr_str += f' translation_std={self.translation_std},'
        repr_str += f' shift_height={self.shift_height})'
zhangwenwei's avatar
zhangwenwei committed
801
802
803
        return repr_str


804
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
805
class PointShuffle(BaseTransform):
806
    """Shuffle input points."""
zhangwenwei's avatar
zhangwenwei committed
807

ZCMax's avatar
ZCMax committed
808
    def transform(self, input_dict: dict) -> dict:
809
810
811
812
813
814
        """Call function to shuffle points.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
815
            dict: Results after filtering, 'points', 'pts_instance_mask'
816
                and 'pts_semantic_mask' keys are updated in the result dict.
817
        """
818
819
820
821
822
823
824
825
826
827
828
829
        idx = input_dict['points'].shuffle()
        idx = idx.numpy()

        pts_instance_mask = input_dict.get('pts_instance_mask', None)
        pts_semantic_mask = input_dict.get('pts_semantic_mask', None)

        if pts_instance_mask is not None:
            input_dict['pts_instance_mask'] = pts_instance_mask[idx]

        if pts_semantic_mask is not None:
            input_dict['pts_semantic_mask'] = pts_semantic_mask[idx]

zhangwenwei's avatar
zhangwenwei committed
830
831
832
        return input_dict

    def __repr__(self):
833
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
834
835
836
        return self.__class__.__name__


837
@TRANSFORMS.register_module()
838
class ObjectRangeFilter(BaseTransform):
839
840
    """Filter objects by the range.

841
842
843
844
845
846
847
848
    Required Keys:

    - gt_bboxes_3d

    Modified Keys:

    - gt_bboxes_3d

849
850
851
    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
852

853
    def __init__(self, point_cloud_range: List[float]):
zhangwenwei's avatar
zhangwenwei committed
854
855
        self.pcd_range = np.array(point_cloud_range, dtype=np.float32)

856
857
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter objects by the range.
858
859
860
861
862

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
863
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d'
864
865
                keys are updated in the result dict.
        """
866
867
868
869
870
871
872
        # Check points instance type and initialise bev_range
        if isinstance(input_dict['gt_bboxes_3d'],
                      (LiDARInstance3DBoxes, DepthInstance3DBoxes)):
            bev_range = self.pcd_range[[0, 1, 3, 4]]
        elif isinstance(input_dict['gt_bboxes_3d'], CameraInstance3DBoxes):
            bev_range = self.pcd_range[[0, 2, 3, 5]]

zhangwenwei's avatar
zhangwenwei committed
873
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
874
        gt_labels_3d = input_dict['gt_labels_3d']
875
        mask = gt_bboxes_3d.in_range_bev(bev_range)
zhangwenwei's avatar
zhangwenwei committed
876
        gt_bboxes_3d = gt_bboxes_3d[mask]
ZwwWayne's avatar
ZwwWayne committed
877
878
879
880
881
        # mask is a torch tensor but gt_labels_3d is still numpy array
        # using mask to index gt_labels_3d will cause bug when
        # len(gt_labels_3d) == 1, where mask=1 will be interpreted
        # as gt_labels_3d[1] and cause out of index error
        gt_labels_3d = gt_labels_3d[mask.numpy().astype(np.bool)]
zhangwenwei's avatar
zhangwenwei committed
882
883

        # limit rad to [-pi, pi]
884
885
        gt_bboxes_3d.limit_yaw(offset=0.5, period=2 * np.pi)
        input_dict['gt_bboxes_3d'] = gt_bboxes_3d
zhangwenwei's avatar
zhangwenwei committed
886
887
        input_dict['gt_labels_3d'] = gt_labels_3d

zhangwenwei's avatar
zhangwenwei committed
888
889
890
        return input_dict

    def __repr__(self):
891
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
892
        repr_str = self.__class__.__name__
893
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
894
895
896
        return repr_str


897
@TRANSFORMS.register_module()
898
class PointsRangeFilter(BaseTransform):
899
900
    """Filter points by the range.

901
902
903
904
905
906
907
908
909
910
    Required Keys:

    - points
    - pts_instance_mask (optional)

    Modified Keys:

    - points
    - pts_instance_mask (optional)

911
912
913
    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
914

915
    def __init__(self, point_cloud_range: List[float]) -> None:
916
        self.pcd_range = np.array(point_cloud_range, dtype=np.float32)
zhangwenwei's avatar
zhangwenwei committed
917

918
919
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter points by the range.
920
921
922
923
924

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
925
            dict: Results after filtering, 'points', 'pts_instance_mask'
926
                and 'pts_semantic_mask' keys are updated in the result dict.
927
        """
zhangwenwei's avatar
zhangwenwei committed
928
        points = input_dict['points']
929
930
        points_mask = points.in_range_3d(self.pcd_range)
        clean_points = points[points_mask]
zhangwenwei's avatar
zhangwenwei committed
931
        input_dict['points'] = clean_points
932
933
934
935
936
937
938
939
940
941
942
        points_mask = points_mask.numpy()

        pts_instance_mask = input_dict.get('pts_instance_mask', None)
        pts_semantic_mask = input_dict.get('pts_semantic_mask', None)

        if pts_instance_mask is not None:
            input_dict['pts_instance_mask'] = pts_instance_mask[points_mask]

        if pts_semantic_mask is not None:
            input_dict['pts_semantic_mask'] = pts_semantic_mask[points_mask]

zhangwenwei's avatar
zhangwenwei committed
943
944
945
        return input_dict

    def __repr__(self):
946
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
947
        repr_str = self.__class__.__name__
948
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
949
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
950
951


952
@TRANSFORMS.register_module()
953
class ObjectNameFilter(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
954
    """Filter GT objects by their names.
zhangwenwei's avatar
zhangwenwei committed
955

956
957
958
959
960
961
962
963
    Required Keys:

    - gt_labels_3d

    Modified Keys:

    - gt_labels_3d

zhangwenwei's avatar
zhangwenwei committed
964
    Args:
liyinhao's avatar
liyinhao committed
965
        classes (list[str]): List of class names to be kept for training.
zhangwenwei's avatar
zhangwenwei committed
966
967
    """

968
    def __init__(self, classes: List[str]) -> None:
zhangwenwei's avatar
zhangwenwei committed
969
970
971
        self.classes = classes
        self.labels = list(range(len(self.classes)))

972
973
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter objects by their names.
974
975
976
977
978

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
979
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d'
980
981
                keys are updated in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
982
983
984
985
986
987
988
989
990
        gt_labels_3d = input_dict['gt_labels_3d']
        gt_bboxes_mask = np.array([n in self.labels for n in gt_labels_3d],
                                  dtype=np.bool_)
        input_dict['gt_bboxes_3d'] = input_dict['gt_bboxes_3d'][gt_bboxes_mask]
        input_dict['gt_labels_3d'] = input_dict['gt_labels_3d'][gt_bboxes_mask]

        return input_dict

    def __repr__(self):
991
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
992
993
994
        repr_str = self.__class__.__name__
        repr_str += f'(classes={self.classes})'
        return repr_str
wuyuefeng's avatar
wuyuefeng committed
995
996


997
998
@TRANSFORMS.register_module()
class PointSample(BaseTransform):
999
    """Point sample.
wuyuefeng's avatar
wuyuefeng committed
1000
1001
1002

    Sampling data to a certain number.

1003
    Required Keys:
1004

1005
1006
1007
1008
1009
    - points
    - pts_instance_mask (optional)
    - pts_semantic_mask (optional)

    Modified Keys:
1010

1011
1012
1013
1014
    - points
    - pts_instance_mask (optional)
    - pts_semantic_mask (optional)

wuyuefeng's avatar
wuyuefeng committed
1015
1016
    Args:
        num_points (int): Number of points to be sampled.
1017
        sample_range (float, optional): The range where to sample points.
1018
1019
1020
1021
            If not None, the points with depth larger than `sample_range` are
            prior to be sampled. Defaults to None.
        replace (bool, optional): Whether the sampling is with or without
            replacement. Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
1022
1023
    """

1024
1025
    def __init__(self,
                 num_points: int,
1026
1027
                 sample_range: Optional[float] = None,
                 replace: bool = False) -> None:
wuyuefeng's avatar
wuyuefeng committed
1028
        self.num_points = num_points
1029
1030
1031
        self.sample_range = sample_range
        self.replace = replace

1032
1033
1034
1035
1036
1037
1038
1039
    def _points_random_sampling(
        self,
        points: BasePoints,
        num_samples: int,
        sample_range: Optional[float] = None,
        replace: bool = False,
        return_choices: bool = False
    ) -> Union[Tuple[BasePoints, np.ndarray], BasePoints]:
wuyuefeng's avatar
wuyuefeng committed
1040
1041
1042
1043
1044
        """Points random sampling.

        Sample points to a certain number.

        Args:
1045
            points (:obj:`BasePoints`): 3D Points.
wuyuefeng's avatar
wuyuefeng committed
1046
            num_samples (int): Number of samples to be sampled.
1047
            sample_range (float, optional): Indicating the range where the
1048
                points will be sampled. Defaults to None.
1049
            replace (bool, optional): Sampling with or without replacement.
1050
                Defaults to False.
1051
1052
            return_choices (bool, optional): Whether return choice.
                Defaults to False.
1053

wuyuefeng's avatar
wuyuefeng committed
1054
        Returns:
1055
1056
1057
            tuple[:obj:`BasePoints`, np.ndarray] | :obj:`BasePoints`:

                - points (:obj:`BasePoints`): 3D Points.
1058
                - choices (np.ndarray, optional): The generated random samples.
wuyuefeng's avatar
wuyuefeng committed
1059
        """
1060
        if not replace:
wuyuefeng's avatar
wuyuefeng committed
1061
            replace = (points.shape[0] < num_samples)
1062
1063
1064
        point_range = range(len(points))
        if sample_range is not None and not replace:
            # Only sampling the near points when len(points) >= num_samples
1065
            dist = np.linalg.norm(points.coord.numpy(), axis=1)
1066
1067
            far_inds = np.where(dist >= sample_range)[0]
            near_inds = np.where(dist < sample_range)[0]
1068
1069
1070
1071
            # in case there are too many far points
            if len(far_inds) > num_samples:
                far_inds = np.random.choice(
                    far_inds, num_samples, replace=False)
1072
1073
1074
1075
1076
1077
1078
            point_range = near_inds
            num_samples -= len(far_inds)
        choices = np.random.choice(point_range, num_samples, replace=replace)
        if sample_range is not None and not replace:
            choices = np.concatenate((far_inds, choices))
            # Shuffle points after sampling
            np.random.shuffle(choices)
wuyuefeng's avatar
wuyuefeng committed
1079
1080
1081
1082
1083
        if return_choices:
            return points[choices], choices
        else:
            return points[choices]

1084
    def transform(self, input_dict: dict) -> dict:
1085
        """Transform function to sample points to in indoor scenes.
1086
1087
1088

        Args:
            input_dict (dict): Result dict from loading pipeline.
1089

1090
        Returns:
1091
            dict: Results after sampling, 'points', 'pts_instance_mask'
1092
1093
                and 'pts_semantic_mask' keys are updated in the result dict.
        """
1094
        points = input_dict['points']
1095
1096
1097
1098
1099
1100
        points, choices = self._points_random_sampling(
            points,
            self.num_points,
            self.sample_range,
            self.replace,
            return_choices=True)
1101
        input_dict['points'] = points
1102

1103
1104
        pts_instance_mask = input_dict.get('pts_instance_mask', None)
        pts_semantic_mask = input_dict.get('pts_semantic_mask', None)
wuyuefeng's avatar
wuyuefeng committed
1105

1106
        if pts_instance_mask is not None:
wuyuefeng's avatar
wuyuefeng committed
1107
            pts_instance_mask = pts_instance_mask[choices]
1108
            input_dict['pts_instance_mask'] = pts_instance_mask
1109
1110
1111

        if pts_semantic_mask is not None:
            pts_semantic_mask = pts_semantic_mask[choices]
1112
            input_dict['pts_semantic_mask'] = pts_semantic_mask
wuyuefeng's avatar
wuyuefeng committed
1113

1114
        return input_dict
wuyuefeng's avatar
wuyuefeng committed
1115
1116

    def __repr__(self):
1117
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
1118
        repr_str = self.__class__.__name__
1119
        repr_str += f'(num_points={self.num_points},'
1120
1121
        repr_str += f' sample_range={self.sample_range},'
        repr_str += f' replace={self.replace})'
1122

1123
1124
1125
        return repr_str


1126
@TRANSFORMS.register_module()
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
class IndoorPointSample(PointSample):
    """Indoor point sample.

    Sampling data to a certain number.
    NOTE: IndoorPointSample is deprecated in favor of PointSample

    Args:
        num_points (int): Number of points to be sampled.
    """

    def __init__(self, *args, **kwargs):
        warnings.warn(
            'IndoorPointSample is deprecated in favor of PointSample')
        super(IndoorPointSample, self).__init__(*args, **kwargs)


1143
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1144
class IndoorPatchPointSample(BaseTransform):
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
    r"""Indoor point sample within a patch. Modified from `PointNet++ <https://
    github.com/charlesq34/pointnet2/blob/master/scannet/scannet_dataset.py>`_.

    Sampling data to a certain number for semantic segmentation.

    Args:
        num_points (int): Number of points to be sampled.
        block_size (float, optional): Size of a block to sample points from.
            Defaults to 1.5.
        sample_rate (float, optional): Stride used in sliding patch generation.
1155
1156
1157
            This parameter is unused in `IndoorPatchPointSample` and thus has
            been deprecated. We plan to remove it in the future.
            Defaults to None.
1158
1159
        ignore_index (int, optional): Label index that won't be used for the
            segmentation task. This is set in PointSegClassMapping as neg_cls.
1160
            If not None, will be used as a patch selection criterion.
1161
1162
1163
1164
1165
            Defaults to None.
        use_normalized_coord (bool, optional): Whether to use normalized xyz as
            additional features. Defaults to False.
        num_try (int, optional): Number of times to try if the patch selected
            is invalid. Defaults to 10.
1166
        enlarge_size (float, optional): Enlarge the sampled patch to
1167
            [-block_size / 2 - enlarge_size, block_size / 2 + enlarge_size] as
1168
            an augmentation. If None, set it as 0. Defaults to 0.2.
1169
        min_unique_num (int, optional): Minimum number of unique points
1170
1171
            the sampled patch should contain. If None, use PointNet++'s method
            to judge uniqueness. Defaults to None.
1172
1173
        eps (float, optional): A value added to patch boundary to guarantee
            points coverage. Defaults to 1e-2.
1174
1175
1176
1177
1178
1179

    Note:
        This transform should only be used in the training process of point
            cloud segmentation tasks. For the sliding patch generation and
            inference process in testing, please refer to the `slide_inference`
            function of `EncoderDecoder3D` class.
1180
1181
1182
    """

    def __init__(self,
ZCMax's avatar
ZCMax committed
1183
1184
1185
1186
1187
1188
1189
1190
1191
                 num_points: int,
                 block_size: float = 1.5,
                 sample_rate: Optional[float] = None,
                 ignore_index: Optional[int] = None,
                 use_normalized_coord: bool = False,
                 num_try: int = 10,
                 enlarge_size: float = 0.2,
                 min_unique_num: Optional[int] = None,
                 eps: float = 1e-2) -> None:
1192
1193
1194
1195
1196
        self.num_points = num_points
        self.block_size = block_size
        self.ignore_index = ignore_index
        self.use_normalized_coord = use_normalized_coord
        self.num_try = num_try
1197
        self.enlarge_size = enlarge_size if enlarge_size is not None else 0.0
1198
        self.min_unique_num = min_unique_num
1199
        self.eps = eps
1200
1201
1202
1203
1204

        if sample_rate is not None:
            warnings.warn(
                "'sample_rate' has been deprecated and will be removed in "
                'the future. Please remove them from your code.')
1205

ZCMax's avatar
ZCMax committed
1206
1207
1208
1209
    def _input_generation(self, coords: np.ndarray, patch_center: np.ndarray,
                          coord_max: np.ndarray, attributes: np.ndarray,
                          attribute_dims: dict,
                          point_type: type) -> BasePoints:
1210
1211
        """Generating model input.

1212
        Generate input by subtracting patch center and adding additional
1213
1214
1215
1216
1217
1218
1219
1220
1221
            features. Currently support colors and normalized xyz as features.

        Args:
            coords (np.ndarray): Sampled 3D Points.
            patch_center (np.ndarray): Center coordinate of the selected patch.
            coord_max (np.ndarray): Max coordinate of all 3D Points.
            attributes (np.ndarray): features of input points.
            attribute_dims (dict): Dictionary to indicate the meaning of extra
                dimension.
1222
            point_type (type): class of input points inherited from BasePoints.
1223
1224

        Returns:
1225
            :obj:`BasePoints`: The generated input data.
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
        """
        # subtract patch center, the z dimension is not centered
        centered_coords = coords.copy()
        centered_coords[:, 0] -= patch_center[0]
        centered_coords[:, 1] -= patch_center[1]

        if self.use_normalized_coord:
            normalized_coord = coords / coord_max
            attributes = np.concatenate([attributes, normalized_coord], axis=1)
            if attribute_dims is None:
                attribute_dims = dict()
            attribute_dims.update(
                dict(normalized_coord=[
                    attributes.shape[1], attributes.shape[1] +
                    1, attributes.shape[1] + 2
                ]))

        points = np.concatenate([centered_coords, attributes], axis=1)
        points = point_type(
            points, points_dim=points.shape[1], attribute_dims=attribute_dims)

        return points

1249
    def _patch_points_sampling(
1250
1251
            self, points: BasePoints,
            sem_mask: np.ndarray) -> Tuple[BasePoints, np.ndarray]:
1252
1253
1254
1255
1256
1257
        """Patch points sampling.

        First sample a valid patch.
        Then sample points within that patch to a certain number.

        Args:
1258
            points (:obj:`BasePoints`): 3D Points.
1259
1260
1261
            sem_mask (np.ndarray): semantic segmentation mask for input points.

        Returns:
1262
            tuple[:obj:`BasePoints`, np.ndarray]:
1263

1264
                - points (:obj:`BasePoints`): 3D Points.
1265
1266
1267
1268
1269
1270
1271
1272
1273
1274
                - choices (np.ndarray): The generated random samples.
        """
        coords = points.coord.numpy()
        attributes = points.tensor[:, 3:].numpy()
        attribute_dims = points.attribute_dims
        point_type = type(points)

        coord_max = np.amax(coords, axis=0)
        coord_min = np.amin(coords, axis=0)

1275
        for _ in range(self.num_try):
1276
1277
1278
            # random sample a point as patch center
            cur_center = coords[np.random.choice(coords.shape[0])]

1279
1280
            # boundary of a patch, which would be enlarged by
            # `self.enlarge_size` as an augmentation
1281
1282
1283
1284
1285
1286
1287
            cur_max = cur_center + np.array(
                [self.block_size / 2.0, self.block_size / 2.0, 0.0])
            cur_min = cur_center - np.array(
                [self.block_size / 2.0, self.block_size / 2.0, 0.0])
            cur_max[2] = coord_max[2]
            cur_min[2] = coord_min[2]
            cur_choice = np.sum(
1288
1289
                (coords >= (cur_min - self.enlarge_size)) *
                (coords <= (cur_max + self.enlarge_size)),
1290
1291
1292
1293
1294
1295
1296
                axis=1) == 3

            if not cur_choice.any():  # no points in this patch
                continue

            cur_coords = coords[cur_choice, :]
            cur_sem_mask = sem_mask[cur_choice]
1297
            point_idxs = np.where(cur_choice)[0]
1298
            mask = np.sum(
1299
1300
                (cur_coords >= (cur_min - self.eps)) * (cur_coords <=
                                                        (cur_max + self.eps)),
1301
                axis=1) == 3
1302

1303
1304
            # two criteria for patch sampling, adopted from PointNet++
            # 1. selected patch should contain enough unique points
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
            if self.min_unique_num is None:
                # use PointNet++'s method as default
                # [31, 31, 62] are just some big values used to transform
                # coords from 3d array to 1d and then check their uniqueness
                # this is used in all the ScanNet code following PointNet++
                vidx = np.ceil(
                    (cur_coords[mask, :] - cur_min) / (cur_max - cur_min) *
                    np.array([31.0, 31.0, 62.0]))
                vidx = np.unique(vidx[:, 0] * 31.0 * 62.0 + vidx[:, 1] * 62.0 +
                                 vidx[:, 2])
                flag1 = len(vidx) / 31.0 / 31.0 / 62.0 >= 0.02
            else:
1317
                # if `min_unique_num` is provided, directly compare with it
1318
                flag1 = mask.sum() >= self.min_unique_num
1319

1320
            # 2. selected patch should contain enough annotated points
1321
1322
1323
1324
1325
1326
1327
1328
1329
            if self.ignore_index is None:
                flag2 = True
            else:
                flag2 = np.sum(cur_sem_mask != self.ignore_index) / \
                               len(cur_sem_mask) >= 0.7

            if flag1 and flag2:
                break

1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
        # sample idx to `self.num_points`
        if point_idxs.size >= self.num_points:
            # no duplicate in sub-sampling
            choices = np.random.choice(
                point_idxs, self.num_points, replace=False)
        else:
            # do not use random choice here to avoid some points not counted
            dup = np.random.choice(point_idxs.size,
                                   self.num_points - point_idxs.size)
            idx_dup = np.concatenate(
                [np.arange(point_idxs.size),
                 np.array(dup)], 0)
            choices = point_idxs[idx_dup]
1343
1344
1345
1346
1347
1348
1349
1350

        # construct model input
        points = self._input_generation(coords[choices], cur_center, coord_max,
                                        attributes[choices], attribute_dims,
                                        point_type)

        return points, choices

ZCMax's avatar
ZCMax committed
1351
    def transform(self, input_dict: dict) -> dict:
1352
1353
1354
1355
1356
1357
        """Call function to sample points to in indoor scenes.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
1358
            dict: Results after sampling, 'points', 'pts_instance_mask'
1359
1360
                and 'pts_semantic_mask' keys are updated in the result dict.
        """
ZCMax's avatar
ZCMax committed
1361
        points = input_dict['points']
1362

ZCMax's avatar
ZCMax committed
1363
        assert 'pts_semantic_mask' in input_dict.keys(), \
1364
            'semantic mask should be provided in training and evaluation'
ZCMax's avatar
ZCMax committed
1365
        pts_semantic_mask = input_dict['pts_semantic_mask']
1366
1367
1368
1369

        points, choices = self._patch_points_sampling(points,
                                                      pts_semantic_mask)

ZCMax's avatar
ZCMax committed
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
        input_dict['points'] = points
        input_dict['pts_semantic_mask'] = pts_semantic_mask[choices]

        # 'eval_ann_info' will be passed to evaluator
        if 'eval_ann_info' in input_dict:
            input_dict['eval_ann_info']['pts_semantic_mask'] = \
                pts_semantic_mask[choices]

        pts_instance_mask = input_dict.get('pts_instance_mask', None)

1380
        if pts_instance_mask is not None:
ZCMax's avatar
ZCMax committed
1381
1382
1383
1384
1385
            input_dict['pts_instance_mask'] = pts_instance_mask[choices]
            # 'eval_ann_info' will be passed to evaluator
            if 'eval_ann_info' in input_dict:
                input_dict['eval_ann_info']['pts_instance_mask'] = \
                    pts_instance_mask[choices]
1386

ZCMax's avatar
ZCMax committed
1387
        return input_dict
1388
1389
1390
1391
1392
1393
1394
1395

    def __repr__(self):
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
        repr_str += f'(num_points={self.num_points},'
        repr_str += f' block_size={self.block_size},'
        repr_str += f' ignore_index={self.ignore_index},'
        repr_str += f' use_normalized_coord={self.use_normalized_coord},'
1396
1397
        repr_str += f' num_try={self.num_try},'
        repr_str += f' enlarge_size={self.enlarge_size},'
1398
1399
        repr_str += f' min_unique_num={self.min_unique_num},'
        repr_str += f' eps={self.eps})'
wuyuefeng's avatar
wuyuefeng committed
1400
        return repr_str
1401
1402


1403
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1404
class BackgroundPointsFilter(BaseTransform):
1405
1406
1407
1408
1409
1410
    """Filter background points near the bounding box.

    Args:
        bbox_enlarge_range (tuple[float], float): Bbox enlarge range.
    """

ZCMax's avatar
ZCMax committed
1411
    def __init__(self, bbox_enlarge_range: Union[Tuple[float], float]) -> None:
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
        assert (is_tuple_of(bbox_enlarge_range, float)
                and len(bbox_enlarge_range) == 3) \
            or isinstance(bbox_enlarge_range, float), \
            f'Invalid arguments bbox_enlarge_range {bbox_enlarge_range}'

        if isinstance(bbox_enlarge_range, float):
            bbox_enlarge_range = [bbox_enlarge_range] * 3
        self.bbox_enlarge_range = np.array(
            bbox_enlarge_range, dtype=np.float32)[np.newaxis, :]

ZCMax's avatar
ZCMax committed
1422
    def transform(self, input_dict: dict) -> dict:
1423
1424
1425
1426
1427
1428
        """Call function to filter points by the range.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
1429
            dict: Results after filtering, 'points', 'pts_instance_mask'
1430
                and 'pts_semantic_mask' keys are updated in the result dict.
1431
1432
1433
1434
        """
        points = input_dict['points']
        gt_bboxes_3d = input_dict['gt_bboxes_3d']

xiliu8006's avatar
xiliu8006 committed
1435
1436
1437
1438
        # avoid groundtruth being modified
        gt_bboxes_3d_np = gt_bboxes_3d.tensor.clone().numpy()
        gt_bboxes_3d_np[:, :3] = gt_bboxes_3d.gravity_center.clone().numpy()

1439
1440
        enlarged_gt_bboxes_3d = gt_bboxes_3d_np.copy()
        enlarged_gt_bboxes_3d[:, 3:6] += self.bbox_enlarge_range
xiliu8006's avatar
xiliu8006 committed
1441
        points_numpy = points.tensor.clone().numpy()
1442
1443
        foreground_masks = box_np_ops.points_in_rbbox(
            points_numpy, gt_bboxes_3d_np, origin=(0.5, 0.5, 0.5))
1444
        enlarge_foreground_masks = box_np_ops.points_in_rbbox(
1445
            points_numpy, enlarged_gt_bboxes_3d, origin=(0.5, 0.5, 0.5))
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
        foreground_masks = foreground_masks.max(1)
        enlarge_foreground_masks = enlarge_foreground_masks.max(1)
        valid_masks = ~np.logical_and(~foreground_masks,
                                      enlarge_foreground_masks)

        input_dict['points'] = points[valid_masks]
        pts_instance_mask = input_dict.get('pts_instance_mask', None)
        if pts_instance_mask is not None:
            input_dict['pts_instance_mask'] = pts_instance_mask[valid_masks]

        pts_semantic_mask = input_dict.get('pts_semantic_mask', None)
        if pts_semantic_mask is not None:
            input_dict['pts_semantic_mask'] = pts_semantic_mask[valid_masks]
        return input_dict

    def __repr__(self):
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
1464
        repr_str += f'(bbox_enlarge_range={self.bbox_enlarge_range.tolist()})'
1465
        return repr_str
1466
1467


1468
@TRANSFORMS.register_module()
1469
class VoxelBasedPointSampler(BaseTransform):
1470
1471
1472
1473
1474
1475
1476
    """Voxel based point sampler.

    Apply voxel sampling to multiple sweep points.

    Args:
        cur_sweep_cfg (dict): Config for sampling current points.
        prev_sweep_cfg (dict): Config for sampling previous points.
1477
        time_dim (int): Index that indicate the time dimension
1478
1479
1480
            for input points.
    """

1481
1482
1483
1484
    def __init__(self,
                 cur_sweep_cfg: dict,
                 prev_sweep_cfg: Optional[dict] = None,
                 time_dim: int = 3) -> None:
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
        self.cur_voxel_generator = VoxelGenerator(**cur_sweep_cfg)
        self.cur_voxel_num = self.cur_voxel_generator._max_voxels
        self.time_dim = time_dim
        if prev_sweep_cfg is not None:
            assert prev_sweep_cfg['max_num_points'] == \
                cur_sweep_cfg['max_num_points']
            self.prev_voxel_generator = VoxelGenerator(**prev_sweep_cfg)
            self.prev_voxel_num = self.prev_voxel_generator._max_voxels
        else:
            self.prev_voxel_generator = None
            self.prev_voxel_num = 0

1497
    def _sample_points(self, points: np.ndarray, sampler: VoxelGenerator,
1498
                       point_dim: int) -> np.ndarray:
1499
1500
1501
1502
1503
1504
        """Sample points for each points subset.

        Args:
            points (np.ndarray): Points subset to be sampled.
            sampler (VoxelGenerator): Voxel based sampler for
                each points subset.
1505
            point_dim (int): The dimension of each points
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523

        Returns:
            np.ndarray: Sampled points.
        """
        voxels, coors, num_points_per_voxel = sampler.generate(points)
        if voxels.shape[0] < sampler._max_voxels:
            padding_points = np.zeros([
                sampler._max_voxels - voxels.shape[0], sampler._max_num_points,
                point_dim
            ],
                                      dtype=points.dtype)
            padding_points[:] = voxels[0]
            sample_points = np.concatenate([voxels, padding_points], axis=0)
        else:
            sample_points = voxels

        return sample_points

1524
    def transform(self, results: dict) -> dict:
1525
1526
1527
1528
1529
1530
        """Call function to sample points from multiple sweeps.

        Args:
            input_dict (dict): Result dict from loading pipeline.

        Returns:
1531
            dict: Results after sampling, 'points', 'pts_instance_mask'
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
                and 'pts_semantic_mask' keys are updated in the result dict.
        """
        points = results['points']
        original_dim = points.shape[1]

        # TODO: process instance and semantic mask while _max_num_points
        # is larger than 1
        # Extend points with seg and mask fields
        map_fields2dim = []
        start_dim = original_dim
1542
1543
        points_numpy = points.tensor.numpy()
        extra_channel = [points_numpy]
1544
1545
1546
1547
1548
1549
1550
1551
1552
        for idx, key in enumerate(results['pts_mask_fields']):
            map_fields2dim.append((key, idx + start_dim))
            extra_channel.append(results[key][..., None])

        start_dim += len(results['pts_mask_fields'])
        for idx, key in enumerate(results['pts_seg_fields']):
            map_fields2dim.append((key, idx + start_dim))
            extra_channel.append(results[key][..., None])

1553
        points_numpy = np.concatenate(extra_channel, axis=-1)
1554
1555
1556
1557
1558

        # Split points into two part, current sweep points and
        # previous sweeps points.
        # TODO: support different sampling methods for next sweeps points
        # and previous sweeps points.
1559
1560
1561
        cur_points_flag = (points_numpy[:, self.time_dim] == 0)
        cur_sweep_points = points_numpy[cur_points_flag]
        prev_sweeps_points = points_numpy[~cur_points_flag]
1562
1563
1564
1565
1566
1567
1568
1569
1570
        if prev_sweeps_points.shape[0] == 0:
            prev_sweeps_points = cur_sweep_points

        # Shuffle points before sampling
        np.random.shuffle(cur_sweep_points)
        np.random.shuffle(prev_sweeps_points)

        cur_sweep_points = self._sample_points(cur_sweep_points,
                                               self.cur_voxel_generator,
1571
                                               points_numpy.shape[1])
1572
1573
1574
        if self.prev_voxel_generator is not None:
            prev_sweeps_points = self._sample_points(prev_sweeps_points,
                                                     self.prev_voxel_generator,
1575
                                                     points_numpy.shape[1])
1576

1577
1578
            points_numpy = np.concatenate(
                [cur_sweep_points, prev_sweeps_points], 0)
1579
        else:
1580
            points_numpy = cur_sweep_points
1581
1582

        if self.cur_voxel_generator._max_num_points == 1:
1583
1584
            points_numpy = points_numpy.squeeze(1)
        results['points'] = points.new_point(points_numpy[..., :original_dim])
1585

1586
        # Restore the corresponding seg and mask fields
1587
        for key, dim_index in map_fields2dim:
1588
            results[key] = points_numpy[..., dim_index]
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611

        return results

    def __repr__(self):
        """str: Return a string that describes the module."""

        def _auto_indent(repr_str, indent):
            repr_str = repr_str.split('\n')
            repr_str = [' ' * indent + t + '\n' for t in repr_str]
            repr_str = ''.join(repr_str)[:-1]
            return repr_str

        repr_str = self.__class__.__name__
        indent = 4
        repr_str += '(\n'
        repr_str += ' ' * indent + f'num_cur_sweep={self.cur_voxel_num},\n'
        repr_str += ' ' * indent + f'num_prev_sweep={self.prev_voxel_num},\n'
        repr_str += ' ' * indent + f'time_dim={self.time_dim},\n'
        repr_str += ' ' * indent + 'cur_voxel_generator=\n'
        repr_str += f'{_auto_indent(repr(self.cur_voxel_generator), 8)},\n'
        repr_str += ' ' * indent + 'prev_voxel_generator=\n'
        repr_str += f'{_auto_indent(repr(self.prev_voxel_generator), 8)})'
        return repr_str
1612
1613


1614
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1615
class AffineResize(BaseTransform):
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
    """Get the affine transform matrices to the target size.

    Different from :class:`RandomAffine` in MMDetection, this class can
    calculate the affine transform matrices while resizing the input image
    to a fixed size. The affine transform matrices include: 1) matrix
    transforming original image to the network input image size. 2) matrix
    transforming original image to the network output feature map size.

    Args:
        img_scale (tuple): Images scales for resizing.
        down_ratio (int): The down ratio of feature map.
            Actually the arg should be >= 1.
        bbox_clip_border (bool, optional): Whether clip the objects
            outside the border of the image. Defaults to True.
    """

ZCMax's avatar
ZCMax committed
1632
1633
1634
1635
    def __init__(self,
                 img_scale: Tuple,
                 down_ratio: int,
                 bbox_clip_border: bool = True) -> None:
1636
1637
1638
1639
1640

        self.img_scale = img_scale
        self.down_ratio = down_ratio
        self.bbox_clip_border = bbox_clip_border

ZCMax's avatar
ZCMax committed
1641
    def transform(self, results: dict) -> dict:
1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
1681
1682
1683
1684
1685
        """Call function to do affine transform to input image and labels.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: Results after affine resize, 'affine_aug', 'trans_mat'
                keys are added in the result dict.
        """
        # The results have gone through RandomShiftScale before AffineResize
        if 'center' not in results:
            img = results['img']
            height, width = img.shape[:2]
            center = np.array([width / 2, height / 2], dtype=np.float32)
            size = np.array([width, height], dtype=np.float32)
            results['affine_aug'] = False
        else:
            # The results did not go through RandomShiftScale before
            # AffineResize
            img = results['img']
            center = results['center']
            size = results['size']

        trans_affine = self._get_transform_matrix(center, size, self.img_scale)

        img = cv2.warpAffine(img, trans_affine[:2, :], self.img_scale)

        if isinstance(self.down_ratio, tuple):
            trans_mat = [
                self._get_transform_matrix(
                    center, size,
                    (self.img_scale[0] // ratio, self.img_scale[1] // ratio))
                for ratio in self.down_ratio
            ]  # (3, 3)
        else:
            trans_mat = self._get_transform_matrix(
                center, size, (self.img_scale[0] // self.down_ratio,
                               self.img_scale[1] // self.down_ratio))

        results['img'] = img
        results['img_shape'] = img.shape
        results['pad_shape'] = img.shape
        results['trans_mat'] = trans_mat

ZCMax's avatar
ZCMax committed
1686
1687
        if 'gt_bboxes' in results:
            self._affine_bboxes(results, trans_affine)
1688

ZCMax's avatar
ZCMax committed
1689
1690
        if 'centers_2d' in results:
            centers2d = self._affine_transform(results['centers_2d'],
1691
1692
1693
1694
1695
                                               trans_affine)
            valid_index = (centers2d[:, 0] >
                           0) & (centers2d[:, 0] <
                                 self.img_scale[0]) & (centers2d[:, 1] > 0) & (
                                     centers2d[:, 1] < self.img_scale[1])
ZCMax's avatar
ZCMax committed
1696
1697
1698
1699
            results['centers_2d'] = centers2d[valid_index]

            if 'gt_bboxes' in results:
                results['gt_bboxes'] = results['gt_bboxes'][valid_index]
1700
1701
1702
                if 'gt_bboxes_labels' in results:
                    results['gt_bboxes_labels'] = results['gt_bboxes_labels'][
                        valid_index]
ZCMax's avatar
ZCMax committed
1703
1704
1705
1706
1707
1708
1709
1710
1711
1712
                if 'gt_masks' in results:
                    raise NotImplementedError(
                        'AffineResize only supports bbox.')

            if 'gt_bboxes_3d' in results:
                results['gt_bboxes_3d'].tensor = results[
                    'gt_bboxes_3d'].tensor[valid_index]
                if 'gt_labels_3d' in results:
                    results['gt_labels_3d'] = results['gt_labels_3d'][
                        valid_index]
1713
1714
1715
1716
1717

            results['depths'] = results['depths'][valid_index]

        return results

ZCMax's avatar
ZCMax committed
1718
    def _affine_bboxes(self, results: dict, matrix: np.ndarray) -> None:
1719
1720
1721
1722
1723
1724
1725
1726
1727
        """Affine transform bboxes to input image.

        Args:
            results (dict): Result dict from loading pipeline.
            matrix (np.ndarray): Matrix transforming original
                image to the network input image size.
                shape: (3, 3)
        """

ZCMax's avatar
ZCMax committed
1728
1729
1730
1731
1732
1733
1734
1735
1736
1737
1738
1739
        bboxes = results['gt_bboxes']
        bboxes[:, :2] = self._affine_transform(bboxes[:, :2], matrix)
        bboxes[:, 2:] = self._affine_transform(bboxes[:, 2:], matrix)
        if self.bbox_clip_border:
            bboxes[:, [0, 2]] = bboxes[:, [0, 2]].clip(0,
                                                       self.img_scale[0] - 1)
            bboxes[:, [1, 3]] = bboxes[:, [1, 3]].clip(0,
                                                       self.img_scale[1] - 1)
        results['gt_bboxes'] = bboxes

    def _affine_transform(self, points: np.ndarray,
                          matrix: np.ndarray) -> np.ndarray:
1740
        """Affine transform bbox points to input image.
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757

        Args:
            points (np.ndarray): Points to be transformed.
                shape: (N, 2)
            matrix (np.ndarray): Affine transform matrix.
                shape: (3, 3)

        Returns:
            np.ndarray: Transformed points.
        """
        num_points = points.shape[0]
        hom_points_2d = np.concatenate((points, np.ones((num_points, 1))),
                                       axis=1)
        hom_points_2d = hom_points_2d.T
        affined_points = np.matmul(matrix, hom_points_2d).T
        return affined_points[:, :2]

ZCMax's avatar
ZCMax committed
1758
1759
    def _get_transform_matrix(self, center: Tuple, scale: Tuple,
                              output_scale: Tuple[float]) -> np.ndarray:
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
        """Get affine transform matrix.

        Args:
            center (tuple): Center of current image.
            scale (tuple): Scale of current image.
            output_scale (tuple[float]): The transform target image scales.

        Returns:
            np.ndarray: Affine transform matrix.
        """
        # TODO: further add rot and shift here.
        src_w = scale[0]
        dst_w = output_scale[0]
        dst_h = output_scale[1]

        src_dir = np.array([0, src_w * -0.5])
        dst_dir = np.array([0, dst_w * -0.5])

        src = np.zeros((3, 2), dtype=np.float32)
        dst = np.zeros((3, 2), dtype=np.float32)
        src[0, :] = center
        src[1, :] = center + src_dir
        dst[0, :] = np.array([dst_w * 0.5, dst_h * 0.5])
        dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir

        src[2, :] = self._get_ref_point(src[0, :], src[1, :])
        dst[2, :] = self._get_ref_point(dst[0, :], dst[1, :])

        get_matrix = cv2.getAffineTransform(src, dst)

        matrix = np.concatenate((get_matrix, [[0., 0., 1.]]))

        return matrix.astype(np.float32)

ZCMax's avatar
ZCMax committed
1794
1795
    def _get_ref_point(self, ref_point1: np.ndarray,
                       ref_point2: np.ndarray) -> np.ndarray:
1796
        """Get reference point to calculate affine transform matrix.
1797
1798

        While using opencv to calculate the affine matrix, we need at least
1799
        three corresponding points separately on original image and target
1800
1801
1802
1803
1804
1805
1806
        image. Here we use two points to get the the third reference point.
        """
        d = ref_point1 - ref_point2
        ref_point3 = ref_point2 + np.array([-d[1], d[0]])
        return ref_point3

    def __repr__(self):
1807
        """str: Return a string that describes the module."""
1808
1809
1810
1811
1812
1813
        repr_str = self.__class__.__name__
        repr_str += f'(img_scale={self.img_scale}, '
        repr_str += f'down_ratio={self.down_ratio}) '
        return repr_str


1814
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1815
class RandomShiftScale(BaseTransform):
1816
1817
1818
1819
    """Random shift scale.

    Different from the normal shift and scale function, it doesn't
    directly shift or scale image. It can record the shift and scale
1820
    infos into loading TRANSFORMS. It's designed to be used with
1821
1822
1823
1824
1825
1826
1827
    AffineResize together.

    Args:
        shift_scale (tuple[float]): Shift and scale range.
        aug_prob (float): The shifting and scaling probability.
    """

1828
    def __init__(self, shift_scale: Tuple[float], aug_prob: float) -> None:
1829
1830
1831
1832

        self.shift_scale = shift_scale
        self.aug_prob = aug_prob

ZCMax's avatar
ZCMax committed
1833
    def transform(self, results: dict) -> dict:
1834
1835
1836
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
        """Call function to record random shift and scale infos.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: Results after random shift and scale, 'center', 'size'
                and 'affine_aug' keys are added in the result dict.
        """
        img = results['img']

        height, width = img.shape[:2]

        center = np.array([width / 2, height / 2], dtype=np.float32)
        size = np.array([width, height], dtype=np.float32)

        if random.random() < self.aug_prob:
            shift, scale = self.shift_scale[0], self.shift_scale[1]
            shift_ranges = np.arange(-shift, shift + 0.1, 0.1)
            center[0] += size[0] * random.choice(shift_ranges)
            center[1] += size[1] * random.choice(shift_ranges)
            scale_ranges = np.arange(1 - scale, 1 + scale + 0.1, 0.1)
            size *= random.choice(scale_ranges)
            results['affine_aug'] = True
        else:
            results['affine_aug'] = False

        results['center'] = center
        results['size'] = size

        return results

    def __repr__(self):
1867
        """str: Return a string that describes the module."""
1868
1869
1870
1871
        repr_str = self.__class__.__name__
        repr_str += f'(shift_scale={self.shift_scale}, '
        repr_str += f'aug_prob={self.aug_prob}) '
        return repr_str
1872
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
1920


@TRANSFORMS.register_module()
class Resize3D(Resize):

    def _resize_3d(self, results):
        """Resize centers_2d and modify camera intrinisc with
        ``results['scale']``."""
        if 'centers_2d' in results:
            results['centers_2d'] *= results['scale_factor'][:2]
        results['cam2img'][0] *= np.array(results['scale_factor'][0])
        results['cam2img'][1] *= np.array(results['scale_factor'][1])

    def transform(self, results: dict) -> dict:
        """Transform function to resize images, bounding boxes, semantic
        segmentation map and keypoints.

        Args:
            results (dict): Result dict from loading pipeline.
        Returns:
            dict: Resized results, 'img', 'gt_bboxes', 'gt_seg_map',
            'gt_keypoints', 'scale', 'scale_factor', 'img_shape',
            and 'keep_ratio' keys are updated in result dict.
        """

        super(Resize3D, self).transform(results)
        self._resize_3d(results)
        return results


@TRANSFORMS.register_module()
class RandomResize3D(RandomResize):
    """The difference between RandomResize3D and RandomResize:

    1. Compared to RandomResize, this class would further
        check if scale is already set in results.
    2. During resizing, this class would modify the centers_2d
        and cam2img with ``results['scale']``.
    """

    def _resize_3d(self, results):
        """Resize centers_2d and modify camera intrinisc with
        ``results['scale']``."""
        if 'centers_2d' in results:
            results['centers_2d'] *= results['scale_factor'][:2]
        results['cam2img'][0] *= np.array(results['scale_factor'][0])
        results['cam2img'][1] *= np.array(results['scale_factor'][1])

    def transform(self, results):
1921
1922
        """Transform function to resize images, bounding boxes, masks, semantic
        segmentation map. Compared to RandomResize, this function would further
1923
1924
1925
1926
        check if scale is already set in results.

        Args:
            results (dict): Result dict from loading pipeline.
1927

1928
1929
1930
1931
1932
1933
1934
1935
1936
1937
1938
        Returns:
            dict: Resized results, 'img_shape', 'pad_shape', 'scale_factor', \
                'keep_ratio' keys are added into result dict.
        """
        if 'scale' not in results:
            results['scale'] = self._random_scale()
        self.resize.scale = results['scale']
        results = self.resize(results)
        self._resize_3d(results)

        return results
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987
1988
1989
1990
1991
1992
1993
1994
1995
1996
1997
1998
1999
2000
2001
2002
2003
2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
2037
2038
2039
2040
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247


@TRANSFORMS.register_module()
class RandomCrop3D(RandomCrop):
    """3D version of RandomCrop. RamdomCrop3D supports the modifications of
    camera intrinsic matrix and using predefined randomness variable to do the
    augmentation.

    The absolute ``crop_size`` is sampled based on ``crop_type`` and
    ``image_size``, then the cropped results are generated.

    Required Keys:

    - img
    - gt_bboxes (np.float32) (optional)
    - gt_bboxes_labels (np.int64) (optional)
    - gt_masks (BitmapMasks | PolygonMasks) (optional)
    - gt_ignore_flags (np.bool) (optional)
    - gt_seg_map (np.uint8) (optional)

    Modified Keys:

    - img
    - img_shape
    - gt_bboxes (optional)
    - gt_bboxes_labels (optional)
    - gt_masks (optional)
    - gt_ignore_flags (optional)
    - gt_seg_map (optional)

    Added Keys:

    - homography_matrix

    Args:
        crop_size (tuple): The relative ratio or absolute pixels of
            height and width.
        crop_type (str): One of "relative_range", "relative",
            "absolute", "absolute_range". "relative" randomly crops
            (h * crop_size[0], w * crop_size[1]) part from an input of size
            (h, w). "relative_range" uniformly samples relative crop size from
            range [crop_size[0], 1] and [crop_size[1], 1] for height and width
            respectively. "absolute" crops from an input with absolute size
            (crop_size[0], crop_size[1]). "absolute_range" uniformly samples
            crop_h in range [crop_size[0], min(h, crop_size[1])] and crop_w
            in range [crop_size[0], min(w, crop_size[1])].
            Defaults to "absolute".
        allow_negative_crop (bool): Whether to allow a crop that does
            not contain any bbox area. Defaults to False.
        recompute_bbox (bool): Whether to re-compute the boxes based
            on cropped instance masks. Defaults to False.
        bbox_clip_border (bool): Whether clip the objects outside
            the border of the image. Defaults to True.
        rel_offset_h (tuple): The cropping interval of image height. Default
            to (0., 1.).
        rel_offset_w (tuple): The cropping interval of image width. Default
            to (0., 1.).

    Note:
        - If the image is smaller than the absolute crop size, return the
            original image.
        - The keys for bboxes, labels and masks must be aligned. That is,
          ``gt_bboxes`` corresponds to ``gt_labels`` and ``gt_masks``, and
          ``gt_bboxes_ignore`` corresponds to ``gt_labels_ignore`` and
          ``gt_masks_ignore``.
        - If the crop does not contain any gt-bbox region and
          ``allow_negative_crop`` is set to False, skip this image.
    """

    def __init__(self,
                 crop_size,
                 crop_type='absolute',
                 allow_negative_crop=False,
                 recompute_bbox=False,
                 bbox_clip_border=True,
                 rel_offset_h=(0., 1.),
                 rel_offset_w=(0., 1.)):
        super().__init__(
            crop_size=crop_size,
            crop_type=crop_type,
            allow_negative_crop=allow_negative_crop,
            recompute_bbox=recompute_bbox,
            bbox_clip_border=bbox_clip_border)
        # rel_offset specifies the relative offset range of cropping origin
        # [0., 1.] means starting from 0*margin to 1*margin + 1
        self.rel_offset_h = rel_offset_h
        self.rel_offset_w = rel_offset_w

    def _crop_data(self, results, crop_size, allow_negative_crop):
        """Function to randomly crop images, bounding boxes, masks, semantic
        segmentation maps.

        Args:
            results (dict): Result dict from loading pipeline.
            crop_size (tuple): Expected absolute size after cropping, (h, w).
            allow_negative_crop (bool): Whether to allow a crop that does not
                contain any bbox area. Default to False.

        Returns:
            dict: Randomly cropped results, 'img_shape' key in result dict is
                updated according to crop size.
        """
        assert crop_size[0] > 0 and crop_size[1] > 0
        for key in results.get('img_fields', ['img']):
            img = results[key]
            if 'img_crop_offset' not in results:
                margin_h = max(img.shape[0] - crop_size[0], 0)
                margin_w = max(img.shape[1] - crop_size[1], 0)
                # TOCHECK: a little different from LIGA implementation
                offset_h = np.random.randint(
                    self.rel_offset_h[0] * margin_h,
                    self.rel_offset_h[1] * margin_h + 1)
                offset_w = np.random.randint(
                    self.rel_offset_w[0] * margin_w,
                    self.rel_offset_w[1] * margin_w + 1)
            else:
                offset_w, offset_h = results['img_crop_offset']

            crop_h = min(crop_size[0], img.shape[0])
            crop_w = min(crop_size[1], img.shape[1])
            crop_y1, crop_y2 = offset_h, offset_h + crop_h
            crop_x1, crop_x2 = offset_w, offset_w + crop_w

            # crop the image
            img = img[crop_y1:crop_y2, crop_x1:crop_x2, ...]
            img_shape = img.shape
            results[key] = img
        results['img_shape'] = img_shape

        # crop bboxes accordingly and clip to the image boundary
        for key in results.get('bbox_fields', []):
            # e.g. gt_bboxes and gt_bboxes_ignore
            bbox_offset = np.array([offset_w, offset_h, offset_w, offset_h],
                                   dtype=np.float32)
            bboxes = results[key] - bbox_offset
            if self.bbox_clip_border:
                bboxes[:, 0::2] = np.clip(bboxes[:, 0::2], 0, img_shape[1])
                bboxes[:, 1::2] = np.clip(bboxes[:, 1::2], 0, img_shape[0])
            valid_inds = (bboxes[:, 2] > bboxes[:, 0]) & (
                bboxes[:, 3] > bboxes[:, 1])
            # If the crop does not contain any gt-bbox area and
            # allow_negative_crop is False, skip this image.
            if (key == 'gt_bboxes' and not valid_inds.any()
                    and not allow_negative_crop):
                return None
            results[key] = bboxes[valid_inds, :]
            # label fields. e.g. gt_labels and gt_labels_ignore
            label_key = self.bbox2label.get(key)
            if label_key in results:
                results[label_key] = results[label_key][valid_inds]

            # mask fields, e.g. gt_masks and gt_masks_ignore
            mask_key = self.bbox2mask.get(key)
            if mask_key in results:
                results[mask_key] = results[mask_key][
                    valid_inds.nonzero()[0]].crop(
                        np.asarray([crop_x1, crop_y1, crop_x2, crop_y2]))
                if self.recompute_bbox:
                    results[key] = results[mask_key].get_bboxes()

        # crop semantic seg
        for key in results.get('seg_fields', []):
            results[key] = results[key][crop_y1:crop_y2, crop_x1:crop_x2]

        # manipulate camera intrinsic matrix
        # needs to apply offset to K instead of P2 (on KITTI)
        if isinstance(results['cam2img'], list):
            # TODO ignore this, but should handle it in the future
            pass
        else:
            K = results['cam2img'][:3, :3].copy()
            inv_K = np.linalg.inv(K)
            T = np.matmul(inv_K, results['cam2img'][:3])
            K[0, 2] -= crop_x1
            K[1, 2] -= crop_y1
            offset_cam2img = np.matmul(K, T)
            results['cam2img'][:offset_cam2img.shape[0], :offset_cam2img.
                               shape[1]] = offset_cam2img

        results['img_crop_offset'] = [offset_w, offset_h]

        return results

    def transform(self, results):
        """Transform function to randomly crop images, bounding boxes, masks,
        semantic segmentation maps.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: Randomly cropped results, 'img_shape' key in result dict is
                updated according to crop size.
        """
        image_size = results['img'].shape[:2]
        if 'crop_size' not in results:
            crop_size = self._get_crop_size(image_size)
            results['crop_size'] = crop_size
        else:
            crop_size = results['crop_size']
        results = self._crop_data(results, crop_size, self.allow_negative_crop)
        return results

    def __repr__(self):
        repr_str = self.__class__.__name__
        repr_str += f'(crop_size={self.crop_size}, '
        repr_str += f'crop_type={self.crop_type}, '
        repr_str += f'allow_negative_crop={self.allow_negative_crop}, '
        repr_str += f'bbox_clip_border={self.bbox_clip_border}), '
        repr_str += f'rel_offset_h={self.rel_offset_h}), '
        repr_str += f'rel_offset_w={self.rel_offset_w})'
        return repr_str


@TRANSFORMS.register_module()
class PhotoMetricDistortion3D(PhotoMetricDistortion):
    """Apply photometric distortion to image sequentially, every transformation
    is applied with a probability of 0.5. The position of random contrast is in
    second or second to last.

    PhotoMetricDistortion3D further support using predefined randomness
    variable to do the augmentation.

    1. random brightness
    2. random contrast (mode 0)
    3. convert color from BGR to HSV
    4. random saturation
    5. random hue
    6. convert color from HSV to BGR
    7. random contrast (mode 1)
    8. randomly swap channels

    Required Keys:

    - img (np.uint8)

    Modified Keys:

    - img (np.float32)

    Args:
        brightness_delta (int): delta of brightness.
        contrast_range (sequence): range of contrast.
        saturation_range (sequence): range of saturation.
        hue_delta (int): delta of hue.
    """

    def transform(self, results: dict) -> dict:
        """Transform function to perform photometric distortion on images.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: Result dict with images distorted.
        """
        assert 'img' in results, '`img` is not found in results'
        img = results['img']
        img = img.astype(np.float32)
        if 'photometric_param' not in results:
            photometric_param = self._random_flags()
            results['photometric_param'] = photometric_param
        else:
            photometric_param = results['photometric_param']

        (mode, brightness_flag, contrast_flag, saturation_flag, hue_flag,
         swap_flag, delta_value, alpha_value, saturation_value, hue_value,
         swap_value) = photometric_param

        # random brightness
        if brightness_flag:
            img += delta_value

        # mode == 0 --> do random contrast first
        # mode == 1 --> do random contrast last
        if mode == 1:
            if contrast_flag:
                img *= alpha_value

        # convert color from BGR to HSV
        img = mmcv.bgr2hsv(img)

        # random saturation
        if saturation_flag:
            img[..., 1] *= saturation_value

        # random hue
        if hue_flag:
            img[..., 0] += hue_value
            img[..., 0][img[..., 0] > 360] -= 360
            img[..., 0][img[..., 0] < 0] += 360

        # convert color from HSV to BGR
        img = mmcv.hsv2bgr(img)

        # random contrast
        if mode == 0:
            if contrast_flag:
                img *= alpha_value

        # randomly swap channels
        if swap_flag:
            img = img[..., swap_value]

        results['img'] = img
        return results


@TRANSFORMS.register_module()
2248
class MultiViewWrapper(BaseTransform):
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
2277
2278
2279
2280
2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
    """Wrap transformation from single-view into multi-view.

    The wrapper processes the images from multi-view one by one. For each
    image, it constructs a pseudo dict according to the keys specified by the
    'process_fields' parameter. After the transformation is finished, desired
    information can be collected by specifying the keys in the 'collected_keys'
    parameter. Multi-view images share the same transformation parameters
    but do not share the same magnitude when a random transformation is
    conducted.

    Args:
        transforms (list[dict]): A list of dict specifying the transformations
            for the monocular situation.
        override_aug_config (bool): flag of whether to use the same aug config
            for multiview image. Default to True.
        process_fields (list): Desired keys that the transformations should
            be conducted on. Default to ['img', 'cam2img', 'lidar2cam'],

        collected_keys (list): Collect information in transformation
            like rotate angles, crop roi, and flip state. Default to
                ['scale', 'scale_factor', 'crop',
                 'crop_offset', 'ori_shape',
                 'pad_shape', 'img_shape',
                 'pad_fixed_size', 'pad_size_divisor',
                 'flip', 'flip_direction', 'rotate'],
        randomness_keys (list): The keys that related to the randomness
            in transformation Default to
                    ['scale', 'scale_factor', 'crop_size', 'flip',
                     'flip_direction', 'photometric_param']
    """

    def __init__(self,
                 transforms: dict,
                 override_aug_config: bool = True,
                 process_fields: list = ['img', 'cam2img', 'lidar2cam'],
                 collected_keys: list = [
                     'scale', 'scale_factor', 'crop', 'img_crop_offset',
                     'ori_shape', 'pad_shape', 'img_shape', 'pad_fixed_size',
                     'pad_size_divisor', 'flip', 'flip_direction', 'rotate'
                 ],
                 randomness_keys: list = [
                     'scale', 'scale_factor', 'crop_size', 'img_crop_offset',
                     'flip', 'flip_direction', 'photometric_param'
                 ]):
2293
        self.transforms = Compose(transforms)
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
        self.override_aug_config = override_aug_config
        self.collected_keys = collected_keys
        self.process_fields = process_fields
        self.randomness_keys = randomness_keys

    def transform(self, input_dict):
        """Transform function to do the transform for multiview image.

        Args:
            results (dict): Result dict from loading pipeline.

        Returns:
            dict: output dict after transformtaion
        """
        # store the augmentation related keys for each image.
        for key in self.collected_keys:
            if key not in input_dict or \
                    not isinstance(input_dict[key], list):
                input_dict[key] = []
        prev_process_dict = {}
        for img_id in range(len(input_dict['img'])):
            process_dict = {}

            # override the process dict (e.g. scale in random scale,
            # crop_size in random crop, flip, flip_direction in
            # random flip)
            if img_id != 0 and self.override_aug_config:
                for key in self.randomness_keys:
                    if key in prev_process_dict:
                        process_dict[key] = prev_process_dict[key]

            for key in self.process_fields:
                if key in input_dict:
                    process_dict[key] = input_dict[key][img_id]
2328
            process_dict = self.transforms(process_dict)
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
            # store the randomness variable in transformation.
            prev_process_dict = process_dict

            # store the related results to results_dict
            for key in self.process_fields:
                if key in process_dict:
                    input_dict[key][img_id] = process_dict[key]
            # update the keys
            for key in self.collected_keys:
                if key in process_dict:
                    if len(input_dict[key]) == img_id + 1:
                        input_dict[key][img_id] = process_dict[key]
                    else:
                        input_dict[key].append(process_dict[key])

        for key in self.collected_keys:
            if len(input_dict[key]) == 0:
                input_dict.pop(key)
        return input_dict