85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
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": []}
|