transforms_3d.py 41.1 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
import numpy as np
2
from mmcv import is_tuple_of
3
from mmcv.utils import build_from_cfg
zhangwenwei's avatar
zhangwenwei committed
4

5
from mmdet3d.core import VoxelGenerator
zhangwenwei's avatar
zhangwenwei committed
6
from mmdet3d.core.bbox import box_np_ops
7
from mmdet.datasets.builder import PIPELINES
zhangwenwei's avatar
zhangwenwei committed
8
from mmdet.datasets.pipelines import RandomFlip
9
from ..builder import OBJECTSAMPLERS
zhangwenwei's avatar
zhangwenwei committed
10
11
12
from .data_augment_utils import noise_per_object_v3_


13
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
14
15
16
17
18
19
20
21
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.

    Args:
zhangwenwei's avatar
zhangwenwei committed
22
23
24
        sync_2d (bool, optional): Whether to apply flip according to the 2D
            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
25
            to that of 2D images. Defaults to True.
wuyuefeng's avatar
wuyuefeng committed
26
        flip_ratio_bev_horizontal (float, optional): The flipping probability
liyinhao's avatar
liyinhao committed
27
            in horizontal direction. Defaults to 0.0.
wuyuefeng's avatar
wuyuefeng committed
28
        flip_ratio_bev_vertical (float, optional): The flipping probability
liyinhao's avatar
liyinhao committed
29
            in vertical direction. Defaults to 0.0.
zhangwenwei's avatar
zhangwenwei committed
30
31
    """

wuyuefeng's avatar
wuyuefeng committed
32
33
34
35
36
37
38
    def __init__(self,
                 sync_2d=True,
                 flip_ratio_bev_horizontal=0.0,
                 flip_ratio_bev_vertical=0.0,
                 **kwargs):
        super(RandomFlip3D, self).__init__(
            flip_ratio=flip_ratio_bev_horizontal, **kwargs)
zhangwenwei's avatar
zhangwenwei committed
39
        self.sync_2d = sync_2d
wuyuefeng's avatar
wuyuefeng committed
40
41
42
43
44
45
46
47
48
49
50
        self.flip_ratio_bev_vertical = flip_ratio_bev_vertical
        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

    def random_flip_data_3d(self, input_dict, direction='horizontal'):
51
52
53
54
55
56
57
58
59
60
        """Flip 3D data randomly.

        Args:
            input_dict (dict): Result dict from loading pipeline.
            direction (str): Flip direction. Default: horizontal.

        Returns:
            dict: Flipped results, 'points', 'bbox3d_fields' keys are \
                updated in the result dict.
        """
wuyuefeng's avatar
wuyuefeng committed
61
        assert direction in ['horizontal', 'vertical']
62
63
64
65
66
        if len(input_dict['bbox3d_fields']) == 0:  # test mode
            input_dict['bbox3d_fields'].append('empty_box3d')
            input_dict['empty_box3d'] = input_dict['box_type_3d'](
                np.array([], dtype=np.float32))
        assert len(input_dict['bbox3d_fields']) == 1
zhangwenwei's avatar
zhangwenwei committed
67
        for key in input_dict['bbox3d_fields']:
68
69
70
71
72
73
74
75
76
77
78
            if 'points' in input_dict:
                input_dict['points'] = input_dict[key].flip(
                    direction, points=input_dict['points'])
            else:
                input_dict[key].flip(direction)
        if 'centers2d' in input_dict:
            assert self.sync_2d is True and direction == 'horizontal', \
                'Only support sync_2d=True and horizontal flip with images'
            w = input_dict['img_shape'][1]
            input_dict['centers2d'][..., 0] = \
                w - input_dict['centers2d'][..., 0]
zhangwenwei's avatar
zhangwenwei committed
79
80

    def __call__(self, input_dict):
81
82
83
84
85
86
87
88
89
90
91
        """Call function to flip points, values in the ``bbox3d_fields`` and \
        also flip 2D image and its annotations.

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

        Returns:
            dict: Flipped results, 'flip', 'flip_direction', \
                'pcd_horizontal_flip' and 'pcd_vertical_flip' keys are added \
                into result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
92
        # filp 2D image and its annotations
zhangwenwei's avatar
zhangwenwei committed
93
        super(RandomFlip3D, self).__call__(input_dict)
zhangwenwei's avatar
zhangwenwei committed
94

zhangwenwei's avatar
zhangwenwei committed
95
        if self.sync_2d:
wuyuefeng's avatar
wuyuefeng committed
96
97
            input_dict['pcd_horizontal_flip'] = input_dict['flip']
            input_dict['pcd_vertical_flip'] = False
zhangwenwei's avatar
zhangwenwei committed
98
        else:
wuyuefeng's avatar
wuyuefeng committed
99
100
101
102
103
104
105
106
107
            if 'pcd_horizontal_flip' not in input_dict:
                flip_horizontal = True if np.random.rand(
                ) < self.flip_ratio else False
                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

108
109
110
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

wuyuefeng's avatar
wuyuefeng committed
111
112
        if input_dict['pcd_horizontal_flip']:
            self.random_flip_data_3d(input_dict, 'horizontal')
113
            input_dict['transformation_3d_flow'].extend(['HF'])
wuyuefeng's avatar
wuyuefeng committed
114
115
        if input_dict['pcd_vertical_flip']:
            self.random_flip_data_3d(input_dict, 'vertical')
116
            input_dict['transformation_3d_flow'].extend(['VF'])
zhangwenwei's avatar
zhangwenwei committed
117
118
        return input_dict

zhangwenwei's avatar
zhangwenwei committed
119
    def __repr__(self):
120
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
121
        repr_str = self.__class__.__name__
122
        repr_str += f'(sync_2d={self.sync_2d},'
123
        repr_str += f' flip_ratio_bev_vertical={self.flip_ratio_bev_vertical})'
wuyuefeng's avatar
wuyuefeng committed
124
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
125

zhangwenwei's avatar
zhangwenwei committed
126

127
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
128
class ObjectSample(object):
zhangwenwei's avatar
zhangwenwei committed
129
    """Sample GT objects to the data.
zhangwenwei's avatar
zhangwenwei committed
130
131
132
133
134

    Args:
        db_sampler (dict): Config dict of the database sampler.
        sample_2d (bool): Whether to also paste 2D image patch to the images
            This should be true when applying multi-modality cut-and-paste.
liyinhao's avatar
liyinhao committed
135
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
136
    """
zhangwenwei's avatar
zhangwenwei committed
137
138
139
140
141
142
143
144
145
146

    def __init__(self, db_sampler, sample_2d=False):
        self.sampler_cfg = db_sampler
        self.sample_2d = sample_2d
        if 'type' not in db_sampler.keys():
            db_sampler['type'] = 'DataBaseSampler'
        self.db_sampler = build_from_cfg(db_sampler, OBJECTSAMPLERS)

    @staticmethod
    def remove_points_in_boxes(points, boxes):
147
148
149
        """Remove the points in the sampled bounding boxes.

        Args:
150
            points (:obj:`BasePoints`): Input point cloud array.
151
152
153
154
155
            boxes (np.ndarray): Sampled ground truth boxes.

        Returns:
            np.ndarray: Points with those in the boxes removed.
        """
156
        masks = box_np_ops.points_in_rbbox(points.coord.numpy(), boxes)
zhangwenwei's avatar
zhangwenwei committed
157
158
159
160
        points = points[np.logical_not(masks.any(-1))]
        return points

    def __call__(self, input_dict):
161
162
163
164
165
166
167
168
169
170
        """Call function to sample ground truth objects to the data.

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

        Returns:
            dict: Results after object sampling augmentation, \
                'points', 'gt_bboxes_3d', 'gt_labels_3d' keys are updated \
                in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
171
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
172
173
        gt_labels_3d = input_dict['gt_labels_3d']

zhangwenwei's avatar
zhangwenwei committed
174
175
176
        # change to float for blending operation
        points = input_dict['points']
        if self.sample_2d:
wuyuefeng's avatar
wuyuefeng committed
177
            img = input_dict['img']
zhangwenwei's avatar
zhangwenwei committed
178
179
180
            gt_bboxes_2d = input_dict['gt_bboxes']
            # Assume for now 3D & 2D bboxes are the same
            sampled_dict = self.db_sampler.sample_all(
181
182
183
184
                gt_bboxes_3d.tensor.numpy(),
                gt_labels_3d,
                gt_bboxes_2d=gt_bboxes_2d,
                img=img)
zhangwenwei's avatar
zhangwenwei committed
185
186
        else:
            sampled_dict = self.db_sampler.sample_all(
187
                gt_bboxes_3d.tensor.numpy(), gt_labels_3d, img=None)
zhangwenwei's avatar
zhangwenwei committed
188
189
190
191

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

zhangwenwei's avatar
zhangwenwei committed
194
195
            gt_labels_3d = np.concatenate([gt_labels_3d, sampled_gt_labels],
                                          axis=0)
196
197
198
            gt_bboxes_3d = gt_bboxes_3d.new_box(
                np.concatenate(
                    [gt_bboxes_3d.tensor.numpy(), sampled_gt_bboxes_3d]))
zhangwenwei's avatar
zhangwenwei committed
199

zhangwenwei's avatar
zhangwenwei committed
200
201
            points = self.remove_points_in_boxes(points, sampled_gt_bboxes_3d)
            # check the points dimension
202
            points = points.cat([sampled_points, points])
zhangwenwei's avatar
zhangwenwei committed
203
204
205
206
207

            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
208

zhangwenwei's avatar
zhangwenwei committed
209
                input_dict['gt_bboxes'] = gt_bboxes_2d
wuyuefeng's avatar
wuyuefeng committed
210
                input_dict['img'] = sampled_dict['img']
zhangwenwei's avatar
zhangwenwei committed
211
212

        input_dict['gt_bboxes_3d'] = gt_bboxes_3d
213
        input_dict['gt_labels_3d'] = gt_labels_3d.astype(np.long)
zhangwenwei's avatar
zhangwenwei committed
214
        input_dict['points'] = points
zhangwenwei's avatar
zhangwenwei committed
215

zhangwenwei's avatar
zhangwenwei committed
216
217
218
        return input_dict

    def __repr__(self):
219
        """str: Return a string that describes the module."""
220
221
222
223
224
225
226
227
228
        repr_str = self.__class__.__name__
        repr_str += f' sample_2d={self.sample_2d},'
        repr_str += f' data_root={self.sampler_cfg.data_root},'
        repr_str += f' info_path={self.sampler_cfg.info_path},'
        repr_str += f' rate={self.sampler_cfg.rate},'
        repr_str += f' prepare={self.sampler_cfg.prepare},'
        repr_str += f' classes={self.sampler_cfg.classes},'
        repr_str += f' sample_groups={self.sampler_cfg.sample_groups}'
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
229
230


231
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
232
class ObjectNoise(object):
zhangwenwei's avatar
zhangwenwei committed
233
    """Apply noise to each GT objects in the scene.
zhangwenwei's avatar
zhangwenwei committed
234
235

    Args:
236
        translation_std (list[float], optional): Standard deviation of the
zhangwenwei's avatar
zhangwenwei committed
237
238
            distribution where translation noise are sampled from.
            Defaults to [0.25, 0.25, 0.25].
239
        global_rot_range (list[float], optional): Global rotation to the scene.
zhangwenwei's avatar
zhangwenwei committed
240
            Defaults to [0.0, 0.0].
241
        rot_range (list[float], optional): Object rotation range.
zhangwenwei's avatar
zhangwenwei committed
242
243
244
245
            Defaults to [-0.15707963267, 0.15707963267].
        num_try (int, optional): Number of times to try if the noise applied is
            invalid. Defaults to 100.
    """
zhangwenwei's avatar
zhangwenwei committed
246
247

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
248
                 translation_std=[0.25, 0.25, 0.25],
zhangwenwei's avatar
zhangwenwei committed
249
                 global_rot_range=[0.0, 0.0],
zhangwenwei's avatar
zhangwenwei committed
250
                 rot_range=[-0.15707963267, 0.15707963267],
zhangwenwei's avatar
zhangwenwei committed
251
                 num_try=100):
zhangwenwei's avatar
zhangwenwei committed
252
        self.translation_std = translation_std
zhangwenwei's avatar
zhangwenwei committed
253
        self.global_rot_range = global_rot_range
zhangwenwei's avatar
zhangwenwei committed
254
        self.rot_range = rot_range
zhangwenwei's avatar
zhangwenwei committed
255
256
257
        self.num_try = num_try

    def __call__(self, input_dict):
258
259
260
261
262
263
264
265
266
        """Call function to apply noise to each ground truth in the scene.

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

        Returns:
            dict: Results after adding noise to each object, \
                'points', 'gt_bboxes_3d' keys are updated in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
267
268
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
        points = input_dict['points']
zhangwenwei's avatar
zhangwenwei committed
269

zhangwenwei's avatar
zhangwenwei committed
270
        # TODO: check this inplace function
271
        numpy_box = gt_bboxes_3d.tensor.numpy()
272
273
        numpy_points = points.tensor.numpy()

zhangwenwei's avatar
zhangwenwei committed
274
        noise_per_object_v3_(
275
            numpy_box,
276
            numpy_points,
zhangwenwei's avatar
zhangwenwei committed
277
278
            rotation_perturb=self.rot_range,
            center_noise_std=self.translation_std,
zhangwenwei's avatar
zhangwenwei committed
279
280
            global_random_rot_range=self.global_rot_range,
            num_try=self.num_try)
281
282

        input_dict['gt_bboxes_3d'] = gt_bboxes_3d.new_box(numpy_box)
283
        input_dict['points'] = points.new_point(numpy_points)
zhangwenwei's avatar
zhangwenwei committed
284
285
286
        return input_dict

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


296
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
297
class GlobalRotScaleTrans(object):
zhangwenwei's avatar
zhangwenwei committed
298
    """Apply global rotation, scaling and translation to a 3D scene.
zhangwenwei's avatar
zhangwenwei committed
299
300
301

    Args:
        rot_range (list[float]): Range of rotation angle.
liyinhao's avatar
liyinhao committed
302
            Defaults to [-0.78539816, 0.78539816] (close to [-pi/4, pi/4]).
zhangwenwei's avatar
zhangwenwei committed
303
        scale_ratio_range (list[float]): Range of scale ratio.
liyinhao's avatar
liyinhao committed
304
            Defaults to [0.95, 1.05].
zhangwenwei's avatar
zhangwenwei committed
305
306
307
        translation_std (list[float]): The standard deviation of ranslation
            noise. This apply random translation to a scene by a noise, which
            is sampled from a gaussian distribution whose standard deviation
liyinhao's avatar
liyinhao committed
308
309
            is set by ``translation_std``. Defaults to [0, 0, 0]
        shift_height (bool): Whether to shift height.
wuyuefeng's avatar
wuyuefeng committed
310
            (the fourth dimension of indoor points) when scaling.
liyinhao's avatar
liyinhao committed
311
            Defaults to False.
zhangwenwei's avatar
zhangwenwei committed
312
    """
zhangwenwei's avatar
zhangwenwei committed
313
314

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
315
316
                 rot_range=[-0.78539816, 0.78539816],
                 scale_ratio_range=[0.95, 1.05],
wuyuefeng's avatar
wuyuefeng committed
317
318
                 translation_std=[0, 0, 0],
                 shift_height=False):
zhangwenwei's avatar
zhangwenwei committed
319
320
321
        self.rot_range = rot_range
        self.scale_ratio_range = scale_ratio_range
        self.translation_std = translation_std
wuyuefeng's avatar
wuyuefeng committed
322
        self.shift_height = shift_height
zhangwenwei's avatar
zhangwenwei committed
323
324

    def _trans_bbox_points(self, input_dict):
325
326
327
328
329
330
331
332
333
334
        """Private function to translate bounding boxes and points.

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

        Returns:
            dict: Results after translation, 'points', 'pcd_trans' \
                and keys in input_dict['bbox3d_fields'] are updated \
                in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
335
336
337
338
339
340
341
342
343
344
        if not isinstance(self.translation_std, (list, tuple, np.ndarray)):
            translation_std = [
                self.translation_std, self.translation_std,
                self.translation_std
            ]
        else:
            translation_std = self.translation_std
        translation_std = np.array(translation_std, dtype=np.float32)
        trans_factor = np.random.normal(scale=translation_std, size=3).T

345
        input_dict['points'].translate(trans_factor)
zhangwenwei's avatar
zhangwenwei committed
346
347
348
349
350
        input_dict['pcd_trans'] = trans_factor
        for key in input_dict['bbox3d_fields']:
            input_dict[key].translate(trans_factor)

    def _rot_bbox_points(self, input_dict):
351
352
353
354
355
356
357
358
359
360
        """Private function to rotate bounding boxes and points.

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

        Returns:
            dict: Results after rotation, 'points', 'pcd_rotation' \
                and keys in input_dict['bbox3d_fields'] are updated \
                in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
361
        rotation = self.rot_range
zhangwenwei's avatar
zhangwenwei committed
362
363
364
        if not isinstance(rotation, list):
            rotation = [-rotation, rotation]
        noise_rotation = np.random.uniform(rotation[0], rotation[1])
zhangwenwei's avatar
zhangwenwei committed
365
366

        for key in input_dict['bbox3d_fields']:
wuyuefeng's avatar
wuyuefeng committed
367
368
369
370
371
            if len(input_dict[key].tensor) != 0:
                points, rot_mat_T = input_dict[key].rotate(
                    noise_rotation, input_dict['points'])
                input_dict['points'] = points
                input_dict['pcd_rotation'] = rot_mat_T
372
        # input_dict['points_instance'].rotate(noise_rotation)
373

zhangwenwei's avatar
zhangwenwei committed
374
    def _scale_bbox_points(self, input_dict):
375
376
377
378
379
380
381
382
383
        """Private function to scale bounding boxes and points.

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

        Returns:
            dict: Results after scaling, 'points'and keys in \
                input_dict['bbox3d_fields'] are updated in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
384
        scale = input_dict['pcd_scale_factor']
385
386
        points = input_dict['points']
        points.scale(scale)
wuyuefeng's avatar
wuyuefeng committed
387
        if self.shift_height:
388
389
390
            assert 'height' in points.attribute_dims.keys()
            points.tensor[:, points.attribute_dims['height']] *= scale
        input_dict['points'] = points
wuyuefeng's avatar
wuyuefeng committed
391

zhangwenwei's avatar
zhangwenwei committed
392
393
        for key in input_dict['bbox3d_fields']:
            input_dict[key].scale(scale)
zhangwenwei's avatar
zhangwenwei committed
394

zhangwenwei's avatar
zhangwenwei committed
395
    def _random_scale(self, input_dict):
396
397
398
399
400
401
402
403
404
        """Private function to randomly set the scale factor.

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

        Returns:
            dict: Results after scaling, 'pcd_scale_factor' are updated \
                in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
405
406
407
        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
408
409

    def __call__(self, input_dict):
410
411
412
413
414
415
416
417
418
419
420
        """Private function to rotate, scale and translate bounding boxes and \
        points.

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

        Returns:
            dict: Results after scaling, 'points', 'pcd_rotation',
                'pcd_scale_factor', 'pcd_trans' and keys in \
                input_dict['bbox3d_fields'] are updated in the result dict.
        """
421
422
423
        if 'transformation_3d_flow' not in input_dict:
            input_dict['transformation_3d_flow'] = []

zhangwenwei's avatar
zhangwenwei committed
424
        self._rot_bbox_points(input_dict)
zhangwenwei's avatar
zhangwenwei committed
425

zhangwenwei's avatar
zhangwenwei committed
426
427
428
        if 'pcd_scale_factor' not in input_dict:
            self._random_scale(input_dict)
        self._scale_bbox_points(input_dict)
zhangwenwei's avatar
zhangwenwei committed
429

zhangwenwei's avatar
zhangwenwei committed
430
        self._trans_bbox_points(input_dict)
431
432

        input_dict['transformation_3d_flow'].extend(['R', 'S', 'T'])
zhangwenwei's avatar
zhangwenwei committed
433
434
435
        return input_dict

    def __repr__(self):
436
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
437
        repr_str = self.__class__.__name__
438
439
440
441
        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
442
443
444
        return repr_str


445
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
446
class PointShuffle(object):
447
    """Shuffle input points."""
zhangwenwei's avatar
zhangwenwei committed
448
449

    def __call__(self, input_dict):
450
451
452
453
454
455
        """Call function to shuffle points.

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

        Returns:
456
457
            dict: Results after filtering, 'points', 'pts_instance_mask' \
                and 'pts_semantic_mask' keys are updated in the result dict.
458
        """
459
460
461
462
463
464
465
466
467
468
469
470
        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
471
472
473
474
475
476
        return input_dict

    def __repr__(self):
        return self.__class__.__name__


477
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
478
class ObjectRangeFilter(object):
479
480
481
482
483
    """Filter objects by the range.

    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
484
485
486
487
488
489

    def __init__(self, point_cloud_range):
        self.pcd_range = np.array(point_cloud_range, dtype=np.float32)
        self.bev_range = self.pcd_range[[0, 1, 3, 4]]

    def __call__(self, input_dict):
490
491
492
493
494
495
496
497
498
        """Call function to filter objects by the range.

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

        Returns:
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d' \
                keys are updated in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
499
        gt_bboxes_3d = input_dict['gt_bboxes_3d']
zhangwenwei's avatar
zhangwenwei committed
500
        gt_labels_3d = input_dict['gt_labels_3d']
501
        mask = gt_bboxes_3d.in_range_bev(self.bev_range)
zhangwenwei's avatar
zhangwenwei committed
502
        gt_bboxes_3d = gt_bboxes_3d[mask]
ZwwWayne's avatar
ZwwWayne committed
503
504
505
506
507
        # 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
508
509

        # limit rad to [-pi, pi]
510
511
        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
512
513
        input_dict['gt_labels_3d'] = gt_labels_3d

zhangwenwei's avatar
zhangwenwei committed
514
515
516
        return input_dict

    def __repr__(self):
517
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
518
        repr_str = self.__class__.__name__
519
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
520
521
522
        return repr_str


523
@PIPELINES.register_module()
zhangwenwei's avatar
zhangwenwei committed
524
class PointsRangeFilter(object):
525
526
527
528
529
    """Filter points by the range.

    Args:
        point_cloud_range (list[float]): Point cloud range.
    """
zhangwenwei's avatar
zhangwenwei committed
530
531

    def __init__(self, point_cloud_range):
532
        self.pcd_range = np.array(point_cloud_range, dtype=np.float32)
zhangwenwei's avatar
zhangwenwei committed
533
534

    def __call__(self, input_dict):
535
536
537
538
539
540
        """Call function to filter points by the range.

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

        Returns:
541
542
            dict: Results after filtering, 'points', 'pts_instance_mask' \
                and 'pts_semantic_mask' keys are updated in the result dict.
543
        """
zhangwenwei's avatar
zhangwenwei committed
544
        points = input_dict['points']
545
546
        points_mask = points.in_range_3d(self.pcd_range)
        clean_points = points[points_mask]
zhangwenwei's avatar
zhangwenwei committed
547
        input_dict['points'] = clean_points
548
549
550
551
552
553
554
555
556
557
558
        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
559
560
561
        return input_dict

    def __repr__(self):
562
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
563
        repr_str = self.__class__.__name__
564
        repr_str += f'(point_cloud_range={self.pcd_range.tolist()})'
zhangwenwei's avatar
zhangwenwei committed
565
        return repr_str
zhangwenwei's avatar
zhangwenwei committed
566
567
568
569


@PIPELINES.register_module()
class ObjectNameFilter(object):
zhangwenwei's avatar
zhangwenwei committed
570
    """Filter GT objects by their names.
zhangwenwei's avatar
zhangwenwei committed
571
572

    Args:
liyinhao's avatar
liyinhao committed
573
        classes (list[str]): List of class names to be kept for training.
zhangwenwei's avatar
zhangwenwei committed
574
575
576
577
578
579
580
    """

    def __init__(self, classes):
        self.classes = classes
        self.labels = list(range(len(self.classes)))

    def __call__(self, input_dict):
581
582
583
584
585
586
587
588
589
        """Call function to filter objects by their names.

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

        Returns:
            dict: Results after filtering, 'gt_bboxes_3d', 'gt_labels_3d' \
                keys are updated in the result dict.
        """
zhangwenwei's avatar
zhangwenwei committed
590
591
592
593
594
595
596
597
598
        gt_labels_3d = input_dict['gt_labels_3d']
        gt_bboxes_mask = np.array([n in self.labels for n in gt_labels_3d],
                                  dtype=np.bool_)
        input_dict['gt_bboxes_3d'] = input_dict['gt_bboxes_3d'][gt_bboxes_mask]
        input_dict['gt_labels_3d'] = input_dict['gt_labels_3d'][gt_bboxes_mask]

        return input_dict

    def __repr__(self):
599
        """str: Return a string that describes the module."""
zhangwenwei's avatar
zhangwenwei committed
600
601
602
        repr_str = self.__class__.__name__
        repr_str += f'(classes={self.classes})'
        return repr_str
wuyuefeng's avatar
wuyuefeng committed
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628


@PIPELINES.register_module()
class IndoorPointSample(object):
    """Indoor point sample.

    Sampling data to a certain number.

    Args:
        name (str): Name of the dataset.
        num_points (int): Number of points to be sampled.
    """

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

    def points_random_sampling(self,
                               points,
                               num_samples,
                               replace=None,
                               return_choices=False):
        """Points random sampling.

        Sample points to a certain number.

        Args:
629
            points (np.ndarray | :obj:`BasePoints`): 3D Points.
wuyuefeng's avatar
wuyuefeng committed
630
631
            num_samples (int): Number of samples to be sampled.
            replace (bool): Whether the sample is with or without replacement.
liyinhao's avatar
liyinhao committed
632
633
            Defaults to None.
            return_choices (bool): Whether return choice. Defaults to False.
wuyuefeng's avatar
wuyuefeng committed
634
635

        Returns:
636
637
            tuple[np.ndarray] | np.ndarray:

638
                - points (np.ndarray | :obj:`BasePoints`): 3D Points.
639
                - choices (np.ndarray, optional): The generated random samples.
wuyuefeng's avatar
wuyuefeng committed
640
641
642
643
644
645
646
647
648
649
650
        """
        if replace is None:
            replace = (points.shape[0] < num_samples)
        choices = np.random.choice(
            points.shape[0], num_samples, replace=replace)
        if return_choices:
            return points[choices], choices
        else:
            return points[choices]

    def __call__(self, results):
651
652
653
654
655
656
657
658
659
        """Call function to sample points to in indoor scenes.

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

        Returns:
            dict: Results after sampling, 'points', 'pts_instance_mask' \
                and 'pts_semantic_mask' keys are updated in the result dict.
        """
wuyuefeng's avatar
wuyuefeng committed
660
661
662
        points = results['points']
        points, choices = self.points_random_sampling(
            points, self.num_points, return_choices=True)
663
        results['points'] = points
664

wuyuefeng's avatar
wuyuefeng committed
665
666
667
        pts_instance_mask = results.get('pts_instance_mask', None)
        pts_semantic_mask = results.get('pts_semantic_mask', None)

668
        if pts_instance_mask is not None:
wuyuefeng's avatar
wuyuefeng committed
669
670
            pts_instance_mask = pts_instance_mask[choices]
            results['pts_instance_mask'] = pts_instance_mask
671
672
673

        if pts_semantic_mask is not None:
            pts_semantic_mask = pts_semantic_mask[choices]
wuyuefeng's avatar
wuyuefeng committed
674
675
676
677
678
            results['pts_semantic_mask'] = pts_semantic_mask

        return results

    def __repr__(self):
679
        """str: Return a string that describes the module."""
wuyuefeng's avatar
wuyuefeng committed
680
        repr_str = self.__class__.__name__
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
        repr_str += f'(num_points={self.num_points})'
        return repr_str


@PIPELINES.register_module()
class IndoorPatchPointSample(object):
    r"""Indoor point sample within a patch. Modified from `PointNet++ <https://
    github.com/charlesq34/pointnet2/blob/master/scannet/scannet_dataset.py>`_.

    Sampling data to a certain number for semantic segmentation.

    Args:
        num_points (int): Number of points to be sampled.
        block_size (float, optional): Size of a block to sample points from.
            Defaults to 1.5.
        sample_rate (float, optional): Stride used in sliding patch generation.
            Defaults to 1.0.
        ignore_index (int, optional): Label index that won't be used for the
            segmentation task. This is set in PointSegClassMapping as neg_cls.
            Defaults to None.
        use_normalized_coord (bool, optional): Whether to use normalized xyz as
            additional features. Defaults to False.
        num_try (int, optional): Number of times to try if the patch selected
            is invalid. Defaults to 10.
    """

    def __init__(self,
                 num_points,
                 block_size=1.5,
                 sample_rate=1.0,
                 ignore_index=None,
                 use_normalized_coord=False,
                 num_try=10):
        self.num_points = num_points
        self.block_size = block_size
        self.sample_rate = sample_rate
        self.ignore_index = ignore_index
        self.use_normalized_coord = use_normalized_coord
        self.num_try = num_try

    def _input_generation(self, coords, patch_center, coord_max, attributes,
                          attribute_dims, point_type):
        """Generating model input.

        Generate input by subtracting patch center and adding additional \
            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.
735
            point_type (type): class of input points inherited from BasePoints.
736
737

        Returns:
738
            :obj:`BasePoints`: The generated input data.
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
        """
        # 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

    def _patch_points_sampling(self, points, sem_mask, replace=None):
        """Patch points sampling.

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

        Args:
769
            points (:obj:`BasePoints`): 3D Points.
770
771
772
773
774
            sem_mask (np.ndarray): semantic segmentation mask for input points.
            replace (bool): Whether the sample is with or without replacement.
                Defaults to None.

        Returns:
775
            tuple[:obj:`BasePoints`, np.ndarray] | :obj:`BasePoints`:
776

777
                - points (:obj:`BasePoints`): 3D Points.
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
                - 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)

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

            # boundary of a patch
            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(
                (coords >= (cur_min - 0.2)) * (coords <= (cur_max + 0.2)),
                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]

            # two criterion for patch sampling, adopted from PointNet++
            # points within selected patch shoule be scattered separately
            mask = np.sum(
                (cur_coords >= (cur_min - 0.01)) * (cur_coords <=
                                                    (cur_max + 0.01)),
                axis=1) == 3
            # not sure if 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

            # selected patch should contain enough annotated points
            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

        # random sample idx
        if replace is None:
            replace = (cur_sem_mask.shape[0] < self.num_points)
        choices = np.random.choice(
            np.where(cur_choice)[0], self.num_points, replace=replace)

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

        return points, choices

    def __call__(self, results):
        """Call function to sample points to in indoor scenes.

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

        Returns:
            dict: Results after sampling, 'points', 'pts_instance_mask' \
                and 'pts_semantic_mask' keys are updated in the result dict.
        """
        points = results['points']

        assert 'pts_semantic_mask' in results.keys(), \
            'semantic mask should be provided in training and evaluation'
        pts_semantic_mask = results['pts_semantic_mask']

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

        results['points'] = points
        results['pts_semantic_mask'] = pts_semantic_mask[choices]
        pts_instance_mask = results.get('pts_instance_mask', None)
        if pts_instance_mask is not None:
            results['pts_instance_mask'] = pts_instance_mask[choices]

        return results

    def __repr__(self):
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
        repr_str += f'(num_points={self.num_points},'
        repr_str += f' block_size={self.block_size},'
        repr_str += f' sample_rate={self.sample_rate},'
        repr_str += f' ignore_index={self.ignore_index},'
        repr_str += f' use_normalized_coord={self.use_normalized_coord},'
        repr_str += f' num_try={self.num_try})'
wuyuefeng's avatar
wuyuefeng committed
883
        return repr_str
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911


@PIPELINES.register_module()
class BackgroundPointsFilter(object):
    """Filter background points near the bounding box.

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

    def __init__(self, bbox_enlarge_range):
        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, :]

    def __call__(self, input_dict):
        """Call function to filter points by the range.

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

        Returns:
912
913
            dict: Results after filtering, 'points', 'pts_instance_mask' \
                and 'pts_semantic_mask' keys are updated in the result dict.
914
915
916
917
918
919
920
921
        """
        points = input_dict['points']
        gt_bboxes_3d = input_dict['gt_bboxes_3d']

        gt_bboxes_3d_np = gt_bboxes_3d.tensor.numpy()
        gt_bboxes_3d_np[:, :3] = gt_bboxes_3d.gravity_center.numpy()
        enlarged_gt_bboxes_3d = gt_bboxes_3d_np.copy()
        enlarged_gt_bboxes_3d[:, 3:6] += self.bbox_enlarge_range
922
923
924
        points_numpy = points.tensor.numpy()
        foreground_masks = box_np_ops.points_in_rbbox(points_numpy,
                                                      gt_bboxes_3d_np)
925
        enlarge_foreground_masks = box_np_ops.points_in_rbbox(
926
            points_numpy, enlarged_gt_bboxes_3d)
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
        foreground_masks = foreground_masks.max(1)
        enlarge_foreground_masks = enlarge_foreground_masks.max(1)
        valid_masks = ~np.logical_and(~foreground_masks,
                                      enlarge_foreground_masks)

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

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

    def __repr__(self):
        """str: Return a string that describes the module."""
        repr_str = self.__class__.__name__
945
        repr_str += f'(bbox_enlarge_range={self.bbox_enlarge_range.tolist()})'
946
        return repr_str
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018


@PIPELINES.register_module()
class VoxelBasedPointSampler(object):
    """Voxel based point sampler.

    Apply voxel sampling to multiple sweep points.

    Args:
        cur_sweep_cfg (dict): Config for sampling current points.
        prev_sweep_cfg (dict): Config for sampling previous points.
        time_dim (int): Index that indicate the time dimention
            for input points.
    """

    def __init__(self, cur_sweep_cfg, prev_sweep_cfg=None, time_dim=3):
        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

    def _sample_points(self, points, sampler, point_dim):
        """Sample points for each points subset.

        Args:
            points (np.ndarray): Points subset to be sampled.
            sampler (VoxelGenerator): Voxel based sampler for
                each points subset.
            point_dim (int): The dimention of each points

        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

    def __call__(self, results):
        """Call function to sample points from multiple sweeps.

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

        Returns:
            dict: Results after sampling, 'points', 'pts_instance_mask' \
                and 'pts_semantic_mask' keys are updated in the result dict.
        """
        points = results['points']
        original_dim = points.shape[1]

        # TODO: process instance and semantic mask while _max_num_points
        # is larger than 1
        # Extend points with seg and mask fields
        map_fields2dim = []
        start_dim = original_dim
1019
1020
        points_numpy = points.tensor.numpy()
        extra_channel = [points_numpy]
1021
1022
1023
1024
1025
1026
1027
1028
1029
        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])

1030
        points_numpy = np.concatenate(extra_channel, axis=-1)
1031
1032
1033
1034
1035

        # 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.
1036
1037
1038
        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]
1039
1040
1041
1042
1043
1044
1045
1046
1047
        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,
1048
                                               points_numpy.shape[1])
1049
1050
1051
        if self.prev_voxel_generator is not None:
            prev_sweeps_points = self._sample_points(prev_sweeps_points,
                                                     self.prev_voxel_generator,
1052
                                                     points_numpy.shape[1])
1053

1054
1055
            points_numpy = np.concatenate(
                [cur_sweep_points, prev_sweeps_points], 0)
1056
        else:
1057
            points_numpy = cur_sweep_points
1058
1059

        if self.cur_voxel_generator._max_num_points == 1:
1060
1061
            points_numpy = points_numpy.squeeze(1)
        results['points'] = points.new_point(points_numpy[..., :original_dim])
1062
1063
1064

        # Restore the correspoinding seg and mask fields
        for key, dim_index in map_fields2dim:
1065
            results[key] = points_numpy[..., dim_index]
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088

        return results

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

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

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