first version
This commit is contained in:
206
sam2-backend/adapters/sam3_adapter.py
Normal file
206
sam2-backend/adapters/sam3_adapter.py
Normal 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
|
||||
|
||||
|
||||
Reference in New Issue
Block a user