vote_head.py 32.9 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
jshilong's avatar
jshilong committed
2
3
from typing import Dict, List, Optional, Union

wuyuefeng's avatar
Votenet  
wuyuefeng committed
4
5
import numpy as np
import torch
6
from mmcv.ops import furthest_point_sample
7
from mmcv.runner import BaseModule, force_fp32
jshilong's avatar
jshilong committed
8
from mmengine import ConfigDict, InstanceData
zhangwenwei's avatar
zhangwenwei committed
9
from torch.nn import functional as F
wuyuefeng's avatar
Votenet  
wuyuefeng committed
10
11
12
13

from mmdet3d.core.post_processing import aligned_3d_nms
from mmdet3d.models.losses import chamfer_distance
from mmdet3d.models.model_utils import VoteModule
14
from mmdet3d.ops import build_sa_module
jshilong's avatar
jshilong committed
15
16
17
from mmdet3d.registry import MODELS, TASK_UTILS
from mmdet.core.utils import multi_apply
from ...core import Det3DDataSample
18
from .base_conv_bbox_head import BaseConvBboxHead
wuyuefeng's avatar
Votenet  
wuyuefeng committed
19
20


21
@MODELS.register_module()
22
class VoteHead(BaseModule):
zhangwenwei's avatar
zhangwenwei committed
23
    r"""Bbox head of `Votenet <https://arxiv.org/abs/1904.09664>`_.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
24
25
26

    Args:
        num_classes (int): The number of class.
jshilong's avatar
jshilong committed
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
        bbox_coder (ConfigDict, dict): Bbox coder for encoding and
            decoding boxes. Defaults to None.
        train_cfg (dict, optional): Config for training. Defaults to None.
        test_cfg (dict, optional): Config for testing. Defaults to None.
        vote_module_cfg (dict, optional): Config of VoteModule for
            point-wise votes. Defaults to None.
        vote_aggregation_cfg (dict, optional): Config of vote
            aggregation layer. Defaults to None.
        pred_layer_cfg (dict, optional): Config of classification
            and regression prediction layers. Defaults to None.
        objectness_loss (dict, optional): Config of objectness loss.
            Defaults to None.
        center_loss (dict, optional): Config of center loss.
            Defaults to None.
        dir_class_loss (dict, optional): Config of direction
            classification loss. Defaults to None.
        dir_res_loss (dict, optional): Config of direction
            residual regression loss. Defaults to None.
        size_class_loss (dict, optional): Config of size
            classification loss. Defaults to None.
        size_res_loss (dict, optional): Config of size
            residual regression loss. Defaults to None.
        semantic_loss (dict, optional): Config of point-wise
            semantic segmentation loss. Defaults to None.
        iou_loss (dict, optional): Config of IOU loss for
            regression. Defaults to None.
        init_cfg (dict, optional): Config of model weight
            initialization. Defaults to None.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
55
56
57
    """

    def __init__(self,
jshilong's avatar
jshilong committed
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
                 num_classes: int,
                 bbox_coder: Union[ConfigDict, dict],
                 train_cfg: Optional[dict] = None,
                 test_cfg: Optional[dict] = None,
                 vote_module_cfg: Optional[dict] = None,
                 vote_aggregation_cfg: Optional[dict] = None,
                 pred_layer_cfg: Optional[dict] = None,
                 objectness_loss: Optional[dict] = None,
                 center_loss: Optional[dict] = None,
                 dir_class_loss: Optional[dict] = None,
                 dir_res_loss: Optional[dict] = None,
                 size_class_loss: Optional[dict] = None,
                 size_res_loss: Optional[dict] = None,
                 semantic_loss: Optional[dict] = None,
                 iou_loss: Optional[dict] = None,
                 init_cfg: Optional[dict] = None):
74
        super(VoteHead, self).__init__(init_cfg=init_cfg)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
75
76
77
        self.num_classes = num_classes
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
jshilong's avatar
jshilong committed
78

79
        self.gt_per_seed = vote_module_cfg['gt_per_seed']
wuyuefeng's avatar
Votenet  
wuyuefeng committed
80
81
        self.num_proposal = vote_aggregation_cfg['num_point']

jshilong's avatar
jshilong committed
82
83
84
85
86
        self.loss_objectness = MODELS.build(objectness_loss)
        self.loss_center = MODELS.build(center_loss)
        self.loss_dir_res = MODELS.build(dir_res_loss)
        self.loss_dir_class = MODELS.build(dir_class_loss)
        self.loss_size_res = MODELS.build(size_res_loss)
87
        if size_class_loss is not None:
jshilong's avatar
jshilong committed
88
            self.size_class_loss = MODELS.build(size_class_loss)
89
        if semantic_loss is not None:
jshilong's avatar
jshilong committed
90
            self.semantic_loss = MODELS.build(semantic_loss)
91
        if iou_loss is not None:
jshilong's avatar
jshilong committed
92
            self.iou_loss = MODELS.build(iou_loss)
93
94
        else:
            self.iou_loss = None
wuyuefeng's avatar
Votenet  
wuyuefeng committed
95

jshilong's avatar
jshilong committed
96
        self.bbox_coder = TASK_UTILS.build(bbox_coder)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
97
98
99
        self.num_sizes = self.bbox_coder.num_sizes
        self.num_dir_bins = self.bbox_coder.num_dir_bins

100
        self.vote_module = VoteModule(**vote_module_cfg)
101
        self.vote_aggregation = build_sa_module(vote_aggregation_cfg)
102
        self.fp16_enabled = False
wuyuefeng's avatar
Votenet  
wuyuefeng committed
103

104
105
106
107
108
109
        # Bbox classification and regression
        self.conv_pred = BaseConvBboxHead(
            **pred_layer_cfg,
            num_cls_out_channels=self._get_cls_out_channels(),
            num_reg_out_channels=self._get_reg_out_channels())

jshilong's avatar
jshilong committed
110
111
112
113
114
115
116
117
118
    @property
    def sample_mode(self):
        if self.training:
            sample_mode = self.train_cfg.sample_mode
        else:
            sample_mode = self.test_cfg.sample_mode
        assert sample_mode in ['vote', 'seed', 'random', 'spec']
        return sample_mode

119
120
121
122
123
124
125
    def _get_cls_out_channels(self):
        """Return the channel number of classification outputs."""
        # Class numbers (k) + objectness (2)
        return self.num_classes + 2

    def _get_reg_out_channels(self):
        """Return the channel number of regression outputs."""
wuyuefeng's avatar
Votenet  
wuyuefeng committed
126
127
128
        # Objectness scores (2), center residual (3),
        # heading class+residual (num_dir_bins*2),
        # size class+residual(num_sizes*4)
129
        return 3 + self.num_dir_bins * 2 + self.num_sizes * 4
wuyuefeng's avatar
Votenet  
wuyuefeng committed
130

jshilong's avatar
jshilong committed
131
    def _extract_input(self, feat_dict: dict) -> tuple:
132
133
134
135
136
137
        """Extract inputs from features dictionary.

        Args:
            feat_dict (dict): Feature dict from backbone.

        Returns:
jshilong's avatar
jshilong committed
138
139
140
141
142
            tuple[Tensor]: Arrage as following three tensor.

                - Coordinates of input points.
                - Features of input points.
                - Indices of input points.
143
        """
144
145
146
147
148
149
150
151
152
153
154
155
156

        # for imvotenet
        if 'seed_points' in feat_dict and \
           'seed_features' in feat_dict and \
           'seed_indices' in feat_dict:
            seed_points = feat_dict['seed_points']
            seed_features = feat_dict['seed_features']
            seed_indices = feat_dict['seed_indices']
        # for votenet
        else:
            seed_points = feat_dict['fp_xyz'][-1]
            seed_features = feat_dict['fp_features'][-1]
            seed_indices = feat_dict['fp_indices'][-1]
157
158

        return seed_points, seed_features, seed_indices
wuyuefeng's avatar
Votenet  
wuyuefeng committed
159

jshilong's avatar
jshilong committed
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
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
    def predict(self,
                points: List[torch.Tensor],
                feats_dict: Dict[str, torch.Tensor],
                batch_data_samples: List[Det3DDataSample],
                rescale=True,
                **kwargs) -> List[InstanceData]:
        """
        Args:
            points (list[tensor]): Point clouds of multiple samples.
            feats_dict (dict): Features from FPN or backbone..
            batch_data_samples (List[:obj:`Det3DDataSample`]): The Data
                Samples. It usually includes meta information of data.
            rescale (bool): Whether rescale the resutls to
                the original scale.

        Returns:
            list[:obj:`InstanceData`]: List of processed predictions. Each
            InstanceData contains 3d Bounding boxes and corresponding
            scores and labels.
        """
        preds_dict = self(feats_dict)
        batch_size = len(batch_data_samples)
        batch_input_metas = []
        for batch_index in range(batch_size):
            metainfo = batch_data_samples[batch_index].metainfo
            batch_input_metas.append(metainfo)

        results_list = self.predict_by_feat(
            points, preds_dict, batch_input_metas, rescale=rescale, **kwargs)
        return results_list

    def loss(self, points: List[torch.Tensor], feats_dict: Dict[str,
                                                                torch.Tensor],
             batch_data_samples: List[Det3DDataSample], **kwargs) -> dict:
        """
        Args:
            points (list[tensor]): Points cloud of multiple samples.
            feats_dict (dict): Predictions from backbone or FPN.
            batch_data_samples (list[:obj:`Det3DDataSample`]): Each item
                contains the meta information of each sample and
                corresponding annotations.

        Returns:
            dict:  A dictionary of loss components.
        """
        preds_dict = self.forward(feats_dict)
        batch_gt_instance_3d = []
        batch_gt_instances_ignore = []
        batch_input_metas = []
        batch_pts_semantic_mask = []
        batch_pts_instance_mask = []
        for data_sample in batch_data_samples:
            batch_input_metas.append(data_sample.metainfo)
            batch_gt_instance_3d.append(data_sample.gt_instances_3d)
            batch_gt_instances_ignore.append(
                data_sample.get('ignored_instances', None))
            batch_pts_semantic_mask.append(
                data_sample.seg_data.get('pts_semantic_mask', None))
            batch_pts_instance_mask.append(
                data_sample.seg_data.get('pts_instance_mask', None))

        loss_inputs = (points, preds_dict, batch_gt_instance_3d)
        losses = self.loss_by_feat(
            *loss_inputs,
            batch_pts_semantic_mask=batch_pts_semantic_mask,
            batch_pts_instance_mask=batch_pts_instance_mask,
            batch_input_metas=batch_input_metas,
            batch_gt_instances_ignore=batch_gt_instances_ignore)
        return losses

    def forward(self, feat_dict: dict) -> dict:
wuyuefeng's avatar
Votenet  
wuyuefeng committed
231
232
        """Forward pass.

zhangwenwei's avatar
zhangwenwei committed
233
        Note:
234
            The forward of VoteHead is divided into 4 steps:
zhangwenwei's avatar
zhangwenwei committed
235
236
237
238
239

                1. Generate vote_points from seed_points.
                2. Aggregate vote_points.
                3. Predict bbox and score.
                4. Decode predictions.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
240
241

        Args:
wangtai's avatar
wangtai committed
242
            feat_dict (dict): Feature dict from backbone.
wuyuefeng's avatar
wuyuefeng committed
243
244
245

        Returns:
            dict: Predictions of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
246
247
        """

248
249
        seed_points, seed_features, seed_indices = self._extract_input(
            feat_dict)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
250
251

        # 1. generate vote_points from seed_points
252
253
        vote_points, vote_features, vote_offset = self.vote_module(
            seed_points, seed_features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
254
255
256
257
        results = dict(
            seed_points=seed_points,
            seed_indices=seed_indices,
            vote_points=vote_points,
258
259
            vote_features=vote_features,
            vote_offset=vote_offset)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
260
261

        # 2. aggregate vote_points
jshilong's avatar
jshilong committed
262
        if self.sample_mode == 'vote':
wuyuefeng's avatar
Votenet  
wuyuefeng committed
263
            # use fps in vote_aggregation
264
265
            aggregation_inputs = dict(
                points_xyz=vote_points, features=vote_features)
jshilong's avatar
jshilong committed
266
        elif self.sample_mode == 'seed':
wuyuefeng's avatar
Votenet  
wuyuefeng committed
267
268
269
            # FPS on seed and choose the votes corresponding to the seeds
            sample_indices = furthest_point_sample(seed_points,
                                                   self.num_proposal)
270
271
272
273
            aggregation_inputs = dict(
                points_xyz=vote_points,
                features=vote_features,
                indices=sample_indices)
jshilong's avatar
jshilong committed
274
        elif self.sample_mode == 'random':
wuyuefeng's avatar
Votenet  
wuyuefeng committed
275
276
277
278
279
            # Random sampling from the votes
            batch_size, num_seed = seed_points.shape[:2]
            sample_indices = seed_points.new_tensor(
                torch.randint(0, num_seed, (batch_size, self.num_proposal)),
                dtype=torch.int32)
280
281
282
283
            aggregation_inputs = dict(
                points_xyz=vote_points,
                features=vote_features,
                indices=sample_indices)
jshilong's avatar
jshilong committed
284
        elif self.sample_mode == 'spec':
285
286
287
288
289
            # Specify the new center in vote_aggregation
            aggregation_inputs = dict(
                points_xyz=seed_points,
                features=seed_features,
                target_xyz=vote_points)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
290
        else:
Wenwei Zhang's avatar
Wenwei Zhang committed
291
            raise NotImplementedError(
jshilong's avatar
jshilong committed
292
                f'Sample mode {self.sample_mode} is not supported!')
wuyuefeng's avatar
Votenet  
wuyuefeng committed
293

294
        vote_aggregation_ret = self.vote_aggregation(**aggregation_inputs)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
295
        aggregated_points, features, aggregated_indices = vote_aggregation_ret
296

wuyuefeng's avatar
Votenet  
wuyuefeng committed
297
        results['aggregated_points'] = aggregated_points
encore-zhou's avatar
encore-zhou committed
298
        results['aggregated_features'] = features
wuyuefeng's avatar
Votenet  
wuyuefeng committed
299
300
301
        results['aggregated_indices'] = aggregated_indices

        # 3. predict bbox and score
302
        cls_predictions, reg_predictions = self.conv_pred(features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
303
304

        # 4. decode predictions
305
306
307
        decode_res = self.bbox_coder.split_pred(cls_predictions,
                                                reg_predictions,
                                                aggregated_points)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
308
309
310
        results.update(decode_res)
        return results

311
    @force_fp32(apply_to=('bbox_preds', ))
jshilong's avatar
jshilong committed
312
313
314
315
316
317
318
319
320
    def loss_by_feat(
            self,
            points: List[torch.Tensor],
            bbox_preds_dict: dict,
            batch_gt_instances_3d: List[InstanceData],
            batch_pts_semantic_mask: Optional[List[torch.Tensor]] = None,
            batch_pts_instance_mask: Optional[List[torch.Tensor]] = None,
            ret_target: bool = False,
            **kwargs) -> dict:
wuyuefeng's avatar
wuyuefeng committed
321
322
323
        """Compute loss.

        Args:
liyinhao's avatar
liyinhao committed
324
            points (list[torch.Tensor]): Input points.
jshilong's avatar
jshilong committed
325
326
327
328
329
330
331
332
333
334
            bbox_preds_dict (dict): Predictions from forward of vote head.
            batch_gt_instances_3d (list[:obj:`InstanceData`]): Batch of
                gt_instances. It usually includes ``bboxes`` and ``labels``
                attributes.
            batch_pts_semantic_mask (list[tensor]): Semantic mask
                of points cloud. Defaults to None.
            batch_pts_semantic_mask (list[tensor]): Instance mask
                of points cloud. Defaults to None.
            batch_input_metas (list[dict]): Contain pcd and img's meta info.
            ret_target (bool): Return targets or not.
wuyuefeng's avatar
wuyuefeng committed
335
336
337
338

        Returns:
            dict: Losses of Votenet.
        """
jshilong's avatar
jshilong committed
339
340
341
342
343

        targets = self.get_targets(points, bbox_preds_dict,
                                   batch_gt_instances_3d,
                                   batch_pts_semantic_mask,
                                   batch_pts_instance_mask)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
344
        (vote_targets, vote_target_masks, size_class_targets, size_res_targets,
345
346
347
348
         dir_class_targets, dir_res_targets, center_targets,
         assigned_center_targets, mask_targets, valid_gt_masks,
         objectness_targets, objectness_weights, box_loss_weights,
         valid_gt_weights) = targets
wuyuefeng's avatar
Votenet  
wuyuefeng committed
349
350

        # calculate vote loss
jshilong's avatar
jshilong committed
351
352
353
        vote_loss = self.vote_module.get_loss(bbox_preds_dict['seed_points'],
                                              bbox_preds_dict['vote_points'],
                                              bbox_preds_dict['seed_indices'],
wuyuefeng's avatar
Votenet  
wuyuefeng committed
354
355
356
                                              vote_target_masks, vote_targets)

        # calculate objectness loss
jshilong's avatar
jshilong committed
357
358
        objectness_loss = self.loss_objectness(
            bbox_preds_dict['obj_scores'].transpose(2, 1),
wuyuefeng's avatar
Votenet  
wuyuefeng committed
359
360
361
362
            objectness_targets,
            weight=objectness_weights)

        # calculate center loss
jshilong's avatar
jshilong committed
363
364
        source2target_loss, target2source_loss = self.loss_center(
            bbox_preds_dict['center'],
wuyuefeng's avatar
Votenet  
wuyuefeng committed
365
366
367
368
369
370
            center_targets,
            src_weight=box_loss_weights,
            dst_weight=valid_gt_weights)
        center_loss = source2target_loss + target2source_loss

        # calculate direction class loss
jshilong's avatar
jshilong committed
371
372
        dir_class_loss = self.loss_dir_class(
            bbox_preds_dict['dir_class'].transpose(2, 1),
wuyuefeng's avatar
Votenet  
wuyuefeng committed
373
374
375
376
377
378
379
380
381
            dir_class_targets,
            weight=box_loss_weights)

        # calculate direction residual loss
        batch_size, proposal_num = size_class_targets.shape[:2]
        heading_label_one_hot = vote_targets.new_zeros(
            (batch_size, proposal_num, self.num_dir_bins))
        heading_label_one_hot.scatter_(2, dir_class_targets.unsqueeze(-1), 1)
        dir_res_norm = torch.sum(
jshilong's avatar
jshilong committed
382
383
            bbox_preds_dict['dir_res_norm'] * heading_label_one_hot, -1)
        dir_res_loss = self.loss_dir_res(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
384
385
386
387
            dir_res_norm, dir_res_targets, weight=box_loss_weights)

        # calculate size class loss
        size_class_loss = self.size_class_loss(
jshilong's avatar
jshilong committed
388
            bbox_preds_dict['size_class'].transpose(2, 1),
wuyuefeng's avatar
Votenet  
wuyuefeng committed
389
390
391
392
393
394
395
396
            size_class_targets,
            weight=box_loss_weights)

        # calculate size residual loss
        one_hot_size_targets = vote_targets.new_zeros(
            (batch_size, proposal_num, self.num_sizes))
        one_hot_size_targets.scatter_(2, size_class_targets.unsqueeze(-1), 1)
        one_hot_size_targets_expand = one_hot_size_targets.unsqueeze(
Wenwei Zhang's avatar
Wenwei Zhang committed
397
            -1).repeat(1, 1, 1, 3).contiguous()
wuyuefeng's avatar
Votenet  
wuyuefeng committed
398
        size_residual_norm = torch.sum(
jshilong's avatar
jshilong committed
399
            bbox_preds_dict['size_res_norm'] * one_hot_size_targets_expand, 2)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
400
401
        box_loss_weights_expand = box_loss_weights.unsqueeze(-1).repeat(
            1, 1, 3)
jshilong's avatar
jshilong committed
402
        size_res_loss = self.loss_size_res(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
403
404
405
406
407
408
            size_residual_norm,
            size_res_targets,
            weight=box_loss_weights_expand)

        # calculate semantic loss
        semantic_loss = self.semantic_loss(
jshilong's avatar
jshilong committed
409
            bbox_preds_dict['sem_scores'].transpose(2, 1),
wuyuefeng's avatar
Votenet  
wuyuefeng committed
410
411
412
413
414
415
416
417
418
419
420
421
            mask_targets,
            weight=box_loss_weights)

        losses = dict(
            vote_loss=vote_loss,
            objectness_loss=objectness_loss,
            semantic_loss=semantic_loss,
            center_loss=center_loss,
            dir_class_loss=dir_class_loss,
            dir_res_loss=dir_res_loss,
            size_class_loss=size_class_loss,
            size_res_loss=size_res_loss)
encore-zhou's avatar
encore-zhou committed
422

423
424
        if self.iou_loss:
            corners_pred = self.bbox_coder.decode_corners(
jshilong's avatar
jshilong committed
425
                bbox_preds_dict['center'], size_residual_norm,
426
427
428
429
430
431
432
433
                one_hot_size_targets_expand)
            corners_target = self.bbox_coder.decode_corners(
                assigned_center_targets, size_res_targets,
                one_hot_size_targets_expand)
            iou_loss = self.iou_loss(
                corners_pred, corners_target, weight=box_loss_weights)
            losses['iou_loss'] = iou_loss

encore-zhou's avatar
encore-zhou committed
434
435
436
        if ret_target:
            losses['targets'] = targets

wuyuefeng's avatar
Votenet  
wuyuefeng committed
437
438
        return losses

jshilong's avatar
jshilong committed
439
440
441
442
443
444
445
446
    def get_targets(
        self,
        points,
        bbox_preds: dict = None,
        batch_gt_instances_3d: List[InstanceData] = None,
        batch_pts_semantic_mask: List[torch.Tensor] = None,
        batch_pts_instance_mask: List[torch.Tensor] = None,
    ):
wuyuefeng's avatar
wuyuefeng committed
447
        """Generate targets of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
448
449

        Args:
liyinhao's avatar
liyinhao committed
450
            points (list[torch.Tensor]): Points of each batch.
wangtai's avatar
wangtai committed
451
            bbox_preds (torch.Tensor): Bounding box predictions of vote head.
jshilong's avatar
jshilong committed
452
453
454
455
456
457
458
            batch_gt_instances_3d (list[:obj:`InstanceData`]): Batch of
                gt_instances. It usually includes ``bboxes`` and ``labels``
                attributes.
            batch_pts_semantic_mask (list[tensor]): Semantic gt mask for
                multiple images.
            batch_pts_instance_mask (list[tensor]): Instance gt mask for
                multiple images.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
459
460

        Returns:
461
            tuple[torch.Tensor]: Targets of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
462
463
464
465
        """
        # find empty example
        valid_gt_masks = list()
        gt_num = list()
jshilong's avatar
jshilong committed
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
        batch_gt_labels_3d = [
            gt_instances_3d.labels_3d
            for gt_instances_3d in batch_gt_instances_3d
        ]
        batch_gt_bboxes_3d = [
            gt_instances_3d.bboxes_3d
            for gt_instances_3d in batch_gt_instances_3d
        ]
        for index in range(len(batch_gt_labels_3d)):
            if len(batch_gt_labels_3d[index]) == 0:
                fake_box = batch_gt_bboxes_3d[index].tensor.new_zeros(
                    1, batch_gt_bboxes_3d[index].tensor.shape[-1])
                batch_gt_bboxes_3d[index] = batch_gt_bboxes_3d[index].new_box(
                    fake_box)
                batch_gt_labels_3d[index] = batch_gt_labels_3d[
                    index].new_zeros(1)
                valid_gt_masks.append(batch_gt_labels_3d[index].new_zeros(1))
wuyuefeng's avatar
Votenet  
wuyuefeng committed
483
484
                gt_num.append(1)
            else:
jshilong's avatar
jshilong committed
485
486
487
                valid_gt_masks.append(batch_gt_labels_3d[index].new_ones(
                    batch_gt_labels_3d[index].shape))
                gt_num.append(batch_gt_labels_3d[index].shape[0])
wuyuefeng's avatar
Votenet  
wuyuefeng committed
488
489
490
491
        max_gt_num = max(gt_num)

        aggregated_points = [
            bbox_preds['aggregated_points'][i]
jshilong's avatar
jshilong committed
492
            for i in range(len(batch_gt_labels_3d))
wuyuefeng's avatar
Votenet  
wuyuefeng committed
493
494
495
        ]

        (vote_targets, vote_target_masks, size_class_targets, size_res_targets,
496
         dir_class_targets, dir_res_targets, center_targets,
jshilong's avatar
jshilong committed
497
498
499
500
501
         assigned_center_targets, mask_targets,
         objectness_targets, objectness_masks) = multi_apply(
             self._get_targets_single, points, batch_gt_bboxes_3d,
             batch_gt_labels_3d, batch_pts_semantic_mask,
             batch_pts_instance_mask, aggregated_points)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
502
503

        # pad targets as original code of votenet.
jshilong's avatar
jshilong committed
504
505
        for index in range(len(batch_gt_labels_3d)):
            pad_num = max_gt_num - batch_gt_labels_3d[index].shape[0]
wuyuefeng's avatar
Votenet  
wuyuefeng committed
506
507
508
509
510
511
512
513
514
            center_targets[index] = F.pad(center_targets[index],
                                          (0, 0, 0, pad_num))
            valid_gt_masks[index] = F.pad(valid_gt_masks[index], (0, pad_num))

        vote_targets = torch.stack(vote_targets)
        vote_target_masks = torch.stack(vote_target_masks)
        center_targets = torch.stack(center_targets)
        valid_gt_masks = torch.stack(valid_gt_masks)

515
        assigned_center_targets = torch.stack(assigned_center_targets)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
        objectness_targets = torch.stack(objectness_targets)
        objectness_weights = torch.stack(objectness_masks)
        objectness_weights /= (torch.sum(objectness_weights) + 1e-6)
        box_loss_weights = objectness_targets.float() / (
            torch.sum(objectness_targets).float() + 1e-6)
        valid_gt_weights = valid_gt_masks.float() / (
            torch.sum(valid_gt_masks.float()) + 1e-6)
        dir_class_targets = torch.stack(dir_class_targets)
        dir_res_targets = torch.stack(dir_res_targets)
        size_class_targets = torch.stack(size_class_targets)
        size_res_targets = torch.stack(size_res_targets)
        mask_targets = torch.stack(mask_targets)

        return (vote_targets, vote_target_masks, size_class_targets,
                size_res_targets, dir_class_targets, dir_res_targets,
531
532
533
                center_targets, assigned_center_targets, mask_targets,
                valid_gt_masks, objectness_targets, objectness_weights,
                box_loss_weights, valid_gt_weights)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
534

jshilong's avatar
jshilong committed
535
536
537
538
539
540
541
    def _get_targets_single(self,
                            points,
                            gt_bboxes_3d,
                            gt_labels_3d,
                            pts_semantic_mask=None,
                            pts_instance_mask=None,
                            aggregated_points=None):
wuyuefeng's avatar
wuyuefeng committed
542
543
544
        """Generate targets of vote head for single batch.

        Args:
liyinhao's avatar
liyinhao committed
545
            points (torch.Tensor): Points of each batch.
546
            gt_bboxes_3d (:obj:`BaseInstance3DBoxes`): Ground truth
wangtai's avatar
wangtai committed
547
548
                boxes of each batch.
            gt_labels_3d (torch.Tensor): Labels of each batch.
549
            pts_semantic_mask (torch.Tensor): Point-wise semantic
wuyuefeng's avatar
wuyuefeng committed
550
                label of each batch.
551
            pts_instance_mask (torch.Tensor): Point-wise instance
wuyuefeng's avatar
wuyuefeng committed
552
                label of each batch.
liyinhao's avatar
liyinhao committed
553
            aggregated_points (torch.Tensor): Aggregated points from
wuyuefeng's avatar
wuyuefeng committed
554
555
556
                vote aggregation layer.

        Returns:
557
            tuple[torch.Tensor]: Targets of vote head.
wuyuefeng's avatar
wuyuefeng committed
558
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
559
560
        assert self.bbox_coder.with_rot or pts_semantic_mask is not None

wuyuefeng's avatar
wuyuefeng committed
561
562
        gt_bboxes_3d = gt_bboxes_3d.to(points.device)

wuyuefeng's avatar
Votenet  
wuyuefeng committed
563
564
565
566
567
568
569
        # generate votes target
        num_points = points.shape[0]
        if self.bbox_coder.with_rot:
            vote_targets = points.new_zeros([num_points, 3 * self.gt_per_seed])
            vote_target_masks = points.new_zeros([num_points],
                                                 dtype=torch.long)
            vote_target_idx = points.new_zeros([num_points], dtype=torch.long)
570
            box_indices_all = gt_bboxes_3d.points_in_boxes_all(points)
wuyuefeng's avatar
wuyuefeng committed
571
            for i in range(gt_labels_3d.shape[0]):
wuyuefeng's avatar
Votenet  
wuyuefeng committed
572
                box_indices = box_indices_all[:, i]
573
574
                indices = torch.nonzero(
                    box_indices, as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
575
576
577
                selected_points = points[indices]
                vote_target_masks[indices] = 1
                vote_targets_tmp = vote_targets[indices]
wuyuefeng's avatar
wuyuefeng committed
578
                votes = gt_bboxes_3d.gravity_center[i].unsqueeze(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
579
580
581
582
                    0) - selected_points[:, :3]

                for j in range(self.gt_per_seed):
                    column_indices = torch.nonzero(
583
584
                        vote_target_idx[indices] == j,
                        as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
                    vote_targets_tmp[column_indices,
                                     int(j * 3):int(j * 3 +
                                                    3)] = votes[column_indices]
                    if j == 0:
                        vote_targets_tmp[column_indices] = votes[
                            column_indices].repeat(1, self.gt_per_seed)

                vote_targets[indices] = vote_targets_tmp
                vote_target_idx[indices] = torch.clamp(
                    vote_target_idx[indices] + 1, max=2)
        elif pts_semantic_mask is not None:
            vote_targets = points.new_zeros([num_points, 3])
            vote_target_masks = points.new_zeros([num_points],
                                                 dtype=torch.long)
            for i in torch.unique(pts_instance_mask):
600
601
                indices = torch.nonzero(
                    pts_instance_mask == i, as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
                if pts_semantic_mask[indices[0]] < self.num_classes:
                    selected_points = points[indices, :3]
                    center = 0.5 * (
                        selected_points.min(0)[0] + selected_points.max(0)[0])
                    vote_targets[indices, :] = center - selected_points
                    vote_target_masks[indices] = 1
            vote_targets = vote_targets.repeat((1, self.gt_per_seed))
        else:
            raise NotImplementedError

        (center_targets, size_class_targets, size_res_targets,
         dir_class_targets,
         dir_res_targets) = self.bbox_coder.encode(gt_bboxes_3d, gt_labels_3d)

        proposal_num = aggregated_points.shape[0]
        distance1, _, assignment, _ = chamfer_distance(
            aggregated_points.unsqueeze(0),
            center_targets.unsqueeze(0),
            reduction='none')
        assignment = assignment.squeeze(0)
        euclidean_distance1 = torch.sqrt(distance1.squeeze(0) + 1e-6)

        objectness_targets = points.new_zeros((proposal_num), dtype=torch.long)
        objectness_targets[
            euclidean_distance1 < self.train_cfg['pos_distance_thr']] = 1

        objectness_masks = points.new_zeros((proposal_num))
        objectness_masks[
            euclidean_distance1 < self.train_cfg['pos_distance_thr']] = 1.0
        objectness_masks[
            euclidean_distance1 > self.train_cfg['neg_distance_thr']] = 1.0

        dir_class_targets = dir_class_targets[assignment]
        dir_res_targets = dir_res_targets[assignment]
        dir_res_targets /= (np.pi / self.num_dir_bins)
        size_class_targets = size_class_targets[assignment]
        size_res_targets = size_res_targets[assignment]

wuyuefeng's avatar
wuyuefeng committed
640
        one_hot_size_targets = gt_bboxes_3d.tensor.new_zeros(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
641
642
643
644
645
646
647
648
649
650
            (proposal_num, self.num_sizes))
        one_hot_size_targets.scatter_(1, size_class_targets.unsqueeze(-1), 1)
        one_hot_size_targets = one_hot_size_targets.unsqueeze(-1).repeat(
            1, 1, 3)
        mean_sizes = size_res_targets.new_tensor(
            self.bbox_coder.mean_sizes).unsqueeze(0)
        pos_mean_sizes = torch.sum(one_hot_size_targets * mean_sizes, 1)
        size_res_targets /= pos_mean_sizes

        mask_targets = gt_labels_3d[assignment]
651
        assigned_center_targets = center_targets[assignment]
wuyuefeng's avatar
Votenet  
wuyuefeng committed
652
653

        return (vote_targets, vote_target_masks, size_class_targets,
654
655
                size_res_targets, dir_class_targets,
                dir_res_targets, center_targets, assigned_center_targets,
wuyuefeng's avatar
Votenet  
wuyuefeng committed
656
657
                mask_targets.long(), objectness_targets, objectness_masks)

jshilong's avatar
jshilong committed
658
659
660
661
662
663
    def predict_by_feat(self,
                        points: List[torch.Tensor],
                        bbox_preds_dict: dict,
                        batch_input_metas: List[dict],
                        use_nms: bool = True,
                        **kwargs) -> List[InstanceData]:
wuyuefeng's avatar
wuyuefeng committed
664
665
666
        """Generate bboxes from vote head predictions.

        Args:
jshilong's avatar
jshilong committed
667
668
669
670
            points (List[torch.Tensor]): Input points of multiple samples.
            bbox_preds_dict (dict): Predictions from vote head.
            batch_input_metas (list[dict]): Each item
                contains the meta information of each sample.
encore-zhou's avatar
encore-zhou committed
671
672
            use_nms (bool): Whether to apply NMS, skip nms postprocessing
                while using vote head in rpn stage.
wuyuefeng's avatar
wuyuefeng committed
673
674

        Returns:
jshilong's avatar
jshilong committed
675
676
677
            list[:obj:`InstanceData`]: List of processed predictions. Each
            InstanceData cantains 3d Bounding boxes and corresponding
            scores and labels.
wuyuefeng's avatar
wuyuefeng committed
678
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
679
        # decode boxes
jshilong's avatar
jshilong committed
680
681
682
683
684
685
686
687
688
689
        stack_points = torch.stack(points)
        obj_scores = F.softmax(bbox_preds_dict['obj_scores'], dim=-1)[..., -1]
        sem_scores = F.softmax(bbox_preds_dict['sem_scores'], dim=-1)
        bbox3d = self.bbox_coder.decode(bbox_preds_dict)

        batch_size = bbox3d.shape[0]
        results_list = list()
        for b in range(batch_size):
            temp_results = InstanceData()
            if use_nms:
encore-zhou's avatar
encore-zhou committed
690
                bbox_selected, score_selected, labels = \
jshilong's avatar
jshilong committed
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
                    self.multiclass_nms_single(obj_scores[b],
                                               sem_scores[b],
                                               bbox3d[b],
                                               stack_points[b, ..., :3],
                                               batch_input_metas[b])
                bbox = batch_input_metas[b]['box_type_3d'](
                    bbox_selected,
                    box_dim=bbox_selected.shape[-1],
                    with_yaw=self.bbox_coder.with_rot)
                temp_results.bboxes_3d = bbox
                temp_results.scores_3d = score_selected
                temp_results.labels_3d = labels
                results_list.append(temp_results)
            else:
                bbox = batch_input_metas[b]['box_type_3d'](
encore-zhou's avatar
encore-zhou committed
706
707
708
                    bbox_selected,
                    box_dim=bbox_selected.shape[-1],
                    with_yaw=self.bbox_coder.with_rot)
jshilong's avatar
jshilong committed
709
710
711
712
                temp_results.bboxes_3d = bbox
                temp_results.obj_scores_3d = obj_scores[b]
                temp_results.sem_scores_3d = obj_scores[b]
                results_list.append(temp_results)
encore-zhou's avatar
encore-zhou committed
713

jshilong's avatar
jshilong committed
714
        return results_list
wuyuefeng's avatar
Votenet  
wuyuefeng committed
715

wuyuefeng's avatar
wuyuefeng committed
716
717
    def multiclass_nms_single(self, obj_scores, sem_scores, bbox, points,
                              input_meta):
wangtai's avatar
wangtai committed
718
        """Multi-class nms in single batch.
wuyuefeng's avatar
wuyuefeng committed
719
720

        Args:
wangtai's avatar
wangtai committed
721
722
723
            obj_scores (torch.Tensor): Objectness score of bounding boxes.
            sem_scores (torch.Tensor): semantic class score of bounding boxes.
            bbox (torch.Tensor): Predicted bounding boxes.
liyinhao's avatar
liyinhao committed
724
            points (torch.Tensor): Input points.
wangtai's avatar
wangtai committed
725
            input_meta (dict): Point cloud and image's meta info.
wuyuefeng's avatar
wuyuefeng committed
726
727

        Returns:
wangtai's avatar
wangtai committed
728
            tuple[torch.Tensor]: Bounding boxes, scores and labels.
wuyuefeng's avatar
wuyuefeng committed
729
        """
wuyuefeng's avatar
wuyuefeng committed
730
731
732
733
734
        bbox = input_meta['box_type_3d'](
            bbox,
            box_dim=bbox.shape[-1],
            with_yaw=self.bbox_coder.with_rot,
            origin=(0.5, 0.5, 0.5))
735
        box_indices = bbox.points_in_boxes_all(points)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
736

wuyuefeng's avatar
wuyuefeng committed
737
        corner3d = bbox.corners
wuyuefeng's avatar
Votenet  
wuyuefeng committed
738
739
740
741
        minmax_box3d = corner3d.new(torch.Size((corner3d.shape[0], 6)))
        minmax_box3d[:, :3] = torch.min(corner3d, dim=1)[0]
        minmax_box3d[:, 3:] = torch.max(corner3d, dim=1)[0]

wuyuefeng's avatar
wuyuefeng committed
742
743
744
        nonempty_box_mask = box_indices.T.sum(1) > 5

        bbox_classes = torch.argmax(sem_scores, -1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
745
746
747
748
749
750
751
        nms_selected = aligned_3d_nms(minmax_box3d[nonempty_box_mask],
                                      obj_scores[nonempty_box_mask],
                                      bbox_classes[nonempty_box_mask],
                                      self.test_cfg.nms_thr)

        # filter empty boxes and boxes with low score
        scores_mask = (obj_scores > self.test_cfg.score_thr)
752
753
        nonempty_box_inds = torch.nonzero(
            nonempty_box_mask, as_tuple=False).flatten()
wuyuefeng's avatar
Votenet  
wuyuefeng committed
754
755
756
757
758
759
760
        nonempty_mask = torch.zeros_like(bbox_classes).scatter(
            0, nonempty_box_inds[nms_selected], 1)
        selected = (nonempty_mask.bool() & scores_mask.bool())

        if self.test_cfg.per_class_proposal:
            bbox_selected, score_selected, labels = [], [], []
            for k in range(sem_scores.shape[-1]):
wuyuefeng's avatar
wuyuefeng committed
761
                bbox_selected.append(bbox[selected].tensor)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
762
763
764
765
766
767
768
769
                score_selected.append(obj_scores[selected] *
                                      sem_scores[selected][:, k])
                labels.append(
                    torch.zeros_like(bbox_classes[selected]).fill_(k))
            bbox_selected = torch.cat(bbox_selected, 0)
            score_selected = torch.cat(score_selected, 0)
            labels = torch.cat(labels, 0)
        else:
wuyuefeng's avatar
wuyuefeng committed
770
            bbox_selected = bbox[selected].tensor
wuyuefeng's avatar
Votenet  
wuyuefeng committed
771
772
773
774
            score_selected = obj_scores[selected]
            labels = bbox_classes[selected]

        return bbox_selected, score_selected, labels