update

2026-05-14 13:30:06 +08:00
parent 974946cee4
commit e43171521d
91 changed files with 10485 additions and 3254 deletions
@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+from typing import Any, Dict
+
+import numpy as np
+from fastapi import APIRouter
+from PIL import Image
+
+from ..deps import OUTPUT_DIR
+from ..schemas import SamPromptSegmentRequest, SegmentRequest
+from ..services.image_io import b64_to_pil_image
+from ..services.predictors import get_seg_predictor
+from model.Seg.seg_loader import expand_mask, mask_to_contour_xy, run_sam_prompt
+
+router = APIRouter(tags=["segment"])
+
+
+@router.post("/segment")
+def segment(req: SegmentRequest) -> Dict[str, Any]:
+    try:
+        model_name = req.model_name or "sam"
+        pil = b64_to_pil_image(req.image_b64).convert("RGB")
+        rgb = np.array(pil, dtype=np.uint8)
+
+        predictor = get_seg_predictor(model_name)
+        label_map = predictor(rgb).astype(np.int32)
+
+        out_dir = OUTPUT_DIR / "segment"
+        out_dir.mkdir(parents=True, exist_ok=True)
+        out_path = out_dir / f"{model_name}_label.png"
+        Image.fromarray(np.clip(label_map, 0, 255).astype(np.uint8), mode="L").save(out_path)
+
+        return {"success": True, "label_path": str(out_path)}
+    except Exception as e:
+        return {"success": False, "error": str(e)}
+
+
+@router.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": []}
+
+        mask = expand_mask(mask, int(req.expand_px))
+        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": []}