arc_atart and arc_end

models/line_detect/trainer.py

@@ -8,6 +8,7 @@ import torch
 import torchvision
 from PIL.ImageDraw import ImageDraw
 from matplotlib import pyplot as plt
+from mpmath import polar
 from scipy.ndimage import gaussian_filter
 from torch.optim.lr_scheduler import ReduceLROnPlateau
 from torch.utils.tensorboard import SummaryWriter
@@ -54,20 +55,7 @@ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 
 def draw_ellipses_on_image(image, masks_pred, threshold=0.5, color=(0, 255, 0), thickness=2):
-    """
-    在单张原始图像上绘制从 masks 拟合出的椭圆。
-    自动将 masks resize 到 image 的空间尺寸。
-
-    Args:
-        image: Tensor [3, H_img, W_img] —— 原始图像（如 [3, 2000, 2000]）
-        masks_pred: Tensor [N, 1, H_mask, W_mask] or [N, H_mask, W_mask] —— 模型输出 mask（如 [2, 1, 672, 672]）
-        threshold: 二值化阈值
-        color: BGR color for OpenCV
-        thickness: ellipse line thickness
 
-    Returns:
-        drawn_image: numpy array [H_img, W_img, 3] in RGB
-    """
     # Step 1: 标准化 masks_pred to [N, H, W]
     if masks_pred.ndim == 4:
         if masks_pred.shape[1] == 1:
@@ -97,6 +85,8 @@ def draw_ellipses_on_image(image, masks_pred, threshold=0.5, color=(0, 255, 0),
     img_rgb = np.transpose(img_np, (1, 2, 0))  # [H, W, 3]
     img_out = img_rgb.copy()
 
+    ellipses_info_list = []
+
     # Step 4: Process each mask
     for mask in masks_resized:
         mask_cpu = mask.detach().cpu()
@@ -109,13 +99,36 @@ def draw_ellipses_on_image(image, masks_pred, threshold=0.5, color=(0, 255, 0),
             if len(largest_contour) >= 5:
                 try:
                     ellipse = cv2.fitEllipse(largest_contour)
+                    (center_x, center_y),( a,b),angle = ellipse
+
+                    contour_points=largest_contour.reshape(-1,2)
+                    points_centered=contour_points-np.array([center_x, center_y])
+
+                    angle_red=np.deg2rad(angle)
+                    rot_matrix=np.array([[np.cos(angle_red), -np.sin(angle_red)],[-np.sin(angle_red), np.cos(angle_red)]])
+                    points_rotated=np.dot(points_centered, rot_matrix.T)
+
+                    polar_angles=np.arctan2(points_rotated[:,1],points_rotated[:,1])
+
+                    min_angle_idx=np.argmin(polar_angles)
+                    max_angle_idx=np.argmax(polar_angles)
+
+                    arc_start=tuple(contour_points[min_angle_idx].astype(int))
+                    arc_end=tuple(contour_points[max_angle_idx].astype(int))
+
                     img_bgr = cv2.cvtColor(img_out, cv2.COLOR_RGB2BGR)
                     cv2.ellipse(img_bgr, ellipse, color=color, thickness=thickness)
+
+                    cv2.circle(img_bgr, arc_start, 2, color=(0,0,255), thickness=-1)
+                    cv2.circle(img_bgr, arc_end, 2, color=(255,0,0), thickness=-1)
+
                     img_out = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)
+
+                    ellipses_info_list.append({"ellipse":ellipse,"arc_start":arc_start,"arc_end":arc_end})
                 except cv2.error as e:
                     print(f"Warning: Failed to fit ellipse: {e}")
 
-    return img_out
+    return img_out, ellipses_info_list
 
 
 def fit_circle(points):
@@ -469,7 +482,7 @@ class Trainer(BaseTrainer):
             # keypoint_img = draw_keypoints((img * 255).to(torch.uint8), points, colors='red', width=3)
             self.writer.add_image('z-ins-masks', sum_mask.squeeze(0), global_step=epoch)
 
-            result_imgs = draw_ellipses_on_image(img, ins_masks, threshold=0.5)
+            result_imgs, _ = draw_ellipses_on_image(img, ins_masks, threshold=0.5)
             self.writer.add_image('z-out-ellipses', result_imgs, dataformats='HWC', global_step=epoch)
 
             features = self.apply_gaussian_blur_to_tensor(features, sigma=3)