debug mask heatmap

models/line_detect/heads/head_losses.py

@@ -784,6 +784,7 @@ def lines_point_pair_loss(line_logits, proposals, gt_lines, line_matched_idxs):
 def compute_arc_loss(feature_logits, proposals, gt_, pos_matched_idxs):
     print(f'compute_arc_loss:{feature_logits.shape}')
     N, K, H, W = feature_logits.shape
+
     len_proposals = len(proposals)
 
     empty_count = 0
@@ -808,6 +809,7 @@ def compute_arc_loss(feature_logits, proposals, gt_, pos_matched_idxs):
 
     gs_heatmaps = []
     # print(f'point_matched_idxs:{point_matched_idxs}')
+    print(f'gt_masks:{gt_[0].shape}')
     for proposals_per_image, gt_kp_in_image, midx in zip(proposals, gt_, pos_matched_idxs):
         # [
         #   (Tensor(38, 4), Tensor(1, 57, 2), Tensor(38, 1)),
@@ -815,8 +817,12 @@ def compute_arc_loss(feature_logits, proposals, gt_, pos_matched_idxs):
         # ]
         print(f'proposals_per_image:{proposals_per_image.shape}')
         kp = gt_kp_in_image[midx]
-        # print(f'gt_kp_in_image:{gt_kp_in_image}')
+        t_h, t_w = kp.shape[-2:]
+        print(f't_h:{t_h}, t_w:{t_w}')
+
+        print(f'gt_kp_in_image:{gt_kp_in_image.shape}')
         if proposals_per_image.shape[0] > 0 and gt_kp_in_image.shape[0] > 0:
+
             gs_heatmaps_per_img = arc_points_to_heatmap(kp, proposals_per_image, discretization_size)
             gs_heatmaps.append(gs_heatmaps_per_img)
 
@@ -839,23 +845,31 @@ def compute_arc_loss(feature_logits, proposals, gt_, pos_matched_idxs):
 
     return line_loss
 
+
 def arc_points_to_heatmap(keypoints, rois, heatmap_size):
     print(f'rois:{rois.shape}')
     print(f'heatmap_size:{heatmap_size}')
 
     print(f'keypoints.shape:{keypoints.shape}')
     # batch_size, num_keypoints, _ = keypoints.shape
+    t_h, t_w = keypoints.shape[-2:]
+    scale=heatmap_size/t_w
+    print(f'scale:{scale}')
+    x = keypoints[..., 0]*scale
+    y = keypoints[..., 1]*scale
+
+    x = x.unsqueeze(1)
+    y = y.unsqueeze(1)
 
-    x = keypoints[..., 0].unsqueeze(1)
-    y = keypoints[..., 1].unsqueeze(1)
     num_points=x.shape[2]
     print(f'num_points:{num_points}')
-    gs = generate_mask_gaussian_heatmaps(x, y, num_points=num_points, heatmap_size=heatmap_size, sigma=2.0)
+    gs = generate_mask_gaussian_heatmaps(x, y, num_points=num_points, heatmap_size=heatmap_size, sigma=10)
+    print(f'gs max :{gs.max()},gs.shape:{gs.shape}')
     # show_heatmap(gs[0],'target')
     all_roi_heatmap = []
     for roi, heatmap in zip(rois, gs):
-        # show_heatmap(heatmap, 'target')
-        print(f'heatmap:{heatmap.shape}')
+        show_heatmap(heatmap, 'target')
+        print(f'heatmap.shape:{heatmap.shape}')
         heatmap = heatmap.unsqueeze(0)
         x1, y1, x2, y2 = map(int, roi)
         roi_heatmap = torch.zeros_like(heatmap)
@@ -1182,7 +1196,10 @@ def arc_inference(x, arc_boxes,th):
 
 import torch.nn.functional as F
 
-def heatmaps_to_arc(maps, rois, threshold=0.1, output_size=(128, 128)):
+
+
+
+def heatmaps_to_arc(maps, rois, threshold=0, output_size=(128, 128)):
     """
     Args:
         maps: [N, 3, H, W] - full heatmaps
@@ -1233,7 +1250,7 @@ def heatmaps_to_arc(maps, rois, threshold=0.1, output_size=(128, 128)):
 
         # NMS + threshold
         nms_roi = non_maximum_suppression(roi_map_resized)  # shape: [1, H, W]
-        bin_mask = (nms_roi > threshold).float()  # 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
@@ -1254,6 +1271,13 @@ def heatmaps_to_arc(maps, rois, threshold=0.1, output_size=(128, 128)):
         masks[i, 0, y1:y2, x1:x2] = bin_mask_original_size.squeeze()
         scores[i] = bin_mask_original_size.sum()
 
+
+        # plt.figure(figsize=(6, 6))
+        # plt.imshow(masks[i, 0].cpu().numpy(), cmap='gray')
+        # plt.title(f"Mask {i}, score={scores[i].item():.1f}")
+        # plt.axis('off')
+        # plt.show()
+
         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}")

models/line_detect/line_dataset.py

@@ -1,5 +1,6 @@
 import imageio
 import numpy as np
+from skimage.draw import ellipse
 from torch.utils.data.dataset import T_co
 
 from libs.vision_libs.utils import draw_keypoints
@@ -113,6 +114,8 @@ class LineDataset(BaseDataset):
 
         if circle_4points is not None:
             target['circles']=circle_4points
+            circle_masks=generate_ellipse_mask(shape,points_to_ellipse(circle_4points))
+            target['circle_masks'] =torch.tensor(circle_masks,dtype=torch.float32).unsqueeze(0)
 
         target["boxes"]=boxes
         target["labels"]=labels
@@ -148,9 +151,12 @@ class LineDataset(BaseDataset):
             boxed_image = draw_bounding_boxes((img * 255).to(torch.uint8), target["boxes"],
                                                   colors="yellow", width=1)
             circle=target['circles']
+            circle_mask=target['circle_masks']
             print(f'taget circle:{circle.shape}')
+            print(f'target circle_masks:{circle_mask.shape}')
+            plt.imshow(circle_mask.squeeze(0))
             keypoint_img=draw_keypoints(boxed_image,circle,colors='red',width=3)
-            plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
+            # plt.imshow(keypoint_img.permute(1, 2, 0).numpy())
             plt.show()
 
         # if show_type=='lines':
@@ -177,6 +183,62 @@ class LineDataset(BaseDataset):
         pass
 
 
+def points_to_ellipse(points):
+    """
+    根据提供的四个点估计椭圆参数。
+
+    :param points: Tensor of shape (4, 2) 表示椭圆上的四个点
+    :return: 返回 (cx, cy, r1, r2, orientation) 其中 cx, cy 是中心坐标，r1, r2 分别是长轴和短轴半径，orientation 是椭圆的方向（弧度）
+    """
+    # 转换为numpy数组进行计算
+    pts = points.numpy()
+    pts = pts.reshape(-1, 2)
+
+    # 计算中心点
+    center = np.mean(pts, axis=0)
+
+    # 使用最小二乘法拟合椭圆
+    A = np.hstack(
+        [pts[:, 0:1] ** 2, pts[:, 0:1] * pts[:, 1:2], pts[:, 1:2] ** 2, pts[:, :2], np.ones((pts.shape[0], 1))])
+    b = np.ones(pts.shape[0])
+    x = np.linalg.lstsq(A, b, rcond=None)[0]
+
+    # 解析解参见 https://en.wikipedia.org/wiki/Ellipse#General_ellipse
+    a, b, c, d, f, g = x.ravel()
+    numerator = 2 * (a * f * f + c * d * d + g * b * b - 2 * b * d * f - a * c * g)
+    denominator1 = (b * b - a * c) * ((c - a) * np.sqrt(1 + 4 * b * b / ((a - c) * (a - c))) - (c + a))
+    denominator2 = (b * b - a * c) * ((a - c) * np.sqrt(1 + 4 * b * b / ((a - c) * (a - c))) - (c + a))
+    major_axis = np.sqrt(numerator / denominator1)
+    minor_axis = np.sqrt(numerator / denominator2)
+
+    # 简化处理：直接用点间距离估算长短轴
+    distances = np.linalg.norm(pts - center, axis=1)
+    long_axis_length = np.max(distances) * 2
+    short_axis_length = np.min(distances) * 2
+
+    # 方向可以通过两点之间的连线斜率来近似计算
+    orientation = np.arctan2(pts[1, 1] - pts[0, 1], pts[1, 0] - pts[0, 0])
+
+    return center[0], center[1], long_axis_length / 2, short_axis_length / 2, orientation
+
+
+def generate_ellipse_mask(shape, ellipse_params):
+    """
+    在指定形状的图像上生成椭圆mask。
+
+    :param shape: 输出mask的形状 (HxW)
+    :param ellipse_params: 椭圆参数 (cx, cy, rx, ry, orientation)
+    :return: 椭圆mask
+    """
+    cx, cy, rx, ry, orientation = ellipse_params
+    img = np.zeros(shape, dtype=np.uint8)
+    cx, cy, rx, ry = int(cx), int(cy), int(rx), int(ry)
+
+    # 注意skimage的ellipse函数不直接支持旋转，所以这里简化处理
+    rr, cc = ellipse(cy, cx, ry, rx, shape)
+    img[rr, cc] = 1
+
+    return img
 def sort_points_clockwise(points):
     points = np.array(points)
 
@@ -305,4 +367,4 @@ def get_boxes_lines(objs,shape):
 if __name__ == '__main__':
     path=r'/data/share/zyh/master_dataset/circle/huyan_eclipse/a_dataset'
     dataset= LineDataset(dataset_path=path, dataset_type='train',augmentation=False, data_type='jpg')
-    dataset.show(33,show_type='all')
+    dataset.show(9,show_type='all')

models/line_detect/loi_heads.py

@@ -1365,7 +1365,7 @@ class RoIHeads(nn.Module):
 
 
 
-
+                print(f'features from backbone:{features['0'].shape}')
                 feature_logits = self.circle_forward1(features, image_shapes, circle_proposals)
 
                 loss_circle = None
@@ -1376,7 +1376,7 @@ class RoIHeads(nn.Module):
                     if targets is None or circle_pos_matched_idxs is None:
                         raise ValueError("both targets and pos_matched_idxs should not be None when in training mode")
 
-                    gt_circles = [t["circles"] for t in targets if "circles" in t]
+                    gt_circles = [t["circle_masks"] for t in targets if "circle_masks" in t]
 
                     print(f'gt_circle:{gt_circles[0].shape}')
                     h, w = targets[0]["img_size"]
@@ -1390,7 +1390,7 @@ class RoIHeads(nn.Module):
                     if gt_circles_tensor.shape[0] > 0:
                         print(f'start to compute circle_loss')
 
-                        loss_circle = compute_circle_loss(feature_logits, circle_proposals, gt_circles, circle_pos_matched_idxs)
+                        loss_circle = compute_arc_loss(feature_logits, circle_proposals, gt_circles, circle_pos_matched_idxs)
 
                         # loss_circle_extra=compute_circle_extra_losses(feature_logits, circle_proposals, gt_circles, circle_pos_matched_idxs)
 
@@ -1409,7 +1409,7 @@ class RoIHeads(nn.Module):
                     if targets is not None:
                         h, w = targets[0]["img_size"]
                         img_size = h
-                        gt_circles = [t["circles"] for t in targets if "circles" in t]
+                        gt_circles = [t["circle_masks"] for t in targets if "circle_masks" in t]
 
                         gt_circles_tensor = torch.zeros(0, 0)
                         if len(gt_circles) > 0:
@@ -1419,7 +1419,7 @@ class RoIHeads(nn.Module):
                         if gt_circles_tensor.shape[0] > 0:
                             print(f'start to compute circle_loss')
 
-                            loss_circle = compute_circle_loss(feature_logits, circle_proposals, gt_circles,
+                            loss_circle = compute_arc_loss(feature_logits, circle_proposals, gt_circles,
                                                             circle_pos_matched_idxs)
 
                             # loss_circle_extra = compute_circle_extra_losses(feature_logits, circle_proposals, gt_circles,circle_pos_matched_idxs)
@@ -1447,7 +1447,7 @@ class RoIHeads(nn.Module):
 
                         if feature_logits is not None:
 
-                            circles_probs, circles_scores = circle_inference(feature_logits, circle_proposals)
+                            circles_probs, circles_scores,circle_points = arc_inference(feature_logits, circle_proposals,th=0)
                             # print(f'circles_probs:{circles_probs.shape}, circles_scores:{circles_scores.shape}')
                             proposals_per_image = [box.size(0) for box in  circle_proposals]
                             print(f'circle_proposals_per_image:{proposals_per_image}')

models/line_detect/trainer.py

@@ -290,27 +290,29 @@ class Trainer(BaseTrainer):
             circle_image = img.cpu().numpy().transpose((1, 2, 0))  # CHW -> HWC
             circle_image = (circle_image * 255).clip(0, 255).astype(np.uint8)
 
-
-            if isinstance(points, torch.Tensor):
-                points = points.cpu().numpy()
-
-            for point_set in points:
-                center, radius = fit_circle(point_set)
-                cx, cy = map(int, center)
-
-                circle_image = cv2.circle(circle_image, (cx, cy), int(radius), (0, 0, 255), 2)
-
-                for point in point_set:
-                    px, py = map(int, point)
-                    circle_image = cv2.circle(circle_image, (px, py), 3, (0, 255, 255), -1)
-
-            img_rgb = cv2.cvtColor(circle_image, cv2.COLOR_BGR2RGB)
-            img_tensor = img_rgb.transpose((2, 0, 1))  # HWC -> CHW
-            img_tensor = img_tensor / 255.0  # 归一化到 [0, 1]
+            sum_mask = points.sum(dim=0, keepdim=True)
+            sum_mask = sum_mask / (sum_mask.max() + 1e-8)
+
+            # if isinstance(points, torch.Tensor):
+            #     points = points.cpu().numpy()
+            #
+            # for point_set in points:
+            #     center, radius = fit_circle(point_set)
+            #     cx, cy = map(int, center)
+            #
+            #     circle_image = cv2.circle(circle_image, (cx, cy), int(radius), (0, 0, 255), 2)
+            #
+            #     for point in point_set:
+            #         px, py = map(int, point)
+            #         circle_image = cv2.circle(circle_image, (px, py), 3, (0, 255, 255), -1)
+            #
+            # img_rgb = cv2.cvtColor(circle_image, cv2.COLOR_BGR2RGB)
+            # img_tensor = img_rgb.transpose((2, 0, 1))  # HWC -> CHW
+            # img_tensor = img_tensor / 255.0  # 归一化到 [0, 1]
 
 
             # keypoint_img = draw_keypoints((img * 255).to(torch.uint8), points, colors='red', width=3)
-            self.writer.add_image('z-out-circle', img_tensor, global_step=epoch)
+            self.writer.add_image('z-out-circle', sum_mask.squeeze(0), global_step=epoch)
             features=self.apply_gaussian_blur_to_tensor(features,sigma=3)
             self.writer.add_image('z-feature', features, global_step=epoch)