transforms_3d.py 87.9 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
import random
3
import warnings
4
from typing import List, Optional, Tuple, Union
5
6

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

zhangshilong's avatar
zhangshilong committed
14
from mmdet3d.models.task_modules import VoxelGenerator
15
from mmdet3d.registry import TRANSFORMS
zhangshilong's avatar
zhangshilong committed
16
17
18
19
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
20
21
22
from .data_augment_utils import noise_per_object_v3_


23
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
24
class RandomDropPointsColor(BaseTransform):
25
26
27
28
29
30
31
    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:
32
        drop_ratio (float): The probability of dropping point colors.
33
34
35
            Defaults to 0.2.
    """

ZCMax's avatar
ZCMax committed
36
    def __init__(self, drop_ratio: float = 0.2) -> None:
37
38
39
40
        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
41
    def transform(self, input_dict: dict) -> dict:
42
43
44
45
46
47
        """Call function to drop point colors.

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

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

56
57
58
59
60
61
62
        # 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:
63
64
65
            points.color = points.color * 0.0
        return input_dict

66
    def __repr__(self) -> str:
67
68
69
70
71
72
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
        repr_str += f'(drop_ratio={self.drop_ratio})'
        return repr_str


73
@TRANSFORMS.register_module()
zhangwenwei's avatar
zhangwenwei committed
74
75
76
77
78
79
80
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
    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
99
    Args:
100
        sync_2d (bool): Whether to apply flip according to the 2D
zhangwenwei's avatar
zhangwenwei committed
101
102
            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
103
            to that of 2D images. Defaults to True.
104
        flip_ratio_bev_horizontal (float): The flipping probability
liyinhao's avatar
liyinhao committed
105
            in horizontal direction. Defaults to 0.0.
106
        flip_ratio_bev_vertical (float): The flipping probability
liyinhao's avatar
liyinhao committed
107
            in vertical direction. Defaults to 0.0.
108
109
        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
110
111
            to False, since the flip of 2D images does not influence the 3D
            box. Defaults to True.
zhangwenwei's avatar
zhangwenwei committed
112
113
    """

wuyuefeng's avatar
wuyuefeng committed
114
    def __init__(self,
jshilong's avatar
jshilong committed
115
116
117
                 sync_2d: bool = True,
                 flip_ratio_bev_horizontal: float = 0.0,
                 flip_ratio_bev_vertical: float = 0.0,
118
                 flip_box3d: bool = True,
jshilong's avatar
jshilong committed
119
120
121
122
                 **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
123
        super(RandomFlip3D, self).__init__(
jshilong's avatar
jshilong committed
124
            prob=flip_ratio_bev_horizontal, direction='horizontal', **kwargs)
zhangwenwei's avatar
zhangwenwei committed
125
        self.sync_2d = sync_2d
jshilong's avatar
jshilong committed
126
        self.flip_ratio_bev_horizontal = flip_ratio_bev_horizontal
wuyuefeng's avatar
wuyuefeng committed
127
        self.flip_ratio_bev_vertical = flip_ratio_bev_vertical
128
        self.flip_box3d = flip_box3d
wuyuefeng's avatar
wuyuefeng committed
129
130
131
132
133
134
135
136
137
        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
138
139
140
    def random_flip_data_3d(self,
                            input_dict: dict,
                            direction: str = 'horizontal') -> None:
141
142
        """Flip 3D data randomly.

jshilong's avatar
jshilong committed
143
144
145
146
147
148
149
        `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.

150
151
        Args:
            input_dict (dict): Result dict from loading pipeline.
152
            direction (str): Flip direction. Defaults to 'horizontal'.
153
154

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

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

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

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

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

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

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

jshilong's avatar
jshilong committed
218
        if self.sync_2d and 'img' in input_dict:
wuyuefeng's avatar
wuyuefeng committed
219
220
            input_dict['pcd_horizontal_flip'] = input_dict['flip']
            input_dict['pcd_vertical_flip'] = False
zhangwenwei's avatar
zhangwenwei committed
221
        else:
wuyuefeng's avatar
wuyuefeng committed
222
223
            if 'pcd_horizontal_flip' not in input_dict:
                flip_horizontal = True if np.random.rand(
jshilong's avatar
jshilong committed
224
                ) < self.flip_ratio_bev_horizontal else False
wuyuefeng's avatar
wuyuefeng committed
225
226
227
228
229
230
                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

231
232
233
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

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

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

zhangwenwei's avatar
zhangwenwei committed
249

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

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

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

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

    def __init__(self,
ZCMax's avatar
ZCMax committed
273
274
                 jitter_std: List[float] = [0.01, 0.01, 0.01],
                 clip_range: List[float] = [-0.05, 0.05]) -> None:
275
276
277
278
279
280
281
282
283
284
285
286
287
288
        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
289
    def transform(self, input_dict: dict) -> dict:
290
291
292
293
294
295
        """Call function to jitter all the points in the scene.

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

        Returns:
296
            dict: Results after adding noise to each point,
297
            'points' key is updated in the result dict.
298
299
300
301
302
303
304
305
306
307
308
309
        """
        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

310
    def __repr__(self) -> str:
311
312
313
314
315
316
317
        """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


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

322
323
324
325
326
327
328
329
330
331
    Required Keys:

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

    Modified Keys:
332

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

    Added Keys:

    - plane (optional)

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

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

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

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

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

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

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

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

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

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

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

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

            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
436

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

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

zhangwenwei's avatar
zhangwenwei committed
444
445
        return input_dict

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


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

459
460
461
462
463
464
465
466
467
468
    Required Keys:

    - points
    - gt_bboxes_3d

    Modified Keys:

    - points
    - gt_bboxes_3d

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

    def __init__(self,
482
483
484
485
                 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
486
        self.translation_std = translation_std
zhangwenwei's avatar
zhangwenwei committed
487
        self.global_rot_range = global_rot_range
zhangwenwei's avatar
zhangwenwei committed
488
        self.rot_range = rot_range
zhangwenwei's avatar
zhangwenwei committed
489
490
        self.num_try = num_try

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

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

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

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

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

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

520
    def __repr__(self) -> str:
521
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
522
        repr_str = self.__class__.__name__
523
524
525
526
        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
527
528
529
        return repr_str


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

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

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

545
    def __init__(self, rotation_axis: int) -> None:
546
547
        self.rotation_axis = rotation_axis

548
    def _trans_points(self, results: dict, trans_factor: np.ndarray) -> None:
549
550
551
552
553
554
555
556
557
        """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.
        """
558
        results['points'].translate(trans_factor)
559

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

573
    def _check_rot_mat(self, rot_mat: np.ndarray) -> None:
574
575
576
577
578
579
580
581
582
583
584
585
        """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}'

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

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

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

599
        axis_align_matrix = results['axis_align_matrix']
600
601
602
603
604
605
        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)
606
607
        self._rot_points(results, rot_mat)
        self._trans_points(results, trans_vec)
608

609
        return results
610

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


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

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

    def __init__(self,
jshilong's avatar
jshilong committed
655
656
657
658
                 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:
659
660
661
662
663
        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
664
        self.rot_range = rot_range
665
666
667

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

zhangwenwei's avatar
zhangwenwei committed
669
        self.scale_ratio_range = scale_ratio_range
670
671
672
673
674
675
676

        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
            ]
677
678
        assert all([std >= 0 for std in translation_std]), \
            'translation_std should be positive'
zhangwenwei's avatar
zhangwenwei committed
679
        self.translation_std = translation_std
wuyuefeng's avatar
wuyuefeng committed
680
        self.shift_height = shift_height
zhangwenwei's avatar
zhangwenwei committed
681

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

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

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

695
        input_dict['points'].translate(trans_factor)
zhangwenwei's avatar
zhangwenwei committed
696
        input_dict['pcd_trans'] = trans_factor
jshilong's avatar
jshilong committed
697
698
        if 'gt_bboxes_3d' in input_dict:
            input_dict['gt_bboxes_3d'].translate(trans_factor)
zhangwenwei's avatar
zhangwenwei committed
699

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

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

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

jshilong's avatar
jshilong committed
713
714
715
716
717
718
719
720
        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
721
            rot_mat_T = input_dict['points'].rotate(noise_rotation)
jshilong's avatar
jshilong committed
722
723
724
725
726

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

    def _scale_bbox_points(self, input_dict: dict) -> None:
727
728
729
730
731
732
        """Private function to scale bounding boxes and points.

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

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

jshilong's avatar
jshilong committed
745
746
747
        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
748

jshilong's avatar
jshilong committed
749
    def _random_scale(self, input_dict: dict) -> None:
750
751
752
753
754
755
        """Private function to randomly set the scale factor.

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

        Returns:
jshilong's avatar
jshilong committed
756
757
            dict: Results after scaling, 'pcd_scale_factor'
            are updated in the result dict.
758
        """
zhangwenwei's avatar
zhangwenwei committed
759
760
761
        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
762

jshilong's avatar
jshilong committed
763
    def transform(self, input_dict: dict) -> dict:
764
        """Private function to rotate, scale and translate bounding boxes and
765
766
767
768
769
770
771
        points.

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

        Returns:
            dict: Results after scaling, 'points', 'pcd_rotation',
772
            'pcd_scale_factor', 'pcd_trans' and `gt_bboxes_3d` are updated
jshilong's avatar
jshilong committed
773
            in the result dict.
774
        """
775
776
777
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

zhangwenwei's avatar
zhangwenwei committed
778
        self._rot_bbox_points(input_dict)
zhangwenwei's avatar
zhangwenwei committed
779

zhangwenwei's avatar
zhangwenwei committed
780
781
782
        if 'pcd_scale_factor' not in input_dict:
            self._random_scale(input_dict)
        self._scale_bbox_points(input_dict)
zhangwenwei's avatar
zhangwenwei committed
783

zhangwenwei's avatar
zhangwenwei committed
784
        self._trans_bbox_points(input_dict)
785
786

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

789
    def __repr__(self) -> str:
790
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
791
        repr_str = self.__class__.__name__
792
793
794
795
        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
796
797
798
        return repr_str


799
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
800
class PointShuffle(BaseTransform):
801
    """Shuffle input points."""
zhangwenwei's avatar
zhangwenwei committed
802

ZCMax's avatar
ZCMax committed
803
    def transform(self, input_dict: dict) -> dict:
804
805
806
807
808
809
        """Call function to shuffle points.

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

        Returns:
810
            dict: Results after filtering, 'points', 'pts_instance_mask'
811
            and 'pts_semantic_mask' keys are updated in the result dict.
812
        """
813
814
815
816
817
818
819
820
821
822
823
824
        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
825
826
        return input_dict

827
    def __repr__(self) -> str:
828
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
829
830
831
        return self.__class__.__name__


832
@TRANSFORMS.register_module()
833
class ObjectRangeFilter(BaseTransform):
834
835
    """Filter objects by the range.

836
837
838
839
840
841
842
843
    Required Keys:

    - gt_bboxes_3d

    Modified Keys:

    - gt_bboxes_3d

844
845
846
    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
847

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

851
852
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter objects by the range.
853
854
855
856
857

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

        Returns:
858
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d'
859
            keys are updated in the result dict.
860
        """
861
862
863
864
865
866
867
        # 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
868
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
869
        gt_labels_3d = input_dict['gt_labels_3d']
870
        mask = gt_bboxes_3d.in_range_bev(bev_range)
zhangwenwei's avatar
zhangwenwei committed
871
        gt_bboxes_3d = gt_bboxes_3d[mask]
ZwwWayne's avatar
ZwwWayne committed
872
873
874
875
876
        # mask is a torch tensor but gt_labels_3d is still numpy array
        # using mask to index gt_labels_3d will cause bug when
        # len(gt_labels_3d) == 1, where mask=1 will be interpreted
        # as gt_labels_3d[1] and cause out of index error
        gt_labels_3d = gt_labels_3d[mask.numpy().astype(np.bool)]
zhangwenwei's avatar
zhangwenwei committed
877
878

        # limit rad to [-pi, pi]
879
880
        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
881
882
        input_dict['gt_labels_3d'] = gt_labels_3d

zhangwenwei's avatar
zhangwenwei committed
883
884
        return input_dict

885
    def __repr__(self) -> str:
886
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
887
        repr_str = self.__class__.__name__
888
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
889
890
891
        return repr_str


892
@TRANSFORMS.register_module()
893
class PointsRangeFilter(BaseTransform):
894
895
    """Filter points by the range.

896
897
898
899
900
901
902
903
904
905
    Required Keys:

    - points
    - pts_instance_mask (optional)

    Modified Keys:

    - points
    - pts_instance_mask (optional)

906
907
908
    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
909

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

913
914
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter points by the range.
915
916
917
918
919

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

        Returns:
920
            dict: Results after filtering, 'points', 'pts_instance_mask'
921
            and 'pts_semantic_mask' keys are updated in the result dict.
922
        """
zhangwenwei's avatar
zhangwenwei committed
923
        points = input_dict['points']
924
925
        points_mask = points.in_range_3d(self.pcd_range)
        clean_points = points[points_mask]
zhangwenwei's avatar
zhangwenwei committed
926
        input_dict['points'] = clean_points
927
928
929
930
931
932
933
934
935
936
937
        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
938
939
        return input_dict

940
    def __repr__(self) -> str:
941
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
942
        repr_str = self.__class__.__name__
943
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
944
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
945
946


947
@TRANSFORMS.register_module()
948
class ObjectNameFilter(BaseTransform):
zhangwenwei's avatar
zhangwenwei committed
949
    """Filter GT objects by their names.
zhangwenwei's avatar
zhangwenwei committed
950

951
952
953
954
955
956
957
958
    Required Keys:

    - gt_labels_3d

    Modified Keys:

    - gt_labels_3d

zhangwenwei's avatar
zhangwenwei committed
959
    Args:
liyinhao's avatar
liyinhao committed
960
        classes (list[str]): List of class names to be kept for training.
zhangwenwei's avatar
zhangwenwei committed
961
962
    """

963
    def __init__(self, classes: List[str]) -> None:
zhangwenwei's avatar
zhangwenwei committed
964
965
966
        self.classes = classes
        self.labels = list(range(len(self.classes)))

967
968
    def transform(self, input_dict: dict) -> dict:
        """Transform function to filter objects by their names.
969
970
971
972
973

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

        Returns:
974
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d'
975
            keys are updated in the result dict.
976
        """
zhangwenwei's avatar
zhangwenwei committed
977
978
979
980
981
982
983
984
        gt_labels_3d = input_dict['gt_labels_3d']
        gt_bboxes_mask = np.array([n in self.labels for n in gt_labels_3d],
                                  dtype=np.bool_)
        input_dict['gt_bboxes_3d'] = input_dict['gt_bboxes_3d'][gt_bboxes_mask]
        input_dict['gt_labels_3d'] = input_dict['gt_labels_3d'][gt_bboxes_mask]

        return input_dict

985
    def __repr__(self) -> str:
986
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
987
988
989
        repr_str = self.__class__.__name__
        repr_str += f'(classes={self.classes})'
        return repr_str
wuyuefeng's avatar
wuyuefeng committed
990
991


992
993
@TRANSFORMS.register_module()
class PointSample(BaseTransform):
994
    """Point sample.
wuyuefeng's avatar
wuyuefeng committed
995
996
997

    Sampling data to a certain number.

998
    Required Keys:
999

1000
1001
1002
1003
1004
    - points
    - pts_instance_mask (optional)
    - pts_semantic_mask (optional)

    Modified Keys:
1005

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

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

1019
1020
    def __init__(self,
                 num_points: int,
1021
1022
                 sample_range: Optional[float] = None,
                 replace: bool = False) -> None:
wuyuefeng's avatar
wuyuefeng committed
1023
        self.num_points = num_points
1024
1025
1026
        self.sample_range = sample_range
        self.replace = replace

1027
1028
1029
1030
1031
1032
1033
1034
    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
1035
1036
1037
1038
1039
        """Points random sampling.

        Sample points to a certain number.

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

wuyuefeng's avatar
wuyuefeng committed
1048
        Returns:
1049
1050
1051
            tuple[:obj:`BasePoints`, np.ndarray] | :obj:`BasePoints`:

                - points (:obj:`BasePoints`): 3D Points.
1052
                - choices (np.ndarray, optional): The generated random samples.
wuyuefeng's avatar
wuyuefeng committed
1053
        """
1054
        if not replace:
wuyuefeng's avatar
wuyuefeng committed
1055
            replace = (points.shape[0] < num_samples)
1056
1057
1058
        point_range = range(len(points))
        if sample_range is not None and not replace:
            # Only sampling the near points when len(points) >= num_samples
1059
            dist = np.linalg.norm(points.coord.numpy(), axis=1)
1060
1061
            far_inds = np.where(dist >= sample_range)[0]
            near_inds = np.where(dist < sample_range)[0]
1062
1063
1064
1065
            # 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)
1066
1067
1068
1069
1070
1071
1072
            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
1073
1074
1075
1076
1077
        if return_choices:
            return points[choices], choices
        else:
            return points[choices]

1078
    def transform(self, input_dict: dict) -> dict:
1079
        """Transform function to sample points to in indoor scenes.
1080
1081
1082

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

1084
        Returns:
1085
            dict: Results after sampling, 'points', 'pts_instance_mask'
1086
            and 'pts_semantic_mask' keys are updated in the result dict.
1087
        """
1088
        points = input_dict['points']
1089
1090
1091
1092
1093
1094
        points, choices = self._points_random_sampling(
            points,
            self.num_points,
            self.sample_range,
            self.replace,
            return_choices=True)
1095
        input_dict['points'] = points
1096

1097
1098
        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
1099

1100
        if pts_instance_mask is not None:
wuyuefeng's avatar
wuyuefeng committed
1101
            pts_instance_mask = pts_instance_mask[choices]
1102
            input_dict['pts_instance_mask'] = pts_instance_mask
1103
1104
1105

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

1108
        return input_dict
wuyuefeng's avatar
wuyuefeng committed
1109

1110
    def __repr__(self) -> str:
1111
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
1112
        repr_str = self.__class__.__name__
1113
        repr_str += f'(num_points={self.num_points},'
1114
1115
        repr_str += f' sample_range={self.sample_range},'
        repr_str += f' replace={self.replace})'
1116

1117
1118
1119
        return repr_str


1120
@TRANSFORMS.register_module()
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
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)


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

    Note:
        This transform should only be used in the training process of point
1171
1172
1173
        cloud segmentation tasks. For the sliding patch generation and
        inference process in testing, please refer to the `slide_inference`
        function of `EncoderDecoder3D` class.
1174
1175
1176
    """

    def __init__(self,
ZCMax's avatar
ZCMax committed
1177
1178
1179
1180
1181
1182
1183
1184
1185
                 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:
1186
1187
1188
1189
1190
        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
1191
        self.enlarge_size = enlarge_size if enlarge_size is not None else 0.0
1192
        self.min_unique_num = min_unique_num
1193
        self.eps = eps
1194
1195
1196
1197
1198

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

ZCMax's avatar
ZCMax committed
1200
1201
1202
1203
    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:
1204
1205
        """Generating model input.

1206
        Generate input by subtracting patch center and adding additional
1207
1208
1209
1210
1211
1212
1213
1214
1215
            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.
1216
            point_type (type): class of input points inherited from BasePoints.
1217
1218

        Returns:
1219
            :obj:`BasePoints`: The generated input data.
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
        """
        # 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

1243
    def _patch_points_sampling(
1244
1245
            self, points: BasePoints,
            sem_mask: np.ndarray) -> Tuple[BasePoints, np.ndarray]:
1246
1247
1248
1249
1250
1251
        """Patch points sampling.

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

        Args:
1252
            points (:obj:`BasePoints`): 3D Points.
1253
1254
1255
            sem_mask (np.ndarray): semantic segmentation mask for input points.

        Returns:
1256
            tuple[:obj:`BasePoints`, np.ndarray]:
1257

1258
                - points (:obj:`BasePoints`): 3D Points.
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
                - 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)

1269
        for _ in range(self.num_try):
1270
1271
1272
            # random sample a point as patch center
            cur_center = coords[np.random.choice(coords.shape[0])]

1273
1274
            # boundary of a patch, which would be enlarged by
            # `self.enlarge_size` as an augmentation
1275
1276
1277
1278
1279
1280
1281
            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(
1282
1283
                (coords >= (cur_min - self.enlarge_size)) *
                (coords <= (cur_max + self.enlarge_size)),
1284
1285
1286
1287
1288
1289
1290
                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]
1291
            point_idxs = np.where(cur_choice)[0]
1292
            mask = np.sum(
1293
1294
                (cur_coords >= (cur_min - self.eps)) * (cur_coords <=
                                                        (cur_max + self.eps)),
1295
                axis=1) == 3
1296

1297
1298
            # two criteria for patch sampling, adopted from PointNet++
            # 1. selected patch should contain enough unique points
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
            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:
1311
                # if `min_unique_num` is provided, directly compare with it
1312
                flag1 = mask.sum() >= self.min_unique_num
1313

1314
            # 2. selected patch should contain enough annotated points
1315
1316
1317
1318
1319
1320
1321
1322
1323
            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

1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
        # 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]
1337
1338
1339
1340
1341
1342
1343
1344

        # 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
1345
    def transform(self, input_dict: dict) -> dict:
1346
1347
1348
1349
1350
1351
        """Call function to sample points to in indoor scenes.

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

        Returns:
1352
            dict: Results after sampling, 'points', 'pts_instance_mask'
1353
            and 'pts_semantic_mask' keys are updated in the result dict.
1354
        """
ZCMax's avatar
ZCMax committed
1355
        points = input_dict['points']
1356

ZCMax's avatar
ZCMax committed
1357
        assert 'pts_semantic_mask' in input_dict.keys(), \
1358
            'semantic mask should be provided in training and evaluation'
ZCMax's avatar
ZCMax committed
1359
        pts_semantic_mask = input_dict['pts_semantic_mask']
1360
1361
1362
1363

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

ZCMax's avatar
ZCMax committed
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
        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)

1374
        if pts_instance_mask is not None:
ZCMax's avatar
ZCMax committed
1375
1376
1377
1378
1379
            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]
1380

ZCMax's avatar
ZCMax committed
1381
        return input_dict
1382

1383
    def __repr__(self) -> str:
1384
1385
1386
1387
1388
1389
        """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},'
1390
1391
        repr_str += f' num_try={self.num_try},'
        repr_str += f' enlarge_size={self.enlarge_size},'
1392
1393
        repr_str += f' min_unique_num={self.min_unique_num},'
        repr_str += f' eps={self.eps})'
wuyuefeng's avatar
wuyuefeng committed
1394
        return repr_str
1395
1396


1397
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1398
class BackgroundPointsFilter(BaseTransform):
1399
1400
1401
    """Filter background points near the bounding box.

    Args:
1402
        bbox_enlarge_range (tuple[float] | float): Bbox enlarge range.
1403
1404
    """

ZCMax's avatar
ZCMax committed
1405
    def __init__(self, bbox_enlarge_range: Union[Tuple[float], float]) -> None:
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
        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
1416
    def transform(self, input_dict: dict) -> dict:
1417
1418
1419
1420
1421
1422
        """Call function to filter points by the range.

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

        Returns:
1423
            dict: Results after filtering, 'points', 'pts_instance_mask'
1424
            and 'pts_semantic_mask' keys are updated in the result dict.
1425
1426
1427
1428
        """
        points = input_dict['points']
        gt_bboxes_3d = input_dict['gt_bboxes_3d']

xiliu8006's avatar
xiliu8006 committed
1429
1430
1431
1432
        # 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()

1433
1434
        enlarged_gt_bboxes_3d = gt_bboxes_3d_np.copy()
        enlarged_gt_bboxes_3d[:, 3:6] += self.bbox_enlarge_range
xiliu8006's avatar
xiliu8006 committed
1435
        points_numpy = points.tensor.clone().numpy()
1436
1437
        foreground_masks = box_np_ops.points_in_rbbox(
            points_numpy, gt_bboxes_3d_np, origin=(0.5, 0.5, 0.5))
1438
        enlarge_foreground_masks = box_np_ops.points_in_rbbox(
1439
            points_numpy, enlarged_gt_bboxes_3d, origin=(0.5, 0.5, 0.5))
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
        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

1455
    def __repr__(self) -> str:
1456
1457
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
1458
        repr_str += f'(bbox_enlarge_range={self.bbox_enlarge_range.tolist()})'
1459
        return repr_str
1460
1461


1462
@TRANSFORMS.register_module()
1463
class VoxelBasedPointSampler(BaseTransform):
1464
1465
1466
1467
1468
1469
    """Voxel based point sampler.

    Apply voxel sampling to multiple sweep points.

    Args:
        cur_sweep_cfg (dict): Config for sampling current points.
1470
1471
        prev_sweep_cfg (dict, optional): Config for sampling previous points.
            Defaults to None.
1472
        time_dim (int): Index that indicate the time dimension
1473
            for input points. Defaults to 3.
1474
1475
    """

1476
1477
1478
1479
    def __init__(self,
                 cur_sweep_cfg: dict,
                 prev_sweep_cfg: Optional[dict] = None,
                 time_dim: int = 3) -> None:
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
        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

1492
    def _sample_points(self, points: np.ndarray, sampler: VoxelGenerator,
1493
                       point_dim: int) -> np.ndarray:
1494
1495
1496
1497
1498
1499
        """Sample points for each points subset.

        Args:
            points (np.ndarray): Points subset to be sampled.
            sampler (VoxelGenerator): Voxel based sampler for
                each points subset.
1500
            point_dim (int): The dimension of each points.
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518

        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

1519
    def transform(self, results: dict) -> dict:
1520
1521
1522
1523
1524
1525
        """Call function to sample points from multiple sweeps.

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

        Returns:
1526
            dict: Results after sampling, 'points', 'pts_instance_mask'
1527
            and 'pts_semantic_mask' keys are updated in the result dict.
1528
1529
1530
1531
1532
1533
1534
1535
1536
        """
        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
1537
1538
        points_numpy = points.tensor.numpy()
        extra_channel = [points_numpy]
1539
1540
1541
1542
1543
1544
1545
1546
1547
        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])

1548
        points_numpy = np.concatenate(extra_channel, axis=-1)
1549
1550
1551
1552
1553

        # 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.
1554
1555
1556
        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]
1557
1558
1559
1560
1561
1562
1563
1564
1565
        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,
1566
                                               points_numpy.shape[1])
1567
1568
1569
        if self.prev_voxel_generator is not None:
            prev_sweeps_points = self._sample_points(prev_sweeps_points,
                                                     self.prev_voxel_generator,
1570
                                                     points_numpy.shape[1])
1571

1572
1573
            points_numpy = np.concatenate(
                [cur_sweep_points, prev_sweeps_points], 0)
1574
        else:
1575
            points_numpy = cur_sweep_points
1576
1577

        if self.cur_voxel_generator._max_num_points == 1:
1578
1579
            points_numpy = points_numpy.squeeze(1)
        results['points'] = points.new_point(points_numpy[..., :original_dim])
1580

1581
        # Restore the corresponding seg and mask fields
1582
        for key, dim_index in map_fields2dim:
1583
            results[key] = points_numpy[..., dim_index]
1584
1585
1586

        return results

1587
    def __repr__(self) -> str:
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
1606
        """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
1607
1608


1609
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1610
class AffineResize(BaseTransform):
1611
1612
1613
1614
1615
1616
1617
1618
1619
1620
1621
1622
    """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.
1623
        bbox_clip_border (bool): Whether clip the objects
1624
1625
1626
            outside the border of the image. Defaults to True.
    """

ZCMax's avatar
ZCMax committed
1627
1628
1629
1630
    def __init__(self,
                 img_scale: Tuple,
                 down_ratio: int,
                 bbox_clip_border: bool = True) -> None:
1631
1632
1633
1634
1635

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

ZCMax's avatar
ZCMax committed
1636
    def transform(self, results: dict) -> dict:
1637
1638
1639
1640
1641
1642
1643
        """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'
1644
            keys are added in the result dict.
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
1678
1679
1680
        """
        # 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
1681
1682
        if 'gt_bboxes' in results:
            self._affine_bboxes(results, trans_affine)
1683

ZCMax's avatar
ZCMax committed
1684
1685
        if 'centers_2d' in results:
            centers2d = self._affine_transform(results['centers_2d'],
1686
1687
1688
1689
1690
                                               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
1691
1692
1693
1694
            results['centers_2d'] = centers2d[valid_index]

            if 'gt_bboxes' in results:
                results['gt_bboxes'] = results['gt_bboxes'][valid_index]
1695
1696
1697
                if 'gt_bboxes_labels' in results:
                    results['gt_bboxes_labels'] = results['gt_bboxes_labels'][
                        valid_index]
ZCMax's avatar
ZCMax committed
1698
1699
1700
1701
1702
1703
1704
1705
1706
1707
                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]
1708
1709
1710
1711
1712

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

        return results

ZCMax's avatar
ZCMax committed
1713
    def _affine_bboxes(self, results: dict, matrix: np.ndarray) -> None:
1714
1715
1716
1717
1718
1719
1720
1721
1722
        """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
1723
1724
1725
1726
1727
1728
1729
1730
1731
1732
1733
1734
        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:
1735
        """Affine transform bbox points to input image.
1736
1737
1738
1739
1740
1741
1742
1743
1744
1745
1746
1747
1748
1749
1750
1751
1752

        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
1753
1754
    def _get_transform_matrix(self, center: Tuple, scale: Tuple,
                              output_scale: Tuple[float]) -> np.ndarray:
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
        """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
1789
1790
    def _get_ref_point(self, ref_point1: np.ndarray,
                       ref_point2: np.ndarray) -> np.ndarray:
1791
        """Get reference point to calculate affine transform matrix.
1792
1793

        While using opencv to calculate the affine matrix, we need at least
1794
        three corresponding points separately on original image and target
1795
1796
1797
1798
1799
1800
        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

1801
    def __repr__(self) -> str:
1802
        """str: Return a string that describes the module."""
1803
1804
1805
1806
1807
1808
        repr_str = self.__class__.__name__
        repr_str += f'(img_scale={self.img_scale}, '
        repr_str += f'down_ratio={self.down_ratio}) '
        return repr_str


1809
@TRANSFORMS.register_module()
ZCMax's avatar
ZCMax committed
1810
class RandomShiftScale(BaseTransform):
1811
1812
1813
1814
    """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
1815
    infos into loading TRANSFORMS. It's designed to be used with
1816
1817
1818
1819
1820
1821
1822
    AffineResize together.

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

1823
    def __init__(self, shift_scale: Tuple[float], aug_prob: float) -> None:
1824
1825
1826
1827

        self.shift_scale = shift_scale
        self.aug_prob = aug_prob

ZCMax's avatar
ZCMax committed
1828
    def transform(self, results: dict) -> dict:
1829
1830
1831
1832
1833
1834
1835
        """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'
1836
            and 'affine_aug' keys are added in the result dict.
1837
1838
1839
1840
1841
1842
1843
1844
1845
1846
1847
1848
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
1860
        """
        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

1861
    def __repr__(self) -> str:
1862
        """str: Return a string that describes the module."""
1863
1864
1865
1866
        repr_str = self.__class__.__name__
        repr_str += f'(shift_scale={self.shift_scale}, '
        repr_str += f'aug_prob={self.aug_prob}) '
        return repr_str
1867
1868
1869
1870
1871


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

1872
    def _resize_3d(self, results: dict) -> None:
1873
1874
1875
1876
1877
1878
1879
1880
1881
1882
1883
1884
1885
        """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.
1886

1887
1888
1889
1890
1891
1892
1893
1894
1895
1896
1897
1898
1899
1900
1901
1902
1903
1904
1905
1906
1907
        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']``.
    """

1908
    def _resize_3d(self, results: dict) -> None:
1909
1910
1911
1912
1913
1914
1915
        """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])

1916
    def transform(self, results: dict) -> dict:
1917
1918
        """Transform function to resize images, bounding boxes, masks, semantic
        segmentation map. Compared to RandomResize, this function would further
1919
1920
1921
1922
        check if scale is already set in results.

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

1924
        Returns:
1925
1926
            dict: Resized results, 'img_shape', 'pad_shape', 'scale_factor',
            'keep_ratio' keys are added into result dict.
1927
1928
1929
1930
1931
1932
1933
1934
        """
        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
1935
1936
1937
1938
1939
1940
1941
1942
1943
1944
1945
1946
1947
1948
1949
1950
1951
1952
1953
1954
1955
1956
1957
1958
1959
1960
1961
1962
1963
1964
1965
1966
1967
1968
1969
1970
1971
1972
1973
1974
1975
1976
1977
1978
1979
1980
1981
1982
1983
1984
1985
1986
1987


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

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

    Required Keys:

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

    Modified Keys:

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

    Added Keys:

    - homography_matrix

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

    Note:
        - If the image is smaller than the absolute crop size, return the
1995
          original image.
1996
1997
1998
1999
2000
2001
2002
2003
        - 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.
    """

2004
2005
2006
2007
2008
2009
2010
2011
2012
2013
    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:
2014
2015
2016
2017
2018
2019
2020
2021
2022
2023
2024
        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

2025
2026
2027
2028
    def _crop_data(self,
                   results: dict,
                   crop_size: tuple,
                   allow_negative_crop: bool = False) -> dict:
2029
2030
2031
2032
2033
2034
2035
        """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
2036
                contain any bbox area. Defaults to False.
2037
2038
2039

        Returns:
            dict: Randomly cropped results, 'img_shape' key in result dict is
2040
            updated according to crop size.
2041
2042
2043
2044
2045
2046
2047
2048
2049
2050
2051
2052
2053
2054
2055
2056
2057
2058
2059
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070
2071
2072
2073
2074
2075
2076
2077
2078
2079
2080
2081
2082
2083
2084
2085
2086
2087
2088
2089
2090
2091
2092
2093
2094
2095
2096
2097
2098
2099
2100
2101
2102
2103
2104
2105
2106
2107
2108
2109
2110
2111
2112
2113
2114
2115
2116
2117
2118
2119
2120
2121
2122
        """
        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

2123
    def transform(self, results: dict) -> dict:
2124
2125
2126
2127
2128
2129
2130
2131
        """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
2132
            updated according to crop size.
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
        """
        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

2143
2144
    def __repr__(self) -> dict:
        """str: Return a string that describes the module."""
2145
2146
2147
2148
2149
2150
2151
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162
2163
2164
2165
2166
2167
2168
2169
2170
2171
2172
2173
2174
2175
2176
2177
2178
2179
2180
2181
2182
2183
2184
2185
2186
2187
2188
2189
2190
2191
2192
2193
2194
2195
2196
2197
2198
2199
2200
2201
2202
2203
2204
2205
2206
2207
2208
2209
2210
2211
2212
2213
2214
2215
2216
2217
2218
2219
2220
2221
2222
2223
2224
2225
2226
2227
2228
2229
2230
2231
2232
2233
2234
2235
2236
2237
2238
2239
2240
2241
2242
2243
2244
2245
2246
2247
2248
2249
        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()
2250
class MultiViewWrapper(BaseTransform):
2251
2252
2253
2254
2255
2256
2257
2258
2259
2260
2261
2262
2263
2264
    """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
2265
            for multiview image. Defaults to True.
2266
        process_fields (list): Desired keys that the transformations should
2267
            be conducted on. Defaults to ['img', 'cam2img', 'lidar2cam'].
2268
        collected_keys (list): Collect information in transformation
2269
            like rotate angles, crop roi, and flip state. Defaults to
2270
2271
2272
2273
                ['scale', 'scale_factor', 'crop',
                 'crop_offset', 'ori_shape',
                 'pad_shape', 'img_shape',
                 'pad_fixed_size', 'pad_size_divisor',
2274
                 'flip', 'flip_direction', 'rotate'].
2275
        randomness_keys (list): The keys that related to the randomness
2276
            in transformation. Defaults to
2277
2278
2279
2280
                    ['scale', 'scale_factor', 'crop_size', 'flip',
                     'flip_direction', 'photometric_param']
    """

2281
2282
2283
2284
2285
2286
2287
2288
2289
2290
2291
2292
2293
2294
2295
    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:
2296
        self.transforms = Compose(transforms)
2297
2298
2299
2300
2301
        self.override_aug_config = override_aug_config
        self.collected_keys = collected_keys
        self.process_fields = process_fields
        self.randomness_keys = randomness_keys

2302
    def transform(self, input_dict: dict) -> dict:
2303
2304
2305
2306
2307
2308
2309
2310
2311
2312
2313
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
        """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]
2331
            process_dict = self.transforms(process_dict)
2332
2333
2334
2335
2336
2337
2338
2339
2340
2341
2342
2343
2344
2345
2346
2347
2348
2349
2350
            # 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