data_augmentor.py 10.2 KB
Newer Older
1
from functools import partial
Shaoshuai Shi's avatar
Shaoshuai Shi committed
2

3
import numpy as np
Shaoshuai Shi's avatar
Shaoshuai Shi committed
4

5
from ...utils import common_utils
Shaoshuai Shi's avatar
Shaoshuai Shi committed
6
from . import augmentor_utils, database_sampler
7
8
9
10
11
12
13


class DataAugmentor(object):
    def __init__(self, root_path, augmentor_configs, class_names, logger=None):
        self.root_path = root_path
        self.class_names = class_names
        self.logger = logger
14
        
15
        self.data_augmentor_queue = []
16
17
        aug_config_list = augmentor_configs if isinstance(augmentor_configs, list) \
            else augmentor_configs.AUG_CONFIG_LIST
18
        
19
20
21
22
        for cur_cfg in aug_config_list:
            if not isinstance(augmentor_configs, list):
                if cur_cfg.NAME in augmentor_configs.DISABLE_AUG_LIST:
                    continue
23
            cur_augmentor = getattr(self, cur_cfg.NAME)(config=cur_cfg)
24
            self.data_augmentor_queue.append(cur_augmentor)
25
    
26
27
28
29
30
31
32
33
    def gt_sampling(self, config=None):
        db_sampler = database_sampler.DataBaseSampler(
            root_path=self.root_path,
            sampler_cfg=config,
            class_names=self.class_names,
            logger=self.logger
        )
        return db_sampler
34
    
Gus-Guo's avatar
Gus-Guo committed
35
36
37
38
    def __getstate__(self):
        d = dict(self.__dict__)
        del d['logger']
        return d
39
    
Gus-Guo's avatar
Gus-Guo committed
40
41
    def __setstate__(self, d):
        self.__dict__.update(d)
42
    
43
    def random_world_flip(self, data_dict=None, config=None):
Gus-Guo's avatar
Gus-Guo committed
44
45
        if data_dict is None:
            return partial(self.random_world_flip, config=config)
46
47
48
49
50
51
        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
        for cur_axis in config['ALONG_AXIS_LIST']:
            assert cur_axis in ['x', 'y']
            gt_boxes, points = getattr(augmentor_utils, 'random_flip_along_%s' % cur_axis)(
                gt_boxes, points,
            )
52
        
53
54
55
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
56
    
57
    def random_world_rotation(self, data_dict=None, config=None):
Gus-Guo's avatar
Gus-Guo committed
58
59
        if data_dict is None:
            return partial(self.random_world_rotation, config=config)
60
61
62
63
64
65
        rot_range = config['WORLD_ROT_ANGLE']
        if not isinstance(rot_range, list):
            rot_range = [-rot_range, rot_range]
        gt_boxes, points = augmentor_utils.global_rotation(
            data_dict['gt_boxes'], data_dict['points'], rot_range=rot_range
        )
66
        
67
68
69
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
70
    
71
    def random_world_scaling(self, data_dict=None, config=None):
Gus-Guo's avatar
Gus-Guo committed
72
73
        if data_dict is None:
            return partial(self.random_world_scaling, config=config)
74
75
76
        gt_boxes, points = augmentor_utils.global_scaling(
            data_dict['gt_boxes'], data_dict['points'], config['WORLD_SCALE_RANGE']
        )
77
        
78
79
80
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
81
    
82
83
84
85
86
87
88
89
90
91
92
93
94
    def random_image_flip(self, data_dict=None, config=None):
        if data_dict is None:
            return partial(self.random_image_flip, config=config)
        images = data_dict["images"]
        depth_maps = data_dict["depth_maps"]
        gt_boxes = data_dict['gt_boxes']
        gt_boxes2d = data_dict["gt_boxes2d"]
        calib = data_dict["calib"]
        for cur_axis in config['ALONG_AXIS_LIST']:
            assert cur_axis in ['horizontal']
            images, depth_maps, gt_boxes = getattr(augmentor_utils, 'random_image_flip_%s' % cur_axis)(
                images, depth_maps, gt_boxes, calib,
            )
95
        
96
97
98
99
        data_dict['images'] = images
        data_dict['depth_maps'] = depth_maps
        data_dict['gt_boxes'] = gt_boxes
        return data_dict
100
    
101
102
103
104
105
106
107
108
109
110
111
112
113
    def random_world_translation(self, data_dict=None, config=None):
        """
        Please check the correctness of it before using.
        """
        if data_dict is None:
            return partial(self.random_world_translation, config=config)
        offset_range = config['WORLD_TRANSLATION_RANGE']
        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
        for cur_axis in config['ALONG_AXIS_LIST']:
            assert cur_axis in ['x', 'y', 'z']
            gt_boxes, points = getattr(augmentor_utils, 'random_translation_along_%s' % cur_axis)(
                gt_boxes, points, offset_range,
            )
114
        
115
116
117
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
118
    
119
120
121
122
123
124
125
126
127
128
129
130
131
    def random_local_translation(self, data_dict=None, config=None):
        """
        Please check the correctness of it before using.
        """
        if data_dict is None:
            return partial(self.random_local_translation, config=config)
        offset_range = config['LOCAL_TRANSLATION_RANGE']
        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
        for cur_axis in config['ALONG_AXIS_LIST']:
            assert cur_axis in ['x', 'y', 'z']
            gt_boxes, points = getattr(augmentor_utils, 'random_local_translation_along_%s' % cur_axis)(
                gt_boxes, points, offset_range,
            )
132
        
133
134
135
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
136
    
137
138
139
140
141
142
143
144
145
146
147
148
    def random_local_rotation(self, data_dict=None, config=None):
        """
        Please check the correctness of it before using.
        """
        if data_dict is None:
            return partial(self.random_local_rotation, config=config)
        rot_range = config['LOCAL_ROT_ANGLE']
        if not isinstance(rot_range, list):
            rot_range = [-rot_range, rot_range]
        gt_boxes, points = augmentor_utils.local_rotation(
            data_dict['gt_boxes'], data_dict['points'], rot_range=rot_range
        )
149
        
150
151
152
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
153
    
154
155
156
157
158
159
160
161
162
    def random_local_scaling(self, data_dict=None, config=None):
        """
        Please check the correctness of it before using.
        """
        if data_dict is None:
            return partial(self.random_local_scaling, config=config)
        gt_boxes, points = augmentor_utils.local_scaling(
            data_dict['gt_boxes'], data_dict['points'], config['LOCAL_SCALE_RANGE']
        )
163
        
164
165
166
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
167
    
168
169
170
171
172
173
    def random_world_frustum_dropout(self, data_dict=None, config=None):
        """
        Please check the correctness of it before using.
        """
        if data_dict is None:
            return partial(self.random_world_frustum_dropout, config=config)
174
        
175
176
177
178
179
180
181
        intensity_range = config['INTENSITY_RANGE']
        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
        for direction in config['DIRECTION']:
            assert direction in ['top', 'bottom', 'left', 'right']
            gt_boxes, points = getattr(augmentor_utils, 'global_frustum_dropout_%s' % direction)(
                gt_boxes, points, intensity_range,
            )
182
        
183
184
185
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
186
    
187
188
189
190
191
192
    def random_local_frustum_dropout(self, data_dict=None, config=None):
        """
        Please check the correctness of it before using.
        """
        if data_dict is None:
            return partial(self.random_local_frustum_dropout, config=config)
193
        
194
195
196
197
198
199
200
        intensity_range = config['INTENSITY_RANGE']
        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
        for direction in config['DIRECTION']:
            assert direction in ['top', 'bottom', 'left', 'right']
            gt_boxes, points = getattr(augmentor_utils, 'local_frustum_dropout_%s' % direction)(
                gt_boxes, points, intensity_range,
            )
201
        
202
203
204
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
    
    def random_local_pyramid_aug(self, data_dict=None, config=None):
        """
        Refer to the paper: 
            SE-SSD: Self-Ensembling Single-Stage Object Detector From Point Cloud
        """
        if data_dict is None:
            return partial(self.random_local_pyramid_aug, config=config)
        
        gt_boxes, points = data_dict['gt_boxes'], data_dict['points']
        
        gt_boxes, points, pyramids = augmentor_utils.local_pyramid_dropout(gt_boxes, points, config['DROP_PROB'])
        gt_boxes, points, pyramids = augmentor_utils.local_pyramid_sparsify(gt_boxes, points,
                                                                            config['SPARSIFY_PROB'],
                                                                            config['SPARSIFY_MAX_NUM'],
                                                                            pyramids)
        gt_boxes, points = augmentor_utils.local_pyramid_swap(gt_boxes, points,
                                                                 config['SWAP_PROB'],
                                                                 config['SWAP_MAX_NUM'],
                                                                 pyramids)
        data_dict['gt_boxes'] = gt_boxes
        data_dict['points'] = points
        return data_dict
    
229
230
231
232
233
234
235
236
237
238
239
240
241
    def forward(self, data_dict):
        """
        Args:
            data_dict:
                points: (N, 3 + C_in)
                gt_boxes: optional, (N, 7) [x, y, z, dx, dy, dz, heading]
                gt_names: optional, (N), string
                ...

        Returns:
        """
        for cur_augmentor in self.data_augmentor_queue:
            data_dict = cur_augmentor(data_dict=data_dict)
242
        
243
244
245
246
247
248
249
250
251
252
253
        data_dict['gt_boxes'][:, 6] = common_utils.limit_period(
            data_dict['gt_boxes'][:, 6], offset=0.5, period=2 * np.pi
        )
        if 'calib' in data_dict:
            data_dict.pop('calib')
        if 'road_plane' in data_dict:
            data_dict.pop('road_plane')
        if 'gt_boxes_mask' in data_dict:
            gt_boxes_mask = data_dict['gt_boxes_mask']
            data_dict['gt_boxes'] = data_dict['gt_boxes'][gt_boxes_mask]
            data_dict['gt_names'] = data_dict['gt_names'][gt_boxes_mask]
254
255
            if 'gt_boxes2d' in data_dict:
                data_dict['gt_boxes2d'] = data_dict['gt_boxes2d'][gt_boxes_mask]
256
            
257
            data_dict.pop('gt_boxes_mask')
258
        return data_dict