Ravi/SR-Label-SamTool
1
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

first version

Browse Source
This commit is contained in:
ravi
2026-04-25 16:23:09 +05:30
commit 7d3c36050b
73 changed files with 19033 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
sam2-backend/.gitignore vendored Normal file
View File

@@ -0,0 +1,7 @@
.venv/
__pycache__/
*.pyc
*.pt
*.pth
.vscode/
.idea/

191
sam2-backend/README.md Normal file
View File

@@ -0,0 +1,191 @@
# SAM2 Backend
Standalone HTTP server that runs Meta's SAM-family image predictors and
speaks the Label-Studio-ML `/predict` protocol the SAM-Tool Tauri app
already uses. Can host locally or on any remote PC on your network — point
the Tauri app at `http://<host>:9090` from Settings.
Ships with adapters for both **SAM 2 / 2.1** and **SAM 3**. The registry
(`models.yaml`) is pluggable — new families just drop an adapter file into
`adapters/` and register in `adapters/__init__.py`.
---
## Layout
```
sam2-backend/
├── server.py FastAPI app — /, /models, /predict
├── registry.py Model spec loader + lazy-load cache
├── adapters/
│ ├── base.py Adapter contract
│ ├── sam2_adapter.py SAM2 implementation
│ └── __init__.py Family → adapter class map
├── models.yaml Which models exist + where their checkpoints live
├── requirements.txt HTTP server deps (Torch + SAM2 installed separately)
├── run.sh Convenience launcher
└── README.md
```
---
## Install (one-time, on the host PC)
### 1. Python environment
```bash
cd sam2-backend
python -m venv .venv
source .venv/bin/activate
pip install --upgrade pip
```
### 2. PyTorch (choose one)
Pick the build that matches your GPU/CPU — do NOT rely on the generic
`torch` pulled in transitively by SAM2.
```bash
# CUDA 12.1 (typical modern NVIDIA)
pip install torch==2.4.* torchvision==0.19.* --index-url https://download.pytorch.org/whl/cu121
# CPU-only (works but inference is ~50× slower)
pip install torch==2.4.* torchvision==0.19.* --index-url https://download.pytorch.org/whl/cpu
```
### 3. SAM 2 package
The SAM2 source is already in this repo at `../sam2/sam2/`. Install it as
an editable package:
```bash
pip install -e ../sam2/sam2
```
> If you see `RuntimeError: You're likely running Python from the parent
> directory of the sam2 repository`, you started Python from the *wrong*
> directory — always run `server.py` from `sam2-backend/` (not from the
> repo root).
### 4. SAM 3 package (optional — only if you want SAM3 models)
```bash
# Upstream repo (adjust if you're using a fork):
pip install "sam3 @ git+https://github.com/facebookresearch/sam3"
# For the HuggingFace tracker path (optional, used by `config: tracker`):
pip install transformers
```
The SAM3 image-predictor path uses the `.pt` file in `../sam3_weights/sam3.pt`
— already referenced by `models.yaml`. No additional downloads needed for
the default SAM3 entry.
### 5. The rest
```bash
pip install -r requirements.txt
```
### 6. Checkpoints
**SAM 2 / 2.1** — present at `../sam2/sam2/checkpoints/sam2.1_hiera_{tiny,small,base_plus,large}.pt`.
If any are missing, download them with:
```bash
cd ../sam2/sam2/checkpoints
./download_ckpts.sh
```
**SAM 3** — present at `../files/sam3_weights/sam3.pt`.
`models.yaml` points at all of these by default. Edit the file to add,
remove, or relocate entries.
---
## Run
```bash
# From sam2-backend/
./run.sh # defaults: 0.0.0.0:9090, auto device
HOST=127.0.0.1 PORT=9090 ./run.sh # localhost only
DEVICE=cuda:1 ./run.sh # pin a specific GPU
# or directly
python server.py --host 0.0.0.0 --port 9090 --device auto
```
First request for a given model triggers a one-time load (few seconds). By
default the previously-loaded model is evicted on switch to keep GPU memory
bounded. Set `SAM_BACKEND_KEEP_ALL=1` to keep them all resident.
---
## Smoke-test
```bash
curl -s http://localhost:9090/ # {"status":"ok",...}
curl -s http://localhost:9090/models # {"default":"...","models":[...]}
```
From the Tauri app: open **Settings**, set the SAM URL to
`http://<host>:9090`, enable SAM, pick a model, save. Drawing a bbox should
now round-trip through SAM2 and come back as a polygon.
---
## API
### `GET /`
```json
{"status":"ok","device":"cuda","models":4,"default":"sam2.1_hiera_small"}
```
### `GET /models`
```json
{
"default": "sam2.1_hiera_small",
"models": [
{"id":"sam2.1_hiera_tiny","name":"SAM 2.1 · Hiera Tiny",
"family":"sam2","available":true,
"checkpoint":"/abs/path/sam2.1_hiera_tiny.pt"},
...
]
}
```
### `POST /predict?model=<id>`
Body: Label-Studio-ML task + bbox prompt (see `server.py` docstring).
Optional `?model=<id>` query selects which model. Omit to use the default.
Response: LS-ML `results[].result[].value.points` as percentages.
---
## SAM 3 notes
Two `config` modes are supported per entry in `models.yaml`:
| `config` | Loader | Checkpoint | Notes |
|-------------------|-----------------------------------------------------------------|--------------------------------------------|-----------------------------------------------------------------------|
| `image_predictor` | `sam3.model_builder.build_sam3_image_model` | Local `.pt` file | Default. Box = visual exemplar; picks mask with highest IoU to user box. |
| `tracker` | `transformers.Sam3TrackerModel.from_pretrained("facebook/sam3")` | HF model id (set `remote: true` in yaml) | Cleaner single-box → single-mask. Requires `transformers` + first-run network access. |
The shipped default SAM3 entry uses `image_predictor` + the local weights.
To enable the tracker path, uncomment the `sam3_tracker` block in
`models.yaml`.
---
## Adding another model family later
1. Drop a new file into `adapters/`, e.g. `adapters/sam4_adapter.py`, that
implements `BaseAdapter`.
2. Register it in `adapters/__init__.py`:
```python
ADAPTERS = {"sam2": Sam2Adapter, "sam3": Sam3Adapter, "sam4": Sam4Adapter}
```
3. Add entries to `models.yaml` with `family: sam4`.
4. Restart the server.
No changes to `server.py` or `registry.py` should be needed.

10
sam2-backend/adapters/__init__.py Normal file
View File

@@ -0,0 +1,10 @@
from .base import BaseAdapter, PredictResult
from .sam2_adapter import Sam2Adapter
from .sam3_adapter import Sam3Adapter
ADAPTERS: dict[str, type[BaseAdapter]] = {
"sam2": Sam2Adapter,
"sam3": Sam3Adapter,
}
__all__ = ["BaseAdapter", "PredictResult", "Sam2Adapter", "Sam3Adapter", "ADAPTERS"]

34
sam2-backend/adapters/base.py Normal file
View File

@@ -0,0 +1,34 @@
from dataclasses import dataclass
from typing import Protocol
import numpy as np
@dataclass
class PredictResult:
"""Result of one bbox-prompted segmentation.
`polygon` is a list of (x, y) pixel coordinates describing the outer
contour of the largest mask blob. `score` is the model's confidence for
the selected mask (in [0, 1]).
"""
polygon: list[tuple[float, float]]
score: float
class BaseAdapter(Protocol):
"""Minimal adapter contract. Each model family implements this."""
def load(self, config: str, checkpoint: str, device: str) -> None:
"""Load weights into memory. Called once, lazily on first predict."""
...
def predict_bbox(
self, image_rgb: np.ndarray, bbox_xyxy: tuple[int, int, int, int]
) -> PredictResult:
"""Run segmentation with a bbox prompt, return the best polygon."""
...
def unload(self) -> None:
"""Free GPU memory. Called when the registry evicts this model."""
...

82
sam2-backend/adapters/sam2_adapter.py Normal file
View File

@@ -0,0 +1,82 @@
"""SAM2 adapter — wraps `sam2.sam2_image_predictor.SAM2ImagePredictor`.
The SAM2 repo uses Hydra to resolve config files relative to the `sam2`
package. We rely on that resolution: pass `config="configs/sam2.1/..."` and
it works as long as the `sam2` package is importable (pip-installed or on
PYTHONPATH).
"""
from __future__ import annotations
import logging
from typing import Optional
import numpy as np
from .base import BaseAdapter, PredictResult
from .utils import mask_to_polygon as _mask_to_polygon
log = logging.getLogger("sam2-backend.sam2")
class Sam2Adapter(BaseAdapter):
def __init__(self) -> None:
self._predictor: Optional[object] = None
self._device: str = "cpu"
def load(self, config: str, checkpoint: str, device: str) -> None:
# Imports are deferred so the server can start even without the SAM2
# deps installed (the model stays in `available: false` until loaded).
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
log.info("loading SAM2: config=%s ckpt=%s device=%s",
config, checkpoint, device)
model = build_sam2(config, checkpoint, device=device)
self._predictor = SAM2ImagePredictor(model)
self._device = device
log.info("SAM2 loaded")
def unload(self) -> None:
if self._predictor is None:
return
self._predictor = None
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
log.info("SAM2 unloaded")
def predict_bbox(
self, image_rgb: np.ndarray, bbox_xyxy: tuple[int, int, int, int]
) -> PredictResult:
if self._predictor is None:
raise RuntimeError("SAM2 adapter not loaded")
import torch
# set_image re-embeds on every call; that's fine for an interactive
# one-box-per-request workflow. Batch prompts would want caching.
inference_ctx = (
torch.inference_mode()
if self._device.startswith("cuda") or self._device == "cpu"
else torch.no_grad()
)
with inference_ctx:
self._predictor.set_image(image_rgb) # type: ignore[attr-defined]
box = np.array(list(bbox_xyxy), dtype=np.float32)
masks, scores, _ = self._predictor.predict( # type: ignore[attr-defined]
box=box,
multimask_output=True,
)
# multimask_output returns 3 candidates — take the highest-scoring one.
best_idx = int(np.argmax(scores))
best_mask = masks[best_idx]
best_score = float(scores[best_idx])
polygon = _mask_to_polygon(best_mask)
return PredictResult(polygon=polygon, score=best_score)

206
sam2-backend/adapters/sam3_adapter.py Normal file
View File

@@ -0,0 +1,206 @@
"""SAM3 adapter.
Mirrors the loading pattern from `files/sam3_server.py` in this repo. Two
modes are supported, selected via the model spec's `config` field in
models.yaml:
config: image_predictor — uses `sam3.model_builder.build_sam3_image_model`
with a local `.pt` checkpoint. Treats the user's
bbox as a visual exemplar and returns the mask
whose predicted box has highest IoU with the
user box (single-object annotation workflow).
config: tracker — uses HuggingFace `transformers.Sam3TrackerModel`
(SAM2-style PVS on images). `checkpoint` here is
a HF model ID like "facebook/sam3" rather than
a local .pt. Requires `transformers` installed
and network access for first download.
Both paths return a single polygon + confidence for the best mask, matching
the adapter contract the server's `/predict` route expects.
"""
from __future__ import annotations
import logging
from typing import Optional
import numpy as np
from .base import BaseAdapter, PredictResult
from .utils import mask_to_polygon as _mask_to_polygon
log = logging.getLogger("sam2-backend.sam3")
class Sam3Adapter(BaseAdapter):
def __init__(self) -> None:
self._mode: str = "image_predictor"
self._device: str = "cpu"
# image_predictor path
self._image_model: Optional[object] = None
self._image_processor: Optional[object] = None
# tracker path
self._tracker_model: Optional[object] = None
self._tracker_processor: Optional[object] = None
def load(self, config: str, checkpoint: str, device: str) -> None:
self._device = device
mode = (config or "image_predictor").strip().lower()
self._mode = mode
if mode == "tracker":
self._load_tracker(checkpoint, device)
elif mode in ("image_predictor", "image", "exemplar"):
self._load_image_predictor(checkpoint, device)
else:
raise ValueError(
f"unknown SAM3 config mode '{mode}'. "
f"expected 'image_predictor' or 'tracker'."
)
def _load_image_predictor(self, checkpoint: str, device: str) -> None:
# Deferred imports so the server can start even without sam3 installed
# (the model stays unavailable until first use).
from sam3.model_builder import build_sam3_image_model
from sam3.model.sam3_image_processor import Sam3Processor
log.info("loading SAM3 image predictor: ckpt=%s device=%s",
checkpoint, device)
model = (
build_sam3_image_model(checkpoint_path=checkpoint)
if checkpoint
else build_sam3_image_model()
)
self._image_model = model.to(device).eval()
self._image_processor = Sam3Processor(self._image_model)
log.info("SAM3 image predictor loaded")
def _load_tracker(self, hf_id: str, device: str) -> None:
from transformers import Sam3TrackerModel, Sam3TrackerProcessor # type: ignore
model_id = hf_id or "facebook/sam3"
log.info("loading SAM3 tracker: hf_id=%s device=%s", model_id, device)
self._tracker_processor = Sam3TrackerProcessor.from_pretrained(model_id)
self._tracker_model = (
Sam3TrackerModel.from_pretrained(model_id).to(device).eval()
)
log.info("SAM3 tracker loaded")
def unload(self) -> None:
self._image_model = None
self._image_processor = None
self._tracker_model = None
self._tracker_processor = None
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except Exception:
pass
log.info("SAM3 unloaded")
def predict_bbox(
self, image_rgb: np.ndarray, bbox_xyxy: tuple[int, int, int, int]
) -> PredictResult:
import torch
from PIL import Image
pil_image = Image.fromarray(image_rgb)
# SAM3 weights are loaded in BFloat16 by `build_sam3_image_model`, but
# the processor hands the model Float32 image tensors → matmul dtype
# mismatch. Autocast does the input cast for us. Tracker path goes
# through HF transformers which already handles dtype internally, but
# autocasting is harmless there too.
device_type = "cuda" if self._device.startswith("cuda") else "cpu"
with torch.inference_mode():
with torch.autocast(device_type=device_type, dtype=torch.bfloat16):
if self._mode == "tracker":
return self._predict_tracker(pil_image, bbox_xyxy)
return self._predict_image_predictor(pil_image, bbox_xyxy)
# --------------------------------------------------------------------
# Tracker path — HF transformers `Sam3Tracker`
# --------------------------------------------------------------------
def _predict_tracker(self, image, bbox_xyxy) -> PredictResult:
if self._tracker_model is None or self._tracker_processor is None:
raise RuntimeError("SAM3 tracker not loaded")
x1, y1, x2, y2 = [float(v) for v in bbox_xyxy]
inputs = self._tracker_processor(
images=image,
input_boxes=[[[x1, y1, x2, y2]]], # batch[1] × obj[1] × 4
return_tensors="pt",
).to(self._device)
outputs = self._tracker_model(**inputs, multimask_output=True)
# `post_process_masks` lost `reshaped_input_sizes` in newer
# transformers — only `masks` and `original_sizes` are accepted now.
# Cast to float32 before .numpy() — autocast leaves outputs in
# bfloat16, which NumPy can't ingest.
masks = self._tracker_processor.post_process_masks(
outputs.pred_masks.float().cpu(),
original_sizes=inputs["original_sizes"].cpu(),
)[0] # (num_masks, H, W)
scores = outputs.iou_scores[0].float().cpu().numpy().ravel()
masks_np = masks.float().numpy().astype(bool) if masks.dtype != bool else masks.numpy()
n_masks = masks_np.shape[0]
if n_masks == 0:
return PredictResult(polygon=[], score=0.0)
# Mask and score counts can disagree (`multimask_output=True` returns 3
# IoU scores even when the head only emits 1 mask, e.g. when a
# `sam3_video` checkpoint is loaded into `Sam3TrackerModel`). Pick the
# best within whichever count is smaller.
valid = scores[:n_masks] if scores.size >= n_masks else scores
best = int(np.argmax(valid)) if valid.size > 0 else 0
best_score = float(valid[best]) if valid.size > 0 else 0.0
return PredictResult(
polygon=_mask_to_polygon(masks_np[best]),
score=best_score,
)
# --------------------------------------------------------------------
# Image-predictor path — box-as-exemplar, pick highest-IoU result
# --------------------------------------------------------------------
def _predict_image_predictor(self, image, bbox_xyxy) -> PredictResult:
if self._image_processor is None:
raise RuntimeError("SAM3 image predictor not loaded")
x1, y1, x2, y2 = [float(v) for v in bbox_xyxy]
state = self._image_processor.set_image(image) # type: ignore[attr-defined]
out = self._image_processor.set_box_prompt( # type: ignore[attr-defined]
state=state,
boxes=[[x1, y1, x2 - x1, y2 - y1]], # xywh per SAM3 API
box_labels=[1],
)
masks = out.get("masks", [])
boxes = out.get("boxes", [])
scores = out.get("scores", [])
if len(masks) == 0:
return PredictResult(polygon=[], score=0.0)
# Pick the proposal whose predicted bbox overlaps the user bbox most.
user = (x1, y1, x2, y2)
ious = [_iou(user, tuple(float(v) for v in b)) for b in boxes]
idx = int(np.argmax(ious))
mask_np = np.asarray(masks[idx]).astype(bool)
return PredictResult(
polygon=_mask_to_polygon(mask_np),
score=float(scores[idx]),
)
# --------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------
def _iou(
a: tuple[float, float, float, float],
b: tuple[float, float, float, float],
) -> float:
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 = max(0.0, ix2 - ix1)
ih = max(0.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

141
sam2-backend/adapters/utils.py Normal file
View File

@@ -0,0 +1,141 @@
"""Shared helpers for the SAM-family adapters."""
from __future__ import annotations
import cv2
import numpy as np
def mask_to_polygon(
mask: np.ndarray,
*,
min_points: int = 4,
max_points: int = 10,
smooth: bool = True,
) -> list[tuple[float, float]]:
"""Largest external contour, simplified to a vertex count in
[min_points, max_points].
Pipeline:
1. Squeeze + binarise the mask.
2. Optional morphological close+open + Gaussian smoothing — this is
what turns staircase-jagged SAM mask edges into curves the polygon
can actually trace cleanly.
3. Pull the full contour (`CHAIN_APPROX_NONE`) so the simplifier has
every pixel-edge sample to choose from.
4. Binary-search Douglas-Peucker epsilon to hit a vertex count in
[min_points, max_points], biased toward more points (more detail).
Why this matters: the user has to drag every vertex by hand in the
annotation UI. Raw `findContours` output is often hundreds of points
along a single edge — unusable. Without step 2, even a 15-point output
looks choppy because the chosen vertices sit on stair-step pixels.
"""
m = np.asarray(mask)
# Tolerate stray leading singleton dims (some processors return
# [1, H, W] or [1, 1, H, W]).
m = np.squeeze(m)
while m.ndim > 2:
m = m[0]
if m.ndim != 2 or m.size == 0:
return []
m = m.astype(np.uint8)
if m.max() == 1:
m = m * 255
if smooth:
m = _smooth_mask(m)
contours, _ = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
if not contours:
return []
cnt = max(contours, key=cv2.contourArea)
if len(cnt) < 3:
return []
simplified = _simplify_to_band(cnt, min_points=min_points, max_points=max_points)
if simplified is None or len(simplified) < 3:
return []
pts = simplified.reshape(-1, 2)
return [(float(x), float(y)) for x, y in pts]
def _smooth_mask(mask_u8: np.ndarray) -> np.ndarray:
"""Clean up a binary mask before contour extraction.
Kernel size scales mildly with the mask's bounding box so a tiny object
doesn't get erased and a huge object gets enough smoothing to matter.
"""
h, w = mask_u8.shape[:2]
# Find the actual blob size (not the canvas size) for kernel sizing.
ys, xs = np.where(mask_u8 > 0)
if xs.size == 0:
return mask_u8
bw = xs.max() - xs.min() + 1
bh = ys.max() - ys.min() + 1
blob_diag = max(1.0, float(np.hypot(bw, bh)))
# Kernel ~1.5% of blob diagonal, clamped to a usable range.
k = int(round(blob_diag * 0.015))
k = max(3, min(k | 1, 11)) # odd, between 3 and 11
kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (k, k))
# Close fills tiny holes / single-pixel notches; open removes specks.
m = cv2.morphologyEx(mask_u8, cv2.MORPH_CLOSE, kernel)
m = cv2.morphologyEx(m, cv2.MORPH_OPEN, kernel)
# A light Gaussian + re-threshold smooths the edge by a sub-pixel amount,
# which makes the contour follow curves cleanly instead of staircasing.
blur_k = max(3, min(k | 1, 9))
m = cv2.GaussianBlur(m, (blur_k, blur_k), 0)
_, m = cv2.threshold(m, 127, 255, cv2.THRESH_BINARY)
return m
def _simplify_to_band(
contour: np.ndarray,
*,
min_points: int,
max_points: int,
iterations: int = 24,
) -> np.ndarray | None:
"""Binary-search epsilon for `cv2.approxPolyDP` so vertex count lands in
[min_points, max_points], biased toward the larger end of the band.
Strategy:
- epsilon=0 → original (many) points
- epsilon=arc → degenerates to a single point
- smaller epsilon ⇒ more points; larger ⇒ fewer.
We track the in-band result with the highest vertex count seen
(most-detailed pick within the cap) and return it.
"""
arc = cv2.arcLength(contour, True)
if arc == 0:
return contour
if len(contour) <= max_points:
return contour if len(contour) >= min_points else None
lo, hi = 0.0, arc
best: np.ndarray | None = None
best_n = -1
for _ in range(iterations):
mid = 0.5 * (lo + hi)
approx = cv2.approxPolyDP(contour, mid, True)
n = len(approx)
if n > max_points:
lo = mid
elif n < min_points:
hi = mid
else:
if n > best_n:
best = approx
best_n = n
hi = mid
if hi - lo < 1e-6 * arc:
break
if best is not None:
return best
return cv2.approxPolyDP(contour, lo, True)

69
sam2-backend/models.yaml Normal file
View File

@@ -0,0 +1,69 @@
# Model registry for the SAM backend.
#
# Add/remove entries freely. Paths may be absolute or relative to this file.
# Checkpoints marked as missing on disk are still listed by /models but flagged
# `available: false` so the frontend can grey them out.
#
# `family` selects the adapter code-path. Currently supported:
# sam2 — Meta SAM 2 / SAM 2.1 (image predictor, bbox prompts)
# sam3 — Meta SAM 3 (image predictor with box-as-exemplar, or HF tracker)
#
# Per-entry fields:
# id — stable key used by /predict?model=<id>
# name — human-readable, shown in the Tauri model dropdown
# family — adapter to use
# config — family-specific (SAM2: hydra yaml path; SAM3: "image_predictor"
# or "tracker")
# checkpoint — file path (relative paths resolved against this yaml's
# directory) OR an opaque identifier when `remote: true`
# remote — optional; true = treat checkpoint as a HuggingFace model id
# (no on-disk existence check)
default: sam2.1_hiera_small
models:
- id: sam2.1_hiera_tiny
name: "SAM 2.1 · Hiera Tiny"
family: sam2
config: configs/sam2.1/sam2.1_hiera_t.yaml
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_tiny.pt
- id: sam2.1_hiera_small
name: "SAM 2.1 · Hiera Small"
family: sam2
config: configs/sam2.1/sam2.1_hiera_s.yaml
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_small.pt
- id: sam2.1_hiera_base_plus
name: "SAM 2.1 · Hiera Base+"
family: sam2
config: configs/sam2.1/sam2.1_hiera_b+.yaml
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_base_plus.pt
- id: sam2.1_hiera_large
name: "SAM 2.1 · Hiera Large"
family: sam2
config: configs/sam2.1/sam2.1_hiera_l.yaml
checkpoint: ../sam2/sam2/checkpoints/sam2.1_hiera_large.pt
# ---- SAM 3 -----------------------------------------------------------
# Local image-predictor path — uses the .pt weights in this repo. Box is
# treated as a visual exemplar; the adapter picks the mask whose predicted
# bbox best overlaps the user bbox (classic single-instance annotation).
- id: sam3
name: "SAM 3 · image predictor"
family: sam3
config: image_predictor
checkpoint: ../files/sam3_weights/sam3.pt
# Optional HuggingFace tracker path. Requires `pip install transformers`
# and pulls weights from HF hub on first use. Gives cleaner single-box
# → single-mask results than the exemplar path but needs a network round
# trip the first time. Disabled by leaving `remote: false` — flip to
# `remote: true` once your machine has transformers installed.
# - id: sam3_tracker
# name: "SAM 3 · tracker (HuggingFace)"
# family: sam3
# config: tracker
# checkpoint: facebook/sam3
# remote: true

141
sam2-backend/registry.py Normal file
View File

@@ -0,0 +1,141 @@
"""Model registry. Reads models.yaml, lazy-loads adapters, caches instances.
Only one adapter is resident at a time by default — switching models unloads
the previous one to keep GPU memory bounded. Set SAM_BACKEND_KEEP_ALL=1 to
disable eviction (useful if you have VRAM headroom and switch often).
"""
from __future__ import annotations
import logging
import os
import threading
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional
import yaml
from adapters import ADAPTERS, BaseAdapter, PredictResult
log = logging.getLogger("sam2-backend.registry")
@dataclass
class ModelSpec:
id: str
name: str
family: str
config: str
checkpoint: str
# If true, `checkpoint` is treated as an opaque identifier (e.g. a
# HuggingFace model id like "facebook/sam3") rather than a file on disk.
# Availability check is skipped for remote entries.
remote: bool = False
available: bool = False
# Resolved absolute paths for logging / responses.
checkpoint_abs: str = ""
@dataclass
class Registry:
specs: dict[str, ModelSpec] = field(default_factory=dict)
default_id: Optional[str] = None
device: str = "cpu"
keep_all: bool = False
_loaded: dict[str, BaseAdapter] = field(default_factory=dict)
_lock: threading.Lock = field(default_factory=threading.Lock)
def list_specs(self) -> list[ModelSpec]:
return list(self.specs.values())
def resolve(self, model_id: Optional[str]) -> ModelSpec:
if not model_id:
if self.default_id is None:
raise KeyError("no default model configured")
model_id = self.default_id
if model_id not in self.specs:
raise KeyError(f"unknown model: {model_id}")
return self.specs[model_id]
def get_adapter(self, model_id: str) -> BaseAdapter:
"""Returns a loaded adapter for `model_id`, loading on first use."""
with self._lock:
if model_id in self._loaded:
return self._loaded[model_id]
spec = self.resolve(model_id)
if not spec.available:
raise FileNotFoundError(
f"checkpoint missing: {spec.checkpoint_abs}"
)
cls = ADAPTERS.get(spec.family)
if cls is None:
raise ValueError(
f"no adapter registered for family '{spec.family}'"
)
if not self.keep_all:
# Evict whichever model is currently resident.
for old_id, adapter in list(self._loaded.items()):
log.info("evicting %s to make room for %s", old_id, model_id)
adapter.unload()
del self._loaded[old_id]
adapter = cls()
adapter.load(spec.config, spec.checkpoint_abs, self.device)
self._loaded[model_id] = adapter
return adapter
def predict_bbox(self, model_id: Optional[str], image_rgb, bbox_xyxy) -> tuple[ModelSpec, PredictResult]:
spec = self.resolve(model_id)
adapter = self.get_adapter(spec.id)
# Serialize inference: most of these adapters aren't reentrant.
with self._lock:
result = adapter.predict_bbox(image_rgb, bbox_xyxy)
return spec, result
def load_registry(models_yaml: Path, device: str) -> Registry:
with open(models_yaml, "r") as f:
data = yaml.safe_load(f)
base_dir = models_yaml.parent
keep_all = os.environ.get("SAM_BACKEND_KEEP_ALL", "0") == "1"
reg = Registry(device=device, keep_all=keep_all)
reg.default_id = data.get("default")
for entry in data.get("models", []):
remote = bool(entry.get("remote", False))
ckpt_str = entry["checkpoint"]
if remote:
# Opaque identifier — leave as-is, mark available so the UI
# shows it. The adapter resolves/downloads on first use.
ckpt_abs = ckpt_str
available = True
else:
ckpt = Path(ckpt_str)
ckpt_path = (ckpt if ckpt.is_absolute() else (base_dir / ckpt)).resolve()
ckpt_abs = str(ckpt_path)
available = ckpt_path.is_file()
spec = ModelSpec(
id=entry["id"],
name=entry["name"],
family=entry["family"],
config=entry["config"],
checkpoint=ckpt_str,
remote=remote,
available=available,
checkpoint_abs=ckpt_abs,
)
reg.specs[spec.id] = spec
log.info(
"registered model %s (available=%s, ckpt=%s)",
spec.id, spec.available, spec.checkpoint_abs,
)
# Fall back the default to the first *available* model if the configured
# default's checkpoint is missing.
if reg.default_id and not reg.specs.get(reg.default_id, ModelSpec("", "", "", "", "")).available:
for s in reg.specs.values():
if s.available:
reg.default_id = s.id
break
return reg

13
sam2-backend/requirements.txt Normal file
View File

@@ -0,0 +1,13 @@
# Runtime deps for the SAM2 backend HTTP server.
# Torch is intentionally NOT pinned here — install the right build for your
# CUDA/CPU first (see README.md), then `pip install -r requirements.txt`.
fastapi>=0.110
uvicorn[standard]>=0.27
pydantic>=2.6
pyyaml>=6.0
pillow>=10.0
numpy>=1.24,<2.0
opencv-python-headless>=4.8
hydra-core>=1.3
iopath>=0.1.10

32
sam2-backend/run.sh Executable file
View File

@@ -0,0 +1,32 @@
#!/usr/bin/env bash
# Convenience launcher for the SAM2 backend.
# Assumes a venv at ./.venv with requirements.txt + torch + sam2 installed.
# See README.md for install steps.
set -euo pipefail
cd "$(dirname "$0")"
if [ -d ".venv" ]; then
# shellcheck disable=SC1091
source .venv/bin/activate
fi
HOST="${HOST:-0.0.0.0}"
PORT="${PORT:-9090}"
DEVICE="${DEVICE:-auto}"
# Privacy: disable HF/transformers telemetry, and (after weights are cached)
# refuse any outbound calls to huggingface.co. Inference runs fully local
# regardless; these just stop the cosmetic pings on startup.
export HF_HUB_DISABLE_TELEMETRY="${HF_HUB_DISABLE_TELEMETRY:-1}"
export DO_NOT_TRACK="${DO_NOT_TRACK:-1}"
# Flip both OFFLINE flags to 1 once you've done the first `from_pretrained`
# and the weights are cached under ~/.cache/huggingface/. Leaving them 0
# the first time lets the model download. After that, set them 1 to fully
# air-gap and silence the HEAD-request revalidation pings on every startup.
# HF_HUB_OFFLINE — used by huggingface_hub
# TRANSFORMERS_OFFLINE — used by transformers itself (some legacy paths)
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-0}"
export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-0}"
exec python server.py --host "$HOST" --port "$PORT" --device "$DEVICE" "$@"

345
sam2-backend/server.py Normal file
View File

@@ -0,0 +1,345 @@
"""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()