first version
This commit is contained in:
381
files/sam3_server.py
Normal file
381
files/sam3_server.py
Normal file
@@ -0,0 +1,381 @@
|
||||
"""
|
||||
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()
|
||||
Reference in New Issue
Block a user