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,200pred_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_dictself.forward_ret_dict['pred_dicts'] = pred_dictsif 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'] = roisdata_dict['roi_scores'] = roi_scoresdata_dict['roi_labels'] = roi_labelsdata_dict['has_class_labels'] = Trueelse:data_dict['final_box_dicts'] = pred_dictsreturn 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 #firstpost_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 Nonefinal_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_scoresfinal_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) + 1return ret_dictpost_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 Nonemask = (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_xstopk_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_scoresfinal_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的办法
本文来自互联网用户投稿，文章观点仅代表作者本人，不代表本站立场，不承担相关法律责任。如若转载，请注明出处。 如若内容造成侵权/违法违规/事实不符，请点击【内容举报】进行投诉反馈！