Alirezakzt's picture
init
dd85ec3
Raw
History Blame Contribute Delete
24.2 kB
"""
SAM 3.1 — Promptable Concept Segmentation demo
================================================
A live, language-driven segmentation demo built on Meta's Segment Anything Model 3.1.
Type a short noun phrase (e.g. "horse", "saddle"); the model finds and segments
*every* matching instance in the image. No boxes, no clicks, no retraining.
Model: facebook/sam3.1 (image Promptable Concept Segmentation path)
Runtime: Hugging Face Spaces — works on a standard GPU Space or on ZeroGPU.
Deploy notes are in README.md (gated-model access + HF_TOKEN + hardware).
"""
import os
import time
import colorsys
from contextlib import nullcontext
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import torch
import gradio as gr
# --------------------------------------------------------------------------------------
# ZeroGPU support (optional). On a standard GPU Space this becomes a transparent no-op.
# --------------------------------------------------------------------------------------
try:
import spaces # provided by the ZeroGPU runtime
GPU = spaces.GPU
except Exception: # not on Spaces / package missing → identity decorator
def GPU(*args, **kwargs):
# Supports both `@GPU` and `@GPU(duration=...)`
if len(args) == 1 and callable(args[0]) and not kwargs:
return args[0]
def _deco(fn):
return fn
return _deco
# --------------------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------------------
MODEL_ID = os.environ.get("MODEL_ID", "facebook/sam3.1")
FALLBACK_MODEL_ID = os.environ.get("FALLBACK_MODEL_ID", "facebook/sam3")
HF_TOKEN = (
os.environ.get("HF_TOKEN")
or os.environ.get("HUGGING_FACE_HUB_TOKEN")
or os.environ.get("HUGGINGFACE_TOKEN")
)
# ZeroGPU sets SPACES_ZERO_GPU; in that case CUDA is attached only inside @GPU calls,
# but we can still target "cuda" because the `spaces` runtime patches device placement.
_ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU"))
DEVICE = "cuda" if (torch.cuda.is_available() or _ZERO_GPU) else "cpu"
EXAMPLE_PROMPTS = [
"horse",
"saddle",
"person",
"object used for riding control",
]
KEY_MESSAGE = "Segmentation is fully driven by language prompts — no retraining required."
# Lazily-loaded singletons
_MODEL = None
_PROCESSOR = None
_LOADED_ID = None
# --------------------------------------------------------------------------------------
# Model loading
# --------------------------------------------------------------------------------------
def load_model():
"""Load SAM 3.1 (image PCS) once. Falls back to SAM 3 if 3.1 is unavailable."""
global _MODEL, _PROCESSOR, _LOADED_ID
if _MODEL is not None:
return
from transformers import Sam3Model, Sam3Processor
candidates = [MODEL_ID]
if FALLBACK_MODEL_ID and FALLBACK_MODEL_ID != MODEL_ID:
candidates.append(FALLBACK_MODEL_ID)
last_err = None
for mid in candidates:
try:
processor = Sam3Processor.from_pretrained(mid, token=HF_TOKEN)
model = Sam3Model.from_pretrained(mid, token=HF_TOKEN)
model.eval()
try:
model.to(DEVICE)
except Exception:
# On ZeroGPU the move is handled when the GPU is attached; ignore here.
pass
_MODEL, _PROCESSOR, _LOADED_ID = model, processor, mid
if mid != MODEL_ID:
print(f"[sam3.1-demo] '{MODEL_ID}' unavailable; loaded fallback '{mid}'.")
else:
print(f"[sam3.1-demo] Loaded '{mid}' on {DEVICE}.")
return
except Exception as e: # try next candidate
last_err = e
print(f"[sam3.1-demo] Could not load '{mid}': {e}")
raise RuntimeError(
f"Failed to load any SAM 3 model from {candidates}. Last error: {last_err}"
)
def _friendly_error(err: Exception) -> str:
"""Turn a load/inference exception into actionable guidance."""
text = str(err).lower()
gated = any(k in text for k in ["401", "403", "gated", "access", "token", "authorized"])
if gated:
return (
"Couldn't access the model weights. SAM 3 / 3.1 are gated: request access on the "
"Hugging Face model page, then add your token as a Space secret named "
"<b>HF_TOKEN</b> (Settings → Variables and secrets), and restart the Space."
)
return f"Something went wrong while running the model: {err}"
# --------------------------------------------------------------------------------------
# Inference (GPU-scoped). Everything returned here is CPU/NumPy so it stays valid
# after the GPU is released (important for ZeroGPU).
# --------------------------------------------------------------------------------------
def _amp_ctx():
"""bfloat16 autocast on CUDA for speed; pass-through elsewhere."""
if DEVICE == "cuda":
return torch.autocast("cuda", dtype=torch.bfloat16)
return nullcontext()
def _to_np(x):
if x is None:
return None
if hasattr(x, "detach"):
return x.detach().to("cpu").float().numpy()
if isinstance(x, np.ndarray):
return x
if isinstance(x, (list, tuple)):
if len(x) == 0:
return np.zeros((0,))
if hasattr(x[0], "detach"):
return np.stack([t.detach().to("cpu").float().numpy() for t in x])
return np.asarray(x)
return np.asarray(x)
def _postprocess(outputs, target_sizes, threshold):
res = _PROCESSOR.post_process_instance_segmentation(
outputs,
threshold=float(threshold),
mask_threshold=0.5,
target_sizes=target_sizes,
)[0]
masks = _to_np(res.get("masks"))
boxes = _to_np(res.get("boxes"))
scores = _to_np(res.get("scores"))
return masks, boxes, scores
@GPU(duration=120)
def _infer_single(image: Image.Image, prompt: str, threshold: float):
"""Segment one text prompt on one image. Returns CPU arrays + metadata."""
load_model()
inputs = _PROCESSOR(images=image, text=prompt, return_tensors="pt").to(_MODEL.device)
target_sizes = inputs["original_sizes"].tolist()
t0 = time.perf_counter()
with torch.no_grad():
try:
with _amp_ctx():
outputs = _MODEL(**inputs)
except RuntimeError:
outputs = _MODEL(**inputs) # rare: fall back to full precision
if DEVICE == "cuda":
torch.cuda.synchronize()
ms = (time.perf_counter() - t0) * 1000.0
masks, boxes, scores = _postprocess(outputs, target_sizes, threshold)
return masks, boxes, scores, ms, _LOADED_ID, _MODEL.device.type
@GPU(duration=180)
def _infer_many(image: Image.Image, prompts, threshold: float):
"""Segment several prompts on one image, reusing vision features for speed.
The whole batch runs inside a single GPU call, so the cached vision embeddings
stay valid (safe on ZeroGPU).
"""
load_model()
img_inputs = _PROCESSOR(images=image, return_tensors="pt").to(_MODEL.device)
target_sizes = img_inputs["original_sizes"].tolist()
results = []
t0 = time.perf_counter()
with torch.no_grad():
with _amp_ctx():
vision_embeds = _MODEL.get_vision_features(pixel_values=img_inputs.pixel_values)
for prompt in prompts:
text_inputs = _PROCESSOR(text=prompt, return_tensors="pt").to(_MODEL.device)
with _amp_ctx():
outputs = _MODEL(vision_embeds=vision_embeds, **text_inputs)
masks, boxes, scores = _postprocess(outputs, target_sizes, threshold)
results.append((prompt, masks, boxes, scores))
if DEVICE == "cuda":
torch.cuda.synchronize()
ms = (time.perf_counter() - t0) * 1000.0
return results, ms, _LOADED_ID, _MODEL.device.type
# --------------------------------------------------------------------------------------
# Rendering (CPU). Builds the semi-transparent overlay and the mask-only view.
# --------------------------------------------------------------------------------------
def _palette(n: int):
"""Evenly-spaced, vivid colors (golden-ratio hue spacing) — one per instance."""
cols = []
for i in range(max(n, 1)):
h = (i * 0.61803398875) % 1.0
r, g, b = colorsys.hsv_to_rgb(h, 0.72, 1.0)
cols.append((int(r * 255), int(g * 255), int(b * 255)))
return cols
def _mask_list(masks, h, w):
"""Normalize whatever the model returned into a list of HxW bool arrays."""
out = []
if masks is None:
return out
arr = masks
if arr.ndim == 2:
arr = arr[None, ...]
for i in range(arr.shape[0]):
m = arr[i]
if m.ndim == 3:
m = m[0]
m = m > 0.5
if m.shape[:2] != (h, w):
m = (
np.asarray(
Image.fromarray((m.astype(np.uint8) * 255)).resize(
(w, h), Image.NEAREST
)
)
> 127
)
out.append(m)
return out
def _boundary(mask: np.ndarray) -> np.ndarray:
"""1px boundary via 4-neighbour erosion (no SciPy/OpenCV dependency)."""
e = mask.copy()
e[1:, :] &= mask[:-1, :]
e[:-1, :] &= mask[1:, :]
e[:, 1:] &= mask[:, :-1]
e[:, :-1] &= mask[:, 1:]
return mask & ~e
def _font(size: int):
for path in (
"DejaVuSans-Bold.ttf",
"/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
"DejaVuSans.ttf",
"/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
):
try:
return ImageFont.truetype(path, size)
except Exception:
continue
return ImageFont.load_default()
def render(image: Image.Image, masks, boxes, scores, prompt: str,
alpha: float = 0.5, show_boxes: bool = False):
"""Return (overlay_image, mask_only_image)."""
base = np.asarray(image.convert("RGB")).astype(np.float32)
h, w = base.shape[:2]
mlist = _mask_list(masks, h, w)
cols = _palette(len(mlist))
overlay = base.copy()
mask_only = np.zeros_like(base)
for i, m in enumerate(mlist):
c = np.array(cols[i], dtype=np.float32)
overlay[m] = overlay[m] * (1.0 - alpha) + c * alpha
edge = _boundary(m)
overlay[edge] = c # crisp instance outline
mask_only[m] = c
mask_only[edge] = np.minimum(c + 70, 255)
overlay_img = Image.fromarray(overlay.clip(0, 255).astype(np.uint8))
mask_only_img = Image.fromarray(mask_only.astype(np.uint8))
if show_boxes and boxes is not None and len(boxes) and scores is not None:
draw = ImageDraw.Draw(overlay_img, "RGBA")
fsize = max(13, int(w / 55))
font = _font(fsize)
line_w = max(2, int(w / 480))
for i in range(min(len(boxes), len(mlist) or len(boxes))):
x1, y1, x2, y2 = [float(v) for v in boxes[i][:4]]
c = cols[i % len(cols)]
draw.rectangle([x1, y1, x2, y2], outline=c + (255,), width=line_w)
label = f"{prompt} · {float(scores[i]):.2f}"
tb = draw.textbbox((0, 0), label, font=font)
tw, th = tb[2] - tb[0], tb[3] - tb[1]
ty = max(0, y1 - th - 6)
draw.rectangle([x1, ty, x1 + tw + 10, ty + th + 6], fill=c + (235,))
draw.text((x1 + 5, ty + 3), label, fill=(20, 24, 31, 255), font=font)
return overlay_img, mask_only_img
# --------------------------------------------------------------------------------------
# Status banner HTML
# --------------------------------------------------------------------------------------
def status_html(count: int, ms: float, model_id: str, device: str) -> str:
plural = "" if count == 1 else "es"
return (
"<div class='status'>"
f"<span class='chip chip-count'>{count} match{plural}</span>"
f"<span class='chip chip-ms'>{ms:.0f} ms</span>"
f"<span class='chip chip-dim'>{model_id} · {device}</span>"
"</div>"
)
def empty_status_html(prompt: str) -> str:
return (
"<div class='status'>"
f"<span class='chip chip-empty'>No matches for &ldquo;{prompt}&rdquo;</span>"
"<span class='chip chip-dim'>Try a simpler noun, or lower the threshold</span>"
"</div>"
)
def info_status_html(message: str) -> str:
return f"<div class='status'><span class='chip chip-empty'>{message}</span></div>"
IDLE_STATUS = (
"<div class='status'><span class='chip chip-dim'>"
"Upload an image, type a prompt, then run.</span></div>"
)
# --------------------------------------------------------------------------------------
# Gradio callbacks
# --------------------------------------------------------------------------------------
def _noop(status_md, history):
# leave images & gallery untouched
return (gr.update(), gr.update(), gr.update(), status_md, gr.update(), history)
def run_single(image, prompt, threshold, show_boxes, history):
history = history or []
if image is None:
return _noop(info_status_html("Upload an image to start."), history)
prompt = (prompt or "").strip()
if not prompt:
return _noop(info_status_html("Type a prompt or pick an example."), history)
try:
masks, boxes, scores, ms, mid, dev = _infer_single(image, prompt, threshold)
except Exception as e:
return _noop(info_status_html(_friendly_error(e)), history)
overlay, mask_only = render(image, masks, boxes, scores, prompt,
show_boxes=show_boxes)
count = 0 if scores is None else int(len(scores))
status = status_html(count, ms, mid, dev) if count else empty_status_html(prompt)
history = ([(overlay, f"{prompt} · {count}")] + history)[:12]
return overlay, mask_only, image, status, history, history
def run_many(image, multi_text, threshold, show_boxes, history):
history = history or []
if image is None:
return _noop(info_status_html("Upload an image to start."), history)
prompts, seen = [], set()
for chunk in (multi_text or "").replace(",", "\n").splitlines():
p = chunk.strip()
if p and p.lower() not in seen:
prompts.append(p)
seen.add(p.lower())
prompts = prompts[:6]
if not prompts:
return _noop(info_status_html("Add one prompt per line first."), history)
try:
results, ms, mid, dev = _infer_many(image, prompts, threshold)
except Exception as e:
return _noop(info_status_html(_friendly_error(e)), history)
first_overlay = first_mask = None
total = 0
new_entries = []
for idx, (prompt, masks, boxes, scores) in enumerate(results):
overlay, mask_only = render(image, masks, boxes, scores, prompt,
show_boxes=show_boxes)
count = 0 if scores is None else int(len(scores))
total += count
new_entries.append((overlay, f"{prompt} · {count}"))
if idx == 0:
first_overlay, first_mask = overlay, mask_only
history = (new_entries + history)[:12]
status = (
"<div class='status'>"
f"<span class='chip chip-count'>{len(prompts)} prompts · {total} matches</span>"
f"<span class='chip chip-ms'>{ms:.0f} ms total</span>"
f"<span class='chip chip-dim'>{mid} · {dev} · vision features reused</span>"
"</div>"
)
return first_overlay, first_mask, image, status, history, history
def fill_prompt(choice):
return choice or ""
def reset_all():
return (
None, # image
"", # prompt
None, # example dropdown
None, # overlay
None, # mask only
None, # original
IDLE_STATUS, # status
)
def clear_history():
return [], []
# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@500;600;700&display=swap');
.gradio-container { max-width: 1280px !important; }
#app-header { display:flex; align-items:center; gap:.65rem; margin:.2rem 0 0; }
#app-header .logo {
width:34px; height:34px; border-radius:9px;
background:linear-gradient(135deg,#5145E5,#00B3A4);
box-shadow:0 2px 10px rgba(81,69,229,.35);
}
#app-header h1 {
font-family:'Space Grotesk', Inter, system-ui, sans-serif;
font-size:1.55rem; font-weight:700; letter-spacing:-0.015em; margin:0;
}
#app-sub { color:#5B6472; margin:.15rem 0 0; font-size:.96rem; }
#key-banner {
margin:.5rem 0 1rem; padding:.7rem 1rem; border-radius:12px; color:#fff;
background:linear-gradient(90deg,#5145E5,#00B3A4); font-weight:600;
display:flex; gap:.6rem; align-items:center; line-height:1.3;
}
#key-banner .dot {
width:8px; height:8px; border-radius:50%; background:#fff;
box-shadow:0 0 0 4px rgba(255,255,255,.28); flex:none;
}
.status { display:flex; gap:.4rem; flex-wrap:wrap; align-items:center; min-height:34px; }
.chip { font-size:.8rem; padding:.2rem .6rem; border-radius:999px; font-weight:600;
white-space:nowrap; }
.chip-count { background:#ECEAFE; color:#3F33CF; }
.chip-ms { background:#E1F6F2; color:#00897B; }
.chip-dim { background:#F0F2F5; color:#5B6472; font-weight:500; }
.chip-empty { background:#FFF4E5; color:#B26A00; }
.fade img { animation: sam-fade .45s ease both; }
@keyframes sam-fade { from { opacity:0; transform:scale(.992); } to { opacity:1; transform:none; } }
@media (prefers-reduced-motion: reduce) { .fade img { animation:none; } }
"""
THEME = gr.themes.Soft(
primary_hue=gr.themes.colors.indigo,
secondary_hue=gr.themes.colors.teal,
neutral_hue=gr.themes.colors.slate,
font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
)
def build_demo():
with gr.Blocks(theme=THEME, css=CSS, title="SAM 3.1 · Concept Segmentation") as demo:
history_state = gr.State([])
gr.HTML(
'<div id="app-header"><div class="logo"></div>'
"<div><h1>SAM 3.1 · Concept Segmentation</h1></div></div>"
'<p id="app-sub">Type what you want to find — the model segments every '
"matching instance. No boxes, no clicks, no retraining.</p>"
)
gr.HTML(
f'<div id="key-banner"><span class="dot"></span><span>{KEY_MESSAGE}</span></div>'
)
with gr.Row(equal_height=False):
# ---------------- Inputs ----------------
with gr.Column(scale=5, min_width=360):
image_in = gr.Image(
type="pil",
label="Image",
sources=["upload", "clipboard"],
height=420,
elem_classes=["fade"],
)
prompt_tb = gr.Textbox(
label="Prompt",
placeholder="e.g. horse",
info="Short noun phrases work best, e.g. \u201chorse\u201d or \u201csaddle\u201d.",
autofocus=True,
)
example_dd = gr.Dropdown(
choices=EXAMPLE_PROMPTS,
label="Example prompts",
value=None,
interactive=True,
)
with gr.Row():
run_btn = gr.Button("Run segmentation", variant="primary", scale=3)
reset_btn = gr.Button("Reset", variant="secondary", scale=1)
status = gr.HTML(IDLE_STATUS)
with gr.Accordion("Advanced", open=False):
threshold = gr.Slider(
minimum=0.05, maximum=0.95, value=0.5, step=0.05,
label="Confidence threshold",
info="Lower to reveal more instances; higher to keep only strong matches.",
)
show_boxes = gr.Checkbox(
value=False, label="Show boxes & confidence scores"
)
gr.Markdown(
"**Multiple prompts** — one per line. They share a single vision "
"pass, so adding prompts is fast."
)
multi_tb = gr.Textbox(
label="Prompts (one per line)",
placeholder="horse\nsaddle\nperson",
lines=3,
)
run_many_btn = gr.Button("Run all prompts", variant="secondary")
# ---------------- Outputs (the hero) ----------------
with gr.Column(scale=7, min_width=420):
with gr.Tabs():
with gr.Tab("Overlay"):
overlay_out = gr.Image(
label=None, height=540, interactive=False,
show_label=False, elem_classes=["fade"],
)
with gr.Tab("Mask only"):
mask_out = gr.Image(
label=None, height=540, interactive=False,
show_label=False, elem_classes=["fade"],
)
with gr.Tab("Original"):
original_out = gr.Image(
label=None, height=540, interactive=False,
show_label=False, elem_classes=["fade"],
)
with gr.Accordion("Prompt history", open=False):
history_gallery = gr.Gallery(
label=None, show_label=False, columns=4, height=240,
object_fit="cover", preview=False,
)
clear_btn = gr.Button("Clear history", variant="secondary", size="sm")
# Optional bundled examples (image + prompt). Lights up only if files exist,
# so the Space runs fine without any image assets checked in.
ex_dir = "examples"
ex_pairs = []
if os.path.isdir(ex_dir):
for fn, pr in [("horse.jpg", "horse"), ("street.jpg", "person"),
("kitchen.jpg", "handle")]:
p = os.path.join(ex_dir, fn)
if os.path.exists(p):
ex_pairs.append([p, pr])
if ex_pairs:
gr.Examples(examples=ex_pairs, inputs=[image_in, prompt_tb],
label="Try an example")
gr.Markdown(
"<sub>Built on Meta's Segment Anything Model 3.1 (Promptable Concept "
"Segmentation). SAM 3 / 3.1 weights are gated on Hugging Face. "
"Very descriptive phrases are less reliable than short nouns — for the "
"reins, \u201cbridle\u201d or \u201creins\u201d will usually beat "
"\u201cobject used for riding control\u201d.</sub>"
)
# ----- wiring -----
out_targets = [overlay_out, mask_out, original_out, status,
history_gallery, history_state]
run_btn.click(
run_single,
inputs=[image_in, prompt_tb, threshold, show_boxes, history_state],
outputs=out_targets,
)
prompt_tb.submit(
run_single,
inputs=[image_in, prompt_tb, threshold, show_boxes, history_state],
outputs=out_targets,
)
run_many_btn.click(
run_many,
inputs=[image_in, multi_tb, threshold, show_boxes, history_state],
outputs=out_targets,
)
example_dd.change(fill_prompt, inputs=example_dd, outputs=prompt_tb)
reset_btn.click(
reset_all,
outputs=[image_in, prompt_tb, example_dd, overlay_out, mask_out,
original_out, status],
)
clear_btn.click(clear_history, outputs=[history_gallery, history_state])
return demo
if __name__ == "__main__":
demo = build_demo()
demo.queue(max_size=20).launch()