207 lines
8.8 KiB
Python
207 lines
8.8 KiB
Python
"""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
|
||
|
||
|