vote_head.py 27.7 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
2
3
import numpy as np
import torch
4
from mmcv.runner import BaseModule, force_fp32
zhangwenwei's avatar
zhangwenwei committed
5
from torch.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


@HEADS.register_module()
18
class VoteHead(BaseModule):
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
                 semantic_loss=None,
59
60
61
                 iou_loss=None,
                 init_cfg=None):
        super(VoteHead, self).__init__(init_cfg=init_cfg)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
62
63
64
        self.num_classes = num_classes
        self.train_cfg = train_cfg
        self.test_cfg = test_cfg
65
        self.gt_per_seed = vote_module_cfg['gt_per_seed']
wuyuefeng's avatar
Votenet  
wuyuefeng committed
66
67
68
69
70
        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)
71
        self.dir_class_loss = build_loss(dir_class_loss)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
72
        self.size_res_loss = build_loss(size_res_loss)
73
74
75
76
        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)
77
78
79
80
        if iou_loss is not None:
            self.iou_loss = build_loss(iou_loss)
        else:
            self.iou_loss = None
wuyuefeng's avatar
Votenet  
wuyuefeng committed
81
82
83
84
85

        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

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

90
91
92
93
94
95
96
97
98
99
100
101
102
        # 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 _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
103
104
105
        # Objectness scores (2), center residual (3),
        # heading class+residual (num_dir_bins*2),
        # size class+residual(num_sizes*4)
106
        return 3 + self.num_dir_bins * 2 + self.num_sizes * 4
wuyuefeng's avatar
Votenet  
wuyuefeng committed
107

108
109
110
111
112
113
114
115
116
117
118
    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.
        """
119
120
121
122
123
124
125
126
127
128
129
130
131

        # 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]
132
133

        return seed_points, seed_features, seed_indices
wuyuefeng's avatar
Votenet  
wuyuefeng committed
134
135
136
137

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

zhangwenwei's avatar
zhangwenwei committed
138
139
140
141
142
143
144
        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
145
146

        Args:
wangtai's avatar
wangtai committed
147
148
            feat_dict (dict): Feature dict from backbone.
            sample_mod (str): Sample mode for vote aggregation layer.
149
                valid modes are "vote", "seed", "random" and "spec".
wuyuefeng's avatar
wuyuefeng committed
150
151
152

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

156
157
        seed_points, seed_features, seed_indices = self._extract_input(
            feat_dict)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
158
159

        # 1. generate vote_points from seed_points
160
161
        vote_points, vote_features, vote_offset = self.vote_module(
            seed_points, seed_features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
162
163
164
165
        results = dict(
            seed_points=seed_points,
            seed_indices=seed_indices,
            vote_points=vote_points,
166
167
            vote_features=vote_features,
            vote_offset=vote_offset)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
168
169
170
171

        # 2. aggregate vote_points
        if sample_mod == 'vote':
            # use fps in vote_aggregation
172
173
            aggregation_inputs = dict(
                points_xyz=vote_points, features=vote_features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
174
175
176
177
        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)
178
179
180
181
            aggregation_inputs = dict(
                points_xyz=vote_points,
                features=vote_features,
                indices=sample_indices)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
182
183
184
185
186
187
        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)
188
189
190
191
192
193
194
195
196
197
            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
198
        else:
Wenwei Zhang's avatar
Wenwei Zhang committed
199
200
            raise NotImplementedError(
                f'Sample mode {sample_mod} is not supported!')
wuyuefeng's avatar
Votenet  
wuyuefeng committed
201

202
        vote_aggregation_ret = self.vote_aggregation(**aggregation_inputs)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
203
        aggregated_points, features, aggregated_indices = vote_aggregation_ret
204

wuyuefeng's avatar
Votenet  
wuyuefeng committed
205
        results['aggregated_points'] = aggregated_points
encore-zhou's avatar
encore-zhou committed
206
        results['aggregated_features'] = features
wuyuefeng's avatar
Votenet  
wuyuefeng committed
207
208
209
        results['aggregated_indices'] = aggregated_indices

        # 3. predict bbox and score
210
        cls_predictions, reg_predictions = self.conv_pred(features)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
211
212

        # 4. decode predictions
213
214
215
216
        decode_res = self.bbox_coder.split_pred(cls_predictions,
                                                reg_predictions,
                                                aggregated_points)

wuyuefeng's avatar
Votenet  
wuyuefeng committed
217
218
219
220
        results.update(decode_res)

        return results

221
    @force_fp32(apply_to=('bbox_preds', ))
wuyuefeng's avatar
Votenet  
wuyuefeng committed
222
223
224
225
226
227
228
    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
229
             img_metas=None,
encore-zhou's avatar
encore-zhou committed
230
231
             gt_bboxes_ignore=None,
             ret_target=False):
wuyuefeng's avatar
wuyuefeng committed
232
233
234
235
        """Compute loss.

        Args:
            bbox_preds (dict): Predictions from forward of vote head.
liyinhao's avatar
liyinhao committed
236
            points (list[torch.Tensor]): Input points.
wangtai's avatar
wangtai committed
237
238
239
            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
240
241
242
243
            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
244
            img_metas (list[dict]): Contain pcd and img's meta info.
liyinhao's avatar
liyinhao committed
245
246
            gt_bboxes_ignore (None | list[torch.Tensor]): Specify
                which bounding.
encore-zhou's avatar
encore-zhou committed
247
            ret_target (Bool): Return targets or not.
wuyuefeng's avatar
wuyuefeng committed
248
249
250
251

        Returns:
            dict: Losses of Votenet.
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
252
253
254
255
        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,
256
257
258
259
         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
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
301
302
303
304
305
306
307

        # 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
308
            -1).repeat(1, 1, 1, 3).contiguous()
wuyuefeng's avatar
Votenet  
wuyuefeng committed
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
        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
333

334
335
336
337
338
339
340
341
342
343
344
        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
345
346
347
        if ret_target:
            losses['targets'] = targets

wuyuefeng's avatar
Votenet  
wuyuefeng committed
348
349
350
351
352
353
354
355
356
        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
357
        """Generate targets of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
358
359

        Args:
liyinhao's avatar
liyinhao committed
360
            points (list[torch.Tensor]): Points of each batch.
wangtai's avatar
wangtai committed
361
362
363
364
            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
365
                label of each batch.
wangtai's avatar
wangtai committed
366
            pts_instance_mask (None | list[torch.Tensor]): Point-wise instance
wuyuefeng's avatar
Votenet  
wuyuefeng committed
367
                label of each batch.
wangtai's avatar
wangtai committed
368
            bbox_preds (torch.Tensor): Bounding box predictions of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
369
370

        Returns:
371
            tuple[torch.Tensor]: Targets of vote head.
wuyuefeng's avatar
Votenet  
wuyuefeng committed
372
373
374
375
376
377
        """
        # 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
378
379
380
                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
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
                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,
400
401
402
403
404
405
         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
406
407
408
409
410
411
412
413
414
415
416
417
418

        # 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)

419
        assigned_center_targets = torch.stack(assigned_center_targets)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
        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,
435
436
437
                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
438
439
440
441
442
443
444
445

    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
446
447
448
        """Generate targets of vote head for single batch.

        Args:
liyinhao's avatar
liyinhao committed
449
            points (torch.Tensor): Points of each batch.
wangtai's avatar
wangtai committed
450
451
452
453
            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
454
                label of each batch.
wangtai's avatar
wangtai committed
455
            pts_instance_mask (None | torch.Tensor): Point-wise instance
wuyuefeng's avatar
wuyuefeng committed
456
                label of each batch.
liyinhao's avatar
liyinhao committed
457
            aggregated_points (torch.Tensor): Aggregated points from
wuyuefeng's avatar
wuyuefeng committed
458
459
460
                vote aggregation layer.

        Returns:
461
            tuple[torch.Tensor]: Targets of vote head.
wuyuefeng's avatar
wuyuefeng committed
462
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
463
464
        assert self.bbox_coder.with_rot or pts_semantic_mask is not None

wuyuefeng's avatar
wuyuefeng committed
465
466
        gt_bboxes_3d = gt_bboxes_3d.to(points.device)

wuyuefeng's avatar
Votenet  
wuyuefeng committed
467
468
469
470
471
472
473
        # 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)
474
            box_indices_all = gt_bboxes_3d.points_in_boxes_batch(points)
wuyuefeng's avatar
wuyuefeng committed
475
            for i in range(gt_labels_3d.shape[0]):
wuyuefeng's avatar
Votenet  
wuyuefeng committed
476
                box_indices = box_indices_all[:, i]
477
478
                indices = torch.nonzero(
                    box_indices, as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
479
480
481
                selected_points = points[indices]
                vote_target_masks[indices] = 1
                vote_targets_tmp = vote_targets[indices]
wuyuefeng's avatar
wuyuefeng committed
482
                votes = gt_bboxes_3d.gravity_center[i].unsqueeze(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
483
484
485
486
                    0) - selected_points[:, :3]

                for j in range(self.gt_per_seed):
                    column_indices = torch.nonzero(
487
488
                        vote_target_idx[indices] == j,
                        as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
                    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):
505
506
                indices = torch.nonzero(
                    pts_instance_mask == i, as_tuple=False).squeeze(-1)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
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
538
539
540
541
542
543
544
                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
545
        one_hot_size_targets = gt_bboxes_3d.tensor.new_zeros(
wuyuefeng's avatar
Votenet  
wuyuefeng committed
546
547
548
549
550
551
552
553
554
555
            (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]
556
        assigned_center_targets = center_targets[assignment]
wuyuefeng's avatar
Votenet  
wuyuefeng committed
557
558

        return (vote_targets, vote_target_masks, size_class_targets,
559
560
                size_res_targets, dir_class_targets,
                dir_res_targets, center_targets, assigned_center_targets,
wuyuefeng's avatar
Votenet  
wuyuefeng committed
561
562
                mask_targets.long(), objectness_targets, objectness_masks)

encore-zhou's avatar
encore-zhou committed
563
564
565
566
567
568
    def get_bboxes(self,
                   points,
                   bbox_preds,
                   input_metas,
                   rescale=False,
                   use_nms=True):
wuyuefeng's avatar
wuyuefeng committed
569
570
571
        """Generate bboxes from vote head predictions.

        Args:
liyinhao's avatar
liyinhao committed
572
            points (torch.Tensor): Input points.
wuyuefeng's avatar
wuyuefeng committed
573
            bbox_preds (dict): Predictions from vote head.
wangtai's avatar
wangtai committed
574
            input_metas (list[dict]): Point cloud and image's meta info.
wuyuefeng's avatar
wuyuefeng committed
575
            rescale (bool): Whether to rescale bboxes.
encore-zhou's avatar
encore-zhou committed
576
577
            use_nms (bool): Whether to apply NMS, skip nms postprocessing
                while using vote head in rpn stage.
wuyuefeng's avatar
wuyuefeng committed
578
579

        Returns:
wangtai's avatar
wangtai committed
580
            list[tuple[torch.Tensor]]: Bounding boxes, scores and labels.
wuyuefeng's avatar
wuyuefeng committed
581
        """
wuyuefeng's avatar
Votenet  
wuyuefeng committed
582
583
584
        # 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
585
        bbox3d = self.bbox_coder.decode(bbox_preds)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
586

encore-zhou's avatar
encore-zhou committed
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
        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
604

wuyuefeng's avatar
wuyuefeng committed
605
606
    def multiclass_nms_single(self, obj_scores, sem_scores, bbox, points,
                              input_meta):
wangtai's avatar
wangtai committed
607
        """Multi-class nms in single batch.
wuyuefeng's avatar
wuyuefeng committed
608
609

        Args:
wangtai's avatar
wangtai committed
610
611
612
            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
613
            points (torch.Tensor): Input points.
wangtai's avatar
wangtai committed
614
            input_meta (dict): Point cloud and image's meta info.
wuyuefeng's avatar
wuyuefeng committed
615
616

        Returns:
wangtai's avatar
wangtai committed
617
            tuple[torch.Tensor]: Bounding boxes, scores and labels.
wuyuefeng's avatar
wuyuefeng committed
618
        """
wuyuefeng's avatar
wuyuefeng committed
619
620
621
622
623
        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))
624
        box_indices = bbox.points_in_boxes_batch(points)
wuyuefeng's avatar
Votenet  
wuyuefeng committed
625

wuyuefeng's avatar
wuyuefeng committed
626
        corner3d = bbox.corners
wuyuefeng's avatar
Votenet  
wuyuefeng committed
627
628
629
630
        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
631
632
633
        nonempty_box_mask = box_indices.T.sum(1) > 5

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

        return bbox_selected, score_selected, labels