train_mixins.py 16.9 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
zhangwenwei's avatar
zhangwenwei committed
2
3
import numpy as np
import torch
4
from mmdet.models.utils import images_to_levels, multi_apply
5
from mmengine.structures import InstanceData
zhangwenwei's avatar
zhangwenwei committed
6

zhangshilong's avatar
zhangshilong committed
7
from mmdet3d.structures import limit_period
zhangwenwei's avatar
zhangwenwei committed
8
9
10


class AnchorTrainMixin(object):
11
    """Mixin class for target assigning of dense heads."""
zhangwenwei's avatar
zhangwenwei committed
12
13
14

    def anchor_target_3d(self,
                         anchor_list,
15
16
17
                         batch_gt_instances_3d,
                         batch_input_metas,
                         batch_gt_instances_ignore=None,
zhangwenwei's avatar
zhangwenwei committed
18
19
20
21
22
23
24
                         label_channels=1,
                         num_classes=1,
                         sampling=True):
        """Compute regression and classification targets for anchors.

        Args:
            anchor_list (list[list]): Multi level anchors of each image.
25
            batch_gt_instances_3d (list[:obj:`InstanceData`]): Ground truth
wuyuefeng's avatar
wuyuefeng committed
26
                bboxes of each image.
27
28
            batch_input_metas (list[dict]): Meta info of each image.
            batch_gt_instances_ignore (list): Ignore list of gt bboxes.
wuyuefeng's avatar
wuyuefeng committed
29
30
31
            label_channels (int): The channel of labels.
            num_classes (int): The number of classes.
            sampling (bool): Whether to sample anchors.
zhangwenwei's avatar
zhangwenwei committed
32
33

        Returns:
34
35
36
            tuple (list, list, list, list, list, list, int, int):
                Anchor targets, including labels, label weights,
                bbox targets, bbox weights, direction targets,
37
                direction weights, number of positive anchors and
38
                number of negative anchors.
zhangwenwei's avatar
zhangwenwei committed
39
        """
40
41
        num_inputs = len(batch_input_metas)
        assert len(anchor_list) == num_inputs
zhangwenwei's avatar
zhangwenwei committed
42

43
44
45
46
47
48
49
        if isinstance(anchor_list[0][0], list):
            # sizes of anchors are different
            # anchor number of a single level
            num_level_anchors = [
                sum([anchor.size(0) for anchor in anchors])
                for anchors in anchor_list[0]
            ]
50
            for i in range(num_inputs):
51
52
53
54
55
56
57
58
                anchor_list[i] = anchor_list[i][0]
        else:
            # anchor number of multi levels
            num_level_anchors = [
                anchors.view(-1, self.box_code_size).size(0)
                for anchors in anchor_list[0]
            ]
            # concat all level anchors and flags to a single tensor
59
            for i in range(num_inputs):
60
                anchor_list[i] = torch.cat(anchor_list[i])
zhangwenwei's avatar
zhangwenwei committed
61
62

        # compute targets for each image
63
64
        if batch_gt_instances_ignore is None:
            batch_gt_instances_ignore = [None for _ in range(num_inputs)]
zhangwenwei's avatar
zhangwenwei committed
65
66
67
68
69
70

        (all_labels, all_label_weights, all_bbox_targets, all_bbox_weights,
         all_dir_targets, all_dir_weights, pos_inds_list,
         neg_inds_list) = multi_apply(
             self.anchor_target_3d_single,
             anchor_list,
71
72
73
             batch_gt_instances_3d,
             batch_gt_instances_ignore,
             batch_input_metas,
zhangwenwei's avatar
zhangwenwei committed
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
             label_channels=label_channels,
             num_classes=num_classes,
             sampling=sampling)

        # no valid anchors
        if any([labels is None for labels in all_labels]):
            return None
        # sampled anchors of all images
        num_total_pos = sum([max(inds.numel(), 1) for inds in pos_inds_list])
        num_total_neg = sum([max(inds.numel(), 1) for inds in neg_inds_list])
        # split targets to a list w.r.t. multiple levels
        labels_list = images_to_levels(all_labels, num_level_anchors)
        label_weights_list = images_to_levels(all_label_weights,
                                              num_level_anchors)
        bbox_targets_list = images_to_levels(all_bbox_targets,
                                             num_level_anchors)
        bbox_weights_list = images_to_levels(all_bbox_weights,
                                             num_level_anchors)
        dir_targets_list = images_to_levels(all_dir_targets, num_level_anchors)
        dir_weights_list = images_to_levels(all_dir_weights, num_level_anchors)
        return (labels_list, label_weights_list, bbox_targets_list,
                bbox_weights_list, dir_targets_list, dir_weights_list,
                num_total_pos, num_total_neg)

    def anchor_target_3d_single(self,
                                anchors,
100
101
                                gt_instance_3d,
                                gt_instance_ignore,
zhangwenwei's avatar
zhangwenwei committed
102
103
104
105
                                input_meta,
                                label_channels=1,
                                num_classes=1,
                                sampling=True):
wuyuefeng's avatar
wuyuefeng committed
106
107
108
        """Compute targets of anchors in single batch.

        Args:
liyinhao's avatar
liyinhao committed
109
            anchors (torch.Tensor): Concatenated multi-level anchor.
110
111
            gt_instance_3d (:obj:`InstanceData`): Gt bboxes.
            gt_instance_ignore (:obj:`InstanceData`): Ignored gt bboxes.
wuyuefeng's avatar
wuyuefeng committed
112
113
114
115
116
117
            input_meta (dict): Meta info of each image.
            label_channels (int): The channel of labels.
            num_classes (int): The number of classes.
            sampling (bool): Whether to sample anchors.

        Returns:
118
            tuple[torch.Tensor]: Anchor targets.
wuyuefeng's avatar
wuyuefeng committed
119
        """
120
121
        if isinstance(self.bbox_assigner,
                      list) and (not isinstance(anchors, list)):
zhangwenwei's avatar
zhangwenwei committed
122
123
124
125
126
127
128
129
130
131
132
133
            feat_size = anchors.size(0) * anchors.size(1) * anchors.size(2)
            rot_angles = anchors.size(-2)
            assert len(self.bbox_assigner) == anchors.size(-3)
            (total_labels, total_label_weights, total_bbox_targets,
             total_bbox_weights, total_dir_targets, total_dir_weights,
             total_pos_inds, total_neg_inds) = [], [], [], [], [], [], [], []
            current_anchor_num = 0
            for i, assigner in enumerate(self.bbox_assigner):
                current_anchors = anchors[..., i, :, :].reshape(
                    -1, self.box_code_size)
                current_anchor_num += current_anchors.size(0)
                if self.assign_per_class:
134
135
136
137
138
139
                    gt_per_cls = (gt_instance_3d.labels_3d == i)
                    gt_per_cls_instance = InstanceData()
                    gt_per_cls_instance.labels_3d = gt_instance_3d.labels_3d[
                        gt_per_cls]
                    gt_per_cls_instance.bboxes_3d = gt_instance_3d.bboxes_3d[
                        gt_per_cls, :]
zhangwenwei's avatar
zhangwenwei committed
140
                    anchor_targets = self.anchor_target_single_assigner(
141
142
                        assigner, current_anchors, gt_per_cls_instance,
                        gt_instance_ignore, input_meta, num_classes, sampling)
zhangwenwei's avatar
zhangwenwei committed
143
144
                else:
                    anchor_targets = self.anchor_target_single_assigner(
145
146
                        assigner, current_anchors, gt_instance_3d,
                        gt_instance_ignore, input_meta, num_classes, sampling)
zhangwenwei's avatar
zhangwenwei committed
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181

                (labels, label_weights, bbox_targets, bbox_weights,
                 dir_targets, dir_weights, pos_inds, neg_inds) = anchor_targets
                total_labels.append(labels.reshape(feat_size, 1, rot_angles))
                total_label_weights.append(
                    label_weights.reshape(feat_size, 1, rot_angles))
                total_bbox_targets.append(
                    bbox_targets.reshape(feat_size, 1, rot_angles,
                                         anchors.size(-1)))
                total_bbox_weights.append(
                    bbox_weights.reshape(feat_size, 1, rot_angles,
                                         anchors.size(-1)))
                total_dir_targets.append(
                    dir_targets.reshape(feat_size, 1, rot_angles))
                total_dir_weights.append(
                    dir_weights.reshape(feat_size, 1, rot_angles))
                total_pos_inds.append(pos_inds)
                total_neg_inds.append(neg_inds)

            total_labels = torch.cat(total_labels, dim=-2).reshape(-1)
            total_label_weights = torch.cat(
                total_label_weights, dim=-2).reshape(-1)
            total_bbox_targets = torch.cat(
                total_bbox_targets, dim=-3).reshape(-1, anchors.size(-1))
            total_bbox_weights = torch.cat(
                total_bbox_weights, dim=-3).reshape(-1, anchors.size(-1))
            total_dir_targets = torch.cat(
                total_dir_targets, dim=-2).reshape(-1)
            total_dir_weights = torch.cat(
                total_dir_weights, dim=-2).reshape(-1)
            total_pos_inds = torch.cat(total_pos_inds, dim=0).reshape(-1)
            total_neg_inds = torch.cat(total_neg_inds, dim=0).reshape(-1)
            return (total_labels, total_label_weights, total_bbox_targets,
                    total_bbox_weights, total_dir_targets, total_dir_weights,
                    total_pos_inds, total_neg_inds)
182
183
184
185
186
187
188
189
190
191
192
193
194
        elif isinstance(self.bbox_assigner, list) and isinstance(
                anchors, list):
            # class-aware anchors with different feature map sizes
            assert len(self.bbox_assigner) == len(anchors), \
                'The number of bbox assigners and anchors should be the same.'
            (total_labels, total_label_weights, total_bbox_targets,
             total_bbox_weights, total_dir_targets, total_dir_weights,
             total_pos_inds, total_neg_inds) = [], [], [], [], [], [], [], []
            current_anchor_num = 0
            for i, assigner in enumerate(self.bbox_assigner):
                current_anchors = anchors[i]
                current_anchor_num += current_anchors.size(0)
                if self.assign_per_class:
195
196
197
198
199
200
                    gt_per_cls = (gt_instance_3d.labels_3d == i)
                    gt_per_cls_instance = InstanceData()
                    gt_per_cls_instance.labels_3d = gt_instance_3d.labels_3d[
                        gt_per_cls]
                    gt_per_cls_instance.bboxes_3d = gt_instance_3d.bboxes_3d[
                        gt_per_cls, :]
201
                    anchor_targets = self.anchor_target_single_assigner(
202
203
                        assigner, current_anchors, gt_per_cls_instance,
                        gt_instance_ignore, input_meta, num_classes, sampling)
204
205
                else:
                    anchor_targets = self.anchor_target_single_assigner(
206
207
                        assigner, current_anchors, gt_instance_3d,
                        gt_instance_ignore, input_meta, num_classes, sampling)
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232

                (labels, label_weights, bbox_targets, bbox_weights,
                 dir_targets, dir_weights, pos_inds, neg_inds) = anchor_targets
                total_labels.append(labels)
                total_label_weights.append(label_weights)
                total_bbox_targets.append(
                    bbox_targets.reshape(-1, anchors[i].size(-1)))
                total_bbox_weights.append(
                    bbox_weights.reshape(-1, anchors[i].size(-1)))
                total_dir_targets.append(dir_targets)
                total_dir_weights.append(dir_weights)
                total_pos_inds.append(pos_inds)
                total_neg_inds.append(neg_inds)

            total_labels = torch.cat(total_labels, dim=0)
            total_label_weights = torch.cat(total_label_weights, dim=0)
            total_bbox_targets = torch.cat(total_bbox_targets, dim=0)
            total_bbox_weights = torch.cat(total_bbox_weights, dim=0)
            total_dir_targets = torch.cat(total_dir_targets, dim=0)
            total_dir_weights = torch.cat(total_dir_weights, dim=0)
            total_pos_inds = torch.cat(total_pos_inds, dim=0)
            total_neg_inds = torch.cat(total_neg_inds, dim=0)
            return (total_labels, total_label_weights, total_bbox_targets,
                    total_bbox_weights, total_dir_targets, total_dir_weights,
                    total_pos_inds, total_neg_inds)
zhangwenwei's avatar
zhangwenwei committed
233
        else:
234
            return self.anchor_target_single_assigner(self.bbox_assigner,
235
236
237
238
                                                      anchors, gt_instance_3d,
                                                      gt_instance_ignore,
                                                      input_meta, num_classes,
                                                      sampling)
zhangwenwei's avatar
zhangwenwei committed
239
240
241
242

    def anchor_target_single_assigner(self,
                                      bbox_assigner,
                                      anchors,
243
244
                                      gt_instance_3d,
                                      gt_instance_ignore,
zhangwenwei's avatar
zhangwenwei committed
245
246
247
                                      input_meta,
                                      num_classes=1,
                                      sampling=True):
wuyuefeng's avatar
wuyuefeng committed
248
249
250
251
        """Assign anchors and encode positive anchors.

        Args:
            bbox_assigner (BaseAssigner): assign positive and negative boxes.
liyinhao's avatar
liyinhao committed
252
            anchors (torch.Tensor): Concatenated multi-level anchor.
253
254
            gt_instance_3d (:obj:`InstanceData`): Gt bboxes.
            gt_instance_ignore (torch.Tensor): Ignored gt bboxes.
wuyuefeng's avatar
wuyuefeng committed
255
256
257
258
259
            input_meta (dict): Meta info of each image.
            num_classes (int): The number of classes.
            sampling (bool): Whether to sample anchors.

        Returns:
260
            tuple[torch.Tensor]: Anchor targets.
wuyuefeng's avatar
wuyuefeng committed
261
        """
zhangwenwei's avatar
zhangwenwei committed
262
263
264
265
266
267
268
269
        anchors = anchors.reshape(-1, anchors.size(-1))
        num_valid_anchors = anchors.shape[0]
        bbox_targets = torch.zeros_like(anchors)
        bbox_weights = torch.zeros_like(anchors)
        dir_targets = anchors.new_zeros((anchors.shape[0]), dtype=torch.long)
        dir_weights = anchors.new_zeros((anchors.shape[0]), dtype=torch.float)
        labels = anchors.new_zeros(num_valid_anchors, dtype=torch.long)
        label_weights = anchors.new_zeros(num_valid_anchors, dtype=torch.float)
270
271
272
273
274
275
276
277
278
279
280
        if len(gt_instance_3d.bboxes_3d) > 0:
            if not isinstance(gt_instance_3d.bboxes_3d, torch.Tensor):
                gt_instance_3d.bboxes_3d = gt_instance_3d.bboxes_3d.tensor.to(
                    anchors.device)
            pred_instance_3d = InstanceData(priors=anchors)
            assign_result = bbox_assigner.assign(pred_instance_3d,
                                                 gt_instance_3d,
                                                 gt_instance_ignore)
            sampling_result = self.bbox_sampler.sample(assign_result,
                                                       pred_instance_3d,
                                                       gt_instance_3d)
zhangwenwei's avatar
zhangwenwei committed
281
282
283
284
            pos_inds = sampling_result.pos_inds
            neg_inds = sampling_result.neg_inds
        else:
            pos_inds = torch.nonzero(
Wenhao Wu's avatar
Wenhao Wu committed
285
286
                anchors.new_zeros((anchors.shape[0], ), dtype=torch.bool) > 0,
                as_tuple=False).squeeze(-1).unique()
zhangwenwei's avatar
zhangwenwei committed
287
            neg_inds = torch.nonzero(
Wenhao Wu's avatar
Wenhao Wu committed
288
289
                anchors.new_zeros((anchors.shape[0], ), dtype=torch.bool) == 0,
                as_tuple=False).squeeze(-1).unique()
zhangwenwei's avatar
zhangwenwei committed
290

291
        if gt_instance_3d.labels_3d is not None:
zhangwenwei's avatar
zhangwenwei committed
292
293
            labels += num_classes
        if len(pos_inds) > 0:
294
295
            pos_bbox_targets = self.bbox_coder.encode(
                sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes)
zhangwenwei's avatar
zhangwenwei committed
296
297
298
299
            pos_dir_targets = get_direction_target(
                sampling_result.pos_bboxes,
                pos_bbox_targets,
                self.dir_offset,
300
                self.dir_limit_offset,
zhangwenwei's avatar
zhangwenwei committed
301
302
303
304
305
306
                one_hot=False)
            bbox_targets[pos_inds, :] = pos_bbox_targets
            bbox_weights[pos_inds, :] = 1.0
            dir_targets[pos_inds] = pos_dir_targets
            dir_weights[pos_inds] = 1.0

307
            if gt_instance_3d.labels_3d is None:
zhangwenwei's avatar
zhangwenwei committed
308
309
                labels[pos_inds] = 1
            else:
310
                labels[pos_inds] = gt_instance_3d.labels_3d[
zhangwenwei's avatar
zhangwenwei committed
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
                    sampling_result.pos_assigned_gt_inds]
            if self.train_cfg.pos_weight <= 0:
                label_weights[pos_inds] = 1.0
            else:
                label_weights[pos_inds] = self.train_cfg.pos_weight

        if len(neg_inds) > 0:
            label_weights[neg_inds] = 1.0
        return (labels, label_weights, bbox_targets, bbox_weights, dir_targets,
                dir_weights, pos_inds, neg_inds)


def get_direction_target(anchors,
                         reg_targets,
                         dir_offset=0,
326
                         dir_limit_offset=0,
zhangwenwei's avatar
zhangwenwei committed
327
328
                         num_bins=2,
                         one_hot=True):
wuyuefeng's avatar
wuyuefeng committed
329
330
331
    """Encode direction to 0 ~ num_bins-1.

    Args:
liyinhao's avatar
liyinhao committed
332
333
        anchors (torch.Tensor): Concatenated multi-level anchor.
        reg_targets (torch.Tensor): Bbox regression targets.
wuyuefeng's avatar
wuyuefeng committed
334
335
336
337
338
        dir_offset (int): Direction offset.
        num_bins (int): Number of bins to divide 2*PI.
        one_hot (bool): Whether to encode as one hot.

    Returns:
liyinhao's avatar
liyinhao committed
339
        torch.Tensor: Encoded direction targets.
wuyuefeng's avatar
wuyuefeng committed
340
    """
zhangwenwei's avatar
zhangwenwei committed
341
    rot_gt = reg_targets[..., 6] + anchors[..., 6]
342
    offset_rot = limit_period(rot_gt - dir_offset, dir_limit_offset, 2 * np.pi)
zhangwenwei's avatar
zhangwenwei committed
343
344
345
346
347
348
349
350
351
352
353
    dir_cls_targets = torch.floor(offset_rot / (2 * np.pi / num_bins)).long()
    dir_cls_targets = torch.clamp(dir_cls_targets, min=0, max=num_bins - 1)
    if one_hot:
        dir_targets = torch.zeros(
            *list(dir_cls_targets.shape),
            num_bins,
            dtype=anchors.dtype,
            device=dir_cls_targets.device)
        dir_targets.scatter_(dir_cls_targets.unsqueeze(dim=-1).long(), 1.0)
        dir_cls_targets = dir_targets
    return dir_cls_targets