vote_head.py 27.3 KB
Newer Older
wuyuefeng's avatar
Votenet  
wuyuefeng committed
1
2
import numpy as np
import torch
3
from mmcv.runner import force_fp32
zhangwenwei's avatar
zhangwenwei committed
4
5
from torch import nn as nn
from torch.nn import functional as F
wuyuefeng's avatar
Votenet  
wuyuefeng committed
6
7
8
9
10

from mmdet3d.core.post_processing import aligned_3d_nms
from mmdet3d.models.builder import build_loss
from mmdet3d.models.losses import chamfer_distance
from mmdet3d.models.model_utils import VoteModule
11
from mmdet3d.ops import build_sa_module, furthest_point_sample
zhangwenwei's avatar
zhangwenwei committed
12
from mmdet.core import build_bbox_coder, multi_apply
wuyuefeng's avatar
Votenet  
wuyuefeng committed
13
from mmdet.models import HEADS
14
from .base_conv_bbox_head import BaseConvBboxHead
wuyuefeng's avatar
Votenet  
wuyuefeng committed
15
16
17
18


@HEADS.register_module()
class VoteHead(nn.Module):
zhangwenwei's avatar
zhangwenwei committed
19
    r"""Bbox head of `Votenet <https://arxiv.org/abs/1904.09664>`_.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
20
21
22

    Args:
        num_classes (int): The number of class.
23
        bbox_coder (:obj:`BaseBBoxCoder`): Bbox coder for encoding and
wuyuefeng's avatar
Votenet  
wuyuefeng committed
24
25
26
            decoding boxes.
        train_cfg (dict): Config for training.
        test_cfg (dict): Config for testing.
27
        vote_module_cfg (dict): Config of VoteModule for point-wise votes.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
28
        vote_aggregation_cfg (dict): Config of vote aggregation layer.
29
30
        pred_layer_cfg (dict): Config of classfication and regression
            prediction layers.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
        conv_cfg (dict): Config of convolution in prediction layer.
        norm_cfg (dict): Config of BN in prediction layer.
        objectness_loss (dict): Config of objectness loss.
        center_loss (dict): Config of center loss.
        dir_class_loss (dict): Config of direction classification loss.
        dir_res_loss (dict): Config of direction residual regression loss.
        size_class_loss (dict): Config of size classification loss.
        size_res_loss (dict): Config of size residual regression loss.
        semantic_loss (dict): Config of point-wise semantic segmentation loss.
    """

    def __init__(self,
                 num_classes,
                 bbox_coder,
                 train_cfg=None,
                 test_cfg=None,
47
                 vote_module_cfg=None,
wuyuefeng's avatar
Votenet  
wuyuefeng committed
48
                 vote_aggregation_cfg=None,
49
                 pred_layer_cfg=None,
wuyuefeng's avatar
Votenet  
wuyuefeng committed
50
51
52
53
54
55
56
57
                 conv_cfg=dict(type='Conv1d'),
                 norm_cfg=dict(type='BN1d'),
                 objectness_loss=None,
                 center_loss=None,
                 dir_class_loss=None,
                 dir_res_loss=None,
                 size_class_loss=None,
                 size_res_loss=None,
58
59
                 semantic_loss=None,
                 iou_loss=None):
wuyuefeng's avatar
Votenet  
wuyuefeng committed
60
61
62
63
        super(VoteHead, self).__init__()
        self.num_classes = num_classes
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
64
        self.gt_per_seed = vote_module_cfg['gt_per_seed']
wuyuefeng's avatar
Votenet  
wuyuefeng committed
65
66
67
68
69
        self.num_proposal = vote_aggregation_cfg['num_point']

        self.objectness_loss = build_loss(objectness_loss)
        self.center_loss = build_loss(center_loss)
        self.dir_res_loss = build_loss(dir_res_loss)
70
        self.dir_class_loss = build_loss(dir_class_loss)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
71
        self.size_res_loss = build_loss(size_res_loss)
72
73
74
75
        if size_class_loss is not None:
            self.size_class_loss = build_loss(size_class_loss)
        if semantic_loss is not None:
            self.semantic_loss = build_loss(semantic_loss)
76
77
78
79
        if iou_loss is not None:
            self.iou_loss = build_loss(iou_loss)
        else:
            self.iou_loss = None
wuyuefeng's avatar
Votenet  
wuyuefeng committed
80
81
82
83
84

        self.bbox_coder = build_bbox_coder(bbox_coder)
        self.num_sizes = self.bbox_coder.num_sizes
        self.num_dir_bins = self.bbox_coder.num_dir_bins

85
        self.vote_module = VoteModule(**vote_module_cfg)
86
        self.vote_aggregation = build_sa_module(vote_aggregation_cfg)
87
        self.fp16_enabled = False
wuyuefeng's avatar
Votenet  
wuyuefeng committed
88

89
90
91
92
93
94
95
96
97
        # 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())

    def init_weights(self):
        """Initialize weights of VoteHead."""
        pass
wuyuefeng's avatar
Votenet  
wuyuefeng committed
98

99
100
101
102
103
104
105
    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
106
107
108
        # Objectness scores (2), center residual (3),
        # heading class+residual (num_dir_bins*2),
        # size class+residual(num_sizes*4)
109
        return 3 + self.num_dir_bins * 2 + self.num_sizes * 4
wuyuefeng's avatar
Votenet  
wuyuefeng committed
110

111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
    def _extract_input(self, feat_dict):
        """Extract inputs from features dictionary.

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

        Returns:
            torch.Tensor: Coordinates of input points.
            torch.Tensor: Features of input points.
            torch.Tensor: Indices of input points.
        """
        seed_points = feat_dict['fp_xyz'][-1]
        seed_features = feat_dict['fp_features'][-1]
        seed_indices = feat_dict['fp_indices'][-1]

        return seed_points, seed_features, seed_indices
wuyuefeng's avatar
Votenet  
wuyuefeng committed
127
128
129
130

    def forward(self, feat_dict, sample_mod):
        """Forward pass.

zhangwenwei's avatar
zhangwenwei committed
131
132
133
134
135
136
137
        Note:
            The forward of VoteHead is devided into 4 steps:

                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
138
139

        Args:
wangtai's avatar
wangtai committed
140
141
            feat_dict (dict): Feature dict from backbone.
            sample_mod (str): Sample mode for vote aggregation layer.
142
                valid modes are "vote", "seed", "random" and "spec".
wuyuefeng's avatar
wuyuefeng committed
143
144
145

        Returns:
            dict: Predictions of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
146
        """
147
        assert sample_mod in ['vote', 'seed', 'random', 'spec']
wuyuefeng's avatar
Votenet  
wuyuefeng committed
148

149
150
        seed_points, seed_features, seed_indices = self._extract_input(
            feat_dict)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
151
152

        # 1. generate vote_points from seed_points
153
154
        vote_points, vote_features, vote_offset = self.vote_module(
            seed_points, seed_features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
155
156
157
158
        results = dict(
            seed_points=seed_points,
            seed_indices=seed_indices,
            vote_points=vote_points,
159
160
            vote_features=vote_features,
            vote_offset=vote_offset)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
161
162
163
164

        # 2. aggregate vote_points
        if sample_mod == 'vote':
            # use fps in vote_aggregation
165
166
            aggregation_inputs = dict(
                points_xyz=vote_points, features=vote_features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
167
168
169
170
        elif sample_mod == 'seed':
            # FPS on seed and choose the votes corresponding to the seeds
            sample_indices = furthest_point_sample(seed_points,
                                                   self.num_proposal)
171
172
173
174
            aggregation_inputs = dict(
                points_xyz=vote_points,
                features=vote_features,
                indices=sample_indices)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
175
176
177
178
179
180
        elif sample_mod == 'random':
            # 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)
181
182
183
184
185
186
187
188
189
190
            aggregation_inputs = dict(
                points_xyz=vote_points,
                features=vote_features,
                indices=sample_indices)
        elif sample_mod == 'spec':
            # 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
191
        else:
Wenwei Zhang's avatar
Wenwei Zhang committed
192
193
            raise NotImplementedError(
                f'Sample mode {sample_mod} is not supported!')
wuyuefeng's avatar
Votenet  
wuyuefeng committed
194

195
        vote_aggregation_ret = self.vote_aggregation(**aggregation_inputs)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
196
        aggregated_points, features, aggregated_indices = vote_aggregation_ret
197

wuyuefeng's avatar
Votenet  
wuyuefeng committed
198
        results['aggregated_points'] = aggregated_points
encore-zhou's avatar
encore-zhou committed
199
        results['aggregated_features'] = features
wuyuefeng's avatar
Votenet  
wuyuefeng committed
200
201
202
        results['aggregated_indices'] = aggregated_indices

        # 3. predict bbox and score
203
        cls_predictions, reg_predictions = self.conv_pred(features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
204
205

        # 4. decode predictions
206
207
208
209
        decode_res = self.bbox_coder.split_pred(cls_predictions,
                                                reg_predictions,
                                                aggregated_points)

wuyuefeng's avatar
Votenet  
wuyuefeng committed
210
211
212
213
        results.update(decode_res)

        return results

214
    @force_fp32(apply_to=('bbox_preds', ))
wuyuefeng's avatar
Votenet  
wuyuefeng committed
215
216
217
218
219
220
221
    def loss(self,
             bbox_preds,
             points,
             gt_bboxes_3d,
             gt_labels_3d,
             pts_semantic_mask=None,
             pts_instance_mask=None,
zhangwenwei's avatar
zhangwenwei committed
222
             img_metas=None,
encore-zhou's avatar
encore-zhou committed
223
224
             gt_bboxes_ignore=None,
             ret_target=False):
wuyuefeng's avatar
wuyuefeng committed
225
226
227
228
        """Compute loss.

        Args:
            bbox_preds (dict): Predictions from forward of vote head.
liyinhao's avatar
liyinhao committed
229
            points (list[torch.Tensor]): Input points.
wangtai's avatar
wangtai committed
230
231
232
            gt_bboxes_3d (list[:obj:`BaseInstance3DBoxes`]): Ground truth \
                bboxes of each sample.
            gt_labels_3d (list[torch.Tensor]): Labels of each sample.
liyinhao's avatar
liyinhao committed
233
234
235
236
            pts_semantic_mask (None | list[torch.Tensor]): Point-wise
                semantic mask.
            pts_instance_mask (None | list[torch.Tensor]): Point-wise
                instance mask.
zhangwenwei's avatar
zhangwenwei committed
237
            img_metas (list[dict]): Contain pcd and img's meta info.
liyinhao's avatar
liyinhao committed
238
239
            gt_bboxes_ignore (None | list[torch.Tensor]): Specify
                which bounding.
encore-zhou's avatar
encore-zhou committed
240
            ret_target (Bool): Return targets or not.
wuyuefeng's avatar
wuyuefeng committed
241
242
243
244

        Returns:
            dict: Losses of Votenet.
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
245
246
247
248
        targets = self.get_targets(points, gt_bboxes_3d, gt_labels_3d,
                                   pts_semantic_mask, pts_instance_mask,
                                   bbox_preds)
        (vote_targets, vote_target_masks, size_class_targets, size_res_targets,
249
250
251
252
         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
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300

        # calculate vote loss
        vote_loss = self.vote_module.get_loss(bbox_preds['seed_points'],
                                              bbox_preds['vote_points'],
                                              bbox_preds['seed_indices'],
                                              vote_target_masks, vote_targets)

        # calculate objectness loss
        objectness_loss = self.objectness_loss(
            bbox_preds['obj_scores'].transpose(2, 1),
            objectness_targets,
            weight=objectness_weights)

        # calculate center loss
        source2target_loss, target2source_loss = self.center_loss(
            bbox_preds['center'],
            center_targets,
            src_weight=box_loss_weights,
            dst_weight=valid_gt_weights)
        center_loss = source2target_loss + target2source_loss

        # calculate direction class loss
        dir_class_loss = self.dir_class_loss(
            bbox_preds['dir_class'].transpose(2, 1),
            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(
            bbox_preds['dir_res_norm'] * heading_label_one_hot, -1)
        dir_res_loss = self.dir_res_loss(
            dir_res_norm, dir_res_targets, weight=box_loss_weights)

        # calculate size class loss
        size_class_loss = self.size_class_loss(
            bbox_preds['size_class'].transpose(2, 1),
            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
301
            -1).repeat(1, 1, 1, 3).contiguous()
wuyuefeng's avatar
Votenet  
wuyuefeng committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
        size_residual_norm = torch.sum(
            bbox_preds['size_res_norm'] * one_hot_size_targets_expand, 2)
        box_loss_weights_expand = box_loss_weights.unsqueeze(-1).repeat(
            1, 1, 3)
        size_res_loss = self.size_res_loss(
            size_residual_norm,
            size_res_targets,
            weight=box_loss_weights_expand)

        # calculate semantic loss
        semantic_loss = self.semantic_loss(
            bbox_preds['sem_scores'].transpose(2, 1),
            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
326

327
328
329
330
331
332
333
334
335
336
337
        if self.iou_loss:
            corners_pred = self.bbox_coder.decode_corners(
                bbox_preds['center'], size_residual_norm,
                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
338
339
340
        if ret_target:
            losses['targets'] = targets

wuyuefeng's avatar
Votenet  
wuyuefeng committed
341
342
343
344
345
346
347
348
349
        return losses

    def get_targets(self,
                    points,
                    gt_bboxes_3d,
                    gt_labels_3d,
                    pts_semantic_mask=None,
                    pts_instance_mask=None,
                    bbox_preds=None):
wuyuefeng's avatar
wuyuefeng committed
350
        """Generate targets of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
351
352

        Args:
liyinhao's avatar
liyinhao committed
353
            points (list[torch.Tensor]): Points of each batch.
wangtai's avatar
wangtai committed
354
355
356
357
            gt_bboxes_3d (list[:obj:`BaseInstance3DBoxes`]): Ground truth \
                bboxes of each batch.
            gt_labels_3d (list[torch.Tensor]): Labels of each batch.
            pts_semantic_mask (None | list[torch.Tensor]): Point-wise semantic
wuyuefeng's avatar
Votenet  
wuyuefeng committed
358
                label of each batch.
wangtai's avatar
wangtai committed
359
            pts_instance_mask (None | list[torch.Tensor]): Point-wise instance
wuyuefeng's avatar
Votenet  
wuyuefeng committed
360
                label of each batch.
wangtai's avatar
wangtai committed
361
            bbox_preds (torch.Tensor): Bounding box predictions of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
362
363

        Returns:
364
            tuple[torch.Tensor]: Targets of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
365
366
367
368
369
370
        """
        # find empty example
        valid_gt_masks = list()
        gt_num = list()
        for index in range(len(gt_labels_3d)):
            if len(gt_labels_3d[index]) == 0:
wuyuefeng's avatar
wuyuefeng committed
371
372
373
                fake_box = gt_bboxes_3d[index].tensor.new_zeros(
                    1, gt_bboxes_3d[index].tensor.shape[-1])
                gt_bboxes_3d[index] = gt_bboxes_3d[index].new_box(fake_box)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
                gt_labels_3d[index] = gt_labels_3d[index].new_zeros(1)
                valid_gt_masks.append(gt_labels_3d[index].new_zeros(1))
                gt_num.append(1)
            else:
                valid_gt_masks.append(gt_labels_3d[index].new_ones(
                    gt_labels_3d[index].shape))
                gt_num.append(gt_labels_3d[index].shape[0])
        max_gt_num = max(gt_num)

        if pts_semantic_mask is None:
            pts_semantic_mask = [None for i in range(len(gt_labels_3d))]
            pts_instance_mask = [None for i in range(len(gt_labels_3d))]

        aggregated_points = [
            bbox_preds['aggregated_points'][i]
            for i in range(len(gt_labels_3d))
        ]

        (vote_targets, vote_target_masks, size_class_targets, size_res_targets,
393
394
395
396
397
398
         dir_class_targets, dir_res_targets, center_targets,
         assigned_center_targets, mask_targets, objectness_targets,
         objectness_masks) = multi_apply(self.get_targets_single, points,
                                         gt_bboxes_3d, gt_labels_3d,
                                         pts_semantic_mask, pts_instance_mask,
                                         aggregated_points)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
399
400
401
402
403
404
405
406
407
408
409
410
411

        # pad targets as original code of votenet.
        for index in range(len(gt_labels_3d)):
            pad_num = max_gt_num - gt_labels_3d[index].shape[0]
            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)

412
        assigned_center_targets = torch.stack(assigned_center_targets)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
        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,
428
429
430
                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
431
432
433
434
435
436
437
438

    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
439
440
441
        """Generate targets of vote head for single batch.

        Args:
liyinhao's avatar
liyinhao committed
442
            points (torch.Tensor): Points of each batch.
wangtai's avatar
wangtai committed
443
444
445
446
            gt_bboxes_3d (:obj:`BaseInstance3DBoxes`): Ground truth \
                boxes of each batch.
            gt_labels_3d (torch.Tensor): Labels of each batch.
            pts_semantic_mask (None | torch.Tensor): Point-wise semantic
wuyuefeng's avatar
wuyuefeng committed
447
                label of each batch.
wangtai's avatar
wangtai committed
448
            pts_instance_mask (None | torch.Tensor): Point-wise instance
wuyuefeng's avatar
wuyuefeng committed
449
                label of each batch.
liyinhao's avatar
liyinhao committed
450
            aggregated_points (torch.Tensor): Aggregated points from
wuyuefeng's avatar
wuyuefeng committed
451
452
453
                vote aggregation layer.

        Returns:
454
            tuple[torch.Tensor]: Targets of vote head.
wuyuefeng's avatar
wuyuefeng committed
455
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
456
457
        assert self.bbox_coder.with_rot or pts_semantic_mask is not None

wuyuefeng's avatar
wuyuefeng committed
458
459
        gt_bboxes_3d = gt_bboxes_3d.to(points.device)

wuyuefeng's avatar
Votenet  
wuyuefeng committed
460
461
462
463
464
465
466
        # 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)
wuyuefeng's avatar
wuyuefeng committed
467
468
            box_indices_all = gt_bboxes_3d.points_in_boxes(points)
            for i in range(gt_labels_3d.shape[0]):
wuyuefeng's avatar
Votenet  
wuyuefeng committed
469
                box_indices = box_indices_all[:, i]
470
471
                indices = torch.nonzero(
                    box_indices, as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
472
473
474
                selected_points = points[indices]
                vote_target_masks[indices] = 1
                vote_targets_tmp = vote_targets[indices]
wuyuefeng's avatar
wuyuefeng committed
475
                votes = gt_bboxes_3d.gravity_center[i].unsqueeze(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
476
477
478
479
                    0) - selected_points[:, :3]

                for j in range(self.gt_per_seed):
                    column_indices = torch.nonzero(
480
481
                        vote_target_idx[indices] == j,
                        as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
                    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):
498
499
                indices = torch.nonzero(
                    pts_instance_mask == i, as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
                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
538
        one_hot_size_targets = gt_bboxes_3d.tensor.new_zeros(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
539
540
541
542
543
544
545
546
547
548
            (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]
549
        assigned_center_targets = center_targets[assignment]
wuyuefeng's avatar
Votenet  
wuyuefeng committed
550
551

        return (vote_targets, vote_target_masks, size_class_targets,
552
553
                size_res_targets, dir_class_targets,
                dir_res_targets, center_targets, assigned_center_targets,
wuyuefeng's avatar
Votenet  
wuyuefeng committed
554
555
                mask_targets.long(), objectness_targets, objectness_masks)

encore-zhou's avatar
encore-zhou committed
556
557
558
559
560
561
    def get_bboxes(self,
                   points,
                   bbox_preds,
                   input_metas,
                   rescale=False,
                   use_nms=True):
wuyuefeng's avatar
wuyuefeng committed
562
563
564
        """Generate bboxes from vote head predictions.

        Args:
liyinhao's avatar
liyinhao committed
565
            points (torch.Tensor): Input points.
wuyuefeng's avatar
wuyuefeng committed
566
            bbox_preds (dict): Predictions from vote head.
wangtai's avatar
wangtai committed
567
            input_metas (list[dict]): Point cloud and image's meta info.
wuyuefeng's avatar
wuyuefeng committed
568
            rescale (bool): Whether to rescale bboxes.
encore-zhou's avatar
encore-zhou committed
569
570
            use_nms (bool): Whether to apply NMS, skip nms postprocessing
                while using vote head in rpn stage.
wuyuefeng's avatar
wuyuefeng committed
571
572

        Returns:
wangtai's avatar
wangtai committed
573
            list[tuple[torch.Tensor]]: Bounding boxes, scores and labels.
wuyuefeng's avatar
wuyuefeng committed
574
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
575
576
577
        # decode boxes
        obj_scores = F.softmax(bbox_preds['obj_scores'], dim=-1)[..., -1]
        sem_scores = F.softmax(bbox_preds['sem_scores'], dim=-1)
wuyuefeng's avatar
wuyuefeng committed
578
        bbox3d = self.bbox_coder.decode(bbox_preds)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
579

encore-zhou's avatar
encore-zhou committed
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
        if use_nms:
            batch_size = bbox3d.shape[0]
            results = list()
            for b in range(batch_size):
                bbox_selected, score_selected, labels = \
                    self.multiclass_nms_single(obj_scores[b], sem_scores[b],
                                               bbox3d[b], points[b, ..., :3],
                                               input_metas[b])
                bbox = input_metas[b]['box_type_3d'](
                    bbox_selected,
                    box_dim=bbox_selected.shape[-1],
                    with_yaw=self.bbox_coder.with_rot)
                results.append((bbox, score_selected, labels))

            return results
        else:
            return bbox3d
wuyuefeng's avatar
Votenet  
wuyuefeng committed
597

wuyuefeng's avatar
wuyuefeng committed
598
599
    def multiclass_nms_single(self, obj_scores, sem_scores, bbox, points,
                              input_meta):
wangtai's avatar
wangtai committed
600
        """Multi-class nms in single batch.
wuyuefeng's avatar
wuyuefeng committed
601
602

        Args:
wangtai's avatar
wangtai committed
603
604
605
            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
606
            points (torch.Tensor): Input points.
wangtai's avatar
wangtai committed
607
            input_meta (dict): Point cloud and image's meta info.
wuyuefeng's avatar
wuyuefeng committed
608
609

        Returns:
wangtai's avatar
wangtai committed
610
            tuple[torch.Tensor]: Bounding boxes, scores and labels.
wuyuefeng's avatar
wuyuefeng committed
611
        """
wuyuefeng's avatar
wuyuefeng committed
612
613
614
615
616
617
        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))
        box_indices = bbox.points_in_boxes(points)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
618

wuyuefeng's avatar
wuyuefeng committed
619
        corner3d = bbox.corners
wuyuefeng's avatar
Votenet  
wuyuefeng committed
620
621
622
623
        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
624
625
626
        nonempty_box_mask = box_indices.T.sum(1) > 5

        bbox_classes = torch.argmax(sem_scores, -1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
627
628
629
630
631
632
633
        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)
634
635
        nonempty_box_inds = torch.nonzero(
            nonempty_box_mask, as_tuple=False).flatten()
wuyuefeng's avatar
Votenet  
wuyuefeng committed
636
637
638
639
640
641
642
        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
643
                bbox_selected.append(bbox[selected].tensor)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
644
645
646
647
648
649
650
651
                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
652
            bbox_selected = bbox[selected].tensor
wuyuefeng's avatar
Votenet  
wuyuefeng committed
653
654
655
656
            score_selected = obj_scores[selected]
            labels = bbox_classes[selected]

        return bbox_selected, score_selected, labels