CenterPoint 后处理部分（基于OpenPcdet）

centernet检测代码在pcdet.models.dense_heads.center_head.py

def forward(self, data_dict):
    spatial_features_2d = data_dict['spatial_features_2d']   #之前得到的鸟瞰图特征[1,128,200,200]
    x = self.shared_conv(spatial_features_2d)                #128,200,200->64,200,200

    pred_dicts = []
    for head in self.heads_list:
        pred_dicts.append(head(x))                            #输出各个检测头

    if self.training:
        target_dict = self.assign_targets(
            data_dict['gt_boxes'], feature_map_size=spatial_features_2d.size()[2:],
            feature_map_stride=data_dict.get('spatial_features_2d_strides', None)
        )
        self.forward_ret_dict['target_dicts'] = target_dict

    self.forward_ret_dict['pred_dicts'] = pred_dicts

    if not self.training or self.predict_boxes_when_training:
        pred_dicts = self.generate_predicted_boxes(             # 生成预测框
            data_dict['batch_size'], pred_dicts
        )

        if self.predict_boxes_when_training:
            rois, roi_scores, roi_labels = self.reorder_rois_for_refining(data_dict['batch_size'], pred_dicts)
            data_dict['rois'] = rois
            data_dict['roi_scores'] = roi_scores
            data_dict['roi_labels'] = roi_labels
            data_dict['has_class_labels'] = True
        else:
            data_dict['final_box_dicts'] = pred_dicts

    return data_dict

各个检测头

'center': {'out_channels': 2, 'num_conv': 2},                         [2，200，200]

'center_z': {'out_channels': 1, 'num_conv': 2},                    [1，200，200]

'dim': {'out_channels': 3, 'num_conv': 2},                              [3，200，200]

'rot': {'out_channels': 2, 'num_conv': 2},                                [2，200，200]

'hm': {'out_channels': 3, 'num_conv': 1}                                 [3，200，200]                    

test模式：

pred_dicts = self.generate_predicted_boxes(  
    data_dict['batch_size'], pred_dicts
)
生成检测框

详解generate_predicted_boxes()函数：

    def generate_predicted_boxes(self, batch_size, pred_dicts):
        post_process_cfg = self.model_cfg.POST_PROCESSING      #first
        post_center_limit_range = torch.tensor(post_process_cfg.POST_CENTER_LIMIT_RANGE).cuda().float()

        ret_dict = [{
            'pred_boxes': [],
            'pred_scores': [],
            'pred_labels': [],
        } for k in range(batch_size)]
        for idx, pred_dict in enumerate(pred_dicts):
            batch_hm = pred_dict['hm'].sigmoid()
            batch_center = pred_dict['center']
            batch_center_z = pred_dict['center_z']
            batch_dim = pred_dict['dim'].exp()
            batch_rot_cos = pred_dict['rot'][:, 0].unsqueeze(dim=1)
            batch_rot_sin = pred_dict['rot'][:, 1].unsqueeze(dim=1)
            batch_vel = pred_dict['vel'] if 'vel' in self.separate_head_cfg.HEAD_ORDER else None

            final_pred_dicts = centernet_utils.decode_bbox_from_heatmap(     # 生成预测框，根据MAX_OBJ_PER_SAMPLE = 500 选取前500个热力值高的点，再根据SCORE_THRESH=0.1 选出热力值不低于0.1的点
                heatmap=batch_hm, rot_cos=batch_rot_cos, rot_sin=batch_rot_sin,
                center=batch_center, center_z=batch_center_z, dim=batch_dim, vel=batch_vel,
                point_cloud_range=self.point_cloud_range, voxel_size=self.voxel_size,
                feature_map_stride=self.feature_map_stride,
                K=post_process_cfg.MAX_OBJ_PER_SAMPLE,
                circle_nms=(post_process_cfg.NMS_CONFIG.NMS_TYPE == 'circle_nms'),
                score_thresh=post_process_cfg.SCORE_THRESH,
                post_center_limit_range=post_center_limit_range
            )

            for k, final_dict in enumerate(final_pred_dicts):
                final_dict['pred_labels'] = self.class_id_mapping_each_head[idx][final_dict['pred_labels'].long()]
                if post_process_cfg.NMS_CONFIG.NMS_TYPE != 'circle_nms':       # 如果不是 circle_nms ，使用nms_gpu方法
                    selected, selected_scores = model_nms_utils.class_agnostic_nms(
                        box_scores=final_dict['pred_scores'], box_preds=final_dict['pred_boxes'],
                        nms_config=post_process_cfg.NMS_CONFIG,
                        score_thresh=None
                    )

                    final_dict['pred_boxes'] = final_dict['pred_boxes'][selected]
                    final_dict['pred_scores'] = selected_scores
                    final_dict['pred_labels'] = final_dict['pred_labels'][selected]

                ret_dict[k]['pred_boxes'].append(final_dict['pred_boxes'])
                ret_dict[k]['pred_scores'].append(final_dict['pred_scores'])
                ret_dict[k]['pred_labels'].append(final_dict['pred_labels'])

        for k in range(batch_size):
            ret_dict[k]['pred_boxes'] = torch.cat(ret_dict[k]['pred_boxes'], dim=0)
            ret_dict[k]['pred_scores'] = torch.cat(ret_dict[k]['pred_scores'], dim=0)
            ret_dict[k]['pred_labels'] = torch.cat(ret_dict[k]['pred_labels'], dim=0) + 1

        return ret_dict

post_center_limit_range = tensor([-75.2000, -75.2000,  -2.0000,  75.2000,  75.2000,   4.0000]）

for idx, pred_dict in enumerate(pred_dicts):
    batch_hm = pred_dict['hm'].sigmoid()   
    batch_center = pred_dict['center']      
    batch_center_z = pred_dict['center_z']  
    batch_dim = pred_dict['dim'].exp()
    batch_rot_cos = pred_dict['rot'][:, 0].unsqueeze(dim=1)
    batch_rot_sin = pred_dict['rot'][:, 1].unsqueeze(dim=1)
    batch_vel = pred_dict['vel'] if 'vel' in self.separate_head_cfg.HEAD_ORDER else None

热力图的输出做sigmoid函数处理

dim角度的输出做指数函数处理

其余保留，送入 centernet_utils.decode_bbox_from_heatmap()函数还原出预测框

详解decode_bbox_from_heatmap()函数：

def decode_bbox_from_heatmap(heatmap, rot_cos, rot_sin, center, center_z, dim,
                             point_cloud_range=None, voxel_size=None, feature_map_stride=None, vel=None, K=100,
                             circle_nms=False, score_thresh=None, post_center_limit_range=None):
    batch_size, num_class, _, _ = heatmap.size()

    if circle_nms:
        # TODO: not checked yet
        # assert False, 'not checked yet'
        heatmap = _nms(heatmap)

    scores, inds, class_ids, ys, xs = _topk(heatmap, K=K)
    center = _transpose_and_gather_feat(center, inds).view(batch_size, K, 2)  #根据热力值前500的标记获得对应的角度，旋转等信息
    rot_sin = _transpose_and_gather_feat(rot_sin, inds).view(batch_size, K, 1)
    rot_cos = _transpose_and_gather_feat(rot_cos, inds).view(batch_size, K, 1)
    center_z = _transpose_and_gather_feat(center_z, inds).view(batch_size, K, 1)
    dim = _transpose_and_gather_feat(dim, inds).view(batch_size, K, 3)

    angle = torch.atan2(rot_sin, rot_cos)  # arctan(rot_sin/rot_cos)
    xs = xs.view(batch_size, K, 1) + center[:, :, 0:1]
    ys = ys.view(batch_size, K, 1) + center[:, :, 1:2]

    xs = xs * feature_map_stride * voxel_size[0] + point_cloud_range[0]
    ys = ys * feature_map_stride * voxel_size[1] + point_cloud_range[1]

    box_part_list = [xs, ys, center_z, dim, angle]
    if vel is not None:
        vel = _transpose_and_gather_feat(vel, inds).view(batch_size, K, 2)
        box_part_list.append(vel)

    final_box_preds = torch.cat((box_part_list), dim=-1)
    final_scores = scores.view(batch_size, K)
    final_class_ids = class_ids.view(batch_size, K)

    assert post_center_limit_range is not None
    mask = (final_box_preds[..., :3] >= post_center_limit_range[:3]).all(2)
    mask &= (final_box_preds[..., :3] <= post_center_limit_range[3:]).all(2)

    if score_thresh is not None:
        mask &= (final_scores > score_thresh)

    ret_pred_dicts = []
    for k in range(batch_size):
        cur_mask = mask[k]
        cur_boxes = final_box_preds[k, cur_mask]
        cur_scores = final_scores[k, cur_mask]
        cur_labels = final_class_ids[k, cur_mask]

#TODO：  find what this model want to do
        # if circle_nms:
        #     # assert False, 'not checked yet'
        #     centers = cur_boxes[:, [0, 1]]
        #     boxes = torch.cat((centers, scores.view(-1, 1)), dim=1)
        #     keep = _circle_nms(boxes, min_radius=min_radius, post_max_size=nms_post_max_size)
        #
        #     cur_boxes = cur_boxes[keep]
        #     cur_scores = cur_scores[keep]
        #     cur_labels = cur_labels[keep]

        ret_pred_dicts.append({
            'pred_boxes': cur_boxes,
            'pred_scores': cur_scores,
            'pred_labels': cur_labels
        })
    return ret_pred_dicts

_topk()函数

def _topk(scores, K=40):
    batch, num_class, height, width = scores.size()
    # 按特征图每个点展开找到热力值最大的前500个点
    topk_scores, topk_inds = torch.topk(scores.flatten(2, 3), K)

    topk_inds = topk_inds % (height * width)   #本来就比后面的小，取余还是他本身
    topk_ys = (topk_inds // width).float()
    topk_xs = (topk_inds % width).int().float()

    topk_score, topk_ind = torch.topk(topk_scores.view(batch, -1), K)
    topk_classes = (topk_ind // K).int()
    topk_inds = _gather_feat(topk_inds.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_ys = _gather_feat(topk_ys.view(batch, -1, 1), topk_ind).view(batch, K)
    topk_xs = _gather_feat(topk_xs.view(batch, -1, 1), topk_ind).view(batch, K)

    return topk_score, topk_inds, topk_classes, topk_ys, topk_xs

topk_scores, topk_inds = torch.topk(scores.flatten(2, 3), K) 展平长宽后选前K=500个

topk_inds：1，3，500

topk_inds = topk_inds % (height * width)   #本来就比后面的小，取余还是他本身
topk_ys = (topk_inds // width).float()
topk_xs = (topk_inds % width).int().float()

计算出对应长宽

final_box_preds = torch.cat((box_part_list), dim=-1)  1，500，7

1+1+1+3+1 = 7

if score_thresh is not None:
    mask &= (final_scores > score_thresh)

过滤阈值低于0.1的预测框，得到189个框

            for k, final_dict in enumerate(final_pred_dicts):
                final_dict['pred_labels'] = self.class_id_mapping_each_head[idx][final_dict['pred_labels'].long()]
                if post_process_cfg.NMS_CONFIG.NMS_TYPE != 'circle_nms':       # 如果不是 circle_nms ，使用nms_gpu方法
                    selected, selected_scores = model_nms_utils.class_agnostic_nms(
                        box_scores=final_dict['pred_scores'], box_preds=final_dict['pred_boxes'],
                        nms_config=post_process_cfg.NMS_CONFIG,
                        score_thresh=None
                    )

                    final_dict['pred_boxes'] = final_dict['pred_boxes'][selected]
                    final_dict['pred_scores'] = selected_scores
                    final_dict['pred_labels'] = final_dict['pred_labels'][selected]

                ret_dict[k]['pred_boxes'].append(final_dict['pred_boxes'])
                ret_dict[k]['pred_scores'].append(final_dict['pred_scores'])
                ret_dict[k]['pred_labels'].append(final_dict['pred_labels'])

这里还是使用了nms_gpu的办法