346 lines
11 KiB
Python
346 lines
11 KiB
Python
"""FastAPI server exposing SAM-family predictors over HTTP.
|
|
|
|
Implements the Label-Studio-ML `/predict` contract the SAM-Tool Tauri app
|
|
already speaks, and an LS-orthogonal `/models` endpoint so the frontend can
|
|
offer a model picker.
|
|
|
|
Run:
|
|
python server.py --host 0.0.0.0 --port 9090 --device cuda
|
|
|
|
Endpoints:
|
|
GET / → liveness + backend info
|
|
GET /models → registered models + availability
|
|
POST /predict → LS-ML protocol, accepts optional `?model=<id>` query
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
import io
|
|
import logging
|
|
import os
|
|
import sys
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import uvicorn
|
|
from fastapi import FastAPI, HTTPException, Query
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from PIL import Image
|
|
from pydantic import BaseModel
|
|
|
|
from registry import Registry, load_registry
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Logging
|
|
# --------------------------------------------------------------------------
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format="%(asctime)s %(levelname)s %(name)s: %(message)s",
|
|
stream=sys.stdout,
|
|
)
|
|
log = logging.getLogger("sam2-backend.server")
|
|
|
|
# --------------------------------------------------------------------------
|
|
# App + registry. Registry is loaded inside a FastAPI lifespan so it works
|
|
# regardless of how uvicorn ends up importing this module — `python
|
|
# server.py` runs us as `__main__` and uvicorn then re-imports as `server`,
|
|
# so initialising at module top-level (or only in main()) leaves the live
|
|
# copy uninitialised.
|
|
# --------------------------------------------------------------------------
|
|
REGISTRY: Registry | None = None
|
|
|
|
|
|
@asynccontextmanager
|
|
async def lifespan(_app: FastAPI):
|
|
global REGISTRY
|
|
if REGISTRY is None:
|
|
models_path = os.environ.get(
|
|
"SAM_BACKEND_MODELS",
|
|
str(Path(__file__).parent / "models.yaml"),
|
|
)
|
|
device = _pick_device(os.environ.get("SAM_BACKEND_DEVICE", "auto"))
|
|
log.info("lifespan: loading registry (models=%s device=%s)",
|
|
models_path, device)
|
|
REGISTRY = load_registry(Path(models_path), device=device)
|
|
log.info(
|
|
"lifespan: %d model specs (%d available), default=%s",
|
|
len(REGISTRY.specs),
|
|
sum(1 for s in REGISTRY.specs.values() if s.available),
|
|
REGISTRY.default_id,
|
|
)
|
|
yield
|
|
|
|
|
|
app = FastAPI(title="SAM Backend", version="0.1.0", lifespan=lifespan)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
def _registry() -> Registry:
|
|
if REGISTRY is None:
|
|
raise RuntimeError("registry not initialised")
|
|
return REGISTRY
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Schemas
|
|
# --------------------------------------------------------------------------
|
|
class ModelInfo(BaseModel):
|
|
id: str
|
|
name: str
|
|
family: str
|
|
available: bool
|
|
checkpoint: str
|
|
|
|
|
|
class ModelListResponse(BaseModel):
|
|
default: str | None
|
|
models: list[ModelInfo]
|
|
|
|
|
|
class HealthResponse(BaseModel):
|
|
status: str
|
|
device: str
|
|
models: int
|
|
default: str | None
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Routes
|
|
# --------------------------------------------------------------------------
|
|
@app.get("/", response_model=HealthResponse)
|
|
def health() -> HealthResponse:
|
|
reg = _registry()
|
|
return HealthResponse(
|
|
status="ok",
|
|
device=reg.device,
|
|
models=len(reg.specs),
|
|
default=reg.default_id,
|
|
)
|
|
|
|
|
|
@app.get("/models", response_model=ModelListResponse)
|
|
def list_models() -> ModelListResponse:
|
|
reg = _registry()
|
|
return ModelListResponse(
|
|
default=reg.default_id,
|
|
models=[
|
|
ModelInfo(
|
|
id=s.id,
|
|
name=s.name,
|
|
family=s.family,
|
|
available=s.available,
|
|
checkpoint=s.checkpoint_abs,
|
|
)
|
|
for s in reg.list_specs()
|
|
],
|
|
)
|
|
|
|
|
|
@app.post("/predict")
|
|
def predict(payload: dict, model: str | None = Query(default=None)) -> dict:
|
|
"""Label-Studio-ML protocol.
|
|
|
|
Body shape (simplified — we ignore fields we don't need):
|
|
{
|
|
"tasks": [{ "data": { "image": "data:image/jpeg;base64,..." } }],
|
|
"params": {
|
|
"context": {
|
|
"result": [{
|
|
"original_width": W, "original_height": H,
|
|
"value": { "x": %, "y": %, "width": %, "height": %, ... }
|
|
}]
|
|
}
|
|
}
|
|
}
|
|
|
|
Response shape (what sam-tool-tauri / LS Studio expects):
|
|
{
|
|
"results": [{
|
|
"model_version": "<id>",
|
|
"result": [{
|
|
"original_width": W, "original_height": H,
|
|
"value": {
|
|
"points": [[x%, y%], ...],
|
|
"closed": true,
|
|
"polygonlabels": ["Object"]
|
|
},
|
|
"from_name": "polygon", "to_name": "image",
|
|
"type": "polygonlabels"
|
|
}]
|
|
}]
|
|
}
|
|
"""
|
|
reg = _registry()
|
|
|
|
try:
|
|
tasks = payload.get("tasks") or []
|
|
if not tasks:
|
|
raise HTTPException(400, "tasks[] missing or empty")
|
|
data = tasks[0].get("data") or {}
|
|
image_field = data.get("image")
|
|
if not isinstance(image_field, str):
|
|
raise HTTPException(400, "tasks[0].data.image missing")
|
|
image_rgb = _decode_image(image_field)
|
|
|
|
ctx_result = (
|
|
payload.get("params", {})
|
|
.get("context", {})
|
|
.get("result", [])
|
|
)
|
|
if not ctx_result:
|
|
raise HTTPException(400, "params.context.result[] missing (bbox prompt)")
|
|
prompt = ctx_result[0]
|
|
value = prompt.get("value", {})
|
|
orig_w = int(prompt.get("original_width") or image_rgb.shape[1])
|
|
orig_h = int(prompt.get("original_height") or image_rgb.shape[0])
|
|
x_pct = float(value.get("x", 0.0))
|
|
y_pct = float(value.get("y", 0.0))
|
|
w_pct = float(value.get("width", 0.0))
|
|
h_pct = float(value.get("height", 0.0))
|
|
|
|
# Percentages → absolute pixel xyxy, clamped to image bounds.
|
|
x1 = max(0, int(round(x_pct / 100.0 * orig_w)))
|
|
y1 = max(0, int(round(y_pct / 100.0 * orig_h)))
|
|
x2 = min(orig_w, int(round((x_pct + w_pct) / 100.0 * orig_w)))
|
|
y2 = min(orig_h, int(round((y_pct + h_pct) / 100.0 * orig_h)))
|
|
if x2 <= x1 or y2 <= y1:
|
|
raise HTTPException(400, "degenerate bbox after clamping")
|
|
|
|
# Model selection: query param wins, then payload.params.model, then default.
|
|
model_id = (
|
|
model
|
|
or payload.get("params", {}).get("model")
|
|
or payload.get("params", {}).get("model_id")
|
|
)
|
|
# Resolve up-front so an unknown id is a clean 400; any error from
|
|
# the actual predict path falls into the generic 500 with traceback
|
|
# below (and is no longer mis-reported as "unknown model").
|
|
try:
|
|
reg.resolve(model_id)
|
|
except KeyError as e:
|
|
raise HTTPException(400, f"unknown model: {e}")
|
|
|
|
spec, result = reg.predict_bbox(model_id, image_rgb, (x1, y1, x2, y2))
|
|
|
|
except HTTPException:
|
|
raise
|
|
except FileNotFoundError as e:
|
|
raise HTTPException(503, str(e))
|
|
except Exception as e:
|
|
log.exception("predict failed")
|
|
raise HTTPException(500, f"{type(e).__name__}: {e}")
|
|
|
|
if not result.polygon:
|
|
raise HTTPException(500, "no polygon returned from model")
|
|
|
|
# Convert pixel polygon back to percentages for LS-ML round-trip.
|
|
points_pct = [
|
|
[px / orig_w * 100.0, py / orig_h * 100.0]
|
|
for (px, py) in result.polygon
|
|
]
|
|
|
|
return {
|
|
"results": [{
|
|
"model_version": spec.id,
|
|
"score": result.score,
|
|
"result": [{
|
|
"original_width": orig_w,
|
|
"original_height": orig_h,
|
|
"image_rotation": 0,
|
|
"value": {
|
|
"points": points_pct,
|
|
"closed": True,
|
|
"polygonlabels": ["Object"],
|
|
},
|
|
"id": "sam_poly_0",
|
|
"from_name": "polygon",
|
|
"to_name": "image",
|
|
"type": "polygonlabels",
|
|
}],
|
|
}]
|
|
}
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Helpers
|
|
# --------------------------------------------------------------------------
|
|
def _decode_image(field: str) -> np.ndarray:
|
|
"""Accept `data:<mime>;base64,...` or a bare base64 string."""
|
|
b64 = field
|
|
if field.startswith("data:"):
|
|
comma = field.find(",")
|
|
if comma == -1:
|
|
raise HTTPException(400, "malformed data URL")
|
|
b64 = field[comma + 1:]
|
|
try:
|
|
raw = base64.b64decode(b64)
|
|
except Exception as e:
|
|
raise HTTPException(400, f"base64 decode failed: {e}")
|
|
try:
|
|
img = Image.open(io.BytesIO(raw)).convert("RGB")
|
|
except Exception as e:
|
|
raise HTTPException(400, f"image decode failed: {e}")
|
|
return np.array(img)
|
|
|
|
|
|
# --------------------------------------------------------------------------
|
|
# Entrypoint
|
|
# --------------------------------------------------------------------------
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--host", default="0.0.0.0")
|
|
parser.add_argument("--port", type=int, default=9090)
|
|
parser.add_argument(
|
|
"--models",
|
|
default=str(Path(__file__).parent / "models.yaml"),
|
|
help="Path to models.yaml registry file",
|
|
)
|
|
parser.add_argument(
|
|
"--device",
|
|
default="auto",
|
|
help="'auto' (prefer cuda), 'cuda', 'cpu', or specific cuda:N",
|
|
)
|
|
parser.add_argument("--reload", action="store_true",
|
|
help="Enable uvicorn auto-reload (dev only)")
|
|
args = parser.parse_args()
|
|
|
|
# Hand config to the worker module via env — the lifespan reads these
|
|
# whether uvicorn imports this file as `__main__` or as `server`.
|
|
os.environ["SAM_BACKEND_MODELS"] = args.models
|
|
os.environ["SAM_BACKEND_DEVICE"] = args.device
|
|
|
|
log.info("starting uvicorn on %s:%d (device=%s, models=%s)",
|
|
args.host, args.port, args.device, args.models)
|
|
uvicorn.run(
|
|
"server:app",
|
|
host=args.host,
|
|
port=args.port,
|
|
reload=args.reload,
|
|
log_level="info",
|
|
)
|
|
|
|
|
|
def _pick_device(preference: str) -> str:
|
|
if preference != "auto":
|
|
return preference
|
|
try:
|
|
import torch
|
|
if torch.cuda.is_available():
|
|
return "cuda"
|
|
except ImportError:
|
|
pass
|
|
return "cpu"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|