添加模型分割

python_server/server.py

@@ -18,7 +18,13 @@ from PIL import Image, ImageDraw
 from config_loader import load_app_config, get_depth_backend_from_app
 from model.Depth.depth_loader import UnifiedDepthConfig, DepthBackend, build_depth_predictor
 
-from model.Seg.seg_loader import UnifiedSegConfig, SegBackend, build_seg_predictor
+from model.Seg.seg_loader import (
+    UnifiedSegConfig,
+    SegBackend,
+    build_seg_predictor,
+    mask_to_contour_xy,
+    run_sam_prompt,
+)
 from model.Inpaint.inpaint_loader import UnifiedInpaintConfig, InpaintBackend, build_inpaint_predictor
 from model.Animation.animation_loader import (
     UnifiedAnimationConfig,
@@ -47,6 +53,28 @@ class SegmentRequest(ImageInput):
     pass
 
 
+class SamPromptSegmentRequest(BaseModel):
+    image_b64: str = Field(..., description="裁剪后的 RGB 图 base64（PNG/JPG）")
+    overlay_b64: Optional[str] = Field(
+        None,
+        description="与裁剪同尺寸的标记叠加 PNG base64（可选；当前用于校验尺寸一致）",
+    )
+    point_coords: list[list[float]] = Field(
+        ...,
+        description="裁剪坐标系下的提示点 [[x,y], ...]",
+    )
+    point_labels: list[int] = Field(
+        ...,
+        description="与 point_coords 等长：1=前景，0=背景",
+    )
+    box_xyxy: list[float] = Field(
+        ...,
+        description="裁剪内笔画紧包围盒 [x1,y1,x2,y2]（像素）",
+        min_length=4,
+        max_length=4,
+    )
+
+
 class InpaintRequest(ImageInput):
     prompt: Optional[str] = Field("", description="补全 prompt")
     strength: float = Field(0.8, ge=0.0, le=1.0)
@@ -287,6 +315,54 @@ def segment(req: SegmentRequest) -> Dict[str, Any]:
         return {"success": False, "error": str(e)}
 
 
+@app.post("/segment/sam_prompt")
+def segment_sam_prompt(req: SamPromptSegmentRequest) -> Dict[str, Any]:
+    """
+    交互式 SAM：裁剪图 + 点/框提示，返回掩膜外轮廓点列（裁剪像素坐标）。
+    """
+    try:
+        pil = _b64_to_pil_image(req.image_b64).convert("RGB")
+        rgb = np.array(pil, dtype=np.uint8)
+        h, w = rgb.shape[0], rgb.shape[1]
+
+        if req.overlay_b64:
+            ov = _b64_to_pil_image(req.overlay_b64)
+            if ov.size != (w, h):
+                return {
+                    "success": False,
+                    "error": f"overlay 尺寸 {ov.size} 与 image {w}x{h} 不一致",
+                    "contour": [],
+                }
+
+        if len(req.point_coords) != len(req.point_labels):
+            return {
+                "success": False,
+                "error": "point_coords 与 point_labels 长度不一致",
+                "contour": [],
+            }
+        if len(req.point_coords) < 1:
+            return {"success": False, "error": "至少需要一个提示点", "contour": []}
+
+        pc = np.array(req.point_coords, dtype=np.float32)
+        if pc.ndim != 2 or pc.shape[1] != 2:
+            return {"success": False, "error": "point_coords 每项须为 [x,y]", "contour": []}
+
+        pl = np.array(req.point_labels, dtype=np.int64)
+        box = np.array(req.box_xyxy, dtype=np.float32)
+
+        mask = run_sam_prompt(rgb, pc, pl, box_xyxy=box)
+        if not np.any(mask):
+            return {"success": False, "error": "SAM 未产生有效掩膜", "contour": []}
+
+        contour = mask_to_contour_xy(mask, epsilon_px=2.0)
+        if len(contour) < 3:
+            return {"success": False, "error": "轮廓点数不足", "contour": []}
+
+        return {"success": True, "contour": contour, "error": None}
+    except Exception as e:
+        return {"success": False, "error": str(e), "contour": []}
+
+
 # -----------------------------
 # Inpaint
 # -----------------------------