dbsampler.py 12.9 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangwenwei's avatar
zhangwenwei committed
2
import copy
3
import os
4
from typing import List, Optional
5

6
import mmengine
zhangwenwei's avatar
zhangwenwei committed
7
import numpy as np
zhangwenwei's avatar
zhangwenwei committed
8

zhangshilong's avatar
zhangshilong committed
9
from mmdet3d.datasets.transforms import data_augment_utils
10
from mmdet3d.registry import TRANSFORMS
zhangshilong's avatar
zhangshilong committed
11
from mmdet3d.structures.ops import box_np_ops
zhangwenwei's avatar
zhangwenwei committed
12
13
14


class BatchSampler:
wangtai's avatar
wangtai committed
15
16
17
18
    """Class for sampling specific category of ground truths.

    Args:
        sample_list (list[dict]): List of samples.
19
20
21
22
23
        name (str, optional): The category of samples. Defaults to None.
        epoch (int, optional): Sampling epoch. Defaults to None.
        shuffle (bool, optional): Whether to shuffle indices.
            Defaults to False.
        drop_reminder (bool, optional): Drop reminder. Defaults to False.
wangtai's avatar
wangtai committed
24
    """
zhangwenwei's avatar
zhangwenwei committed
25
26

    def __init__(self,
27
28
29
30
31
                 sampled_list: List[dict],
                 name: Optional[str] = None,
                 epoch: Optional[int] = None,
                 shuffle: bool = True,
                 drop_reminder: bool = False) -> None:
zhangwenwei's avatar
zhangwenwei committed
32
33
34
35
36
37
38
39
40
41
42
43
        self._sampled_list = sampled_list
        self._indices = np.arange(len(sampled_list))
        if shuffle:
            np.random.shuffle(self._indices)
        self._idx = 0
        self._example_num = len(sampled_list)
        self._name = name
        self._shuffle = shuffle
        self._epoch = epoch
        self._epoch_counter = 0
        self._drop_reminder = drop_reminder

44
    def _sample(self, num: int) -> List[int]:
wangtai's avatar
wangtai committed
45
46
47
48
49
50
51
52
        """Sample specific number of ground truths and return indices.

        Args:
            num (int): Sampled number.

        Returns:
            list[int]: Indices of sampled ground truths.
        """
zhangwenwei's avatar
zhangwenwei committed
53
54
55
56
57
58
59
60
        if self._idx + num >= self._example_num:
            ret = self._indices[self._idx:].copy()
            self._reset()
        else:
            ret = self._indices[self._idx:self._idx + num]
            self._idx += num
        return ret

61
    def _reset(self) -> None:
wangtai's avatar
wangtai committed
62
        """Reset the index of batchsampler to zero."""
zhangwenwei's avatar
zhangwenwei committed
63
64
65
66
67
68
        assert self._name is not None
        # print("reset", self._name)
        if self._shuffle:
            np.random.shuffle(self._indices)
        self._idx = 0

69
    def sample(self, num: int) -> List[dict]:
wangtai's avatar
wangtai committed
70
71
72
73
74
75
76
77
        """Sample specific number of ground truths.

        Args:
            num (int): Sampled number.

        Returns:
            list[dict]: Sampled ground truths.
        """
zhangwenwei's avatar
zhangwenwei committed
78
79
80
81
        indices = self._sample(num)
        return [self._sampled_list[i] for i in indices]


82
@TRANSFORMS.register_module()
zhangwenwei's avatar
zhangwenwei committed
83
class DataBaseSampler(object):
84
85
86
87
88
89
90
91
    """Class for sampling data from the ground truth database.

    Args:
        info_path (str): Path of groundtruth database info.
        data_root (str): Path of groundtruth database.
        rate (float): Rate of actual sampled over maximum sampled number.
        prepare (dict): Name of preparation functions and the input value.
        sample_groups (dict): Sampled classes and numbers.
92
93
94
95
96
97
98
        classes (list[str], optional): List of classes. Defaults to None.
        points_loader(dict, optional): Config of points loader. Defaults to
            dict(type='LoadPointsFromFile', load_dim=4, use_dim=[0, 1, 2, 3]).
        file_client_args (dict, optional): Config dict of file clients,
            refer to
            https://github.com/open-mmlab/mmengine/blob/main/mmengine/fileio/file_client.py
            for more details. Defaults to dict(backend='disk').
99
    """
zhangwenwei's avatar
zhangwenwei committed
100

101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
    def __init__(
        self,
        info_path: str,
        data_root: str,
        rate: float,
        prepare: dict,
        sample_groups: dict,
        classes: Optional[List[str]] = None,
        points_loader: dict = dict(
            type='LoadPointsFromFile',
            coord_type='LIDAR',
            load_dim=4,
            use_dim=[0, 1, 2, 3]),
        file_client_args: dict = dict(backend='disk')
    ) -> None:
zhangwenwei's avatar
zhangwenwei committed
116
        super().__init__()
zhangwenwei's avatar
zhangwenwei committed
117
        self.data_root = data_root
zhangwenwei's avatar
zhangwenwei committed
118
119
120
        self.info_path = info_path
        self.rate = rate
        self.prepare = prepare
zhangwenwei's avatar
zhangwenwei committed
121
122
123
        self.classes = classes
        self.cat2label = {name: i for i, name in enumerate(classes)}
        self.label2cat = {i: name for i, name in enumerate(classes)}
124
        self.points_loader = TRANSFORMS.build(points_loader)
125
        self.file_client = mmengine.FileClient(**file_client_args)
zhangwenwei's avatar
zhangwenwei committed
126

127
        # load data base infos
128
129
130
        with self.file_client.get_local_path(info_path) as local_path:
            # loading data from a file-like object needs file format
            db_infos = mmengine.load(open(local_path, 'rb'), file_format='pkl')
zhangwenwei's avatar
zhangwenwei committed
131
132

        # filter database infos
133
134
        from mmengine.logging import MMLogger
        logger: MMLogger = MMLogger.get_current_instance()
zhangwenwei's avatar
zhangwenwei committed
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
        for k, v in db_infos.items():
            logger.info(f'load {len(v)} {k} database infos')
        for prep_func, val in prepare.items():
            db_infos = getattr(self, prep_func)(db_infos, val)
        logger.info('After filter database:')
        for k, v in db_infos.items():
            logger.info(f'load {len(v)} {k} database infos')

        self.db_infos = db_infos

        # load sample groups
        # TODO: more elegant way to load sample groups
        self.sample_groups = []
        for name, num in sample_groups.items():
            self.sample_groups.append({name: int(num)})

        self.group_db_infos = self.db_infos  # just use db_infos
        self.sample_classes = []
        self.sample_max_nums = []
        for group_info in self.sample_groups:
            self.sample_classes += list(group_info.keys())
            self.sample_max_nums += list(group_info.values())

        self.sampler_dict = {}
        for k, v in self.group_db_infos.items():
            self.sampler_dict[k] = BatchSampler(v, k, shuffle=True)
        # TODO: No group_sampling currently

    @staticmethod
164
    def filter_by_difficulty(db_infos: dict, removed_difficulty: list) -> dict:
165
166
167
168
169
170
171
172
173
        """Filter ground truths by difficulties.

        Args:
            db_infos (dict): Info of groundtruth database.
            removed_difficulty (list): Difficulties that are not qualified.

        Returns:
            dict: Info of database after filtering.
        """
zhangwenwei's avatar
zhangwenwei committed
174
175
176
177
178
179
180
181
182
        new_db_infos = {}
        for key, dinfos in db_infos.items():
            new_db_infos[key] = [
                info for info in dinfos
                if info['difficulty'] not in removed_difficulty
            ]
        return new_db_infos

    @staticmethod
183
    def filter_by_min_points(db_infos: dict, min_gt_points_dict: dict) -> dict:
184
185
186
187
188
189
190
191
192
193
        """Filter ground truths by number of points in the bbox.

        Args:
            db_infos (dict): Info of groundtruth database.
            min_gt_points_dict (dict): Different number of minimum points
                needed for different categories of ground truths.

        Returns:
            dict: Info of database after filtering.
        """
zhangwenwei's avatar
zhangwenwei committed
194
195
196
197
198
199
200
201
202
203
        for name, min_num in min_gt_points_dict.items():
            min_num = int(min_num)
            if min_num > 0:
                filtered_infos = []
                for info in db_infos[name]:
                    if info['num_points_in_gt'] >= min_num:
                        filtered_infos.append(info)
                db_infos[name] = filtered_infos
        return db_infos

204
205
206
207
208
    def sample_all(self,
                   gt_bboxes: np.ndarray,
                   gt_labels: np.ndarray,
                   img: Optional[np.ndarray] = None,
                   ground_plane: Optional[np.ndarray] = None) -> dict:
209
210
211
212
        """Sampling all categories of bboxes.

        Args:
            gt_bboxes (np.ndarray): Ground truth bounding boxes.
liyinhao's avatar
liyinhao committed
213
            gt_labels (np.ndarray): Ground truth labels of boxes.
214
215
216
            img (np.ndarray, optional): Image array. Defaults to None.
            ground_plane (np.ndarray, optional): Ground plane information.
                Defaults to None.
217
218
219
220

        Returns:
            dict: Dict of sampled 'pseudo ground truths'.

221
                - gt_labels_3d (np.ndarray): ground truths labels
zhangwenwei's avatar
zhangwenwei committed
222
                    of sampled objects.
223
                - gt_bboxes_3d (:obj:`BaseInstance3DBoxes`):
liyinhao's avatar
liyinhao committed
224
                    sampled ground truth 3D bounding boxes
225
226
227
                - points (np.ndarray): sampled points
                - group_ids (np.ndarray): ids of sampled ground truths
        """
zhangwenwei's avatar
zhangwenwei committed
228
229
230
231
        sampled_num_dict = {}
        sample_num_per_class = []
        for class_name, max_sample_num in zip(self.sample_classes,
                                              self.sample_max_nums):
zhangwenwei's avatar
zhangwenwei committed
232
233
234
            class_label = self.cat2label[class_name]
            # sampled_num = int(max_sample_num -
            #                   np.sum([n == class_name for n in gt_names]))
zhangwenwei's avatar
zhangwenwei committed
235
            sampled_num = int(max_sample_num -
zhangwenwei's avatar
zhangwenwei committed
236
                              np.sum([n == class_label for n in gt_labels]))
zhangwenwei's avatar
zhangwenwei committed
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
            sampled_num = np.round(self.rate * sampled_num).astype(np.int64)
            sampled_num_dict[class_name] = sampled_num
            sample_num_per_class.append(sampled_num)

        sampled = []
        sampled_gt_bboxes = []
        avoid_coll_boxes = gt_bboxes

        for class_name, sampled_num in zip(self.sample_classes,
                                           sample_num_per_class):
            if sampled_num > 0:
                sampled_cls = self.sample_class_v2(class_name, sampled_num,
                                                   avoid_coll_boxes)

                sampled += sampled_cls
                if len(sampled_cls) > 0:
                    if len(sampled_cls) == 1:
                        sampled_gt_box = sampled_cls[0]['box3d_lidar'][
                            np.newaxis, ...]
                    else:
                        sampled_gt_box = np.stack(
                            [s['box3d_lidar'] for s in sampled_cls], axis=0)

                    sampled_gt_bboxes += [sampled_gt_box]
                    avoid_coll_boxes = np.concatenate(
                        [avoid_coll_boxes, sampled_gt_box], axis=0)

        ret = None
        if len(sampled) > 0:
            sampled_gt_bboxes = np.concatenate(sampled_gt_bboxes, axis=0)
            # center = sampled_gt_bboxes[:, 0:3]

zhangwenwei's avatar
zhangwenwei committed
269
            # num_sampled = len(sampled)
zhangwenwei's avatar
zhangwenwei committed
270
271
272
273
            s_points_list = []
            count = 0
            for info in sampled:
                file_path = os.path.join(
zhangwenwei's avatar
zhangwenwei committed
274
275
                    self.data_root,
                    info['path']) if self.data_root else info['path']
276
                results = dict(lidar_points=dict(lidar_path=file_path))
277
                s_points = self.points_loader(results)['points']
278
                s_points.translate(info['box3d_lidar'][:3])
zhangwenwei's avatar
zhangwenwei committed
279
280
281
282

                count += 1

                s_points_list.append(s_points)
283
284
285

            gt_labels = np.array([self.cat2label[s['name']] for s in sampled],
                                 dtype=np.long)
286
287
288
289
290
291
292
293
294

            if ground_plane is not None:
                xyz = sampled_gt_bboxes[:, :3]
                dz = (ground_plane[:3][None, :] *
                      xyz).sum(-1) + ground_plane[3]
                sampled_gt_bboxes[:, 2] -= dz
                for i, s_points in enumerate(s_points_list):
                    s_points.tensor[:, 2].sub_(dz[i])

zhangwenwei's avatar
zhangwenwei committed
295
            ret = {
zhangwenwei's avatar
zhangwenwei committed
296
297
                'gt_labels_3d':
                gt_labels,
zhangwenwei's avatar
zhangwenwei committed
298
299
300
                'gt_bboxes_3d':
                sampled_gt_bboxes,
                'points':
301
                s_points_list[0].cat(s_points_list),
zhangwenwei's avatar
zhangwenwei committed
302
303
304
305
306
307
308
                'group_ids':
                np.arange(gt_bboxes.shape[0],
                          gt_bboxes.shape[0] + len(sampled))
            }

        return ret

309
    def sample_class_v2(self, name: str, num: int,
310
                        gt_bboxes: np.ndarray) -> List[dict]:
311
312
313
314
315
316
317
318
319
320
        """Sampling specific categories of bounding boxes.

        Args:
            name (str): Class of objects to be sampled.
            num (int): Number of sampled bboxes.
            gt_bboxes (np.ndarray): Ground truth boxes.

        Returns:
            list[dict]: Valid samples after collision test.
        """
zhangwenwei's avatar
zhangwenwei committed
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
        sampled = self.sampler_dict[name].sample(num)
        sampled = copy.deepcopy(sampled)
        num_gt = gt_bboxes.shape[0]
        num_sampled = len(sampled)
        gt_bboxes_bv = box_np_ops.center_to_corner_box2d(
            gt_bboxes[:, 0:2], gt_bboxes[:, 3:5], gt_bboxes[:, 6])

        sp_boxes = np.stack([i['box3d_lidar'] for i in sampled], axis=0)
        boxes = np.concatenate([gt_bboxes, sp_boxes], axis=0).copy()

        sp_boxes_new = boxes[gt_bboxes.shape[0]:]
        sp_boxes_bv = box_np_ops.center_to_corner_box2d(
            sp_boxes_new[:, 0:2], sp_boxes_new[:, 3:5], sp_boxes_new[:, 6])

        total_bv = np.concatenate([gt_bboxes_bv, sp_boxes_bv], axis=0)
        coll_mat = data_augment_utils.box_collision_test(total_bv, total_bv)
        diag = np.arange(total_bv.shape[0])
        coll_mat[diag, diag] = False

        valid_samples = []
        for i in range(num_gt, num_gt + num_sampled):
            if coll_mat[i].any():
                coll_mat[i] = False
                coll_mat[:, i] = False
            else:
                valid_samples.append(sampled[i - num_gt])
        return valid_samples