anchor3d_head.py 21 KB
Newer Older
zhangwenwei's avatar
zhangwenwei committed
1
2
import numpy as np
import torch
3
from mmcv.runner import BaseModule, force_fp32
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()
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
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
                 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
        self.fp16_enabled = False
zhangwenwei's avatar
zhangwenwei committed
84
85

        # build anchor generator
86
        self.anchor_generator = build_anchor_generator(anchor_generator)
zhangwenwei's avatar
zhangwenwei committed
87
        # In 3D detection, the anchor stride is connected with anchor size
88
        self.num_anchors = self.anchor_generator.num_base_anchors
zhangwenwei's avatar
zhangwenwei committed
89
90
91
        # 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
92

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

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

        Args:
liyinhao's avatar
liyinhao committed
144
            x (torch.Tensor): Input features.
wuyuefeng's avatar
wuyuefeng committed
145
146

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

        Args:
liyinhao's avatar
liyinhao committed
161
            feats (list[torch.Tensor]): Multi-level features, e.g.,
wuyuefeng's avatar
wuyuefeng committed
162
163
164
                features produced by FPN.

        Returns:
zhangwenwei's avatar
zhangwenwei committed
165
            tuple[list[torch.Tensor]]: Multi-level class score, bbox \
wuyuefeng's avatar
wuyuefeng committed
166
167
                and direction predictions.
        """
zhangwenwei's avatar
zhangwenwei committed
168
169
        return multi_apply(self.forward_single, feats)

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

zhangwenwei's avatar
zhangwenwei committed
173
174
175
        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
176
            device (str): device of current module.
zhangwenwei's avatar
zhangwenwei committed
177

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

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

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

        # regression loss
223
224
        bbox_pred = bbox_pred.permute(0, 2, 3,
                                      1).reshape(-1, self.box_code_size)
zhangwenwei's avatar
zhangwenwei committed
225
226
227
        bbox_targets = bbox_targets.reshape(-1, self.box_code_size)
        bbox_weights = bbox_weights.reshape(-1, self.box_code_size)

228
229
        bg_class_ind = self.num_classes
        pos_inds = ((labels >= 0)
Wenhao Wu's avatar
Wenhao Wu committed
230
231
                    & (labels < bg_class_ind)).nonzero(
                        as_tuple=False).reshape(-1)
232
233
234
235
236
237
238
        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
239
240
241
242
        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)
243
244
245
246
247
248
249
            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:
250
                pos_bbox_weights = pos_bbox_weights * bbox_weights.new_tensor(
251
252
253
254
255
256
257
258
                    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
259
260
                avg_factor=num_total_samples)

261
262
263
264
265
266
267
268
269
270
271
272
273
            # 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
274
275
276
277
        return loss_cls, loss_bbox, loss_dir

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

        Args:
zhangwenwei's avatar
zhangwenwei committed
281
282
283
284
            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
285
286

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

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

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

        Returns:
zhangwenwei's avatar
zhangwenwei committed
324
325
            dict[str, list[torch.Tensor]]: Classification, bbox, and \
                direction losses of each level.
326

327
328
                - loss_cls (list[torch.Tensor]): Classification losses.
                - loss_bbox (list[torch.Tensor]): Box regression losses.
zhangwenwei's avatar
zhangwenwei committed
329
                - loss_dir (list[torch.Tensor]): Direction classification \
330
                    losses.
wuyuefeng's avatar
wuyuefeng committed
331
        """
zhangwenwei's avatar
zhangwenwei committed
332
        featmap_sizes = [featmap.size()[-2:] for featmap in cls_scores]
333
334
335
336
        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
337
338
339
340
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
        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
370
            loss_cls=losses_cls, loss_bbox=losses_bbox, loss_dir=losses_dir)
zhangwenwei's avatar
zhangwenwei committed
371
372
373
374
375
376

    def get_bboxes(self,
                   cls_scores,
                   bbox_preds,
                   dir_cls_preds,
                   input_metas,
zhangwenwei's avatar
zhangwenwei committed
377
                   cfg=None,
zhangwenwei's avatar
zhangwenwei committed
378
                   rescale=False):
wuyuefeng's avatar
wuyuefeng committed
379
380
381
        """Get bboxes of anchor head.

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

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

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

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

        Returns:
            tuple: Contain predictions of single batch.
446

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

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

484
            bboxes = self.bbox_coder.decode(anchors, bbox_pred)
zhangwenwei's avatar
zhangwenwei committed
485
486
            mlvl_bboxes.append(bboxes)
            mlvl_scores.append(scores)
zhangwenwei's avatar
zhangwenwei committed
487
            mlvl_dir_scores.append(dir_cls_score)
zhangwenwei's avatar
zhangwenwei committed
488
489

        mlvl_bboxes = torch.cat(mlvl_bboxes)
zhangwenwei's avatar
zhangwenwei committed
490
491
        mlvl_bboxes_for_nms = xywhr2xyxyr(input_meta['box_type_3d'](
            mlvl_bboxes, box_dim=self.box_code_size).bev)
zhangwenwei's avatar
zhangwenwei committed
492
493
494
        mlvl_scores = torch.cat(mlvl_scores)
        mlvl_dir_scores = torch.cat(mlvl_dir_scores)

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