|
|
@@ -204,7 +204,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}')
|
|
|
#
|
|
|
@@ -546,13 +546,13 @@ def arc_points_to_heatmap(keypoints, rois, heatmap_size):
|
|
|
# show_heatmap(gs[0],'target')
|
|
|
all_roi_heatmap = []
|
|
|
for roi, heatmap in zip(rois, gs):
|
|
|
- show_heatmap(heatmap, 'target')
|
|
|
+ # show_heatmap(heatmap, 'target')
|
|
|
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[0],'roi_heatmap')
|
|
|
+ # show_heatmap(roi_heatmap[0],'roi_heatmap')
|
|
|
all_roi_heatmap.append(roi_heatmap)
|
|
|
|
|
|
all_roi_heatmap = torch.cat(all_roi_heatmap)
|
|
|
@@ -742,10 +742,110 @@ def line_inference(x, line_boxes):
|
|
|
|
|
|
boxes_per_image = [box.size(0) for box in line_boxes]
|
|
|
x2 = x.split(boxes_per_image, dim=0)
|
|
|
-
|
|
|
+ # x2:tuple 2 x2[0]:[1,3,1024,1024]
|
|
|
+ # line_box: list:2 [1,4] [1.4] fasterrcnn kuang
|
|
|
for xx, bb in zip(x2, line_boxes):
|
|
|
line_prob, line_scores, = heatmaps_to_lines(xx, bb)
|
|
|
lines_probs.append(line_prob)
|
|
|
lines_scores.append(line_scores)
|
|
|
|
|
|
return lines_probs, lines_scores
|
|
|
+
|
|
|
+def arc_inference(x, point_boxes):
|
|
|
+ # type: (Tensor, List[Tensor]) -> Tuple[List[Tensor], List[Tensor]]
|
|
|
+
|
|
|
+ points_probs = []
|
|
|
+ points_scores = []
|
|
|
+
|
|
|
+ boxes_per_image = [box.size(0) for box in point_boxes]
|
|
|
+ x2 = x.split(boxes_per_image, dim=0)
|
|
|
+
|
|
|
+ for xx, bb in zip(x2, point_boxes):
|
|
|
+ point_prob,point_scores = heatmaps_to_arc(xx, bb)
|
|
|
+
|
|
|
+ points_probs.append(point_prob.unsqueeze(1))
|
|
|
+ points_scores.append(point_scores)
|
|
|
+
|
|
|
+ return points_probs,points_scores
|
|
|
+
|
|
|
+import torch.nn.functional as F
|
|
|
+
|
|
|
+import torch.nn.functional as F
|
|
|
+
|
|
|
+def heatmaps_to_arc(maps, rois, threshold=0.1, output_size=(128, 128)):
|
|
|
+ """
|
|
|
+ Args:
|
|
|
+ maps: [N, 3, H, W] - full heatmaps
|
|
|
+ rois: [N, 4] - bounding boxes
|
|
|
+ threshold: float - binarization threshold
|
|
|
+ output_size: resized size for uniform NMS
|
|
|
+
|
|
|
+ Returns:
|
|
|
+ masks: [N, 1, H, W] - binary mask aligned with input map
|
|
|
+ scores: [N, 1] - count of non-zero pixels in each mask
|
|
|
+ """
|
|
|
+ N, _, H, W = maps.shape
|
|
|
+ masks = torch.zeros((N, 1, H, W), dtype=torch.float32, device=maps.device)
|
|
|
+ scores = torch.zeros((N, 1), dtype=torch.float32, device=maps.device)
|
|
|
+
|
|
|
+ point_maps = maps[:, 0] # È¡µÚÒ»¸öͨµÀ [N, H, W]
|
|
|
+
|
|
|
+ print(f"==> heatmaps_to_arc: maps.shape = {maps.shape}, rois.shape = {rois.shape}")
|
|
|
+
|
|
|
+ for i in range(N):
|
|
|
+ x1, y1, x2, y2 = rois[i].long()
|
|
|
+ x1 = x1.clamp(0, W - 1)
|
|
|
+ x2 = x2.clamp(0, W - 1)
|
|
|
+ y1 = y1.clamp(0, H - 1)
|
|
|
+ y2 = y2.clamp(0, H - 1)
|
|
|
+
|
|
|
+ print(f"[{i}] roi: ({x1.item()}, {y1.item()}, {x2.item()}, {y2.item()})")
|
|
|
+
|
|
|
+ if x2 <= x1 or y2 <= y1:
|
|
|
+ print(f" Skipped invalid ROI at index {i}")
|
|
|
+ continue
|
|
|
+
|
|
|
+ roi_map = point_maps[i, y1:y2, x1:x2] # [h, w]
|
|
|
+ print(f" roi_map.shape: {roi_map.shape}")
|
|
|
+
|
|
|
+ if roi_map.numel() == 0:
|
|
|
+ print(f" Skipped empty ROI at index {i}")
|
|
|
+ continue
|
|
|
+
|
|
|
+ # resize to uniform size
|
|
|
+ roi_map_resized = F.interpolate(
|
|
|
+ roi_map.unsqueeze(0).unsqueeze(0),
|
|
|
+ size=output_size,
|
|
|
+ mode='bilinear',
|
|
|
+ align_corners=False
|
|
|
+ ) # [1, 1, H, W]
|
|
|
+ print(f" roi_map_resized.shape: {roi_map_resized.shape}")
|
|
|
+
|
|
|
+ # NMS + threshold
|
|
|
+ nms_roi = non_maximum_suppression(roi_map_resized) # shape: [1, H, W]
|
|
|
+ bin_mask = (nms_roi > threshold).float() # shape: [1, H, W]
|
|
|
+ print(f" bin_mask.sum(): {bin_mask.sum().item()}")
|
|
|
+
|
|
|
+ # resize back to original roi size
|
|
|
+ h = int((y2 - y1).item())
|
|
|
+ w = int((x2 - x1).item())
|
|
|
+ # È·±£ bin_mask ÊÇ [1, 128, 128]
|
|
|
+ assert bin_mask.dim() == 4, f"Expected 3D tensor [1, H, W], got {bin_mask.shape}"
|
|
|
+
|
|
|
+ # ÉϲÉÑù»Ø ROI Ôʼ´óС
|
|
|
+ bin_mask_original_size = F.interpolate(
|
|
|
+ # bin_mask.unsqueeze(0), # ? [1, 1, 128, 128]
|
|
|
+ bin_mask, # ? [1, 1, 128, 128]
|
|
|
+ size=(h, w),
|
|
|
+ mode='bilinear',
|
|
|
+ align_corners=False
|
|
|
+ )[0] # ? [1, h, w]
|
|
|
+
|
|
|
+ masks[i, 0, y1:y2, x1:x2] = bin_mask_original_size.squeeze()
|
|
|
+ scores[i] = bin_mask_original_size.sum()
|
|
|
+
|
|
|
+ print(f" bin_mask_original_size.shape: {bin_mask_original_size.shape}, sum: {scores[i].item()}")
|
|
|
+
|
|
|
+ print(f"==> Done. Total valid masks: {(scores > 0).sum().item()} / {N}")
|
|
|
+
|
|
|
+ return masks, scores
|
zyhhsss