382 lines
14 KiB
Python
382 lines
14 KiB
Python
"""
|
|
SAM 3 local inference server for annotation tools.
|
|
|
|
Exposes a small FastAPI service that takes an image + bbox and returns a
|
|
segmentation mask. Designed for the classic annotation workflow:
|
|
user draws a box around ONE object -> server returns the mask for that object.
|
|
|
|
Two endpoints are provided:
|
|
|
|
POST /segment/box
|
|
Uses Sam3Tracker (the SAM2-style Promptable Visual Segmentation head).
|
|
Given ONE bbox, returns ONE mask for the object inside that box.
|
|
This is what you want for 99% of annotation-tool use cases.
|
|
|
|
POST /segment/exemplar
|
|
Uses the full SAM 3 image predictor with a box as a visual "exemplar".
|
|
Given one bbox, returns masks for ALL similar objects in the image.
|
|
Useful when you want to label one instance and auto-propagate.
|
|
|
|
Run:
|
|
python sam3_server.py --host 0.0.0.0 --port 8000
|
|
|
|
Request format (JSON):
|
|
{
|
|
"image": "<base64-encoded image>" # OR
|
|
"image_url": "http://..." # OR
|
|
"image_path": "/absolute/path.jpg",
|
|
"bbox": [x1, y1, x2, y2], # absolute pixel coords (xyxy)
|
|
"multimask_output": false # optional, /segment/box only
|
|
}
|
|
|
|
Response:
|
|
{
|
|
"mask": "<base64 PNG, single-channel 0/255>",
|
|
"rle": {"size":[H,W], "counts":"..."}, # COCO RLE, handy for COCO-format tools
|
|
"polygon": [[x,y], [x,y], ...], # largest contour, optional
|
|
"score": 0.97,
|
|
"bbox_out": [x1,y1,x2,y2]
|
|
}
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import io
|
|
import logging
|
|
from typing import List, Optional
|
|
|
|
import numpy as np
|
|
import torch
|
|
import uvicorn
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from PIL import Image
|
|
from pydantic import BaseModel, Field
|
|
|
|
# ----- SAM 3 imports ---------------------------------------------------------
|
|
# The SAM 3 image predictor (concept / exemplar segmentation):
|
|
from sam3.model_builder import build_sam3_image_model
|
|
from sam3.model.sam3_image_processor import Sam3Processor
|
|
|
|
# The SAM 3 tracker (SAM2-style single-instance PVS on images):
|
|
# This is exposed in the HF `transformers` integration as Sam3TrackerModel,
|
|
# and in the reference repo via the image predictor for single-box prompts.
|
|
# We use HuggingFace transformers for the tracker because its API is clean
|
|
# and documented for bbox -> single mask. Fall back to the repo's predictor
|
|
# if transformers isn't available.
|
|
try:
|
|
from transformers import Sam3TrackerModel, Sam3TrackerProcessor # type: ignore
|
|
HAS_TRACKER = True
|
|
except Exception:
|
|
HAS_TRACKER = False
|
|
|
|
log = logging.getLogger("sam3-server")
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
|
|
|
|
|
|
# ============================================================================
|
|
# Request / response schemas
|
|
# ============================================================================
|
|
|
|
class SegmentRequest(BaseModel):
|
|
# Exactly one of these three must be provided.
|
|
image: Optional[str] = Field(None, description="Base64-encoded image bytes.")
|
|
image_url: Optional[str] = None
|
|
image_path: Optional[str] = None
|
|
|
|
bbox: List[float] = Field(..., description="Box in absolute pixels, [x1,y1,x2,y2].")
|
|
multimask_output: bool = Field(
|
|
False,
|
|
description="If true (tracker only), return the 3 candidate masks instead of picking the highest-scoring one.",
|
|
)
|
|
return_polygon: bool = Field(True, description="Also return the largest contour as [[x,y], ...].")
|
|
return_rle: bool = Field(True, description="Also return COCO RLE encoding of the mask.")
|
|
|
|
|
|
class SegmentResponse(BaseModel):
|
|
mask: str # base64 PNG
|
|
score: float
|
|
bbox_out: List[float]
|
|
rle: Optional[dict] = None
|
|
polygon: Optional[List[List[float]]] = None
|
|
extra_masks: Optional[List[str]] = None # only populated when multimask_output=True
|
|
|
|
|
|
# ============================================================================
|
|
# Model holder (loaded once at startup)
|
|
# ============================================================================
|
|
|
|
class SAM3Models:
|
|
def __init__(self, device: str, checkpoint_path: Optional[str] = None):
|
|
self.device = device
|
|
|
|
log.info("Loading SAM 3 image model (concept/exemplar predictor)...")
|
|
if checkpoint_path:
|
|
self.image_model = build_sam3_image_model(checkpoint_path=checkpoint_path)
|
|
else:
|
|
self.image_model = build_sam3_image_model()
|
|
self.image_model.to(device).eval()
|
|
self.image_processor = Sam3Processor(self.image_model)
|
|
|
|
if HAS_TRACKER:
|
|
log.info("Loading SAM 3 tracker (single-instance PVS) from HuggingFace...")
|
|
# facebook/sam3 exposes the tracker weights under this name
|
|
self.tracker_processor = Sam3TrackerProcessor.from_pretrained("facebook/sam3")
|
|
self.tracker_model = Sam3TrackerModel.from_pretrained("facebook/sam3").to(device).eval()
|
|
else:
|
|
log.warning(
|
|
"transformers not available or missing Sam3Tracker classes. "
|
|
"The /segment/box endpoint will fall back to the image predictor "
|
|
"(still works, but treats the box as an exemplar)."
|
|
)
|
|
self.tracker_processor = None
|
|
self.tracker_model = None
|
|
|
|
|
|
MODELS: Optional[SAM3Models] = None
|
|
|
|
|
|
# ============================================================================
|
|
# Helpers
|
|
# ============================================================================
|
|
|
|
def _load_image(req: SegmentRequest) -> Image.Image:
|
|
if req.image:
|
|
raw = base64.b64decode(req.image)
|
|
return Image.open(io.BytesIO(raw)).convert("RGB")
|
|
if req.image_path:
|
|
return Image.open(req.image_path).convert("RGB")
|
|
if req.image_url:
|
|
import requests
|
|
r = requests.get(req.image_url, timeout=20)
|
|
r.raise_for_status()
|
|
return Image.open(io.BytesIO(r.content)).convert("RGB")
|
|
raise HTTPException(400, "Provide one of `image` (base64), `image_path`, or `image_url`.")
|
|
|
|
|
|
def _mask_to_png_b64(mask: np.ndarray) -> str:
|
|
"""mask: HxW bool/uint8 -> base64-encoded PNG string."""
|
|
if mask.dtype != np.uint8:
|
|
mask = (mask.astype(np.uint8) * 255)
|
|
img = Image.fromarray(mask, mode="L")
|
|
buf = io.BytesIO()
|
|
img.save(buf, format="PNG", optimize=True)
|
|
return base64.b64encode(buf.getvalue()).decode("ascii")
|
|
|
|
|
|
def _mask_to_rle(mask: np.ndarray) -> dict:
|
|
"""COCO-style RLE. Uses pycocotools if available, else a pure-python fallback."""
|
|
mask = (mask > 0).astype(np.uint8)
|
|
try:
|
|
from pycocotools import mask as mutil # type: ignore
|
|
rle = mutil.encode(np.asfortranarray(mask))
|
|
rle["counts"] = rle["counts"].decode("ascii")
|
|
return rle
|
|
except Exception:
|
|
# Pure python fallback (uncompressed)
|
|
h, w = mask.shape
|
|
flat = mask.T.flatten() # column-major, like COCO
|
|
counts: List[int] = []
|
|
last = 0
|
|
run = 0
|
|
for v in flat:
|
|
if v == last:
|
|
run += 1
|
|
else:
|
|
counts.append(run)
|
|
run = 1
|
|
last = v
|
|
counts.append(run)
|
|
return {"size": [h, w], "counts": counts}
|
|
|
|
|
|
def _mask_to_polygon(mask: np.ndarray) -> Optional[List[List[float]]]:
|
|
"""Largest-contour polygon, approximated. Returns None if mask is empty."""
|
|
try:
|
|
import cv2 # type: ignore
|
|
except Exception:
|
|
return None
|
|
m = (mask > 0).astype(np.uint8) * 255
|
|
contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_TC89_L1)
|
|
if not contours:
|
|
return None
|
|
c = max(contours, key=cv2.contourArea)
|
|
# Simplify a bit — saves bytes and is friendlier for annotation tools
|
|
epsilon = 0.0015 * cv2.arcLength(c, True)
|
|
c = cv2.approxPolyDP(c, epsilon, True)
|
|
return [[float(pt[0][0]), float(pt[0][1])] for pt in c]
|
|
|
|
|
|
def _mask_bbox(mask: np.ndarray) -> List[float]:
|
|
ys, xs = np.where(mask > 0)
|
|
if len(xs) == 0:
|
|
return [0.0, 0.0, 0.0, 0.0]
|
|
return [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())]
|
|
|
|
|
|
def _build_response(mask: np.ndarray, score: float, req: SegmentRequest,
|
|
extra_masks: Optional[List[np.ndarray]] = None) -> SegmentResponse:
|
|
return SegmentResponse(
|
|
mask=_mask_to_png_b64(mask),
|
|
score=float(score),
|
|
bbox_out=_mask_bbox(mask),
|
|
rle=_mask_to_rle(mask) if req.return_rle else None,
|
|
polygon=_mask_to_polygon(mask) if req.return_polygon else None,
|
|
extra_masks=[_mask_to_png_b64(m) for m in extra_masks] if extra_masks else None,
|
|
)
|
|
|
|
|
|
# ============================================================================
|
|
# FastAPI app
|
|
# ============================================================================
|
|
|
|
app = FastAPI(title="SAM 3 local server", version="1.0")
|
|
|
|
# Allow local annotation tool UIs to call us from the browser
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
@app.get("/health")
|
|
def health():
|
|
return {
|
|
"ok": True,
|
|
"device": MODELS.device if MODELS else "unloaded",
|
|
"tracker_available": HAS_TRACKER,
|
|
}
|
|
|
|
|
|
@app.post("/segment/box", response_model=SegmentResponse)
|
|
@torch.inference_mode()
|
|
def segment_box(req: SegmentRequest):
|
|
"""
|
|
Annotation-tool workflow:
|
|
user draws a bbox around ONE object -> return ONE mask for that object.
|
|
|
|
Uses the SAM3 tracker (SAM2-style PVS) for a clean single-instance result.
|
|
"""
|
|
if MODELS is None:
|
|
raise HTTPException(500, "Models not loaded.")
|
|
image = _load_image(req)
|
|
x1, y1, x2, y2 = req.bbox
|
|
|
|
# ---- Preferred path: HuggingFace Sam3Tracker -------------------------
|
|
if HAS_TRACKER and MODELS.tracker_model is not None:
|
|
inputs = MODELS.tracker_processor(
|
|
images=image,
|
|
input_boxes=[[[x1, y1, x2, y2]]], # batch[1] x obj[1] x 4
|
|
return_tensors="pt",
|
|
).to(MODELS.device)
|
|
|
|
outputs = MODELS.tracker_model(**inputs, multimask_output=req.multimask_output)
|
|
|
|
# Post-process to image resolution, binary masks
|
|
masks = MODELS.tracker_processor.post_process_masks(
|
|
outputs.pred_masks.cpu(),
|
|
original_sizes=inputs["original_sizes"].cpu(),
|
|
reshaped_input_sizes=inputs["reshaped_input_sizes"].cpu(),
|
|
)[0] # shape: (num_masks, H, W) tensor
|
|
scores = outputs.iou_scores[0].cpu().numpy().ravel()
|
|
|
|
masks_np = masks.numpy().astype(bool)
|
|
if req.multimask_output and masks_np.shape[0] > 1:
|
|
best = int(np.argmax(scores))
|
|
extras = [m for i, m in enumerate(masks_np) if i != best]
|
|
return _build_response(masks_np[best], scores[best], req, extra_masks=extras)
|
|
|
|
best = int(np.argmax(scores)) if masks_np.shape[0] > 1 else 0
|
|
return _build_response(masks_np[best], float(scores[best]), req)
|
|
|
|
# ---- Fallback: use image predictor with box as exemplar, pick the
|
|
# detection whose predicted box has highest IoU with the user bbox.
|
|
state = MODELS.image_processor.set_image(image)
|
|
out = MODELS.image_processor.set_box_prompt(
|
|
state=state,
|
|
boxes=[[x1, y1, x2 - x1, y2 - y1]], # xywh as documented
|
|
box_labels=[1],
|
|
)
|
|
masks = out["masks"]
|
|
boxes = out["boxes"]
|
|
scores = out["scores"]
|
|
if len(masks) == 0:
|
|
raise HTTPException(404, "No segmentation produced for the given bbox.")
|
|
|
|
# Pick mask whose predicted bbox most overlaps user bbox
|
|
def _iou(a, b):
|
|
ax1, ay1, ax2, ay2 = a
|
|
bx1, by1, bx2, by2 = b
|
|
ix1, iy1 = max(ax1, bx1), max(ay1, by1)
|
|
ix2, iy2 = min(ax2, bx2), min(ay2, by2)
|
|
iw, ih = max(0, ix2 - ix1), max(0, iy2 - iy1)
|
|
inter = iw * ih
|
|
ua = (ax2 - ax1) * (ay2 - ay1) + (bx2 - bx1) * (by2 - by1) - inter
|
|
return inter / ua if ua > 0 else 0.0
|
|
|
|
user_box = [x1, y1, x2, y2]
|
|
ious = [_iou(user_box, list(map(float, b))) for b in boxes]
|
|
idx = int(np.argmax(ious))
|
|
mask_np = np.asarray(masks[idx]).astype(bool)
|
|
return _build_response(mask_np, float(scores[idx]), req)
|
|
|
|
|
|
@app.post("/segment/exemplar", response_model=SegmentResponse)
|
|
@torch.inference_mode()
|
|
def segment_exemplar(req: SegmentRequest):
|
|
"""
|
|
Auto-propagation workflow:
|
|
user draws a box around ONE instance -> return a merged mask of ALL
|
|
similar instances found in the image. `extra_masks` contains each
|
|
individual instance separately.
|
|
"""
|
|
if MODELS is None:
|
|
raise HTTPException(500, "Models not loaded.")
|
|
image = _load_image(req)
|
|
x1, y1, x2, y2 = req.bbox
|
|
|
|
state = MODELS.image_processor.set_image(image)
|
|
out = MODELS.image_processor.set_box_prompt(
|
|
state=state,
|
|
boxes=[[x1, y1, x2 - x1, y2 - y1]],
|
|
box_labels=[1],
|
|
)
|
|
masks = out["masks"]
|
|
scores = out["scores"]
|
|
if len(masks) == 0:
|
|
raise HTTPException(404, "No instances found.")
|
|
|
|
masks_np = [np.asarray(m).astype(bool) for m in masks]
|
|
merged = np.zeros_like(masks_np[0], dtype=bool)
|
|
for m in masks_np:
|
|
merged |= m
|
|
best_score = float(np.max(scores))
|
|
return _build_response(merged, best_score, req, extra_masks=masks_np)
|
|
|
|
|
|
# ============================================================================
|
|
# Entrypoint
|
|
# ============================================================================
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="0.0.0.0")
|
|
parser.add_argument("--port", type=int, default=8000)
|
|
parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
|
|
parser.add_argument("--checkpoint", default=None,
|
|
help="Optional path to sam3 checkpoint (otherwise uses HF default).")
|
|
args = parser.parse_args()
|
|
|
|
global MODELS
|
|
MODELS = SAM3Models(device=args.device, checkpoint_path=args.checkpoint)
|
|
log.info("SAM 3 server ready on %s:%d (device=%s)", args.host, args.port, args.device)
|
|
uvicorn.run(app, host=args.host, port=args.port, log_level="info")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|