Upload app.py

app.py

@@ -128,7 +128,7 @@ def reset_image(predictor, img):
       - 设置 predictor 的输入图像，
       - 返回原图
     """
-    predictor.set_image(img)
+    # predictor.set_image(img)
     original_img = img.copy()
     # 返回predictor，visible occlusion mask初始化, 原始图像
     return predictor, img, img, original_img
@@ -140,22 +140,25 @@ def button_clickable(selected_points):
         return gr.Button.update(interactive=False)
 
 
-def run_sam(predictor: SamPredictor, selected_points):
+def run_sam(image, predictor, selected_points):
     """
     调用 SAM 模型进行分割。
     """
-    # 确保图像为 RGB 模式
+    predictor.set_image(image)
     if len(selected_points) == 0:
         return [], None
     input_points = [p for p in selected_points]
     input_labels = [1 for _ in range(len(selected_points))]
+    print(input_points)
+    print(input_labels)
     masks, _, _ = predictor.predict(
         point_coords=np.array(input_points),
-        point_labels=input_labels,
+        point_labels=np.array(input_labels),
         multimask_output=False,  # 单对象输出
     )
+    print(masks.shape)
     visible_mask = 255 * np.squeeze(masks).astype(np.uint8)
-    return visible_mask, None
+    return visible_mask
 
 def apply_mask_overlay(image, mask):
     """
@@ -176,7 +179,7 @@ def segment_and_overlay(image, points, sam_predictor):
     """
     调用 run_sam 获得 mask，然后叠加显示分割结果。
     """
-    mask, _ = run_sam(sam_predictor, points)
+    mask, _ = run_sam(image, sam_predictor, points)
     overlaid = apply_mask_overlay(image, mask)
     return overlaid, mask