anchor3d_head.py 21.2 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 mmcv.runner import BaseModule, force_fp32
zhangwenwei's avatar
zhangwenwei committed
5
from torch import nn as nn
zhangwenwei's avatar
zhangwenwei committed
6

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


15
@HEADS.register_module()
16
class Anchor3DHead(BaseModule, AnchorTrainMixin):
zhangwenwei's avatar
zhangwenwei committed
17
    """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
                 anchor_generator=dict(
                     type='Anchor3DRangeGenerator',
                     range=[0, -39.68, -1.78, 69.12, 39.68, -1.78],
                     strides=[2],
53
                     sizes=[[3.9, 1.6, 1.56]],
54
55
56
                     rotations=[0, 1.57],
                     custom_values=[],
                     reshape_out=False),
zhangwenwei's avatar
zhangwenwei committed
57
58
59
                 assigner_per_size=False,
                 assign_per_class=False,
                 diff_rad_by_sin=True,
60
61
                 dir_offset=-np.pi / 2,
                 dir_limit_offset=0,
62
                 bbox_coder=dict(type='DeltaXYZWLHRBBoxCoder'),
zhangwenwei's avatar
zhangwenwei committed
63
64
65
66
67
68
                 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),
69
70
71
                 loss_dir=dict(type='CrossEntropyLoss', loss_weight=0.2),
                 init_cfg=None):
        super().__init__(init_cfg=init_cfg)
zhangwenwei's avatar
zhangwenwei committed
72
        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
        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
83
84
85
86
        import warnings
        warnings.warn(
            'dir_offset and dir_limit_offset will be depressed and be '
            'incorporated into box coder in the future')
87
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
88
89

        # build anchor generator
90
        self.anchor_generator = build_prior_generator(anchor_generator)
zhangwenwei's avatar
zhangwenwei committed
91
        # In 3D detection, the anchor stride is connected with anchor size
92
        self.num_anchors = self.anchor_generator.num_base_anchors
zhangwenwei's avatar
zhangwenwei committed
93
94
95
        # 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
96

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

110
111
112
113
114
115
116
117
        if init_cfg is None:
            self.init_cfg = dict(
                type='Normal',
                layer='Conv2d',
                std=0.01,
                override=dict(
                    type='Normal', name='conv_cls', std=0.01, bias_prob=0.01))

zhangwenwei's avatar
zhangwenwei committed
118
    def _init_assigner_sampler(self):
119
        """Initialize the target assigner and sampler of the head."""
zhangwenwei's avatar
zhangwenwei committed
120
121
122
123
124
125
126
127
128
129
130
131
132
133
        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
134
    def _init_layers(self):
135
        """Initialize neural network layers of the head."""
zhangwenwei's avatar
zhangwenwei committed
136
137
138
139
140
141
142
143
144
        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 forward_single(self, x):
wuyuefeng's avatar
wuyuefeng committed
145
146
147
        """Forward function on a single-scale feature map.

        Args:
liyinhao's avatar
liyinhao committed
148
            x (torch.Tensor): Input features.
wuyuefeng's avatar
wuyuefeng committed
149
150

        Returns:
151
            tuple[torch.Tensor]: Contain score of each class, bbox
zhangwenwei's avatar
zhangwenwei committed
152
                regression and direction classification predictions.
wuyuefeng's avatar
wuyuefeng committed
153
        """
zhangwenwei's avatar
zhangwenwei committed
154
155
156
157
158
159
160
161
        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
162
163
164
        """Forward pass.

        Args:
liyinhao's avatar
liyinhao committed
165
            feats (list[torch.Tensor]): Multi-level features, e.g.,
wuyuefeng's avatar
wuyuefeng committed
166
167
168
                features produced by FPN.

        Returns:
169
            tuple[list[torch.Tensor]]: Multi-level class score, bbox
wuyuefeng's avatar
wuyuefeng committed
170
171
                and direction predictions.
        """
zhangwenwei's avatar
zhangwenwei committed
172
173
        return multi_apply(self.forward_single, feats)

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

zhangwenwei's avatar
zhangwenwei committed
177
178
179
        Args:
            featmap_sizes (list[tuple]): Multi-level feature map sizes.
            input_metas (list[dict]): contain pcd and img's meta info.
wangtai's avatar
wangtai committed
180
            device (str): device of current module.
zhangwenwei's avatar
zhangwenwei committed
181

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

        Args:
liyinhao's avatar
liyinhao committed
200
201
202
            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
203
                in single-level.
liyinhao's avatar
liyinhao committed
204
205
206
207
208
209
            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
210
211
212
            num_total_samples (int): The number of valid samples.

        Returns:
213
            tuple[torch.Tensor]: Losses of class, bbox
liyinhao's avatar
liyinhao committed
214
                and direction, respectively.
wuyuefeng's avatar
wuyuefeng committed
215
        """
zhangwenwei's avatar
zhangwenwei committed
216
217
218
219
220
221
        # 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)
222
        assert labels.max().item() <= self.num_classes
zhangwenwei's avatar
zhangwenwei committed
223
224
225
226
        loss_cls = self.loss_cls(
            cls_score, labels, label_weights, avg_factor=num_total_samples)

        # regression loss
227
228
        bbox_pred = bbox_pred.permute(0, 2, 3,
                                      1).reshape(-1, self.box_code_size)
zhangwenwei's avatar
zhangwenwei committed
229
230
231
        bbox_targets = bbox_targets.reshape(-1, self.box_code_size)
        bbox_weights = bbox_weights.reshape(-1, self.box_code_size)

232
233
        bg_class_ind = self.num_classes
        pos_inds = ((labels >= 0)
Wenhao Wu's avatar
Wenhao Wu committed
234
235
                    & (labels < bg_class_ind)).nonzero(
                        as_tuple=False).reshape(-1)
236
237
238
239
240
241
242
        num_pos = len(pos_inds)

        pos_bbox_pred = bbox_pred[pos_inds]
        pos_bbox_targets = bbox_targets[pos_inds]
        pos_bbox_weights = bbox_weights[pos_inds]

        # dir loss
zhangwenwei's avatar
zhangwenwei committed
243
244
245
246
        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)
247
248
249
250
251
252
253
            pos_dir_cls_preds = dir_cls_preds[pos_inds]
            pos_dir_targets = dir_targets[pos_inds]
            pos_dir_weights = dir_weights[pos_inds]

        if num_pos > 0:
            code_weight = self.train_cfg.get('code_weight', None)
            if code_weight:
254
                pos_bbox_weights = pos_bbox_weights * bbox_weights.new_tensor(
255
256
257
258
259
260
261
262
                    code_weight)
            if self.diff_rad_by_sin:
                pos_bbox_pred, pos_bbox_targets = self.add_sin_difference(
                    pos_bbox_pred, pos_bbox_targets)
            loss_bbox = self.loss_bbox(
                pos_bbox_pred,
                pos_bbox_targets,
                pos_bbox_weights,
zhangwenwei's avatar
zhangwenwei committed
263
264
                avg_factor=num_total_samples)

265
266
267
268
269
270
271
272
273
274
275
276
277
            # direction classification loss
            loss_dir = None
            if self.use_direction_classifier:
                loss_dir = self.loss_dir(
                    pos_dir_cls_preds,
                    pos_dir_targets,
                    pos_dir_weights,
                    avg_factor=num_total_samples)
        else:
            loss_bbox = pos_bbox_pred.sum()
            if self.use_direction_classifier:
                loss_dir = pos_dir_cls_preds.sum()

zhangwenwei's avatar
zhangwenwei committed
278
279
280
281
        return loss_cls, loss_bbox, loss_dir

    @staticmethod
    def add_sin_difference(boxes1, boxes2):
zhangwenwei's avatar
zhangwenwei committed
282
        """Convert the rotation difference to difference in sine function.
zhangwenwei's avatar
zhangwenwei committed
283
284

        Args:
zhangwenwei's avatar
zhangwenwei committed
285
286
287
288
            boxes1 (torch.Tensor): Original Boxes in shape (NxC), where C>=7
                and the 7th dimension is rotation dimension.
            boxes2 (torch.Tensor): Target boxes in shape (NxC), where C>=7 and
                the 7th dimension is rotation dimension.
zhangwenwei's avatar
zhangwenwei committed
289
290

        Returns:
291
            tuple[torch.Tensor]: ``boxes1`` and ``boxes2`` whose 7th
zhangwenwei's avatar
zhangwenwei committed
292
                dimensions are changed.
zhangwenwei's avatar
zhangwenwei committed
293
294
295
296
297
298
299
300
301
        """
        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
302
303
        return boxes1, boxes2

304
    @force_fp32(apply_to=('cls_scores', 'bbox_preds', 'dir_cls_preds'))
zhangwenwei's avatar
zhangwenwei committed
305
306
307
308
309
310
311
312
    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
313
314
315
        """Calculate losses.

        Args:
liyinhao's avatar
liyinhao committed
316
317
318
            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
319
                class predictions.
zhangwenwei's avatar
zhangwenwei committed
320
            gt_bboxes (list[:obj:`BaseInstance3DBoxes`]): Gt bboxes
wuyuefeng's avatar
wuyuefeng committed
321
                of each sample.
liyinhao's avatar
liyinhao committed
322
            gt_labels (list[torch.Tensor]): Gt labels of each sample.
wuyuefeng's avatar
wuyuefeng committed
323
            input_metas (list[dict]): Contain pcd and img's meta info.
324
325
            gt_bboxes_ignore (list[torch.Tensor]): Specify
                which bounding boxes to ignore.
wuyuefeng's avatar
wuyuefeng committed
326
327

        Returns:
328
            dict[str, list[torch.Tensor]]: Classification, bbox, and
zhangwenwei's avatar
zhangwenwei committed
329
                direction losses of each level.
330

331
332
                - loss_cls (list[torch.Tensor]): Classification losses.
                - loss_bbox (list[torch.Tensor]): Box regression losses.
333
                - loss_dir (list[torch.Tensor]): Direction classification
334
                    losses.
wuyuefeng's avatar
wuyuefeng committed
335
        """
zhangwenwei's avatar
zhangwenwei committed
336
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
337
338
339
340
        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
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
        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
374
            loss_cls=losses_cls, loss_bbox=losses_bbox, loss_dir=losses_dir)
zhangwenwei's avatar
zhangwenwei committed
375
376
377
378
379
380

    def get_bboxes(self,
                   cls_scores,
                   bbox_preds,
                   dir_cls_preds,
                   input_metas,
zhangwenwei's avatar
zhangwenwei committed
381
                   cfg=None,
zhangwenwei's avatar
zhangwenwei committed
382
                   rescale=False):
wuyuefeng's avatar
wuyuefeng committed
383
384
385
        """Get bboxes of anchor head.

        Args:
liyinhao's avatar
liyinhao committed
386
387
388
            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
389
390
                class predictions.
            input_metas (list[dict]): Contain pcd and img's meta info.
391
            cfg (:obj:`ConfigDict`): Training or testing config.
wangtai's avatar
wangtai committed
392
            rescale (list[torch.Tensor]): Whether th rescale bbox.
wuyuefeng's avatar
wuyuefeng committed
393
394

        Returns:
wangtai's avatar
wangtai committed
395
            list[tuple]: Prediction resultes of batches.
wuyuefeng's avatar
wuyuefeng committed
396
        """
zhangwenwei's avatar
zhangwenwei committed
397
398
399
        assert len(cls_scores) == len(bbox_preds)
        assert len(cls_scores) == len(dir_cls_preds)
        num_levels = len(cls_scores)
400
401
        featmap_sizes = [cls_scores[i].shape[-2:] for i in range(num_levels)]
        device = cls_scores[0].device
402
        mlvl_anchors = self.anchor_generator.grid_anchors(
403
            featmap_sizes, device=device)
zhangwenwei's avatar
zhangwenwei committed
404
        mlvl_anchors = [
405
            anchor.reshape(-1, self.box_code_size) for anchor in mlvl_anchors
zhangwenwei's avatar
zhangwenwei committed
406
        ]
407

zhangwenwei's avatar
zhangwenwei committed
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
        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
423
                                               input_meta, cfg, rescale)
zhangwenwei's avatar
zhangwenwei committed
424
425
426
427
428
429
430
431
432
            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
433
                          cfg=None,
zhangwenwei's avatar
zhangwenwei committed
434
                          rescale=False):
wuyuefeng's avatar
wuyuefeng committed
435
436
437
        """Get bboxes of single branch.

        Args:
liyinhao's avatar
liyinhao committed
438
439
440
441
442
            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
443
444
                in single batch.
            input_meta (list[dict]): Contain pcd and img's meta info.
445
            cfg (:obj:`ConfigDict`): Training or testing config.
liyinhao's avatar
liyinhao committed
446
            rescale (list[torch.Tensor]): whether th rescale bbox.
wuyuefeng's avatar
wuyuefeng committed
447
448
449

        Returns:
            tuple: Contain predictions of single batch.
450

zhangwenwei's avatar
zhangwenwei committed
451
                - bboxes (:obj:`BaseInstance3DBoxes`): Predicted 3d bboxes.
liyinhao's avatar
liyinhao committed
452
453
                - scores (torch.Tensor): Class score of each bbox.
                - labels (torch.Tensor): Label of each bbox.
wuyuefeng's avatar
wuyuefeng committed
454
        """
zhangwenwei's avatar
zhangwenwei committed
455
        cfg = self.test_cfg if cfg is None else cfg
zhangwenwei's avatar
zhangwenwei committed
456
457
458
459
460
461
462
        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
463
464
465
            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
466
467
468
469
470
471
472
473
474
475

            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
476
477
            nms_pre = cfg.get('nms_pre', -1)
            if nms_pre > 0 and scores.shape[0] > nms_pre:
zhangwenwei's avatar
zhangwenwei committed
478
479
480
                if self.use_sigmoid_cls:
                    max_scores, _ = scores.max(dim=1)
                else:
zhangwenwei's avatar
zhangwenwei committed
481
482
483
484
485
486
487
                    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]

488
            bboxes = self.bbox_coder.decode(anchors, bbox_pred)
zhangwenwei's avatar
zhangwenwei committed
489
490
            mlvl_bboxes.append(bboxes)
            mlvl_scores.append(scores)
zhangwenwei's avatar
zhangwenwei committed
491
            mlvl_dir_scores.append(dir_cls_score)
zhangwenwei's avatar
zhangwenwei committed
492
493

        mlvl_bboxes = torch.cat(mlvl_bboxes)
zhangwenwei's avatar
zhangwenwei committed
494
495
        mlvl_bboxes_for_nms = xywhr2xyxyr(input_meta['box_type_3d'](
            mlvl_bboxes, box_dim=self.box_code_size).bev)
zhangwenwei's avatar
zhangwenwei committed
496
497
498
        mlvl_scores = torch.cat(mlvl_scores)
        mlvl_dir_scores = torch.cat(mlvl_dir_scores)

zhangwenwei's avatar
zhangwenwei committed
499
500
501
502
503
504
505
506
507
508
509
        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
510
511
            dir_rot = limit_period(bboxes[..., 6] - self.dir_offset,
                                   self.dir_limit_offset, np.pi)
zhangwenwei's avatar
zhangwenwei committed
512
            bboxes[..., 6] = (
zhangwenwei's avatar
zhangwenwei committed
513
                dir_rot + self.dir_offset +
zhangwenwei's avatar
zhangwenwei committed
514
                np.pi * dir_scores.to(bboxes.dtype))
515
        bboxes = input_meta['box_type_3d'](bboxes, box_dim=self.box_code_size)
zhangwenwei's avatar
zhangwenwei committed
516
        return bboxes, scores, labels