def get_roi_from_rpn(box_score,box,nms_cfg):
"""
box_score:(num of anchor,)
box:(num of anchor, 7)
nms_cfg:nms的配置文件
"""
selected = []
if box.shape[0]>0:
rank_score,indices = torch.topk(box_score,k=min(nms_cfg.NMS_PRE_MAXSIZE,box.shape[0]))
for_nms_box = box[indices]
nms_indices, _ = iou3d_nms_utils.nms_gpu(for_nms_box,rank_score,nms_cfg.NMS_THRESH,**nms_cfg)
selected = indices[nms_indices[:nms_cfg.NMS_POST_MAXSIZE]]
return selected.view(-1),box_score[selected]
def class_final_nms(self,box_scores,box_preds,nms_config,score_thresh=None):
"""
box_scores:(N,)
box_preds:(N,7)
score thresh:(1,)
"""
box_scores_src = box_scores
score_mask = (box_scores>score_thresh).squeeze(dim=-1)
box_scores = box_scores[score_mask]
box_preds = box_preds[score_mask]
selected =[]
if box_scores.shape[0]>0:
score_for_nms,indices = torch.topk(
box_scores.squeeze(dim=-1),k=min(nms_config.NMS_PRE_MAXSIZE,box_scores.shape[0]))
box_for_nms = box_preds[indices]
keep_id,_ = iou3d_nms_utils.nms_gpu(box_for_nms,score_for_nms,nms_config.NMS_THRESH)
selected = indices[keep_id[:nms_config.NMS_POST_MAXSIZE]]
origin_idx = torch.nonzero(score_mask).view(-1)
selected = origin_idx[selected].view(-1)
return selected,box_scores_src[selected].view(-1)
def forward(self, inputs):
x = self.conv_input(inputs["input_sp_tensor"])
x_conv1 = self.conv1(x)
x_conv2 = self.conv2(x_conv1)
x_conv3 = self.conv3(x_conv2)
x_conv4 = self.conv4(x_conv3)
out = self.conv_out(x_conv4)
spatial_features = out.dense()
N, C, D, H, W = spatial_features.shape
spatial_features = spatial_features.view(N, C * D, H, W) #压缩到二维特征
#######################--------VOXEL_SA----------############################################################
#对原始点云采集keypoints作为new_xyz
keypoints = self.get_sampled_points(inputs["points"])
point_features_list = []
point_bev_features = self.interpolate_from_bev_features(
keypoints, spatial_features)
point_features_list.append(point_bev_features)
batch_size, num_keypoints, _ = keypoints.shape
new_xyz = keypoints.view(-1, 3)
new_xyz_batch_cnt = new_xyz.new_zeros(batch_size).int().fill_(
num_keypoints)
#对原始点云进行聚合
xyz = inputs["points"][:, :3]
xyz_batch_cnt = xyz.new_zeros(batch_size).int()
xyz_batch_cnt[0] = xyz.shape[0]
point_features = inputs["points"][:, 3:]
pooled_points, pooled_features = self.raw_sa_layer(
xyz=xyz.contiguous(),
xyz_batch_cnt=xyz_batch_cnt,
new_xyz=new_xyz,
new_xyz_batch_cnt=new_xyz_batch_cnt,
features=point_features.contiguous(),
)
point_features_list.append(
pooled_features.view(batch_size, num_keypoints, -1))
#对conv1进行聚合
cur_coords = x_conv1.indices
xyz = self.get_voxel_centers(
cur_coords[:, 1:4],
downsample_times=self.downsample_times_map[0],
voxel_size=self.voxel_size,
point_cloud_range=self.point_cloud_range)
xyz_batch_cnt = xyz.new_zeros(batch_size).int()
xyz_batch_cnt[0] = (cur_coords[:, 0] == 0).sum()
pooled_points, pooled_features = self.conv1_sa_layer(
xyz=xyz.contiguous(),
xyz_batch_cnt=xyz_batch_cnt,
new_xyz=new_xyz,
new_xyz_batch_cnt=new_xyz_batch_cnt,
features=x_conv1.features.contiguous(),
)
point_features_list.append(
pooled_features.view(batch_size, num_keypoints, -1))
#对conv2聚合
cur_coords = x_conv2.indices
xyz = self.get_voxel_centers(
cur_coords[:, 1:4],
downsample_times=self.downsample_times_map[1],
voxel_size=self.voxel_size,
point_cloud_range=self.point_cloud_range)
xyz_batch_cnt = xyz.new_zeros(batch_size).int()
xyz_batch_cnt[0] = (cur_coords[:, 0] == 0).sum()
pooled_points, pooled_features = self.conv2_sa_layer(
xyz=xyz.contiguous(),
xyz_batch_cnt=xyz_batch_cnt,
new_xyz=new_xyz,
new_xyz_batch_cnt=new_xyz_batch_cnt,
features=x_conv2.features.contiguous(),
)
point_features_list.append(
pooled_features.view(batch_size, num_keypoints, -1))
# 对conv3聚合
cur_coords = x_conv3.indices
xyz = self.get_voxel_centers(
cur_coords[:, 1:4],
downsample_times=self.downsample_times_map[2],
voxel_size=self.voxel_size,
point_cloud_range=self.point_cloud_range)
xyz_batch_cnt = xyz.new_zeros(batch_size).int()
xyz_batch_cnt[0] = (cur_coords[:, 0] == 0).sum()
pooled_points, pooled_features = self.conv3_sa_layer(
xyz=xyz.contiguous(),
xyz_batch_cnt=xyz_batch_cnt,
new_xyz=new_xyz,
new_xyz_batch_cnt=new_xyz_batch_cnt,
features=x_conv3.features.contiguous(),
)
point_features_list.append(
pooled_features.view(batch_size, num_keypoints, -1))
# 对conv4聚合
cur_coords = x_conv4.indices
xyz = self.get_voxel_centers(
cur_coords[:, 1:4],
downsample_times=self.downsample_times_map[3],
voxel_size=self.voxel_size,
point_cloud_range=self.point_cloud_range)
xyz_batch_cnt = xyz.new_zeros(batch_size).int()
xyz_batch_cnt[0] = (cur_coords[:, 0] == 0).sum()
pooled_points, pooled_features = self.conv4_sa_layer(
xyz=xyz.contiguous(),
xyz_batch_cnt=xyz_batch_cnt,
new_xyz=new_xyz,
new_xyz_batch_cnt=new_xyz_batch_cnt,
features=x_conv4.features.contiguous(),
)
point_features_list.append(
pooled_features.view(batch_size, num_keypoints, -1))
# 关键点的坐标：point_coords(bxyz)
batch_idx = torch.arange(batch_size, device=keypoints.device).view(
-1, 1).repeat(1, keypoints.shape[1]).view(-1)
point_coords = torch.cat(
(batch_idx.view(-1, 1).float(), keypoints.view(-1, 3)), dim=1)
#关键点的特征 :point_features
point_features = torch.cat(point_features_list, dim=2)
point_features_before_fusion = point_features.view(
-1, point_features.shape[-1]) #原始拼接特征
point_features = self.vsa_point_feature_fusion(
point_features.view(-1, point_features.shape[-1])) #经过线性层后的特征
############################################--RPN-Head----#################################################
ups = []
x = self.block1(spatial_features)
ups.append(self.deblock1(x))
x = self.block2(x)
ups.append(self.deblock2(x))
x = torch.cat(ups, dim=1)
rpn_box_preds = self.rpn_head_conv_box(x)
rpn_cls_preds = self.rpn_head_conv_cls(x)
rpn_box_preds = rpn_box_preds.permute(0, 2, 3, 1).contiguous()
rpn_cls_preds = rpn_cls_preds.permute(0, 2, 3, 1).contiguous()
rpn_dir_preds = self.rpn_head_conv_dir_cls(x)
#使用box回归值和anhcor进行decode,
rpn_box_preds = rpn_box_preds.reshape(batch_size, -1, 7)
rpn_cls_preds = rpn_cls_preds.reshape(batch_size, -1, 3)
rpn_dir_preds_score = rpn_dir_preds.permute(0, 2, 3, 1).reshape(
batch_size, -1, 2)
dir_cls_preds_label = torch.argmax(rpn_dir_preds_score, dim=-1)
rpn_box_preds = self.decode_torch(rpn_box_preds, self.anchors)
rot_angle_preds = self.limit_period(rpn_box_preds[..., 6] - 0.78539,
offset=0.0,
period=np.pi)
rot_angle_preds_final = rot_angle_preds + 0.78539 + np.pi * dir_cls_preds_label.to(
rpn_box_preds.dtype)
rpn_box_preds[..., 6] = rot_angle_preds_final % (np.pi * 2)
##################################Point Head######################################################
point_cls_preds = self.point_head_cls_layer(
point_features_before_fusion)
point_cls_score, _ = torch.sigmoid(point_cls_preds).max(dim=-1)
###################################-----RCNN--------##############################################
#感兴趣区域的筛选
rois = rpn_box_preds.new_zeros(batch_size, 100, 7)
roi_score = rpn_box_preds.new_zeros(batch_size, 100)
roi_labels = rpn_box_preds.new_zeros(batch_size, 100, dtype=torch.long)
cur_box = rpn_box_preds[0]
cur_cls = rpn_cls_preds[0]
cur_roi_score, cur_roi_labels = torch.max(cur_cls, dim=1)
rank_score, indices = torch.topk(cur_roi_score,
k=min(1024, 7)) #根据rpn类别得分得到前1024个box
for_nms_box = cur_box[indices]
nms_indices, _ = iou3d_nms_utils.nms_gpu(for_nms_box, rank_score,
0.7) #对1024个box进行nms
selected = indices[
nms_indices[:100]] #选则rpn类别得分前100个box进行second-stage细化
rois[0, :len(selected), :] = cur_box[selected]
roi_score[0, :len(selected)] = cur_roi_score[selected]
roi_labels[0, :len(
selected)] = cur_roi_labels[selected] + 1 #类别从1-3并非0-2，将0认为是背景应该是
#将感兴趣区域网格化，每个3d-box转化为6*6*6的网格
rois = rois.view(-1, 7)
local_grid = rois.new_ones(6, 6, 6)
local_grid_id = torch.nonzero(local_grid).float() #(6*6*6,3)
local_grid_id = local_grid_id.repeat(100, 1, 1) #(100,216,3)
rois_size = rois[..., 3:6].unsqueeze(dim=1).clone()
gloal_grid_points = (local_grid_id + 0.5) * (
rois_size / 6) - rois_size / 2 #将原点转化到box中心
rois_ry = rois[..., 6].clone()
gloal_grid_points = self.rotate_points_along_z(
gloal_grid_points.clone(),
rois_ry.view(-1)) #将网格化的rois由规定坐标系旋转到真实坐标系
roi_centers = rois[..., 0:3].clone()
gloal_grid_points += roi_centers.unsqueeze(dim=1) #平移网格化的roi
gloal_grid_points = gloal_grid_points.view(batch_size, -1,
gloal_grid_points.shape[-1])
#网格化的roi作为关键点，聚合keypoints的特征
xyz = point_coords[:, 1:4]
xyz_batch_count = xyz.new_zeros(batch_size, ).int()
xyz_batch_count[0] = keypoints.shape[1] #keypoints的个数
new_xyz = gloal_grid_points.view(-1, 3)
new_xyz_bn_count = new_xyz.new_zeros(batch_size).int()
new_xyz_bn_count[0] = gloal_grid_points.shape[1]
pooled_xyz, pooled_features = self.roi_pool_layer(
xyz.contiguous(),
xyz_batch_count,
new_xyz,
new_xyz_bn_count,
features=point_features)
pooled_features = pooled_features.view(-1, 6**3,
pooled_features.shape[-1])
pooled_features = pooled_features.permute(
0, 2, 1).contiguous() # (100,128,6*6*6)
share_features = self.rcnn_share_layer(pooled_features.view(
100, -1, 1))
rcnn_cls_preds = self.rcnn_cls_layer(share_features).squeeze(
dim=-1).contiguous()
rcnn_reg_preds = self.rcnn_reg_layer(share_features).squeeze(
dim=-1).contiguous()
rcnn_cls_preds = rcnn_cls_preds.unsqueeze(dim=0)
rcnn_reg_preds = rcnn_reg_preds.unsqueeze(dim=0)
#将second-stage阶段的回归结果根据roi进行解码，得到最终的box,解码顺序就是先不带xyz解码，然后在roi的xyz加上
rois = rois.unsqueeze(dim=0)
rois_y = rois[..., 6].view(-1)
rois_xyz = rois[..., 0:3]
rois_temp = rois.clone().detach()
rois_temp[..., 0:3] = 0
box_ct = self.decode_torch(rcnn_reg_preds, rois_temp).view(-1, 7)
box_preds = self.rotate_points_along_z(box_ct.unsqueeze(dim=1),
rois_y).squeeze(dim=1).view(
batch_size, -1, 7)
box_preds[..., 0:3] += rois_xyz
#最终的输出为：box_preds，rcnn_cls_preds
#############################################----POST-Processing----######################################
cur_box_preds = box_preds[0]
cur_cls_preds = rcnn_cls_preds[0]
cur_normal_cls_preds = torch.sigmoid(cur_cls_preds)
label_preds = roi_labels[0]
#进行后处理
score_mask = (cur_normal_cls_preds > 0.1).squeeze(dim=-1)
pps_box_scores = cur_normal_cls_preds[score_mask] #post processing box
pps_box_preds = cur_box_preds[score_mask]
selected = []
if pps_box_scores.shape[0] > 0:
score_for_nms, indices = torch.topk(pps_box_scores.squeeze(dim=-1),
k=min(4096,
pps_box_scores.shape[0]))
box_for_nms = pps_box_preds[indices]
keep_id, _ = iou3d_nms_utils.nms_gpu(box_for_nms, score_for_nms,
0.1)
selected = indices[keep_id[:500]]
origin_idx = torch.nonzero(score_mask).view(-1)
selected = origin_idx[selected].view(-1)
final_box = cur_box_preds[selected]
final_cls = label_preds[selected]
final_score = cur_normal_cls_preds[selected]
return (final_box, final_cls, final_score)