anchor3d_head.py 19.3 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
2
3
import numpy as np
import torch
import torch.nn as nn
4
from mmcv.cnn import bias_init_with_prob, normal_init
zhangwenwei's avatar
zhangwenwei committed
5

zhangwenwei's avatar
zhangwenwei committed
6
7
8
9
from mmdet3d.core import (PseudoSampler, box3d_multiclass_nms,
                          build_anchor_generator, build_assigner,
                          build_bbox_coder, build_sampler, limit_period,
                          xywhr2xyxyr)
zhangwenwei's avatar
zhangwenwei committed
10
from mmdet.core import multi_apply
zhangwenwei's avatar
zhangwenwei committed
11
from mmdet.models import HEADS
zhangwenwei's avatar
zhangwenwei committed
12
13
14
15
from ..builder import build_loss
from .train_mixins import AnchorTrainMixin


16
@HEADS.register_module()
zhangwenwei's avatar
zhangwenwei committed
17
18
class Anchor3DHead(nn.Module, AnchorTrainMixin):
    """Anchor head for SECOND/PointPillars/MVXNet/PartA2.
19

zhangwenwei's avatar
zhangwenwei committed
20
    Args:
zhangwenwei's avatar
zhangwenwei committed
21
        num_classes (int): Number of classes.
zhangwenwei's avatar
zhangwenwei committed
22
        in_channels (int): Number of channels in the input feature map.
wuyuefeng's avatar
wuyuefeng committed
23
24
        train_cfg (dict): Train configs.
        test_cfg (dict): Test configs.
zhangwenwei's avatar
zhangwenwei committed
25
        feat_channels (int): Number of channels of the feature map.
26
27
28
29
30
31
32
        use_direction_classifier (bool): Whether to add a direction classifier.
        anchor_generator(dict): Config dict of anchor generator.
        assigner_per_size (bool): Whether to do assignment for each separate
            anchor size.
        assign_per_class (bool): Whether to do assignment for each class.
        diff_rad_by_sin (bool): Whether to change the difference into sin
            difference for box regression loss.
wuyuefeng's avatar
wuyuefeng committed
33
        dir_offset (float | int): The offset of BEV rotation angles.
34
            (TODO: may be moved into box coder)
wuyuefeng's avatar
wuyuefeng committed
35
36
37
        dir_limit_offset (float | int): The limited range of BEV
            rotation angles. (TODO: may be moved into box coder)
        bbox_coder (dict): Config dict of box coders.
zhangwenwei's avatar
zhangwenwei committed
38
39
        loss_cls (dict): Config of classification loss.
        loss_bbox (dict): Config of localization loss.
40
        loss_dir (dict): Config of direction classifier loss.
zhangwenwei's avatar
zhangwenwei committed
41
    """
zhangwenwei's avatar
zhangwenwei committed
42
43

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
44
                 num_classes,
zhangwenwei's avatar
zhangwenwei committed
45
46
47
48
49
                 in_channels,
                 train_cfg,
                 test_cfg,
                 feat_channels=256,
                 use_direction_classifier=True,
50
51
52
53
54
55
56
57
                 anchor_generator=dict(
                     type='Anchor3DRangeGenerator',
                     range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
                     strides=[2],
                     sizes=[[1.6, 3.9, 1.56]],
                     rotations=[0, 1.57],
                     custom_values=[],
                     reshape_out=False),
zhangwenwei's avatar
zhangwenwei committed
58
59
60
61
62
                 assigner_per_size=False,
                 assign_per_class=False,
                 diff_rad_by_sin=True,
                 dir_offset=0,
                 dir_limit_offset=1,
63
                 bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
zhangwenwei's avatar
zhangwenwei committed
64
65
66
67
68
69
70
71
72
                 loss_cls=dict(
                     type='CrossEntropyLoss',
                     use_sigmoid=True,
                     loss_weight=1.0),
                 loss_bbox=dict(
                     type='SmoothL1Loss', beta=1.0 / 9.0, loss_weight=2.0),
                 loss_dir=dict(type='CrossEntropyLoss', loss_weight=0.2)):
        super().__init__()
        self.in_channels = in_channels
zhangwenwei's avatar
zhangwenwei committed
73
        self.num_classes = num_classes
zhangwenwei's avatar
zhangwenwei committed
74
75
76
77
78
79
80
81
82
83
84
        self.feat_channels = feat_channels
        self.diff_rad_by_sin = diff_rad_by_sin
        self.use_direction_classifier = use_direction_classifier
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
        self.assigner_per_size = assigner_per_size
        self.assign_per_class = assign_per_class
        self.dir_offset = dir_offset
        self.dir_limit_offset = dir_limit_offset

        # build anchor generator
85
        self.anchor_generator = build_anchor_generator(anchor_generator)
zhangwenwei's avatar
zhangwenwei committed
86
        # In 3D detection, the anchor stride is connected with anchor size
87
        self.num_anchors = self.anchor_generator.num_base_anchors
zhangwenwei's avatar
zhangwenwei committed
88
89
90
        # build box coder
        self.bbox_coder = build_bbox_coder(bbox_coder)
        self.box_code_size = self.bbox_coder.code_size
zhangwenwei's avatar
zhangwenwei committed
91

zhangwenwei's avatar
zhangwenwei committed
92
        # build loss function
zhangwenwei's avatar
zhangwenwei committed
93
        self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
zhangwenwei's avatar
zhangwenwei committed
94
        self.sampling = loss_cls['type'] not in ['FocalLoss', 'GHMC']
zhangwenwei's avatar
zhangwenwei committed
95
96
97
98
99
100
101
        if not self.use_sigmoid_cls:
            self.num_classes += 1
        self.loss_cls = build_loss(loss_cls)
        self.loss_bbox = build_loss(loss_bbox)
        self.loss_dir = build_loss(loss_dir)
        self.fp16_enabled = False

zhangwenwei's avatar
zhangwenwei committed
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
        self._init_layers()
        self._init_assigner_sampler()

    def _init_assigner_sampler(self):
        if self.train_cfg is None:
            return

        if self.sampling:
            self.bbox_sampler = build_sampler(self.train_cfg.sampler)
        else:
            self.bbox_sampler = PseudoSampler()
        if isinstance(self.train_cfg.assigner, dict):
            self.bbox_assigner = build_assigner(self.train_cfg.assigner)
        elif isinstance(self.train_cfg.assigner, list):
            self.bbox_assigner = [
                build_assigner(res) for res in self.train_cfg.assigner
            ]

zhangwenwei's avatar
zhangwenwei committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
    def _init_layers(self):
        self.cls_out_channels = self.num_anchors * self.num_classes
        self.conv_cls = nn.Conv2d(self.feat_channels, self.cls_out_channels, 1)
        self.conv_reg = nn.Conv2d(self.feat_channels,
                                  self.num_anchors * self.box_code_size, 1)
        if self.use_direction_classifier:
            self.conv_dir_cls = nn.Conv2d(self.feat_channels,
                                          self.num_anchors * 2, 1)

    def init_weights(self):
        bias_cls = bias_init_with_prob(0.01)
        normal_init(self.conv_cls, std=0.01, bias=bias_cls)
        normal_init(self.conv_reg, std=0.01)

    def forward_single(self, x):
wuyuefeng's avatar
wuyuefeng committed
135
136
137
        """Forward function on a single-scale feature map.

        Args:
liyinhao's avatar
liyinhao committed
138
            x (torch.Tensor): Input features.
wuyuefeng's avatar
wuyuefeng committed
139
140

        Returns:
liyinhao's avatar
liyinhao committed
141
            tuple[torch.Tensor]: Contain score of each class, bbox predictions
wuyuefeng's avatar
wuyuefeng committed
142
143
                and class predictions of direction.
        """
zhangwenwei's avatar
zhangwenwei committed
144
145
146
147
148
149
150
151
        cls_score = self.conv_cls(x)
        bbox_pred = self.conv_reg(x)
        dir_cls_preds = None
        if self.use_direction_classifier:
            dir_cls_preds = self.conv_dir_cls(x)
        return cls_score, bbox_pred, dir_cls_preds

    def forward(self, feats):
wuyuefeng's avatar
wuyuefeng committed
152
153
154
        """Forward pass.

        Args:
liyinhao's avatar
liyinhao committed
155
            feats (list[torch.Tensor]): Multi-level features, e.g.,
wuyuefeng's avatar
wuyuefeng committed
156
157
158
                features produced by FPN.

        Returns:
liyinhao's avatar
liyinhao committed
159
            tuple[list[torch.Tensor]]: Multi-level class score, bbox
wuyuefeng's avatar
wuyuefeng committed
160
161
                and direction predictions.
        """
zhangwenwei's avatar
zhangwenwei committed
162
163
        return multi_apply(self.forward_single, feats)

164
    def get_anchors(self, featmap_sizes, input_metas, device='cuda'):
zhangwenwei's avatar
zhangwenwei committed
165
        """Get anchors according to feature map sizes.
zhangwenwei's avatar
zhangwenwei committed
166

zhangwenwei's avatar
zhangwenwei committed
167
168
169
        Args:
            featmap_sizes (list[tuple]): Multi-level feature map sizes.
            input_metas (list[dict]): contain pcd and img's meta info.
zhangwenwei's avatar
zhangwenwei committed
170
171
            device (str): device of current module

zhangwenwei's avatar
zhangwenwei committed
172
173
174
175
176
177
        Returns:
            tuple: anchors of each image, valid flags of each image
        """
        num_imgs = len(input_metas)
        # since feature map sizes of all images are the same, we only compute
        # anchors for one time
178
179
        multi_level_anchors = self.anchor_generator.grid_anchors(
            featmap_sizes, device=device)
zhangwenwei's avatar
zhangwenwei committed
180
181
182
183
184
185
        anchor_list = [multi_level_anchors for _ in range(num_imgs)]
        return anchor_list

    def loss_single(self, cls_score, bbox_pred, dir_cls_preds, labels,
                    label_weights, bbox_targets, bbox_weights, dir_targets,
                    dir_weights, num_total_samples):
wuyuefeng's avatar
wuyuefeng committed
186
187
188
        """Calculate loss of Single-level results.

        Args:
liyinhao's avatar
liyinhao committed
189
190
191
            cls_score (torch.Tensor): Class score in single-level.
            bbox_pred (torch.Tensor): Bbox prediction in single-level.
            dir_cls_preds (torch.Tensor): Predictions of direction class
wuyuefeng's avatar
wuyuefeng committed
192
                in single-level.
liyinhao's avatar
liyinhao committed
193
194
195
196
197
198
            labels (torch.Tensor): Labels of class.
            label_weights (torch.Tensor): Weights of class loss.
            bbox_targets (torch.Tensor): Targets of bbox predictions.
            bbox_weights (torch.Tensor): Weights of bbox loss.
            dir_targets (torch.Tensor): Targets of direction predictions.
            dir_weights (torch.Tensor): Weights of direction loss.
wuyuefeng's avatar
wuyuefeng committed
199
200
201
            num_total_samples (int): The number of valid samples.

        Returns:
liyinhao's avatar
liyinhao committed
202
203
            tuple[torch.Tensor]: losses of class, bbox
                and direction, respectively.
wuyuefeng's avatar
wuyuefeng committed
204
        """
zhangwenwei's avatar
zhangwenwei committed
205
206
207
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
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
        # classification loss
        if num_total_samples is None:
            num_total_samples = int(cls_score.shape[0])
        labels = labels.reshape(-1)
        label_weights = label_weights.reshape(-1)
        cls_score = cls_score.permute(0, 2, 3, 1).reshape(-1, self.num_classes)
        loss_cls = self.loss_cls(
            cls_score, labels, label_weights, avg_factor=num_total_samples)

        # regression loss
        bbox_targets = bbox_targets.reshape(-1, self.box_code_size)
        bbox_weights = bbox_weights.reshape(-1, self.box_code_size)
        code_weight = self.train_cfg.get('code_weight', None)

        if code_weight:
            bbox_weights = bbox_weights * bbox_weights.new_tensor(code_weight)
        bbox_pred = bbox_pred.permute(0, 2, 3,
                                      1).reshape(-1, self.box_code_size)
        if self.diff_rad_by_sin:
            bbox_pred, bbox_targets = self.add_sin_difference(
                bbox_pred, bbox_targets)
        loss_bbox = self.loss_bbox(
            bbox_pred,
            bbox_targets,
            bbox_weights,
            avg_factor=num_total_samples)

        # direction classification loss
        loss_dir = None
        if self.use_direction_classifier:
            dir_cls_preds = dir_cls_preds.permute(0, 2, 3, 1).reshape(-1, 2)
            dir_targets = dir_targets.reshape(-1)
            dir_weights = dir_weights.reshape(-1)
            loss_dir = self.loss_dir(
                dir_cls_preds,
                dir_targets,
                dir_weights,
                avg_factor=num_total_samples)

        return loss_cls, loss_bbox, loss_dir

    @staticmethod
    def add_sin_difference(boxes1, boxes2):
zhangwenwei's avatar
zhangwenwei committed
248
249
250
        """Convert the rotation difference to difference in sine function

        Args:
liyinhao's avatar
liyinhao committed
251
252
253
254
            boxes1 (torch.Tensor): shape (NxC), where C>=7 and
                the 7th dimension is rotation dimension
            boxes2 (torch.Tensor): shape (NxC), where C>=7 and the 7th
                dimension is rotation dimension
zhangwenwei's avatar
zhangwenwei committed
255
256
257
258
259
260
261
262
263
264
265
266

        Returns:
            tuple: (boxes1, boxes2) whose 7th dimensions are changed
        """
        rad_pred_encoding = torch.sin(boxes1[..., 6:7]) * torch.cos(
            boxes2[..., 6:7])
        rad_tg_encoding = torch.cos(boxes1[..., 6:7]) * torch.sin(boxes2[...,
                                                                         6:7])
        boxes1 = torch.cat(
            [boxes1[..., :6], rad_pred_encoding, boxes1[..., 7:]], dim=-1)
        boxes2 = torch.cat([boxes2[..., :6], rad_tg_encoding, boxes2[..., 7:]],
                           dim=-1)
zhangwenwei's avatar
zhangwenwei committed
267
268
269
270
271
272
273
274
275
276
        return boxes1, boxes2

    def loss(self,
             cls_scores,
             bbox_preds,
             dir_cls_preds,
             gt_bboxes,
             gt_labels,
             input_metas,
             gt_bboxes_ignore=None):
wuyuefeng's avatar
wuyuefeng committed
277
278
279
        """Calculate losses.

        Args:
liyinhao's avatar
liyinhao committed
280
281
282
            cls_scores (list[torch.Tensor]): Multi-level class scores.
            bbox_preds (list[torch.Tensor]): Multi-level bbox predictions.
            dir_cls_preds (list[torch.Tensor]): Multi-level direction
wuyuefeng's avatar
wuyuefeng committed
283
                class predictions.
zhangwenwei's avatar
zhangwenwei committed
284
            gt_bboxes (list[:obj:`BaseInstance3DBoxes`]): Gt bboxes
wuyuefeng's avatar
wuyuefeng committed
285
                of each sample.
liyinhao's avatar
liyinhao committed
286
            gt_labels (list[torch.Tensor]): Gt labels of each sample.
wuyuefeng's avatar
wuyuefeng committed
287
            input_metas (list[dict]): Contain pcd and img's meta info.
liyinhao's avatar
liyinhao committed
288
289
            gt_bboxes_ignore (None | list[torch.Tensor]): Specify
                which bounding.
wuyuefeng's avatar
wuyuefeng committed
290
291
292
293

        Returns:
            dict: Contain class, bbox and direction losses of each level.
        """
zhangwenwei's avatar
zhangwenwei committed
294
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
295
296
297
298
        assert len(featmap_sizes) == self.anchor_generator.num_levels
        device = cls_scores[0].device
        anchor_list = self.get_anchors(
            featmap_sizes, input_metas, device=device)
zhangwenwei's avatar
zhangwenwei committed
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
        label_channels = self.cls_out_channels if self.use_sigmoid_cls else 1
        cls_reg_targets = self.anchor_target_3d(
            anchor_list,
            gt_bboxes,
            input_metas,
            gt_bboxes_ignore_list=gt_bboxes_ignore,
            gt_labels_list=gt_labels,
            num_classes=self.num_classes,
            label_channels=label_channels,
            sampling=self.sampling)

        if cls_reg_targets is None:
            return None
        (labels_list, label_weights_list, bbox_targets_list, bbox_weights_list,
         dir_targets_list, dir_weights_list, num_total_pos,
         num_total_neg) = cls_reg_targets
        num_total_samples = (
            num_total_pos + num_total_neg if self.sampling else num_total_pos)

        # num_total_samples = None
        losses_cls, losses_bbox, losses_dir = multi_apply(
            self.loss_single,
            cls_scores,
            bbox_preds,
            dir_cls_preds,
            labels_list,
            label_weights_list,
            bbox_targets_list,
            bbox_weights_list,
            dir_targets_list,
            dir_weights_list,
            num_total_samples=num_total_samples)
        return dict(
zhangwenwei's avatar
zhangwenwei committed
332
            loss_cls=losses_cls, loss_bbox=losses_bbox, loss_dir=losses_dir)
zhangwenwei's avatar
zhangwenwei committed
333
334
335
336
337
338

    def get_bboxes(self,
                   cls_scores,
                   bbox_preds,
                   dir_cls_preds,
                   input_metas,
zhangwenwei's avatar
zhangwenwei committed
339
                   cfg=None,
zhangwenwei's avatar
zhangwenwei committed
340
                   rescale=False):
wuyuefeng's avatar
wuyuefeng committed
341
342
343
        """Get bboxes of anchor head.

        Args:
liyinhao's avatar
liyinhao committed
344
345
346
            cls_scores (list[torch.Tensor]): Multi-level class scores.
            bbox_preds (list[torch.Tensor]): Multi-level bbox predictions.
            dir_cls_preds (list[torch.Tensor]): Multi-level direction
wuyuefeng's avatar
wuyuefeng committed
347
348
349
                class predictions.
            input_metas (list[dict]): Contain pcd and img's meta info.
            cfg (None | ConfigDict): Training or testing config.
liyinhao's avatar
liyinhao committed
350
            rescale (list[torch.Tensor]): whether th rescale bbox.
wuyuefeng's avatar
wuyuefeng committed
351
352
353
354

        Returns:
            list[tuple]: prediction resultes of batches.
        """
zhangwenwei's avatar
zhangwenwei committed
355
356
357
        assert len(cls_scores) == len(bbox_preds)
        assert len(cls_scores) == len(dir_cls_preds)
        num_levels = len(cls_scores)
358
359
        featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
        device = cls_scores[0].device
360
        mlvl_anchors = self.anchor_generator.grid_anchors(
361
            featmap_sizes, device=device)
zhangwenwei's avatar
zhangwenwei committed
362
        mlvl_anchors = [
363
            anchor.reshape(-1, self.box_code_size) for anchor in mlvl_anchors
zhangwenwei's avatar
zhangwenwei committed
364
        ]
365

zhangwenwei's avatar
zhangwenwei committed
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
        result_list = []
        for img_id in range(len(input_metas)):
            cls_score_list = [
                cls_scores[i][img_id].detach() for i in range(num_levels)
            ]
            bbox_pred_list = [
                bbox_preds[i][img_id].detach() for i in range(num_levels)
            ]
            dir_cls_pred_list = [
                dir_cls_preds[i][img_id].detach() for i in range(num_levels)
            ]

            input_meta = input_metas[img_id]
            proposals = self.get_bboxes_single(cls_score_list, bbox_pred_list,
                                               dir_cls_pred_list, mlvl_anchors,
zhangwenwei's avatar
zhangwenwei committed
381
                                               input_meta, cfg, rescale)
zhangwenwei's avatar
zhangwenwei committed
382
383
384
385
386
387
388
389
390
            result_list.append(proposals)
        return result_list

    def get_bboxes_single(self,
                          cls_scores,
                          bbox_preds,
                          dir_cls_preds,
                          mlvl_anchors,
                          input_meta,
zhangwenwei's avatar
zhangwenwei committed
391
                          cfg=None,
zhangwenwei's avatar
zhangwenwei committed
392
                          rescale=False):
wuyuefeng's avatar
wuyuefeng committed
393
394
395
        """Get bboxes of single branch.

        Args:
liyinhao's avatar
liyinhao committed
396
397
398
399
400
            cls_scores (torch.Tensor): Class score in single batch.
            bbox_preds (torch.Tensor): Bbox prediction in single batch.
            dir_cls_preds (torch.Tensor): Predictions of direction class
                in single batch.
            mlvl_anchors (List[torch.Tensor]): Multi-level anchors
wuyuefeng's avatar
wuyuefeng committed
401
402
403
                in single batch.
            input_meta (list[dict]): Contain pcd and img's meta info.
            cfg (None | ConfigDict): Training or testing config.
liyinhao's avatar
liyinhao committed
404
            rescale (list[torch.Tensor]): whether th rescale bbox.
wuyuefeng's avatar
wuyuefeng committed
405
406
407

        Returns:
            tuple: Contain predictions of single batch.
zhangwenwei's avatar
zhangwenwei committed
408
                - bboxes (:obj:`BaseInstance3DBoxes`): Predicted 3d bboxes.
liyinhao's avatar
liyinhao committed
409
410
                - scores (torch.Tensor): Class score of each bbox.
                - labels (torch.Tensor): Label of each bbox.
wuyuefeng's avatar
wuyuefeng committed
411
        """
zhangwenwei's avatar
zhangwenwei committed
412
        cfg = self.test_cfg if cfg is None else cfg
zhangwenwei's avatar
zhangwenwei committed
413
414
415
416
417
418
419
        assert len(cls_scores) == len(bbox_preds) == len(mlvl_anchors)
        mlvl_bboxes = []
        mlvl_scores = []
        mlvl_dir_scores = []
        for cls_score, bbox_pred, dir_cls_pred, anchors in zip(
                cls_scores, bbox_preds, dir_cls_preds, mlvl_anchors):
            assert cls_score.size()[-2:] == bbox_pred.size()[-2:]
zhangwenwei's avatar
zhangwenwei committed
420
421
422
            assert cls_score.size()[-2:] == dir_cls_pred.size()[-2:]
            dir_cls_pred = dir_cls_pred.permute(1, 2, 0).reshape(-1, 2)
            dir_cls_score = torch.max(dir_cls_pred, dim=-1)[1]
zhangwenwei's avatar
zhangwenwei committed
423
424
425
426
427
428
429
430
431
432

            cls_score = cls_score.permute(1, 2,
                                          0).reshape(-1, self.num_classes)
            if self.use_sigmoid_cls:
                scores = cls_score.sigmoid()
            else:
                scores = cls_score.softmax(-1)
            bbox_pred = bbox_pred.permute(1, 2,
                                          0).reshape(-1, self.box_code_size)

zhangwenwei's avatar
zhangwenwei committed
433
434
            nms_pre = cfg.get('nms_pre', -1)
            if nms_pre > 0 and scores.shape[0] > nms_pre:
zhangwenwei's avatar
zhangwenwei committed
435
436
437
                if self.use_sigmoid_cls:
                    max_scores, _ = scores.max(dim=1)
                else:
zhangwenwei's avatar
zhangwenwei committed
438
439
440
441
442
443
444
                    max_scores, _ = scores[:, :-1].max(dim=1)
                _, topk_inds = max_scores.topk(nms_pre)
                anchors = anchors[topk_inds, :]
                bbox_pred = bbox_pred[topk_inds, :]
                scores = scores[topk_inds, :]
                dir_cls_score = dir_cls_score[topk_inds]

445
            bboxes = self.bbox_coder.decode(anchors, bbox_pred)
zhangwenwei's avatar
zhangwenwei committed
446
447
            mlvl_bboxes.append(bboxes)
            mlvl_scores.append(scores)
zhangwenwei's avatar
zhangwenwei committed
448
            mlvl_dir_scores.append(dir_cls_score)
zhangwenwei's avatar
zhangwenwei committed
449
450

        mlvl_bboxes = torch.cat(mlvl_bboxes)
zhangwenwei's avatar
zhangwenwei committed
451
452
        mlvl_bboxes_for_nms = xywhr2xyxyr(input_meta['box_type_3d'](
            mlvl_bboxes, box_dim=self.box_code_size).bev)
zhangwenwei's avatar
zhangwenwei committed
453
454
455
        mlvl_scores = torch.cat(mlvl_scores)
        mlvl_dir_scores = torch.cat(mlvl_dir_scores)

zhangwenwei's avatar
zhangwenwei committed
456
457
458
459
460
461
462
463
464
465
466
        if self.use_sigmoid_cls:
            # Add a dummy background class to the front when using sigmoid
            padding = mlvl_scores.new_zeros(mlvl_scores.shape[0], 1)
            mlvl_scores = torch.cat([mlvl_scores, padding], dim=1)

        score_thr = cfg.get('score_thr', 0)
        results = box3d_multiclass_nms(mlvl_bboxes, mlvl_bboxes_for_nms,
                                       mlvl_scores, score_thr, cfg.max_num,
                                       cfg, mlvl_dir_scores)
        bboxes, scores, labels, dir_scores = results
        if bboxes.shape[0] > 0:
zhangwenwei's avatar
zhangwenwei committed
467
468
            dir_rot = limit_period(bboxes[..., 6] - self.dir_offset,
                                   self.dir_limit_offset, np.pi)
zhangwenwei's avatar
zhangwenwei committed
469
            bboxes[..., 6] = (
zhangwenwei's avatar
zhangwenwei committed
470
                dir_rot + self.dir_offset +
zhangwenwei's avatar
zhangwenwei committed
471
                np.pi * dir_scores.to(bboxes.dtype))
472
        bboxes = input_meta['box_type_3d'](bboxes, box_dim=self.box_code_size)
zhangwenwei's avatar
zhangwenwei committed
473
        return bboxes, scores, labels