Add SAM click-to-select with brush refinement in AI Edit tool
Backend: - sam_service.py: auto-downloads SAM ViT-B (~375 MB) on first use with progress tracking; loads model to CUDA/MPS/CPU; predict_points() takes multi-point prompts (include/exclude labels) and returns best mask - POST /api/segment/point: SAM point-prompt endpoint; returns mask PNG - GET /api/segment/install-status: poll download progress - POST /api/segment/install: explicit trigger (also auto on first click) - main.py: pre-download SAM on startup alongside NCNN Frontend (ai_edit.js): - Click mode (default): click object → SAM generates mask instantly Alt+click → subtract (deselect over-selected area) Multiple clicks accumulate for multi-object or refinement - Brush + / Brush − modes: paint to add or erase from SAM mask by hand - If SAM model is still downloading on first click: inline progress bar, user retries the click when done - Unified action bar: Erase | Replace (inline prompt) | Upscale | Expand | Clear - All modes share the same mask canvas; SAM and brush are fully composited https://claude.ai/code/session_01B58MaJCU1R6KwBDJCp8AfN
This commit is contained in:
Claude
@@ -21,6 +21,9 @@ async def lifespan(app: FastAPI):
|
||||
caps = probe_upscale_capabilities()
|
||||
if not caps["realesrgan_pytorch"] and not caps["realesrgan_ncnn"]:
|
||||
asyncio.create_task(ensure_ncnn_installed())
|
||||
# Pre-download SAM model in background so first click is fast
|
||||
from app.services.sam_service import ensure_sam_installed
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
yield
|
||||
|
||||
|
||||
|
||||
@@ -348,3 +348,63 @@ async def get_config():
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
# ─── SAM (Segment Anything) ──────────────────────────────────────────────────
|
||||
|
||||
class SegmentPointRequest(BaseModel):
|
||||
image: str # base64 PNG/JPEG
|
||||
points: list[list[int]] # [[x, y], ...] original image coords
|
||||
labels: list[int] # 1=include, 0=exclude — same length as points
|
||||
|
||||
|
||||
@router.post("/segment/point")
|
||||
async def segment_point(req: SegmentPointRequest):
|
||||
"""
|
||||
Run SAM point-prompt segmentation.
|
||||
Returns a binary mask PNG (white = selected area).
|
||||
Auto-downloads the SAM ViT-B model (~375 MB) on first call.
|
||||
"""
|
||||
if not req.points:
|
||||
raise HTTPException(status_code=400, detail="At least one point required.")
|
||||
if len(req.points) != len(req.labels):
|
||||
raise HTTPException(status_code=400, detail="points and labels must have the same length.")
|
||||
|
||||
try:
|
||||
image_bytes = base64.b64decode(req.image)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=400, detail=f"Could not decode image: {e}")
|
||||
|
||||
from app.services.sam_service import predict_points, get_install_status
|
||||
try:
|
||||
mask_bytes = await predict_points(
|
||||
image_bytes,
|
||||
[tuple(p) for p in req.points],
|
||||
req.labels,
|
||||
)
|
||||
return {
|
||||
"mask": base64.b64encode(mask_bytes).decode(),
|
||||
"sam_install": get_install_status(),
|
||||
}
|
||||
except RuntimeError as e:
|
||||
raise HTTPException(status_code=503, detail=str(e))
|
||||
except Exception as e:
|
||||
import traceback; traceback.print_exc()
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@router.get("/segment/install-status")
|
||||
def segment_install_status():
|
||||
"""Poll SAM model download progress."""
|
||||
from app.services.sam_service import get_install_status, sam_model_available
|
||||
status = get_install_status()
|
||||
status["model_ready"] = sam_model_available()
|
||||
return status
|
||||
|
||||
|
||||
@router.post("/segment/install")
|
||||
async def segment_install():
|
||||
"""Trigger SAM model download explicitly (also auto-triggered on first /segment/point call)."""
|
||||
from app.services.sam_service import ensure_sam_installed, get_install_status
|
||||
asyncio.create_task(ensure_sam_installed())
|
||||
return get_install_status()
|
||||
|
||||
@@ -0,0 +1,184 @@
|
||||
"""
|
||||
SAM (Segment Anything Model) service.
|
||||
|
||||
Auto-downloads the ViT-B checkpoint (~375 MB) on first use.
|
||||
Caches the loaded model in memory; re-uses predictor across calls.
|
||||
|
||||
Prediction API:
|
||||
predict_points(image_bytes, points, labels) -> mask_bytes (PNG, white=selected)
|
||||
points: list of (x, y) in original image pixels
|
||||
labels: list of 1 (include) or 0 (exclude), same length as points
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import io
|
||||
import os
|
||||
import urllib.request
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
# ── Model download ────────────────────────────────────────────────────────────
|
||||
|
||||
SAM_DIR = Path("/app/data/models/sam")
|
||||
SAM_FILENAME = "sam_vit_b_01ec64.pth"
|
||||
SAM_URL = f"https://dl.fbaipublicfiles.com/segment_anything/{SAM_FILENAME}"
|
||||
SAM_PATH = SAM_DIR / SAM_FILENAME
|
||||
|
||||
|
||||
class SamInstallState(str, Enum):
|
||||
idle = "idle"
|
||||
downloading = "downloading"
|
||||
done = "done"
|
||||
failed = "failed"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SamInstallStatus:
|
||||
state: SamInstallState = SamInstallState.idle
|
||||
progress: int = 0
|
||||
message: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
_install_status = SamInstallStatus()
|
||||
_install_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def get_install_status() -> dict:
|
||||
s = _install_status
|
||||
return {"state": s.state.value, "progress": s.progress,
|
||||
"message": s.message, "error": s.error}
|
||||
|
||||
|
||||
def sam_model_available() -> bool:
|
||||
return SAM_PATH.exists() and SAM_PATH.stat().st_size > 100_000_000
|
||||
|
||||
|
||||
async def ensure_sam_installed() -> bool:
|
||||
"""Download SAM ViT-B checkpoint if not present. Returns True on success."""
|
||||
global _install_status
|
||||
|
||||
if sam_model_available():
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
async with _install_lock:
|
||||
if sam_model_available():
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
if _install_status.state == SamInstallState.downloading:
|
||||
return False
|
||||
|
||||
try:
|
||||
SAM_DIR.mkdir(parents=True, exist_ok=True)
|
||||
_install_status = SamInstallStatus(
|
||||
state=SamInstallState.downloading, progress=0,
|
||||
message="Downloading SAM ViT-B model (~375 MB)…",
|
||||
)
|
||||
|
||||
def _download():
|
||||
def _progress(count, block, total):
|
||||
if total > 0:
|
||||
_install_status.progress = min(99, int(count * block * 99 / total))
|
||||
tmp = SAM_PATH.with_suffix(".tmp")
|
||||
urllib.request.urlretrieve(SAM_URL, tmp, _progress)
|
||||
tmp.rename(SAM_PATH)
|
||||
|
||||
loop = asyncio.get_event_loop()
|
||||
await loop.run_in_executor(None, _download)
|
||||
|
||||
_install_status = SamInstallStatus(state=SamInstallState.done, progress=100,
|
||||
message="SAM model ready.")
|
||||
return True
|
||||
|
||||
except Exception as exc:
|
||||
_install_status = SamInstallStatus(
|
||||
state=SamInstallState.failed, error=str(exc),
|
||||
message="SAM download failed.",
|
||||
)
|
||||
print(f"[sam] Download failed: {exc}")
|
||||
return False
|
||||
|
||||
|
||||
# ── Model cache ───────────────────────────────────────────────────────────────
|
||||
|
||||
_predictor = None
|
||||
_predictor_lock = asyncio.Lock()
|
||||
|
||||
|
||||
def _load_predictor():
|
||||
"""Load SAM model and return a SamPredictor. Called in thread pool."""
|
||||
global _predictor
|
||||
if _predictor is not None:
|
||||
return _predictor
|
||||
|
||||
import torch
|
||||
from segment_anything import sam_model_registry, SamPredictor
|
||||
|
||||
if torch.cuda.is_available():
|
||||
device = "cuda"
|
||||
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
|
||||
device = "mps"
|
||||
else:
|
||||
device = "cpu"
|
||||
|
||||
print(f"[sam] Loading SAM ViT-B on {device}…")
|
||||
sam = sam_model_registry["vit_b"](checkpoint=str(SAM_PATH))
|
||||
sam.to(device)
|
||||
_predictor = SamPredictor(sam)
|
||||
print("[sam] Model loaded.")
|
||||
return _predictor
|
||||
|
||||
|
||||
# ── Prediction ────────────────────────────────────────────────────────────────
|
||||
|
||||
def _predict_sync(image_bytes: bytes,
|
||||
points: list[tuple[int, int]],
|
||||
labels: list[int]) -> bytes:
|
||||
"""
|
||||
Run SAM prediction synchronously (call via run_in_executor).
|
||||
Returns PNG bytes: white = selected, black = background.
|
||||
"""
|
||||
predictor = _load_predictor()
|
||||
|
||||
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
|
||||
img_array = np.array(image)
|
||||
|
||||
predictor.set_image(img_array)
|
||||
|
||||
pt_array = np.array(points, dtype=np.float32) # [[x, y], ...]
|
||||
lbl_array = np.array(labels, dtype=np.int32) # [1=fg, 0=bg, ...]
|
||||
|
||||
masks, scores, _ = predictor.predict(
|
||||
point_coords=pt_array,
|
||||
point_labels=lbl_array,
|
||||
multimask_output=True,
|
||||
)
|
||||
|
||||
# Pick the highest-confidence mask
|
||||
best = masks[int(np.argmax(scores))] # bool array H×W
|
||||
|
||||
mask_img = Image.fromarray((best * 255).astype(np.uint8), mode="L")
|
||||
buf = io.BytesIO()
|
||||
mask_img.save(buf, format="PNG")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
async def predict_points(image_bytes: bytes,
|
||||
points: list[tuple[int, int]],
|
||||
labels: list[int]) -> bytes:
|
||||
"""Async wrapper for SAM point prediction."""
|
||||
if not sam_model_available():
|
||||
ok = await ensure_sam_installed()
|
||||
if not ok:
|
||||
raise RuntimeError("SAM model not available.")
|
||||
loop = asyncio.get_event_loop()
|
||||
return await loop.run_in_executor(None, _predict_sync, image_bytes, points, labels)
|
||||
Reference in New Issue
Block a user