transforms_3d.py 100 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, Sequence, Tuple, Union
5
6

import cv2
7
import mmcv
8
import numpy as np
9
import torch
10
from mmcv.transforms import BaseTransform, Compose, RandomResize, Resize
11
12
from mmdet.datasets.transforms import (PhotoMetricDistortion, RandomCrop,
                                       RandomFlip)
13
from mmengine import is_list_of, is_tuple_of
zhangwenwei's avatar
zhangwenwei committed
14

zhangshilong's avatar
zhangshilong committed
15
from mmdet3d.models.task_modules import VoxelGenerator
16
from mmdet3d.registry import TRANSFORMS
zhangshilong's avatar
zhangshilong committed
17
18
19
20
from mmdet3d.structures import (CameraInstance3DBoxes, DepthInstance3DBoxes,
                                LiDARInstance3DBoxes)
from mmdet3d.structures.ops import box_np_ops
from mmdet3d.structures.points import BasePoints
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): 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
50
            dict: Results after color dropping, 'points' key is updated
            in the result dict.
51
52
53
54
55
56
        """
        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
            points.color = points.color * 0.0
        return input_dict

67
    def __repr__(self) -> str:
68
69
70
71
72
73
        """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
        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
111
112
            to False, since the flip of 2D images does not influence the 3D
            box. Defaults 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
            direction (str): Flip direction. Defaults to 'horizontal'.
154
155

        Returns:
156
            dict: Flipped results, 'points', 'bbox3d_fields' keys are
157
            updated in the result dict.
158
        """
wuyuefeng's avatar
wuyuefeng committed
159
        assert direction in ['horizontal', 'vertical']
160
161
162
163
164
165
166
167
        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)
168
            else:
169
                input_dict['points'].flip(direction)
jshilong's avatar
jshilong committed
170
171

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

184
185
186
187
188
189
190
191
192
193
    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:
194
195
196
197
198
199
200
201
202
            # `flip_direction` works only when `flip` is True.
            # For example, in `MultiScaleFlipAug3D`, `flip_direction` is
            # 'horizontal' but `flip` is False.
            if results['flip']:
                assert 'flip_direction' in results, 'flip and flip_direction '
                'must exist simultaneously'
                cur_dir = results['flip_direction']
            else:
                cur_dir = None
203
204
205
206
207
208
209
210
        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
211
    def transform(self, input_dict: dict) -> dict:
212
        """Call function to flip points, values in the ``bbox3d_fields`` and
213
214
215
216
217
218
        also flip 2D image and its annotations.

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

        Returns:
219
            dict: Flipped results, 'flip', 'flip_direction',
220
221
            'pcd_horizontal_flip' and 'pcd_vertical_flip' keys are added
            into result dict.
222
        """
223
        # flip 2D image and its annotations
jshilong's avatar
jshilong committed
224
225
        if 'img' in input_dict:
            super(RandomFlip3D, self).transform(input_dict)
zhangwenwei's avatar
zhangwenwei committed
226

jshilong's avatar
jshilong committed
227
        if self.sync_2d and 'img' in input_dict:
wuyuefeng's avatar
wuyuefeng committed
228
229
            input_dict['pcd_horizontal_flip'] = input_dict['flip']
            input_dict['pcd_vertical_flip'] = False
zhangwenwei's avatar
zhangwenwei committed
230
        else:
wuyuefeng's avatar
wuyuefeng committed
231
232
            if 'pcd_horizontal_flip' not in input_dict:
                flip_horizontal = True if np.random.rand(
jshilong's avatar
jshilong committed
233
                ) < self.flip_ratio_bev_horizontal else False
wuyuefeng's avatar
wuyuefeng committed
234
235
236
237
238
239
                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

240
241
242
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

wuyuefeng's avatar
wuyuefeng committed
243
244
        if input_dict['pcd_horizontal_flip']:
            self.random_flip_data_3d(input_dict, 'horizontal')
245
            input_dict['transformation_3d_flow'].extend(['HF'])
wuyuefeng's avatar
wuyuefeng committed
246
247
        if input_dict['pcd_vertical_flip']:
            self.random_flip_data_3d(input_dict, 'vertical')
248
            input_dict['transformation_3d_flow'].extend(['VF'])
zhangwenwei's avatar
zhangwenwei committed
249
250
        return input_dict

251
    def __repr__(self) -> str:
252
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
253
        repr_str = self.__class__.__name__
254
        repr_str += f'(sync_2d={self.sync_2d},'
255
        repr_str += f' flip_ratio_bev_vertical={self.flip_ratio_bev_vertical})'
wuyuefeng's avatar
wuyuefeng committed
256
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
257

zhangwenwei's avatar
zhangwenwei committed
258

259
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
260
class RandomJitterPoints(BaseTransform):
261
262
    """Randomly jitter point coordinates.

263
    Different from the global translation in ``GlobalRotScaleTrans``, here we
264
    apply different noises to each point in a scene.
265
266
267

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

    Note:
276
        This transform should only be used in point cloud segmentation tasks
277
        because we don't transform ground-truth bboxes accordingly.
278
279
280
281
        For similar transform in detection task, please refer to `ObjectNoise`.
    """

    def __init__(self,
ZCMax's avatar
ZCMax committed
282
283
                 jitter_std: List[float] = [0.01, 0.01, 0.01],
                 clip_range: List[float] = [-0.05, 0.05]) -> None:
284
285
286
287
288
289
290
291
292
293
294
295
296
297
        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
298
    def transform(self, input_dict: dict) -> dict:
299
300
301
302
303
304
        """Call function to jitter all the points in the scene.

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

        Returns:
305
            dict: Results after adding noise to each point,
306
            'points' key is updated in the result dict.
307
308
309
310
311
312
313
314
315
316
317
318
        """
        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

319
    def __repr__(self) -> str:
320
321
322
323
324
325
326
        """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


327
328
@TRANSFORMS.register_module()
class ObjectSample(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
329
    """Sample GT objects to the data.
zhangwenwei's avatar
zhangwenwei committed
330

331
332
333
334
335
336
337
338
339
340
    Required Keys:

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

    Modified Keys:
341

342
343
344
345
346
347
348
349
350
351
    - points
    - gt_bboxes_3d
    - gt_labels_3d
    - img (optional)
    - gt_bboxes (optional)

    Added Keys:

    - plane (optional)

zhangwenwei's avatar
zhangwenwei committed
352
353
    Args:
        db_sampler (dict): Config dict of the database sampler.
354
        sample_2d (bool): Whether to also paste 2D image patch to the images.
zhangwenwei's avatar
zhangwenwei committed
355
            This should be true when applying multi-modality cut-and-paste.
liyinhao's avatar
liyinhao committed
356
            Defaults to False.
357
        use_ground_plane (bool): Whether to use ground plane to adjust the
358
            3D labels. Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
359
    """
zhangwenwei's avatar
zhangwenwei committed
360

361
362
363
    def __init__(self,
                 db_sampler: dict,
                 sample_2d: bool = False,
364
                 use_ground_plane: bool = False) -> None:
zhangwenwei's avatar
zhangwenwei committed
365
366
367
368
        self.sampler_cfg = db_sampler
        self.sample_2d = sample_2d
        if 'type' not in db_sampler.keys():
            db_sampler['type'] = 'DataBaseSampler'
369
        self.db_sampler = TRANSFORMS.build(db_sampler)
370
        self.use_ground_plane = use_ground_plane
371
        self.disabled = False
zhangwenwei's avatar
zhangwenwei committed
372
373

    @staticmethod
374
375
    def remove_points_in_boxes(points: BasePoints,
                               boxes: np.ndarray) -> np.ndarray:
376
377
378
        """Remove the points in the sampled bounding boxes.

        Args:
379
            points (:obj:`BasePoints`): Input point cloud array.
380
381
382
383
384
            boxes (np.ndarray): Sampled ground truth boxes.

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

389
390
    def transform(self, input_dict: dict) -> dict:
        """Transform function to sample ground truth objects to the data.
391
392
393
394
395

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

        Returns:
396
            dict: Results after object sampling augmentation,
397
398
            'points', 'gt_bboxes_3d', 'gt_labels_3d' keys are updated
            in the result dict.
399
        """
400
401
402
        if self.disabled:
            return input_dict

zhangwenwei's avatar
zhangwenwei committed
403
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
404
405
        gt_labels_3d = input_dict['gt_labels_3d']

ChaimZhu's avatar
ChaimZhu committed
406
407
        if self.use_ground_plane:
            ground_plane = input_dict.get('plane', None)
408
409
            assert ground_plane is not None, '`use_ground_plane` is True ' \
                                             'but find plane is None'
410
411
        else:
            ground_plane = None
zhangwenwei's avatar
zhangwenwei committed
412
413
414
        # change to float for blending operation
        points = input_dict['points']
        if self.sample_2d:
wuyuefeng's avatar
wuyuefeng committed
415
            img = input_dict['img']
zhangwenwei's avatar
zhangwenwei committed
416
417
418
            gt_bboxes_2d = input_dict['gt_bboxes']
            # Assume for now 3D & 2D bboxes are the same
            sampled_dict = self.db_sampler.sample_all(
419
                gt_bboxes_3d.numpy(),
420
421
422
                gt_labels_3d,
                gt_bboxes_2d=gt_bboxes_2d,
                img=img)
zhangwenwei's avatar
zhangwenwei committed
423
424
        else:
            sampled_dict = self.db_sampler.sample_all(
425
                gt_bboxes_3d.numpy(),
426
427
428
                gt_labels_3d,
                img=None,
                ground_plane=ground_plane)
zhangwenwei's avatar
zhangwenwei committed
429
430
431
432

        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
433
            sampled_gt_labels = sampled_dict['gt_labels_3d']
zhangwenwei's avatar
zhangwenwei committed
434

zhangwenwei's avatar
zhangwenwei committed
435
436
            gt_labels_3d = np.concatenate([gt_labels_3d, sampled_gt_labels],
                                          axis=0)
437
            gt_bboxes_3d = gt_bboxes_3d.new_box(
438
                np.concatenate([gt_bboxes_3d.numpy(), sampled_gt_bboxes_3d]))
zhangwenwei's avatar
zhangwenwei committed
439

zhangwenwei's avatar
zhangwenwei committed
440
441
            points = self.remove_points_in_boxes(points, sampled_gt_bboxes_3d)
            # check the points dimension
442
            points = points.cat([sampled_points, points])
zhangwenwei's avatar
zhangwenwei committed
443
444
445
446
447

            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
448

zhangwenwei's avatar
zhangwenwei committed
449
                input_dict['gt_bboxes'] = gt_bboxes_2d
wuyuefeng's avatar
wuyuefeng committed
450
                input_dict['img'] = sampled_dict['img']
zhangwenwei's avatar
zhangwenwei committed
451
452

        input_dict['gt_bboxes_3d'] = gt_bboxes_3d
WRH's avatar
WRH committed
453
        input_dict['gt_labels_3d'] = gt_labels_3d.astype(np.int64)
zhangwenwei's avatar
zhangwenwei committed
454
        input_dict['points'] = points
zhangwenwei's avatar
zhangwenwei committed
455

zhangwenwei's avatar
zhangwenwei committed
456
457
        return input_dict

458
    def __repr__(self) -> str:
459
        """str: Return a string that describes the module."""
460
        repr_str = self.__class__.__name__
461
        repr_str += f'(db_sampler={self.db_sampler},'
462
        repr_str += f' sample_2d={self.sample_2d},'
463
        repr_str += f' use_ground_plane={self.use_ground_plane})'
464
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
465
466


467
468
@TRANSFORMS.register_module()
class ObjectNoise(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
469
    """Apply noise to each GT objects in the scene.
zhangwenwei's avatar
zhangwenwei committed
470

471
472
473
474
475
476
477
478
479
480
    Required Keys:

    - points
    - gt_bboxes_3d

    Modified Keys:

    - points
    - gt_bboxes_3d

zhangwenwei's avatar
zhangwenwei committed
481
    Args:
482
        translation_std (list[float]): Standard deviation of the
zhangwenwei's avatar
zhangwenwei committed
483
484
            distribution where translation noise are sampled from.
            Defaults to [0.25, 0.25, 0.25].
485
        global_rot_range (list[float]): Global rotation to the scene.
zhangwenwei's avatar
zhangwenwei committed
486
            Defaults to [0.0, 0.0].
487
        rot_range (list[float]): Object rotation range.
zhangwenwei's avatar
zhangwenwei committed
488
            Defaults to [-0.15707963267, 0.15707963267].
489
490
        num_try (int): Number of times to try if the noise applied is invalid.
            Defaults to 100.
zhangwenwei's avatar
zhangwenwei committed
491
    """
zhangwenwei's avatar
zhangwenwei committed
492
493

    def __init__(self,
494
495
496
497
                 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
498
        self.translation_std = translation_std
zhangwenwei's avatar
zhangwenwei committed
499
        self.global_rot_range = global_rot_range
zhangwenwei's avatar
zhangwenwei committed
500
        self.rot_range = rot_range
zhangwenwei's avatar
zhangwenwei committed
501
502
        self.num_try = num_try

503
504
    def transform(self, input_dict: dict) -> dict:
        """Transform function to apply noise to each ground truth in the scene.
505
506
507
508
509

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

        Returns:
510
            dict: Results after adding noise to each object,
511
            'points', 'gt_bboxes_3d' keys are updated in the result dict.
512
        """
zhangwenwei's avatar
zhangwenwei committed
513
514
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
        points = input_dict['points']
zhangwenwei's avatar
zhangwenwei committed
515

516
        # TODO: this is inplace operation
517
518
        numpy_box = gt_bboxes_3d.numpy()
        numpy_points = points.numpy()
519

zhangwenwei's avatar
zhangwenwei committed
520
        noise_per_object_v3_(
521
            numpy_box,
522
            numpy_points,
zhangwenwei's avatar
zhangwenwei committed
523
524
            rotation_perturb=self.rot_range,
            center_noise_std=self.translation_std,
zhangwenwei's avatar
zhangwenwei committed
525
526
            global_random_rot_range=self.global_rot_range,
            num_try=self.num_try)
527
528

        input_dict['gt_bboxes_3d'] = gt_bboxes_3d.new_box(numpy_box)
529
        input_dict['points'] = points.new_point(numpy_points)
zhangwenwei's avatar
zhangwenwei committed
530
531
        return input_dict

532
    def __repr__(self) -> str:
533
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
534
        repr_str = self.__class__.__name__
535
536
537
538
        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
539
540
541
        return repr_str


542
@TRANSFORMS.register_module()
543
class GlobalAlignment(BaseTransform):
544
545
546
547
548
549
    """Apply global alignment to 3D scene points by rotation and translation.

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

    Note:
550
        We do not record the applied rotation and translation as in
551
552
        GlobalRotScaleTrans. Because usually, we do not need to reverse
        the alignment step.
553
        For example, ScanNet 3D detection task uses aligned ground-truth
554
        bounding boxes for evaluation.
555
556
    """

557
    def __init__(self, rotation_axis: int) -> None:
558
559
        self.rotation_axis = rotation_axis

560
    def _trans_points(self, results: dict, trans_factor: np.ndarray) -> None:
561
562
563
564
565
566
567
568
569
        """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.
        """
570
        results['points'].translate(trans_factor)
571

572
    def _rot_points(self, results: dict, rot_mat: np.ndarray) -> None:
573
574
575
576
577
578
579
580
581
582
        """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
583
        results['points'].rotate(rot_mat.T)
584

585
    def _check_rot_mat(self, rot_mat: np.ndarray) -> None:
586
587
588
589
590
591
592
593
594
595
596
597
        """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}'

598
    def transform(self, results: dict) -> dict:
599
600
601
602
603
604
        """Call function to shuffle points.

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

        Returns:
605
            dict: Results after global alignment, 'points' and keys in
606
            input_dict['bbox3d_fields'] are updated in the result dict.
607
        """
608
        assert 'axis_align_matrix' in results, \
609
610
            'axis_align_matrix is not provided in GlobalAlignment'

611
        axis_align_matrix = results['axis_align_matrix']
612
613
614
615
616
617
        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)
618
619
        self._rot_points(results, rot_mat)
        self._trans_points(results, trans_vec)
620

621
        return results
622

623
    def __repr__(self) -> str:
624
        """str: Return a string that describes the module."""
625
626
627
628
629
        repr_str = self.__class__.__name__
        repr_str += f'(rotation_axis={self.rotation_axis})'
        return repr_str


630
@TRANSFORMS.register_module()
jshilong's avatar
jshilong committed
631
class GlobalRotScaleTrans(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
632
    """Apply global rotation, scaling and translation to a 3D scene.
zhangwenwei's avatar
zhangwenwei committed
633

jshilong's avatar
jshilong committed
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
    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
652
    Args:
653
        rot_range (list[float]): Range of rotation angle.
liyinhao's avatar
liyinhao committed
654
            Defaults to [-0.78539816, 0.78539816] (close to [-pi/4, pi/4]).
655
        scale_ratio_range (list[float]): Range of scale ratio.
liyinhao's avatar
liyinhao committed
656
            Defaults to [0.95, 1.05].
657
        translation_std (list[float]): The standard deviation of
658
            translation noise applied to a scene, which
zhangwenwei's avatar
zhangwenwei committed
659
            is sampled from a gaussian distribution whose standard deviation
660
661
            is set by ``translation_std``. Defaults to [0, 0, 0].
        shift_height (bool): Whether to shift height.
wuyuefeng's avatar
wuyuefeng committed
662
            (the fourth dimension of indoor points) when scaling.
liyinhao's avatar
liyinhao committed
663
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
664
    """
zhangwenwei's avatar
zhangwenwei committed
665
666

    def __init__(self,
jshilong's avatar
jshilong committed
667
668
669
670
                 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:
671
672
673
674
675
        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
676
        self.rot_range = rot_range
677
678
679

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

zhangwenwei's avatar
zhangwenwei committed
681
        self.scale_ratio_range = scale_ratio_range
682
683
684
685
686
687
688

        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
            ]
689
690
        assert all([std >= 0 for std in translation_std]), \
            'translation_std should be positive'
zhangwenwei's avatar
zhangwenwei committed
691
        self.translation_std = translation_std
wuyuefeng's avatar
wuyuefeng committed
692
        self.shift_height = shift_height
zhangwenwei's avatar
zhangwenwei committed
693

jshilong's avatar
jshilong committed
694
    def _trans_bbox_points(self, input_dict: dict) -> None:
695
696
697
698
699
700
        """Private function to translate bounding boxes and points.

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

        Returns:
701
            dict: Results after translation, 'points', 'pcd_trans'
702
            and `gt_bboxes_3d` is updated in the result dict.
703
        """
704
        translation_std = np.array(self.translation_std, dtype=np.float32)
zhangwenwei's avatar
zhangwenwei committed
705
706
        trans_factor = np.random.normal(scale=translation_std, size=3).T

707
        input_dict['points'].translate(trans_factor)
zhangwenwei's avatar
zhangwenwei committed
708
        input_dict['pcd_trans'] = trans_factor
jshilong's avatar
jshilong committed
709
710
        if 'gt_bboxes_3d' in input_dict:
            input_dict['gt_bboxes_3d'].translate(trans_factor)
zhangwenwei's avatar
zhangwenwei committed
711

jshilong's avatar
jshilong committed
712
    def _rot_bbox_points(self, input_dict: dict) -> None:
713
714
715
716
717
718
        """Private function to rotate bounding boxes and points.

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

        Returns:
719
            dict: Results after rotation, 'points', 'pcd_rotation'
720
            and `gt_bboxes_3d` is updated in the result dict.
721
        """
zhangwenwei's avatar
zhangwenwei committed
722
        rotation = self.rot_range
zhangwenwei's avatar
zhangwenwei committed
723
        noise_rotation = np.random.uniform(rotation[0], rotation[1])
zhangwenwei's avatar
zhangwenwei committed
724

jshilong's avatar
jshilong committed
725
726
727
728
729
730
731
732
        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
733
            rot_mat_T = input_dict['points'].rotate(noise_rotation)
jshilong's avatar
jshilong committed
734
735
736
737
738

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

    def _scale_bbox_points(self, input_dict: dict) -> None:
739
740
741
742
743
744
        """Private function to scale bounding boxes and points.

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

        Returns:
jshilong's avatar
jshilong committed
745
            dict: Results after scaling, 'points' and
746
            `gt_bboxes_3d` is updated in the result dict.
747
        """
zhangwenwei's avatar
zhangwenwei committed
748
        scale = input_dict['pcd_scale_factor']
749
750
        points = input_dict['points']
        points.scale(scale)
wuyuefeng's avatar
wuyuefeng committed
751
        if self.shift_height:
752
753
            assert 'height' in points.attribute_dims.keys(), \
                'setting shift_height=True but points have no height attribute'
754
755
            points.tensor[:, points.attribute_dims['height']] *= scale
        input_dict['points'] = points
wuyuefeng's avatar
wuyuefeng committed
756

jshilong's avatar
jshilong committed
757
758
759
        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
760

jshilong's avatar
jshilong committed
761
    def _random_scale(self, input_dict: dict) -> None:
762
763
764
765
766
767
        """Private function to randomly set the scale factor.

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

        Returns:
jshilong's avatar
jshilong committed
768
769
            dict: Results after scaling, 'pcd_scale_factor'
            are updated in the result dict.
770
        """
zhangwenwei's avatar
zhangwenwei committed
771
772
773
        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
774

jshilong's avatar
jshilong committed
775
    def transform(self, input_dict: dict) -> dict:
776
        """Private function to rotate, scale and translate bounding boxes and
777
778
779
780
781
782
783
        points.

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

        Returns:
            dict: Results after scaling, 'points', 'pcd_rotation',
784
            'pcd_scale_factor', 'pcd_trans' and `gt_bboxes_3d` are updated
jshilong's avatar
jshilong committed
785
            in the result dict.
786
        """
787
788
789
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

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

zhangwenwei's avatar
zhangwenwei committed
792
793
794
        if 'pcd_scale_factor' not in input_dict:
            self._random_scale(input_dict)
        self._scale_bbox_points(input_dict)
zhangwenwei's avatar
zhangwenwei committed
795

zhangwenwei's avatar
zhangwenwei committed
796
        self._trans_bbox_points(input_dict)
797
798

        input_dict['transformation_3d_flow'].extend(['R', 'S', 'T'])
zhangwenwei's avatar
zhangwenwei committed
799
800
        return input_dict

801
    def __repr__(self) -> str:
802
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
803
        repr_str = self.__class__.__name__
804
805
806
807
        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
808
809
810
        return repr_str


811
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
812
class PointShuffle(BaseTransform):
813
    """Shuffle input points."""
zhangwenwei's avatar
zhangwenwei committed
814

ZCMax's avatar
ZCMax committed
815
    def transform(self, input_dict: dict) -> dict:
816
817
818
819
820
821
        """Call function to shuffle points.

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

        Returns:
822
            dict: Results after filtering, 'points', 'pts_instance_mask'
823
            and 'pts_semantic_mask' keys are updated in the result dict.
824
        """
825
826
827
828
829
830
831
832
833
834
835
836
        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
837
838
        return input_dict

839
    def __repr__(self) -> str:
840
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
841
842
843
        return self.__class__.__name__


844
@TRANSFORMS.register_module()
845
class ObjectRangeFilter(BaseTransform):
846
847
    """Filter objects by the range.

848
849
850
851
852
853
854
855
    Required Keys:

    - gt_bboxes_3d

    Modified Keys:

    - gt_bboxes_3d

856
857
858
    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
859

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

863
864
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter objects by the range.
865
866
867
868
869

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

        Returns:
870
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d'
871
            keys are updated in the result dict.
872
        """
873
874
875
876
877
878
879
        # 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
880
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
881
        gt_labels_3d = input_dict['gt_labels_3d']
882
        mask = gt_bboxes_3d.in_range_bev(bev_range)
zhangwenwei's avatar
zhangwenwei committed
883
        gt_bboxes_3d = gt_bboxes_3d[mask]
ZwwWayne's avatar
ZwwWayne committed
884
885
886
887
        # 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
888
        gt_labels_3d = gt_labels_3d[mask.numpy().astype(bool)]
zhangwenwei's avatar
zhangwenwei committed
889
890

        # limit rad to [-pi, pi]
891
892
        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
893
894
        input_dict['gt_labels_3d'] = gt_labels_3d

zhangwenwei's avatar
zhangwenwei committed
895
896
        return input_dict

897
    def __repr__(self) -> str:
898
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
899
        repr_str = self.__class__.__name__
900
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
901
902
903
        return repr_str


904
@TRANSFORMS.register_module()
905
class PointsRangeFilter(BaseTransform):
906
907
    """Filter points by the range.

908
909
910
911
912
913
914
915
916
917
    Required Keys:

    - points
    - pts_instance_mask (optional)

    Modified Keys:

    - points
    - pts_instance_mask (optional)

918
919
920
    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
921

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

925
926
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter points by the range.
927
928
929
930
931

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

        Returns:
932
            dict: Results after filtering, 'points', 'pts_instance_mask'
933
            and 'pts_semantic_mask' keys are updated in the result dict.
934
        """
zhangwenwei's avatar
zhangwenwei committed
935
        points = input_dict['points']
936
937
        points_mask = points.in_range_3d(self.pcd_range)
        clean_points = points[points_mask]
zhangwenwei's avatar
zhangwenwei committed
938
        input_dict['points'] = clean_points
939
940
941
942
943
944
945
946
947
948
949
        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
950
951
        return input_dict

952
    def __repr__(self) -> str:
953
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
954
        repr_str = self.__class__.__name__
955
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
956
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
957
958


959
@TRANSFORMS.register_module()
960
class ObjectNameFilter(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
961
    """Filter GT objects by their names.
zhangwenwei's avatar
zhangwenwei committed
962

963
964
965
966
967
968
969
970
    Required Keys:

    - gt_labels_3d

    Modified Keys:

    - gt_labels_3d

zhangwenwei's avatar
zhangwenwei committed
971
    Args:
liyinhao's avatar
liyinhao committed
972
        classes (list[str]): List of class names to be kept for training.
zhangwenwei's avatar
zhangwenwei committed
973
974
    """

975
    def __init__(self, classes: List[str]) -> None:
zhangwenwei's avatar
zhangwenwei committed
976
977
978
        self.classes = classes
        self.labels = list(range(len(self.classes)))

979
980
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter objects by their names.
981
982
983
984
985

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

        Returns:
986
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d'
987
            keys are updated in the result dict.
988
        """
zhangwenwei's avatar
zhangwenwei committed
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],
991
                                  dtype=bool)
zhangwenwei's avatar
zhangwenwei committed
992
993
994
995
996
        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

997
    def __repr__(self) -> str:
998
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
999
1000
1001
        repr_str = self.__class__.__name__
        repr_str += f'(classes={self.classes})'
        return repr_str
wuyuefeng's avatar
wuyuefeng committed
1002
1003


1004
1005
@TRANSFORMS.register_module()
class PointSample(BaseTransform):
1006
    """Point sample.
wuyuefeng's avatar
wuyuefeng committed
1007
1008
1009

    Sampling data to a certain number.

1010
    Required Keys:
1011

1012
1013
1014
1015
1016
    - points
    - pts_instance_mask (optional)
    - pts_semantic_mask (optional)

    Modified Keys:
1017

1018
1019
1020
1021
    - points
    - pts_instance_mask (optional)
    - pts_semantic_mask (optional)

wuyuefeng's avatar
wuyuefeng committed
1022
1023
    Args:
        num_points (int): Number of points to be sampled.
1024
        sample_range (float, optional): The range where to sample points.
1025
1026
            If not None, the points with depth larger than `sample_range` are
            prior to be sampled. Defaults to None.
1027
1028
        replace (bool): Whether the sampling is with or without replacement.
            Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
1029
1030
    """

1031
1032
    def __init__(self,
                 num_points: int,
1033
1034
                 sample_range: Optional[float] = None,
                 replace: bool = False) -> None:
wuyuefeng's avatar
wuyuefeng committed
1035
        self.num_points = num_points
1036
1037
1038
        self.sample_range = sample_range
        self.replace = replace

1039
1040
1041
1042
1043
1044
1045
1046
    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
1047
1048
1049
1050
1051
        """Points random sampling.

        Sample points to a certain number.

        Args:
1052
            points (:obj:`BasePoints`): 3D Points.
wuyuefeng's avatar
wuyuefeng committed
1053
            num_samples (int): Number of samples to be sampled.
1054
            sample_range (float, optional): Indicating the range where the
1055
                points will be sampled. Defaults to None.
1056
            replace (bool): Sampling with or without replacement.
1057
                Defaults to False.
1058
            return_choices (bool): Whether return choice. Defaults to False.
1059

wuyuefeng's avatar
wuyuefeng committed
1060
        Returns:
1061
1062
1063
            tuple[:obj:`BasePoints`, np.ndarray] | :obj:`BasePoints`:

                - points (:obj:`BasePoints`): 3D Points.
1064
                - choices (np.ndarray, optional): The generated random samples.
wuyuefeng's avatar
wuyuefeng committed
1065
        """
1066
        if not replace:
wuyuefeng's avatar
wuyuefeng committed
1067
            replace = (points.shape[0] < num_samples)
1068
1069
1070
        point_range = range(len(points))
        if sample_range is not None and not replace:
            # Only sampling the near points when len(points) >= num_samples
1071
            dist = np.linalg.norm(points.coord.numpy(), axis=1)
1072
1073
            far_inds = np.where(dist >= sample_range)[0]
            near_inds = np.where(dist < sample_range)[0]
1074
1075
1076
1077
            # 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)
1078
1079
1080
1081
1082
1083
1084
            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
1085
1086
1087
1088
1089
        if return_choices:
            return points[choices], choices
        else:
            return points[choices]

1090
    def transform(self, input_dict: dict) -> dict:
1091
        """Transform function to sample points to in indoor scenes.
1092
1093
1094

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

1096
        Returns:
1097
            dict: Results after sampling, 'points', 'pts_instance_mask'
1098
            and 'pts_semantic_mask' keys are updated in the result dict.
1099
        """
1100
        points = input_dict['points']
1101
1102
1103
1104
1105
1106
        points, choices = self._points_random_sampling(
            points,
            self.num_points,
            self.sample_range,
            self.replace,
            return_choices=True)
1107
        input_dict['points'] = points
1108

1109
1110
        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
1111

1112
        if pts_instance_mask is not None:
wuyuefeng's avatar
wuyuefeng committed
1113
            pts_instance_mask = pts_instance_mask[choices]
1114
            input_dict['pts_instance_mask'] = pts_instance_mask
1115
1116
1117

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

1120
        return input_dict
wuyuefeng's avatar
wuyuefeng committed
1121

1122
    def __repr__(self) -> str:
1123
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
1124
        repr_str = self.__class__.__name__
1125
        repr_str += f'(num_points={self.num_points},'
1126
1127
        repr_str += f' sample_range={self.sample_range},'
        repr_str += f' replace={self.replace})'
1128

1129
1130
1131
        return repr_str


1132
@TRANSFORMS.register_module()
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
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)


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

    Note:
        This transform should only be used in the training process of point
1183
1184
1185
        cloud segmentation tasks. For the sliding patch generation and
        inference process in testing, please refer to the `slide_inference`
        function of `EncoderDecoder3D` class.
1186
1187
1188
    """

    def __init__(self,
ZCMax's avatar
ZCMax committed
1189
1190
1191
1192
1193
1194
1195
1196
1197
                 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:
1198
1199
1200
1201
1202
        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
1203
        self.enlarge_size = enlarge_size if enlarge_size is not None else 0.0
1204
        self.min_unique_num = min_unique_num
1205
        self.eps = eps
1206
1207
1208
1209
1210

        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.')
1211

ZCMax's avatar
ZCMax committed
1212
1213
1214
1215
    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:
1216
1217
        """Generating model input.

1218
        Generate input by subtracting patch center and adding additional
1219
1220
1221
1222
1223
1224
1225
1226
1227
            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.
1228
            point_type (type): class of input points inherited from BasePoints.
1229
1230

        Returns:
1231
            :obj:`BasePoints`: The generated input data.
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
1243
1244
1245
1246
1247
1248
1249
1250
1251
1252
1253
1254
        """
        # 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

1255
    def _patch_points_sampling(
1256
1257
            self, points: BasePoints,
            sem_mask: np.ndarray) -> Tuple[BasePoints, np.ndarray]:
1258
1259
1260
1261
1262
1263
        """Patch points sampling.

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

        Args:
1264
            points (:obj:`BasePoints`): 3D Points.
1265
1266
1267
            sem_mask (np.ndarray): semantic segmentation mask for input points.

        Returns:
1268
            tuple[:obj:`BasePoints`, np.ndarray]:
1269

1270
                - points (:obj:`BasePoints`): 3D Points.
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
                - 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)

1281
        for _ in range(self.num_try):
1282
1283
1284
            # random sample a point as patch center
            cur_center = coords[np.random.choice(coords.shape[0])]

1285
1286
            # boundary of a patch, which would be enlarged by
            # `self.enlarge_size` as an augmentation
1287
1288
1289
1290
1291
1292
1293
            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(
1294
1295
                (coords >= (cur_min - self.enlarge_size)) *
                (coords <= (cur_max + self.enlarge_size)),
1296
1297
1298
1299
1300
1301
1302
                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]
1303
            point_idxs = np.where(cur_choice)[0]
1304
            mask = np.sum(
1305
1306
                (cur_coords >= (cur_min - self.eps)) * (cur_coords <=
                                                        (cur_max + self.eps)),
1307
                axis=1) == 3
1308

1309
1310
            # two criteria for patch sampling, adopted from PointNet++
            # 1. selected patch should contain enough unique points
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
            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:
1323
                # if `min_unique_num` is provided, directly compare with it
1324
                flag1 = mask.sum() >= self.min_unique_num
1325

1326
            # 2. selected patch should contain enough annotated points
1327
1328
1329
1330
1331
1332
1333
1334
1335
            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

1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
        # 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]
1349
1350
1351
1352
1353
1354
1355
1356

        # 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
1357
    def transform(self, input_dict: dict) -> dict:
1358
1359
1360
1361
1362
1363
        """Call function to sample points to in indoor scenes.

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

        Returns:
1364
            dict: Results after sampling, 'points', 'pts_instance_mask'
1365
            and 'pts_semantic_mask' keys are updated in the result dict.
1366
        """
ZCMax's avatar
ZCMax committed
1367
        points = input_dict['points']
1368

ZCMax's avatar
ZCMax committed
1369
        assert 'pts_semantic_mask' in input_dict.keys(), \
1370
            'semantic mask should be provided in training and evaluation'
ZCMax's avatar
ZCMax committed
1371
        pts_semantic_mask = input_dict['pts_semantic_mask']
1372
1373
1374
1375

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

ZCMax's avatar
ZCMax committed
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
        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)

1386
        if pts_instance_mask is not None:
ZCMax's avatar
ZCMax committed
1387
1388
1389
1390
1391
            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]
1392

ZCMax's avatar
ZCMax committed
1393
        return input_dict
1394

1395
    def __repr__(self) -> str:
1396
1397
1398
1399
1400
1401
        """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},'
1402
1403
        repr_str += f' num_try={self.num_try},'
        repr_str += f' enlarge_size={self.enlarge_size},'
1404
1405
        repr_str += f' min_unique_num={self.min_unique_num},'
        repr_str += f' eps={self.eps})'
wuyuefeng's avatar
wuyuefeng committed
1406
        return repr_str
1407
1408


1409
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1410
class BackgroundPointsFilter(BaseTransform):
1411
1412
1413
    """Filter background points near the bounding box.

    Args:
1414
        bbox_enlarge_range (tuple[float] | float): Bbox enlarge range.
1415
1416
    """

ZCMax's avatar
ZCMax committed
1417
    def __init__(self, bbox_enlarge_range: Union[Tuple[float], float]) -> None:
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
        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
1428
    def transform(self, input_dict: dict) -> dict:
1429
1430
1431
1432
1433
1434
        """Call function to filter points by the range.

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

        Returns:
1435
            dict: Results after filtering, 'points', 'pts_instance_mask'
1436
            and 'pts_semantic_mask' keys are updated in the result dict.
1437
1438
1439
1440
        """
        points = input_dict['points']
        gt_bboxes_3d = input_dict['gt_bboxes_3d']

xiliu8006's avatar
xiliu8006 committed
1441
1442
1443
1444
        # 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()

1445
1446
        enlarged_gt_bboxes_3d = gt_bboxes_3d_np.copy()
        enlarged_gt_bboxes_3d[:, 3:6] += self.bbox_enlarge_range
xiliu8006's avatar
xiliu8006 committed
1447
        points_numpy = points.tensor.clone().numpy()
1448
1449
        foreground_masks = box_np_ops.points_in_rbbox(
            points_numpy, gt_bboxes_3d_np, origin=(0.5, 0.5, 0.5))
1450
        enlarge_foreground_masks = box_np_ops.points_in_rbbox(
1451
            points_numpy, enlarged_gt_bboxes_3d, origin=(0.5, 0.5, 0.5))
1452
1453
1454
1455
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
        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

1467
    def __repr__(self) -> str:
1468
1469
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
1470
        repr_str += f'(bbox_enlarge_range={self.bbox_enlarge_range.tolist()})'
1471
        return repr_str
1472
1473


1474
@TRANSFORMS.register_module()
1475
class VoxelBasedPointSampler(BaseTransform):
1476
1477
1478
1479
1480
1481
    """Voxel based point sampler.

    Apply voxel sampling to multiple sweep points.

    Args:
        cur_sweep_cfg (dict): Config for sampling current points.
1482
1483
        prev_sweep_cfg (dict, optional): Config for sampling previous points.
            Defaults to None.
1484
        time_dim (int): Index that indicate the time dimension
1485
            for input points. Defaults to 3.
1486
1487
    """

1488
1489
1490
1491
    def __init__(self,
                 cur_sweep_cfg: dict,
                 prev_sweep_cfg: Optional[dict] = None,
                 time_dim: int = 3) -> None:
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
        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

1504
    def _sample_points(self, points: np.ndarray, sampler: VoxelGenerator,
1505
                       point_dim: int) -> np.ndarray:
1506
1507
1508
1509
1510
1511
        """Sample points for each points subset.

        Args:
            points (np.ndarray): Points subset to be sampled.
            sampler (VoxelGenerator): Voxel based sampler for
                each points subset.
1512
            point_dim (int): The dimension of each points.
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530

        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

1531
    def transform(self, results: dict) -> dict:
1532
1533
1534
1535
1536
1537
        """Call function to sample points from multiple sweeps.

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

        Returns:
1538
            dict: Results after sampling, 'points', 'pts_instance_mask'
1539
            and 'pts_semantic_mask' keys are updated in the result dict.
1540
1541
1542
1543
1544
1545
1546
1547
1548
        """
        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
1549
        points_numpy = points.numpy()
1550
        extra_channel = [points_numpy]
1551
1552
1553
1554
1555
1556
1557
1558
1559
        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])

1560
        points_numpy = np.concatenate(extra_channel, axis=-1)
1561
1562
1563
1564
1565

        # 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.
1566
1567
1568
        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]
1569
1570
1571
1572
1573
1574
1575
1576
1577
        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,
1578
                                               points_numpy.shape[1])
1579
1580
1581
        if self.prev_voxel_generator is not None:
            prev_sweeps_points = self._sample_points(prev_sweeps_points,
                                                     self.prev_voxel_generator,
1582
                                                     points_numpy.shape[1])
1583

1584
1585
            points_numpy = np.concatenate(
                [cur_sweep_points, prev_sweeps_points], 0)
1586
        else:
1587
            points_numpy = cur_sweep_points
1588
1589

        if self.cur_voxel_generator._max_num_points == 1:
1590
1591
            points_numpy = points_numpy.squeeze(1)
        results['points'] = points.new_point(points_numpy[..., :original_dim])
1592

1593
        # Restore the corresponding seg and mask fields
1594
        for key, dim_index in map_fields2dim:
1595
            results[key] = points_numpy[..., dim_index]
1596
1597
1598

        return results

1599
    def __repr__(self) -> str:
1600
1601
1602
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618
        """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
1619
1620


1621
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1622
class AffineResize(BaseTransform):
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
    """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.
1635
        bbox_clip_border (bool): Whether clip the objects
1636
1637
1638
            outside the border of the image. Defaults to True.
    """

ZCMax's avatar
ZCMax committed
1639
1640
1641
1642
    def __init__(self,
                 img_scale: Tuple,
                 down_ratio: int,
                 bbox_clip_border: bool = True) -> None:
1643
1644
1645
1646
1647

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

ZCMax's avatar
ZCMax committed
1648
    def transform(self, results: dict) -> dict:
1649
1650
1651
1652
1653
1654
1655
        """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'
1656
            keys are added in the result dict.
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
1686
1687
1688
1689
1690
1691
1692
        """
        # 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
1693
1694
        if 'gt_bboxes' in results:
            self._affine_bboxes(results, trans_affine)
1695

ZCMax's avatar
ZCMax committed
1696
1697
        if 'centers_2d' in results:
            centers2d = self._affine_transform(results['centers_2d'],
1698
1699
1700
1701
1702
                                               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
1703
1704
1705
1706
            results['centers_2d'] = centers2d[valid_index]

            if 'gt_bboxes' in results:
                results['gt_bboxes'] = results['gt_bboxes'][valid_index]
1707
1708
1709
                if 'gt_bboxes_labels' in results:
                    results['gt_bboxes_labels'] = results['gt_bboxes_labels'][
                        valid_index]
ZCMax's avatar
ZCMax committed
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
                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]
1720
1721
1722
1723
1724

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

        return results

ZCMax's avatar
ZCMax committed
1725
    def _affine_bboxes(self, results: dict, matrix: np.ndarray) -> None:
1726
1727
1728
1729
1730
1731
1732
1733
1734
        """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
1735
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
        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:
1747
        """Affine transform bbox points to input image.
1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764

        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
1765
1766
    def _get_transform_matrix(self, center: Tuple, scale: Tuple,
                              output_scale: Tuple[float]) -> np.ndarray:
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
1794
1795
1796
1797
1798
1799
1800
        """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
1801
1802
    def _get_ref_point(self, ref_point1: np.ndarray,
                       ref_point2: np.ndarray) -> np.ndarray:
1803
        """Get reference point to calculate affine transform matrix.
1804
1805

        While using opencv to calculate the affine matrix, we need at least
1806
        three corresponding points separately on original image and target
1807
1808
1809
1810
1811
1812
        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

1813
    def __repr__(self) -> str:
1814
        """str: Return a string that describes the module."""
1815
1816
1817
1818
1819
1820
        repr_str = self.__class__.__name__
        repr_str += f'(img_scale={self.img_scale}, '
        repr_str += f'down_ratio={self.down_ratio}) '
        return repr_str


1821
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1822
class RandomShiftScale(BaseTransform):
1823
1824
1825
1826
    """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
1827
    infos into loading TRANSFORMS. It's designed to be used with
1828
1829
1830
1831
1832
1833
1834
    AffineResize together.

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

1835
    def __init__(self, shift_scale: Tuple[float], aug_prob: float) -> None:
1836
1837
1838
1839

        self.shift_scale = shift_scale
        self.aug_prob = aug_prob

ZCMax's avatar
ZCMax committed
1840
    def transform(self, results: dict) -> dict:
1841
1842
1843
1844
1845
1846
1847
        """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'
1848
            and 'affine_aug' keys are added in the result dict.
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
1868
1869
1870
1871
1872
        """
        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

1873
    def __repr__(self) -> str:
1874
        """str: Return a string that describes the module."""
1875
1876
1877
1878
        repr_str = self.__class__.__name__
        repr_str += f'(shift_scale={self.shift_scale}, '
        repr_str += f'aug_prob={self.aug_prob}) '
        return repr_str
1879
1880
1881
1882
1883


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

1884
    def _resize_3d(self, results: dict) -> None:
1885
1886
1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
        """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.
1898

1899
1900
1901
1902
1903
1904
1905
1906
1907
1908
1909
1910
1911
1912
1913
1914
1915
1916
1917
1918
1919
        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']``.
    """

1920
    def _resize_3d(self, results: dict) -> None:
1921
1922
1923
1924
1925
1926
1927
        """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])

1928
    def transform(self, results: dict) -> dict:
1929
1930
        """Transform function to resize images, bounding boxes, masks, semantic
        segmentation map. Compared to RandomResize, this function would further
1931
1932
1933
1934
        check if scale is already set in results.

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

1936
        Returns:
1937
1938
            dict: Resized results, 'img_shape', 'pad_shape', 'scale_factor',
            'keep_ratio' keys are added into result dict.
1939
1940
1941
1942
1943
1944
1945
1946
        """
        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
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963


@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)
1964
    - gt_ignore_flags (bool) (optional)
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
    - 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.
2000
        rel_offset_h (tuple): The cropping interval of image height. Defaults
2001
            to (0., 1.).
2002
        rel_offset_w (tuple): The cropping interval of image width. Defaults
2003
2004
2005
2006
            to (0., 1.).

    Note:
        - If the image is smaller than the absolute crop size, return the
2007
          original image.
2008
2009
2010
2011
2012
2013
2014
2015
        - 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.
    """

2016
2017
2018
2019
2020
2021
2022
2023
2024
2025
    def __init__(
        self,
        crop_size: tuple,
        crop_type: str = 'absolute',
        allow_negative_crop: bool = False,
        recompute_bbox: bool = False,
        bbox_clip_border: bool = True,
        rel_offset_h: tuple = (0., 1.),
        rel_offset_w: tuple = (0., 1.)
    ) -> None:
2026
2027
2028
2029
2030
2031
2032
2033
2034
2035
2036
        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

2037
2038
2039
2040
    def _crop_data(self,
                   results: dict,
                   crop_size: tuple,
                   allow_negative_crop: bool = False) -> dict:
2041
2042
2043
2044
2045
2046
2047
        """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
2048
                contain any bbox area. Defaults to False.
2049
2050
2051

        Returns:
            dict: Randomly cropped results, 'img_shape' key in result dict is
2052
            updated according to crop size.
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
        """
        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

2135
    def transform(self, results: dict) -> dict:
2136
2137
2138
2139
2140
2141
2142
2143
        """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
2144
            updated according to crop size.
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
        """
        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

2155
2156
    def __repr__(self) -> dict:
        """str: Return a string that describes the module."""
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
2248
2249
2250
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
        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()
2262
class MultiViewWrapper(BaseTransform):
2263
2264
2265
2266
2267
2268
2269
2270
2271
2272
2273
2274
2275
2276
    """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
2277
            for multiview image. Defaults to True.
2278
        process_fields (list): Desired keys that the transformations should
2279
            be conducted on. Defaults to ['img', 'cam2img', 'lidar2cam'].
2280
        collected_keys (list): Collect information in transformation
2281
            like rotate angles, crop roi, and flip state. Defaults to
2282
2283
2284
2285
                ['scale', 'scale_factor', 'crop',
                 'crop_offset', 'ori_shape',
                 'pad_shape', 'img_shape',
                 'pad_fixed_size', 'pad_size_divisor',
2286
                 'flip', 'flip_direction', 'rotate'].
2287
        randomness_keys (list): The keys that related to the randomness
2288
            in transformation. Defaults to
2289
2290
2291
2292
                    ['scale', 'scale_factor', 'crop_size', 'flip',
                     'flip_direction', 'photometric_param']
    """

2293
2294
2295
2296
2297
2298
2299
2300
2301
2302
2303
2304
2305
2306
2307
    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'
        ]
    ) -> None:
2308
        self.transforms = Compose(transforms)
2309
2310
2311
2312
2313
        self.override_aug_config = override_aug_config
        self.collected_keys = collected_keys
        self.process_fields = process_fields
        self.randomness_keys = randomness_keys

2314
    def transform(self, input_dict: dict) -> dict:
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
        """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]
2343
            process_dict = self.transforms(process_dict)
2344
2345
2346
2347
2348
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
            # 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
2363
2364
2365
2366
2367
2368
2369
2370
2371
2372
2373
2374
2375
2376
2377
2378
2379
2380
2381
2382
2383
2384
2385
2386
2387
2388
2389
2390
2391
2392
2393
2394
2395
2396
2397
2398
2399
2400
2401
2402
2403
2404
2405
2406
2407
2408
2409
2410
2411
2412
2413
2414
2415
2416
2417
2418
2419
2420
2421
2422
2423
2424
2425
2426
2427
2428
2429
2430
2431
2432
2433
2434
2435
2436
2437
2438
2439
2440
2441
2442
2443
2444
2445
2446
2447
2448
2449
2450
2451
2452
2453
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464
2465
2466
2467
2468
2469
2470
2471
2472
2473
2474
2475
2476
2477
2478
2479
2480
2481
2482
2483
2484
2485
2486
2487
2488
2489
2490
2491
2492
2493
2494
2495
2496
2497
2498
2499
2500
2501
2502
2503
2504
2505
2506
2507
2508
2509
2510
2511
2512
2513
2514
2515
2516
2517
2518
2519
2520
2521
2522
2523
2524
2525
2526
2527
2528
2529
2530


@TRANSFORMS.register_module()
class PolarMix(BaseTransform):
    """PolarMix data augmentation.

    The polarmix transform steps are as follows:

        1. Another random point cloud is picked by dataset.
        2. Exchange sectors of two point clouds that are cut with certain
           azimuth angles.
        3. Cut point instances from picked point cloud, rotate them by multiple
           azimuth angles, and paste the cut and rotated instances.

    Required Keys:

    - points (:obj:`BasePoints`)
    - pts_semantic_mask (np.int64)
    - dataset (:obj:`BaseDataset`)

    Modified Keys:

    - points (:obj:`BasePoints`)
    - pts_semantic_mask (np.int64)

    Args:
        instance_classes (List[int]): Semantic masks which represent the
            instance.
        swap_ratio (float): Swap ratio of two point cloud. Defaults to 0.5.
        rotate_paste_ratio (float): Rotate paste ratio. Defaults to 1.0.
        pre_transform (Sequence[dict], optional): Sequence of transform object
            or config dict to be composed. Defaults to None.
        prob (float): The transformation probability. Defaults to 1.0.
    """

    def __init__(self,
                 instance_classes: List[int],
                 swap_ratio: float = 0.5,
                 rotate_paste_ratio: float = 1.0,
                 pre_transform: Optional[Sequence[dict]] = None,
                 prob: float = 1.0) -> None:
        assert is_list_of(instance_classes, int), \
            'instance_classes should be a list of int'
        self.instance_classes = instance_classes
        self.swap_ratio = swap_ratio
        self.rotate_paste_ratio = rotate_paste_ratio

        self.prob = prob
        if pre_transform is None:
            self.pre_transform = None
        else:
            self.pre_transform = Compose(pre_transform)

    def polar_mix_transform(self, input_dict: dict, mix_results: dict) -> dict:
        """PolarMix transform function.

        Args:
            input_dict (dict): Result dict from loading pipeline.
            mix_results (dict): Mixed dict picked from dataset.

        Returns:
            dict: output dict after transformation.
        """
        mix_points = mix_results['points']
        mix_pts_semantic_mask = mix_results['pts_semantic_mask']

        points = input_dict['points']
        pts_semantic_mask = input_dict['pts_semantic_mask']

        # 1. swap point cloud
        if np.random.random() < self.swap_ratio:
            start_angle = (np.random.random() - 1) * np.pi  # -pi~0
            end_angle = start_angle + np.pi
            # calculate horizontal angle for each point
            yaw = -torch.atan2(points.coord[:, 1], points.coord[:, 0])
            mix_yaw = -torch.atan2(mix_points.coord[:, 1], mix_points.coord[:,
                                                                            0])

            # select points in sector
            idx = (yaw <= start_angle) | (yaw >= end_angle)
            mix_idx = (mix_yaw > start_angle) & (mix_yaw < end_angle)

            # swap
            points = points.cat([points[idx], mix_points[mix_idx]])
            pts_semantic_mask = np.concatenate(
                (pts_semantic_mask[idx.numpy()],
                 mix_pts_semantic_mask[mix_idx.numpy()]),
                axis=0)

        # 2. rotate-pasting
        if np.random.random() < self.rotate_paste_ratio:
            # extract instance points
            instance_points, instance_pts_semantic_mask = [], []
            for instance_class in self.instance_classes:
                mix_idx = mix_pts_semantic_mask == instance_class
                instance_points.append(mix_points[mix_idx])
                instance_pts_semantic_mask.append(
                    mix_pts_semantic_mask[mix_idx])
            instance_points = mix_points.cat(instance_points)
            instance_pts_semantic_mask = np.concatenate(
                instance_pts_semantic_mask, axis=0)

            # rotate-copy
            copy_points = [instance_points]
            copy_pts_semantic_mask = [instance_pts_semantic_mask]
            angle_list = [
                np.random.random() * np.pi * 2 / 3,
                (np.random.random() + 1) * np.pi * 2 / 3
            ]
            for angle in angle_list:
                new_points = instance_points.clone()
                new_points.rotate(angle)
                copy_points.append(new_points)
                copy_pts_semantic_mask.append(instance_pts_semantic_mask)
            copy_points = instance_points.cat(copy_points)
            copy_pts_semantic_mask = np.concatenate(
                copy_pts_semantic_mask, axis=0)

            points = points.cat([points, copy_points])
            pts_semantic_mask = np.concatenate(
                (pts_semantic_mask, copy_pts_semantic_mask), axis=0)

        input_dict['points'] = points
        input_dict['pts_semantic_mask'] = pts_semantic_mask
        return input_dict

    def transform(self, input_dict: dict) -> dict:
        """PolarMix transform function.

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

        Returns:
            dict: output dict after transformation.
        """
        if np.random.rand() > self.prob:
            return input_dict

        assert 'dataset' in input_dict, \
            '`dataset` is needed to pass through PolarMix, while not found.'
        dataset = input_dict['dataset']

        # get index of other point cloud
        index = np.random.randint(0, len(dataset))

        mix_results = dataset.get_data_info(index)

        if self.pre_transform is not None:
            # pre_transform may also require dataset
            mix_results.update({'dataset': dataset})
            # before polarmix need to go through
            # the necessary pre_transform
            mix_results = self.pre_transform(mix_results)
            mix_results.pop('dataset')

        input_dict = self.polar_mix_transform(input_dict, mix_results)

        return input_dict

    def __repr__(self) -> str:
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
        repr_str += f'(instance_classes={self.instance_classes}, '
        repr_str += f'swap_ratio={self.swap_ratio}, '
        repr_str += f'rotate_paste_ratio={self.rotate_paste_ratio}, '
        repr_str += f'pre_transform={self.pre_transform}, '
        repr_str += f'prob={self.prob})'
        return repr_str
2531
2532
2533
2534
2535
2536
2537
2538
2539
2540
2541
2542
2543
2544
2545
2546
2547
2548
2549
2550
2551
2552
2553
2554
2555
2556
2557
2558
2559
2560
2561
2562
2563
2564
2565
2566
2567
2568
2569
2570
2571
2572
2573
2574
2575
2576
2577
2578
2579
2580
2581
2582
2583
2584
2585
2586
2587
2588
2589
2590
2591
2592
2593
2594
2595
2596
2597
2598
2599
2600
2601


@TRANSFORMS.register_module()
class LaserMix(BaseTransform):
    """LaserMix data augmentation.

    The lasermix transform steps are as follows:

        1. Another random point cloud is picked by dataset.
        2. Divide the point cloud into several regions according to pitch
           angles and combine the areas crossly.

    Required Keys:

    - points (:obj:`BasePoints`)
    - pts_semantic_mask (np.int64)
    - dataset (:obj:`BaseDataset`)

    Modified Keys:

    - points (:obj:`BasePoints`)
    - pts_semantic_mask (np.int64)

    Args:
        num_areas (List[int]): A list of area numbers will be divided into.
        pitch_angles (Sequence[float]): Pitch angles used to divide areas.
        pre_transform (Sequence[dict], optional): Sequence of transform object
            or config dict to be composed. Defaults to None.
        prob (float): The transformation probability. Defaults to 1.0.
    """

    def __init__(self,
                 num_areas: List[int],
                 pitch_angles: Sequence[float],
                 pre_transform: Optional[Sequence[dict]] = None,
                 prob: float = 1.0) -> None:
        assert is_list_of(num_areas, int), \
            'num_areas should be a list of int.'
        self.num_areas = num_areas

        assert len(pitch_angles) == 2, \
            'The length of pitch_angles should be 2, ' \
            f'but got {len(pitch_angles)}.'
        assert pitch_angles[1] > pitch_angles[0], \
            'pitch_angles[1] should be larger than pitch_angles[0].'
        self.pitch_angles = pitch_angles

        self.prob = prob
        if pre_transform is None:
            self.pre_transform = None
        else:
            self.pre_transform = Compose(pre_transform)

    def laser_mix_transform(self, input_dict: dict, mix_results: dict) -> dict:
        """LaserMix transform function.

        Args:
            input_dict (dict): Result dict from loading pipeline.
            mix_results (dict): Mixed dict picked from dataset.

        Returns:
            dict: output dict after transformation.
        """
        mix_points = mix_results['points']
        mix_pts_semantic_mask = mix_results['pts_semantic_mask']

        points = input_dict['points']
        pts_semantic_mask = input_dict['pts_semantic_mask']

        rho = torch.sqrt(points.coord[:, 0]**2 + points.coord[:, 1]**2)
        pitch = torch.atan2(points.coord[:, 2], rho)
Xiang Xu's avatar
Xiang Xu committed
2602
2603
        pitch = torch.clamp(pitch, self.pitch_angles[0] + 1e-5,
                            self.pitch_angles[1] - 1e-5)
2604
2605
2606
2607

        mix_rho = torch.sqrt(mix_points.coord[:, 0]**2 +
                             mix_points.coord[:, 1]**2)
        mix_pitch = torch.atan2(mix_points.coord[:, 2], mix_rho)
Xiang Xu's avatar
Xiang Xu committed
2608
2609
        mix_pitch = torch.clamp(mix_pitch, self.pitch_angles[0] + 1e-5,
                                self.pitch_angles[1] - 1e-5)
2610
2611
2612
2613
2614
2615
2616
2617
2618
2619
2620
2621
2622
2623
2624
2625
2626
2627
2628
2629
2630
2631
2632
2633
2634
2635
2636
2637
2638
2639
2640
2641
2642
2643
2644
2645
2646
2647
2648
2649
2650
2651
2652
2653
2654
2655
2656
2657
2658
2659
2660
2661
2662
2663
2664
2665
2666
2667
2668
2669
2670
2671
2672
2673
2674
2675

        num_areas = np.random.choice(self.num_areas, size=1)[0]
        angle_list = np.linspace(self.pitch_angles[1], self.pitch_angles[0],
                                 num_areas + 1)
        out_points = []
        out_pts_semantic_mask = []
        for i in range(num_areas):
            # convert angle to radian
            start_angle = angle_list[i + 1] / 180 * np.pi
            end_angle = angle_list[i] / 180 * np.pi
            if i % 2 == 0:  # pick from original point cloud
                idx = (pitch > start_angle) & (pitch <= end_angle)
                out_points.append(points[idx])
                out_pts_semantic_mask.append(pts_semantic_mask[idx.numpy()])
            else:  # pickle from mixed point cloud
                idx = (mix_pitch > start_angle) & (mix_pitch <= end_angle)
                out_points.append(mix_points[idx])
                out_pts_semantic_mask.append(
                    mix_pts_semantic_mask[idx.numpy()])
        out_points = points.cat(out_points)
        out_pts_semantic_mask = np.concatenate(out_pts_semantic_mask, axis=0)
        input_dict['points'] = out_points
        input_dict['pts_semantic_mask'] = out_pts_semantic_mask
        return input_dict

    def transform(self, input_dict: dict) -> dict:
        """LaserMix transform function.

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

        Returns:
            dict: output dict after transformation.
        """
        if np.random.rand() > self.prob:
            return input_dict

        assert 'dataset' in input_dict, \
            '`dataset` is needed to pass through LaserMix, while not found.'
        dataset = input_dict['dataset']

        # get index of other point cloud
        index = np.random.randint(0, len(dataset))

        mix_results = dataset.get_data_info(index)

        if self.pre_transform is not None:
            # pre_transform may also require dataset
            mix_results.update({'dataset': dataset})
            # before lasermix need to go through
            # the necessary pre_transform
            mix_results = self.pre_transform(mix_results)
            mix_results.pop('dataset')

        input_dict = self.laser_mix_transform(input_dict, mix_results)

        return input_dict

    def __repr__(self) -> str:
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
        repr_str += f'(num_areas={self.num_areas}, '
        repr_str += f'pitch_angles={self.pitch_angles}, '
        repr_str += f'pre_transform={self.pre_transform}, '
        repr_str += f'prob={self.prob})'
        return repr_str