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

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)