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

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


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

zhangwenwei's avatar
zhangwenwei committed
19
    Args:
zhangwenwei's avatar
zhangwenwei committed
20
        num_classes (int): Number of classes.
zhangwenwei's avatar
zhangwenwei committed
21
        in_channels (int): Number of channels in the input feature map.
wuyuefeng's avatar
wuyuefeng committed
22
23
        train_cfg (dict): Train configs.
        test_cfg (dict): Test configs.
zhangwenwei's avatar
zhangwenwei committed
24
        feat_channels (int): Number of channels of the feature map.
25
26
27
28
29
30
31
        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
32
        dir_offset (float | int): The offset of BEV rotation angles.
33
            (TODO: may be moved into box coder)
wuyuefeng's avatar
wuyuefeng committed
34
35
36
        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
37
38
        loss_cls (dict): Config of classification loss.
        loss_bbox (dict): Config of localization loss.
39
        loss_dir (dict): Config of direction classifier loss.
zhangwenwei's avatar
zhangwenwei committed
40
    """
zhangwenwei's avatar
zhangwenwei committed
41
42

    def __init__(self,
zhangwenwei's avatar
zhangwenwei committed
43
                 num_classes,
zhangwenwei's avatar
zhangwenwei committed
44
45
46
47
48
                 in_channels,
                 train_cfg,
                 test_cfg,
                 feat_channels=256,
                 use_direction_classifier=True,
49
50
51
52
53
54
55
56
                 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
57
58
59
60
61
                 assigner_per_size=False,
                 assign_per_class=False,
                 diff_rad_by_sin=True,
                 dir_offset=0,
                 dir_limit_offset=1,
62
                 bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
zhangwenwei's avatar
zhangwenwei committed
63
64
65
66
67
68
69
70
71
                 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
72
        self.num_classes = num_classes
zhangwenwei's avatar
zhangwenwei committed
73
74
75
76
77
78
79
80
81
82
83
        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
84
        self.anchor_generator = build_anchor_generator(anchor_generator)
zhangwenwei's avatar
zhangwenwei committed
85
        # In 3D detection, the anchor stride is connected with anchor size
86
        self.num_anchors = self.anchor_generator.num_base_anchors
zhangwenwei's avatar
zhangwenwei committed
87
88
89
        # 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
90

zhangwenwei's avatar
zhangwenwei committed
91
        # build loss function
zhangwenwei's avatar
zhangwenwei committed
92
        self.use_sigmoid_cls = loss_cls.get('use_sigmoid', False)
zhangwenwei's avatar
zhangwenwei committed
93
        self.sampling = loss_cls['type'] not in ['FocalLoss', 'GHMC']
zhangwenwei's avatar
zhangwenwei committed
94
95
96
97
98
99
100
        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
101
102
103
104
        self._init_layers()
        self._init_assigner_sampler()

    def _init_assigner_sampler(self):
105
        """Initialize the target assigner and sampler of the head."""
zhangwenwei's avatar
zhangwenwei committed
106
107
108
109
110
111
112
113
114
115
116
117
118
119
        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
    def _init_layers(self):
121
        """Initialize neural network layers of the head."""
zhangwenwei's avatar
zhangwenwei committed
122
123
124
125
126
127
128
129
130
        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):
131
        """Initialize the weights of head."""
zhangwenwei's avatar
zhangwenwei committed
132
133
134
135
136
        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
137
138
139
        """Forward function on a single-scale feature map.

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

        Returns:
liyinhao's avatar
liyinhao committed
143
            tuple[torch.Tensor]: Contain score of each class, bbox predictions
wuyuefeng's avatar
wuyuefeng committed
144
145
                and class predictions of direction.
        """
zhangwenwei's avatar
zhangwenwei committed
146
147
148
149
150
151
152
153
        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
154
155
156
        """Forward pass.

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

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

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

zhangwenwei's avatar
zhangwenwei committed
169
170
171
        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
172
173
            device (str): device of current module

zhangwenwei's avatar
zhangwenwei committed
174
        Returns:
175
176
            list[list[torch.Tensor]]: anchors of each image, valid flags
                of each image
zhangwenwei's avatar
zhangwenwei committed
177
178
179
180
        """
        num_imgs = len(input_metas)
        # since feature map sizes of all images are the same, we only compute
        # anchors for one time
181
182
        multi_level_anchors = self.anchor_generator.grid_anchors(
            featmap_sizes, device=device)
zhangwenwei's avatar
zhangwenwei committed
183
184
185
186
187
188
        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
189
190
191
        """Calculate loss of Single-level results.

        Args:
liyinhao's avatar
liyinhao committed
192
193
194
            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
195
                in single-level.
liyinhao's avatar
liyinhao committed
196
197
198
199
200
201
            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
202
203
204
            num_total_samples (int): The number of valid samples.

        Returns:
liyinhao's avatar
liyinhao committed
205
206
            tuple[torch.Tensor]: losses of class, bbox
                and direction, respectively.
wuyuefeng's avatar
wuyuefeng committed
207
        """
zhangwenwei's avatar
zhangwenwei committed
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
248
249
250
        # 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
251
        """Convert the rotation difference to difference in sine function.
zhangwenwei's avatar
zhangwenwei committed
252
253

        Args:
liyinhao's avatar
liyinhao committed
254
255
256
257
            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
258
259

        Returns:
260
261
            tuple[torch.Tensor]: boxes1 and boxes2 whose 7th dimensions
                are changed
zhangwenwei's avatar
zhangwenwei committed
262
263
264
265
266
267
268
269
270
        """
        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
271
272
273
274
275
276
277
278
279
280
        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
281
282
283
        """Calculate losses.

        Args:
liyinhao's avatar
liyinhao committed
284
285
286
            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
287
                class predictions.
zhangwenwei's avatar
zhangwenwei committed
288
            gt_bboxes (list[:obj:`BaseInstance3DBoxes`]): Gt bboxes
wuyuefeng's avatar
wuyuefeng committed
289
                of each sample.
liyinhao's avatar
liyinhao committed
290
            gt_labels (list[torch.Tensor]): Gt labels of each sample.
wuyuefeng's avatar
wuyuefeng committed
291
            input_metas (list[dict]): Contain pcd and img's meta info.
liyinhao's avatar
liyinhao committed
292
293
            gt_bboxes_ignore (None | list[torch.Tensor]): Specify
                which bounding.
wuyuefeng's avatar
wuyuefeng committed
294
295

        Returns:
296
297
            dict[str, list[torch.Tensor]]: Classification, bbox, and direction
                losses of each level.
298

299
300
301
302
                - loss_cls (list[torch.Tensor]): Classification losses.
                - loss_bbox (list[torch.Tensor]): Box regression losses.
                - loss_dir (list[torch.Tensor]): Direction classification
                    losses.
wuyuefeng's avatar
wuyuefeng committed
303
        """
zhangwenwei's avatar
zhangwenwei committed
304
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
305
306
307
308
        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
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
        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
342
            loss_cls=losses_cls, loss_bbox=losses_bbox, loss_dir=losses_dir)
zhangwenwei's avatar
zhangwenwei committed
343
344
345
346
347
348

    def get_bboxes(self,
                   cls_scores,
                   bbox_preds,
                   dir_cls_preds,
                   input_metas,
zhangwenwei's avatar
zhangwenwei committed
349
                   cfg=None,
zhangwenwei's avatar
zhangwenwei committed
350
                   rescale=False):
wuyuefeng's avatar
wuyuefeng committed
351
352
353
        """Get bboxes of anchor head.

        Args:
liyinhao's avatar
liyinhao committed
354
355
356
            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
357
358
359
                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
360
            rescale (list[torch.Tensor]): whether th rescale bbox.
wuyuefeng's avatar
wuyuefeng committed
361
362
363
364

        Returns:
            list[tuple]: prediction resultes of batches.
        """
zhangwenwei's avatar
zhangwenwei committed
365
366
367
        assert len(cls_scores) == len(bbox_preds)
        assert len(cls_scores) == len(dir_cls_preds)
        num_levels = len(cls_scores)
368
369
        featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
        device = cls_scores[0].device
370
        mlvl_anchors = self.anchor_generator.grid_anchors(
371
            featmap_sizes, device=device)
zhangwenwei's avatar
zhangwenwei committed
372
        mlvl_anchors = [
373
            anchor.reshape(-1, self.box_code_size) for anchor in mlvl_anchors
zhangwenwei's avatar
zhangwenwei committed
374
        ]
375

zhangwenwei's avatar
zhangwenwei committed
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
        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
391
                                               input_meta, cfg, rescale)
zhangwenwei's avatar
zhangwenwei committed
392
393
394
395
396
397
398
399
400
            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
401
                          cfg=None,
zhangwenwei's avatar
zhangwenwei committed
402
                          rescale=False):
wuyuefeng's avatar
wuyuefeng committed
403
404
405
        """Get bboxes of single branch.

        Args:
liyinhao's avatar
liyinhao committed
406
407
408
409
410
            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
411
412
413
                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
414
            rescale (list[torch.Tensor]): whether th rescale bbox.
wuyuefeng's avatar
wuyuefeng committed
415
416
417

        Returns:
            tuple: Contain predictions of single batch.
418

zhangwenwei's avatar
zhangwenwei committed
419
                - bboxes (:obj:`BaseInstance3DBoxes`): Predicted 3d bboxes.
liyinhao's avatar
liyinhao committed
420
421
                - scores (torch.Tensor): Class score of each bbox.
                - labels (torch.Tensor): Label of each bbox.
wuyuefeng's avatar
wuyuefeng committed
422
        """
zhangwenwei's avatar
zhangwenwei committed
423
        cfg = self.test_cfg if cfg is None else cfg
zhangwenwei's avatar
zhangwenwei committed
424
425
426
427
428
429
430
        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
431
432
433
            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
434
435
436
437
438
439
440
441
442
443

            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
444
445
            nms_pre = cfg.get('nms_pre', -1)
            if nms_pre > 0 and scores.shape[0] > nms_pre:
zhangwenwei's avatar
zhangwenwei committed
446
447
448
                if self.use_sigmoid_cls:
                    max_scores, _ = scores.max(dim=1)
                else:
zhangwenwei's avatar
zhangwenwei committed
449
450
451
452
453
454
455
                    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]

456
            bboxes = self.bbox_coder.decode(anchors, bbox_pred)
zhangwenwei's avatar
zhangwenwei committed
457
458
            mlvl_bboxes.append(bboxes)
            mlvl_scores.append(scores)
zhangwenwei's avatar
zhangwenwei committed
459
            mlvl_dir_scores.append(dir_cls_score)
zhangwenwei's avatar
zhangwenwei committed
460
461

        mlvl_bboxes = torch.cat(mlvl_bboxes)
zhangwenwei's avatar
zhangwenwei committed
462
463
        mlvl_bboxes_for_nms = xywhr2xyxyr(input_meta['box_type_3d'](
            mlvl_bboxes, box_dim=self.box_code_size).bev)
zhangwenwei's avatar
zhangwenwei committed
464
465
466
        mlvl_scores = torch.cat(mlvl_scores)
        mlvl_dir_scores = torch.cat(mlvl_dir_scores)

zhangwenwei's avatar
zhangwenwei committed
467
468
469
470
471
472
473
474
475
476
477
        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
478
479
            dir_rot = limit_period(bboxes[..., 6] - self.dir_offset,
                                   self.dir_limit_offset, np.pi)
zhangwenwei's avatar
zhangwenwei committed
480
            bboxes[..., 6] = (
zhangwenwei's avatar
zhangwenwei committed
481
                dir_rot + self.dir_offset +
zhangwenwei's avatar
zhangwenwei committed
482
                np.pi * dir_scores.to(bboxes.dtype))
483
        bboxes = input_meta['box_type_3d'](bboxes, box_dim=self.box_code_size)
zhangwenwei's avatar
zhangwenwei committed
484
        return bboxes, scores, labels