|
|
@@ -129,7 +129,6 @@ def maskrcnn_loss(mask_logits, proposals, gt_masks, gt_labels, mask_matched_idxs
|
|
|
)
|
|
|
return mask_loss
|
|
|
|
|
|
-
|
|
|
def line_points_to_heatmap(keypoints, rois, heatmap_size):
|
|
|
# type: (Tensor, Tensor, int) -> Tuple[Tensor, Tensor]
|
|
|
print(f'rois:{rois.shape}')
|
|
|
@@ -150,10 +149,54 @@ def line_points_to_heatmap(keypoints, rois, heatmap_size):
|
|
|
x = keypoints[..., 0]
|
|
|
y = keypoints[..., 1]
|
|
|
|
|
|
+ gs=generate_gaussian_heatmaps(x,y,heatmap_size,1.0)
|
|
|
+ # show_heatmap(gs[0],'target')
|
|
|
+ all_roi_heatmap=[]
|
|
|
+ for roi ,heatmap in zip(rois,gs):
|
|
|
+ # print(f'heatmap:{heatmap.shape}')
|
|
|
+ heatmap=heatmap.unsqueeze(0)
|
|
|
+ x1, y1, x2, y2 = map(int, roi)
|
|
|
+ roi_heatmap = torch.zeros_like(heatmap)
|
|
|
+ roi_heatmap[..., y1:y2+1, x1:x2+1]=heatmap[..., y1:y2+1, x1:x2+1]
|
|
|
+ # show_heatmap(roi_heatmap,'roi_heatmap')
|
|
|
+ all_roi_heatmap.append(roi_heatmap)
|
|
|
+
|
|
|
+ all_roi_heatmap=torch.cat(all_roi_heatmap)
|
|
|
+ print(f'all_roi_heatmap:{all_roi_heatmap.shape}')
|
|
|
+
|
|
|
+
|
|
|
+ return all_roi_heatmap
|
|
|
+
|
|
|
+
|
|
|
+"""
|
|
|
+修改适配的原结构的点 转热图，适用于带roi_pool版本的
|
|
|
+"""
|
|
|
+def line_points_to_heatmap_(keypoints, rois, heatmap_size):
|
|
|
+ # type: (Tensor, Tensor, int) -> Tuple[Tensor, Tensor]
|
|
|
+ print(f'rois:{rois.shape}')
|
|
|
+ print(f'heatmap_size:{heatmap_size}')
|
|
|
+ offset_x = rois[:, 0]
|
|
|
+ offset_y = rois[:, 1]
|
|
|
+ scale_x = heatmap_size / (rois[:, 2] - rois[:, 0])
|
|
|
+ scale_y = heatmap_size / (rois[:, 3] - rois[:, 1])
|
|
|
+
|
|
|
+ offset_x = offset_x[:, None]
|
|
|
+ offset_y = offset_y[:, None]
|
|
|
+ scale_x = scale_x[:, None]
|
|
|
+ scale_y = scale_y[:, None]
|
|
|
+
|
|
|
+ print(f'keypoints.shape:{keypoints.shape}')
|
|
|
+ # batch_size, num_keypoints, _ = keypoints.shape
|
|
|
+
|
|
|
+ x = keypoints[..., 0]
|
|
|
+ y = keypoints[..., 1]
|
|
|
+
|
|
|
# gs=generate_gaussian_heatmaps(x,y,512,1.0)
|
|
|
+
|
|
|
+
|
|
|
# print(f'gs_heatmap shape:{gs.shape}')
|
|
|
#
|
|
|
- # show_heatmap(gs,'target')
|
|
|
+ # show_heatmap(gs[0],'target')
|
|
|
|
|
|
x_boundary_inds = x == rois[:, 2][:, None]
|
|
|
y_boundary_inds = y == rois[:, 3][:, None]
|
|
|
@@ -174,7 +217,7 @@ def line_points_to_heatmap(keypoints, rois, heatmap_size):
|
|
|
|
|
|
gs_heatmap=generate_gaussian_heatmaps(x,y,heatmap_size,1.0)
|
|
|
|
|
|
- # show_heatmap(gs_heatmap[0],'feature')
|
|
|
+ show_heatmap(gs_heatmap[0],'feature')
|
|
|
|
|
|
# print(f'gs_heatmap:{gs_heatmap.shape}')
|
|
|
#
|
|
|
@@ -1266,6 +1309,7 @@ class RoIHeads(nn.Module):
|
|
|
raise ValueError("both targets and pos_matched_idxs should not be None when in training mode")
|
|
|
|
|
|
gt_lines = [t["lines"] for t in targets]
|
|
|
+ print(f'gt_lines:{gt_lines[0].shape}')
|
|
|
h, w = targets[0]["img_size"]
|
|
|
img_size = h
|
|
|
rcnn_loss_line = lines_point_pair_loss(
|
