Alirezakzt/Concept-Segmentation
4
1"""2SAM 3.1 — Promptable Concept Segmentation demo3================================================4A live, language-driven segmentation demo built on Meta's Segment Anything Model 3.1.5 6Type a short noun phrase (e.g. "horse", "saddle"); the model finds and segments7*every* matching instance in the image. No boxes, no clicks, no retraining.8 9Model: facebook/sam3.1 (image Promptable Concept Segmentation path)10Runtime: Hugging Face Spaces — works on a standard GPU Space or on ZeroGPU.11 12Deploy notes are in README.md (gated-model access + HF_TOKEN + hardware).13"""14 15import os16import time17import colorsys18from contextlib import nullcontext19 20import numpy as np21from PIL import Image, ImageDraw, ImageFont22import torch23import gradio as gr24 25# --------------------------------------------------------------------------------------26# ZeroGPU support (optional). On a standard GPU Space this becomes a transparent no-op.27# --------------------------------------------------------------------------------------28try:29 import spaces # provided by the ZeroGPU runtime30 31 GPU = spaces.GPU32except Exception: # not on Spaces / package missing → identity decorator33 34 def GPU(*args, **kwargs):35 # Supports both `@GPU` and `@GPU(duration=...)`36 if len(args) == 1 and callable(args[0]) and not kwargs:37 return args[0]38 39 def _deco(fn):40 return fn41 42 return _deco43 44 45# --------------------------------------------------------------------------------------46# Configuration47# --------------------------------------------------------------------------------------48MODEL_ID = os.environ.get("MODEL_ID", "facebook/sam3.1")49FALLBACK_MODEL_ID = os.environ.get("FALLBACK_MODEL_ID", "facebook/sam3")50HF_TOKEN = (51 os.environ.get("HF_TOKEN")52 or os.environ.get("HUGGING_FACE_HUB_TOKEN")53 or os.environ.get("HUGGINGFACE_TOKEN")54)55 56# ZeroGPU sets SPACES_ZERO_GPU; in that case CUDA is attached only inside @GPU calls,57# but we can still target "cuda" because the `spaces` runtime patches device placement.58_ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU"))59DEVICE = "cuda" if (torch.cuda.is_available() or _ZERO_GPU) else "cpu"60 61EXAMPLE_PROMPTS = [62 "horse",63 "saddle",64 "person",65 "object used for riding control",66]67 68KEY_MESSAGE = "Segmentation is fully driven by language prompts — no retraining required."69 70# Lazily-loaded singletons71_MODEL = None72_PROCESSOR = None73_LOADED_ID = None74 75 76# --------------------------------------------------------------------------------------77# Model loading78# --------------------------------------------------------------------------------------79def load_model():80 """Load SAM 3.1 (image PCS) once. Falls back to SAM 3 if 3.1 is unavailable."""81 global _MODEL, _PROCESSOR, _LOADED_ID82 if _MODEL is not None:83 return84 85 from transformers import Sam3Model, Sam3Processor86 87 candidates = [MODEL_ID]88 if FALLBACK_MODEL_ID and FALLBACK_MODEL_ID != MODEL_ID:89 candidates.append(FALLBACK_MODEL_ID)90 91 last_err = None92 for mid in candidates:93 try:94 processor = Sam3Processor.from_pretrained(mid, token=HF_TOKEN)95 model = Sam3Model.from_pretrained(mid, token=HF_TOKEN)96 model.eval()97 try:98 model.to(DEVICE)99 except Exception:100 # On ZeroGPU the move is handled when the GPU is attached; ignore here.101 pass102 _MODEL, _PROCESSOR, _LOADED_ID = model, processor, mid103 if mid != MODEL_ID:104 print(f"[sam3.1-demo] '{MODEL_ID}' unavailable; loaded fallback '{mid}'.")105 else:106 print(f"[sam3.1-demo] Loaded '{mid}' on {DEVICE}.")107 return108 except Exception as e: # try next candidate109 last_err = e110 print(f"[sam3.1-demo] Could not load '{mid}': {e}")111 112 raise RuntimeError(113 f"Failed to load any SAM 3 model from {candidates}. Last error: {last_err}"114 )115 116 117def _friendly_error(err: Exception) -> str:118 """Turn a load/inference exception into actionable guidance."""119 text = str(err).lower()120 gated = any(k in text for k in ["401", "403", "gated", "access", "token", "authorized"])121 if gated:122 return (123 "Couldn't access the model weights. SAM 3 / 3.1 are gated: request access on the "124 "Hugging Face model page, then add your token as a Space secret named "125 "<b>HF_TOKEN</b> (Settings → Variables and secrets), and restart the Space."126 )127 return f"Something went wrong while running the model: {err}"128 129 130# --------------------------------------------------------------------------------------131# Inference (GPU-scoped). Everything returned here is CPU/NumPy so it stays valid132# after the GPU is released (important for ZeroGPU).133# --------------------------------------------------------------------------------------134def _amp_ctx():135 """bfloat16 autocast on CUDA for speed; pass-through elsewhere."""136 if DEVICE == "cuda":137 return torch.autocast("cuda", dtype=torch.bfloat16)138 return nullcontext()139 140 141def _to_np(x):142 if x is None:143 return None144 if hasattr(x, "detach"):145 return x.detach().to("cpu").float().numpy()146 if isinstance(x, np.ndarray):147 return x148 if isinstance(x, (list, tuple)):149 if len(x) == 0:150 return np.zeros((0,))151 if hasattr(x[0], "detach"):152 return np.stack([t.detach().to("cpu").float().numpy() for t in x])153 return np.asarray(x)154 return np.asarray(x)155 156 157def _postprocess(outputs, target_sizes, threshold):158 res = _PROCESSOR.post_process_instance_segmentation(159 outputs,160 threshold=float(threshold),161 mask_threshold=0.5,162 target_sizes=target_sizes,163 )[0]164 masks = _to_np(res.get("masks"))165 boxes = _to_np(res.get("boxes"))166 scores = _to_np(res.get("scores"))167 return masks, boxes, scores168 169 170@GPU(duration=120)171def _infer_single(image: Image.Image, prompt: str, threshold: float):172 """Segment one text prompt on one image. Returns CPU arrays + metadata."""173 load_model()174 inputs = _PROCESSOR(images=image, text=prompt, return_tensors="pt").to(_MODEL.device)175 target_sizes = inputs["original_sizes"].tolist()176 177 t0 = time.perf_counter()178 with torch.no_grad():179 try:180 with _amp_ctx():181 outputs = _MODEL(**inputs)182 except RuntimeError:183 outputs = _MODEL(**inputs) # rare: fall back to full precision184 if DEVICE == "cuda":185 torch.cuda.synchronize()186 ms = (time.perf_counter() - t0) * 1000.0187 188 masks, boxes, scores = _postprocess(outputs, target_sizes, threshold)189 return masks, boxes, scores, ms, _LOADED_ID, _MODEL.device.type190 191 192@GPU(duration=180)193def _infer_many(image: Image.Image, prompts, threshold: float):194 """Segment several prompts on one image, reusing vision features for speed.195 196 The whole batch runs inside a single GPU call, so the cached vision embeddings197 stay valid (safe on ZeroGPU).198 """199 load_model()200 img_inputs = _PROCESSOR(images=image, return_tensors="pt").to(_MODEL.device)201 target_sizes = img_inputs["original_sizes"].tolist()202 203 results = []204 t0 = time.perf_counter()205 with torch.no_grad():206 with _amp_ctx():207 vision_embeds = _MODEL.get_vision_features(pixel_values=img_inputs.pixel_values)208 for prompt in prompts:209 text_inputs = _PROCESSOR(text=prompt, return_tensors="pt").to(_MODEL.device)210 with _amp_ctx():211 outputs = _MODEL(vision_embeds=vision_embeds, **text_inputs)212 masks, boxes, scores = _postprocess(outputs, target_sizes, threshold)213 results.append((prompt, masks, boxes, scores))214 if DEVICE == "cuda":215 torch.cuda.synchronize()216 ms = (time.perf_counter() - t0) * 1000.0217 return results, ms, _LOADED_ID, _MODEL.device.type218 219 220# --------------------------------------------------------------------------------------221# Rendering (CPU). Builds the semi-transparent overlay and the mask-only view.222# --------------------------------------------------------------------------------------223def _palette(n: int):224 """Evenly-spaced, vivid colors (golden-ratio hue spacing) — one per instance."""225 cols = []226 for i in range(max(n, 1)):227 h = (i * 0.61803398875) % 1.0228 r, g, b = colorsys.hsv_to_rgb(h, 0.72, 1.0)229 cols.append((int(r * 255), int(g * 255), int(b * 255)))230 return cols231 232 233def _mask_list(masks, h, w):234 """Normalize whatever the model returned into a list of HxW bool arrays."""235 out = []236 if masks is None:237 return out238 arr = masks239 if arr.ndim == 2:240 arr = arr[None, ...]241 for i in range(arr.shape[0]):242 m = arr[i]243 if m.ndim == 3:244 m = m[0]245 m = m > 0.5246 if m.shape[:2] != (h, w):247 m = (248 np.asarray(249 Image.fromarray((m.astype(np.uint8) * 255)).resize(250 (w, h), Image.NEAREST251 )252 )253 > 127254 )255 out.append(m)256 return out257 258 259def _boundary(mask: np.ndarray) -> np.ndarray:260 """1px boundary via 4-neighbour erosion (no SciPy/OpenCV dependency)."""261 e = mask.copy()262 e[1:, :] &= mask[:-1, :]263 e[:-1, :] &= mask[1:, :]264 e[:, 1:] &= mask[:, :-1]265 e[:, :-1] &= mask[:, 1:]266 return mask & ~e267 268 269def _font(size: int):270 for path in (271 "DejaVuSans-Bold.ttf",272 "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",273 "DejaVuSans.ttf",274 "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",275 ):276 try:277 return ImageFont.truetype(path, size)278 except Exception:279 continue280 return ImageFont.load_default()281 282 283def render(image: Image.Image, masks, boxes, scores, prompt: str,284 alpha: float = 0.5, show_boxes: bool = False):285 """Return (overlay_image, mask_only_image)."""286 base = np.asarray(image.convert("RGB")).astype(np.float32)287 h, w = base.shape[:2]288 289 mlist = _mask_list(masks, h, w)290 cols = _palette(len(mlist))291 292 overlay = base.copy()293 mask_only = np.zeros_like(base)294 295 for i, m in enumerate(mlist):296 c = np.array(cols[i], dtype=np.float32)297 overlay[m] = overlay[m] * (1.0 - alpha) + c * alpha298 edge = _boundary(m)299 overlay[edge] = c # crisp instance outline300 mask_only[m] = c301 mask_only[edge] = np.minimum(c + 70, 255)302 303 overlay_img = Image.fromarray(overlay.clip(0, 255).astype(np.uint8))304 mask_only_img = Image.fromarray(mask_only.astype(np.uint8))305 306 if show_boxes and boxes is not None and len(boxes) and scores is not None:307 draw = ImageDraw.Draw(overlay_img, "RGBA")308 fsize = max(13, int(w / 55))309 font = _font(fsize)310 line_w = max(2, int(w / 480))311 for i in range(min(len(boxes), len(mlist) or len(boxes))):312 x1, y1, x2, y2 = [float(v) for v in boxes[i][:4]]313 c = cols[i % len(cols)]314 draw.rectangle([x1, y1, x2, y2], outline=c + (255,), width=line_w)315 label = f"{prompt} · {float(scores[i]):.2f}"316 tb = draw.textbbox((0, 0), label, font=font)317 tw, th = tb[2] - tb[0], tb[3] - tb[1]318 ty = max(0, y1 - th - 6)319 draw.rectangle([x1, ty, x1 + tw + 10, ty + th + 6], fill=c + (235,))320 draw.text((x1 + 5, ty + 3), label, fill=(20, 24, 31, 255), font=font)321 322 return overlay_img, mask_only_img323 324 325# --------------------------------------------------------------------------------------326# Status banner HTML327# --------------------------------------------------------------------------------------328def status_html(count: int, ms: float, model_id: str, device: str) -> str:329 plural = "" if count == 1 else "es"330 return (331 "<div class='status'>"332 f"<span class='chip chip-count'>{count} match{plural}</span>"333 f"<span class='chip chip-ms'>{ms:.0f} ms</span>"334 f"<span class='chip chip-dim'>{model_id} · {device}</span>"335 "</div>"336 )337 338 339def empty_status_html(prompt: str) -> str:340 return (341 "<div class='status'>"342 f"<span class='chip chip-empty'>No matches for “{prompt}”</span>"343 "<span class='chip chip-dim'>Try a simpler noun, or lower the threshold</span>"344 "</div>"345 )346 347 348def info_status_html(message: str) -> str:349 return f"<div class='status'><span class='chip chip-empty'>{message}</span></div>"350 351 352IDLE_STATUS = (353 "<div class='status'><span class='chip chip-dim'>"354 "Upload an image, type a prompt, then run.</span></div>"355)356 357 358# --------------------------------------------------------------------------------------359# Gradio callbacks360# --------------------------------------------------------------------------------------361def _noop(status_md, history):362 # leave images & gallery untouched363 return (gr.update(), gr.update(), gr.update(), status_md, gr.update(), history)364 365 366def run_single(image, prompt, threshold, show_boxes, history):367 history = history or []368 if image is None:369 return _noop(info_status_html("Upload an image to start."), history)370 prompt = (prompt or "").strip()371 if not prompt:372 return _noop(info_status_html("Type a prompt or pick an example."), history)373 374 try:375 masks, boxes, scores, ms, mid, dev = _infer_single(image, prompt, threshold)376 except Exception as e:377 return _noop(info_status_html(_friendly_error(e)), history)378 379 overlay, mask_only = render(image, masks, boxes, scores, prompt,380 show_boxes=show_boxes)381 count = 0 if scores is None else int(len(scores))382 status = status_html(count, ms, mid, dev) if count else empty_status_html(prompt)383 384 history = ([(overlay, f"{prompt} · {count}")] + history)[:12]385 return overlay, mask_only, image, status, history, history386 387 388def run_many(image, multi_text, threshold, show_boxes, history):389 history = history or []390 if image is None:391 return _noop(info_status_html("Upload an image to start."), history)392 393 prompts, seen = [], set()394 for chunk in (multi_text or "").replace(",", "\n").splitlines():395 p = chunk.strip()396 if p and p.lower() not in seen:397 prompts.append(p)398 seen.add(p.lower())399 prompts = prompts[:6]400 if not prompts:401 return _noop(info_status_html("Add one prompt per line first."), history)402 403 try:404 results, ms, mid, dev = _infer_many(image, prompts, threshold)405 except Exception as e:406 return _noop(info_status_html(_friendly_error(e)), history)407 408 first_overlay = first_mask = None409 total = 0410 new_entries = []411 for idx, (prompt, masks, boxes, scores) in enumerate(results):412 overlay, mask_only = render(image, masks, boxes, scores, prompt,413 show_boxes=show_boxes)414 count = 0 if scores is None else int(len(scores))415 total += count416 new_entries.append((overlay, f"{prompt} · {count}"))417 if idx == 0:418 first_overlay, first_mask = overlay, mask_only419 420 history = (new_entries + history)[:12]421 status = (422 "<div class='status'>"423 f"<span class='chip chip-count'>{len(prompts)} prompts · {total} matches</span>"424 f"<span class='chip chip-ms'>{ms:.0f} ms total</span>"425 f"<span class='chip chip-dim'>{mid} · {dev} · vision features reused</span>"426 "</div>"427 )428 return first_overlay, first_mask, image, status, history, history429 430 431def fill_prompt(choice):432 return choice or ""433 434 435def reset_all():436 return (437 None, # image438 "", # prompt439 None, # example dropdown440 None, # overlay441 None, # mask only442 None, # original443 IDLE_STATUS, # status444 )445 446 447def clear_history():448 return [], []449 450 451# --------------------------------------------------------------------------------------452# UI453# --------------------------------------------------------------------------------------454CSS = """455@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@500;600;700&display=swap');456 457.gradio-container { max-width: 1280px !important; }458 459#app-header { display:flex; align-items:center; gap:.65rem; margin:.2rem 0 0; }460#app-header .logo {461 width:34px; height:34px; border-radius:9px;462 background:linear-gradient(135deg,#5145E5,#00B3A4);463 box-shadow:0 2px 10px rgba(81,69,229,.35);464}465#app-header h1 {466 font-family:'Space Grotesk', Inter, system-ui, sans-serif;467 font-size:1.55rem; font-weight:700; letter-spacing:-0.015em; margin:0;468}469#app-sub { color:#5B6472; margin:.15rem 0 0; font-size:.96rem; }470 471#key-banner {472 margin:.5rem 0 1rem; padding:.7rem 1rem; border-radius:12px; color:#fff;473 background:linear-gradient(90deg,#5145E5,#00B3A4); font-weight:600;474 display:flex; gap:.6rem; align-items:center; line-height:1.3;475}476#key-banner .dot {477 width:8px; height:8px; border-radius:50%; background:#fff;478 box-shadow:0 0 0 4px rgba(255,255,255,.28); flex:none;479}480 481.status { display:flex; gap:.4rem; flex-wrap:wrap; align-items:center; min-height:34px; }482.chip { font-size:.8rem; padding:.2rem .6rem; border-radius:999px; font-weight:600;483 white-space:nowrap; }484.chip-count { background:#ECEAFE; color:#3F33CF; }485.chip-ms { background:#E1F6F2; color:#00897B; }486.chip-dim { background:#F0F2F5; color:#5B6472; font-weight:500; }487.chip-empty { background:#FFF4E5; color:#B26A00; }488 489.fade img { animation: sam-fade .45s ease both; }490@keyframes sam-fade { from { opacity:0; transform:scale(.992); } to { opacity:1; transform:none; } }491@media (prefers-reduced-motion: reduce) { .fade img { animation:none; } }492"""493 494THEME = gr.themes.Soft(495 primary_hue=gr.themes.colors.indigo,496 secondary_hue=gr.themes.colors.teal,497 neutral_hue=gr.themes.colors.slate,498 font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],499)500 501 502def build_demo():503 with gr.Blocks(theme=THEME, css=CSS, title="SAM 3.1 · Concept Segmentation") as demo:504 history_state = gr.State([])505 506 gr.HTML(507 '<div id="app-header"><div class="logo"></div>'508 "<div><h1>SAM 3.1 · Concept Segmentation</h1></div></div>"509 '<p id="app-sub">Type what you want to find — the model segments every '510 "matching instance. No boxes, no clicks, no retraining.</p>"511 )512 gr.HTML(513 f'<div id="key-banner"><span class="dot"></span><span>{KEY_MESSAGE}</span></div>'514 )515 516 with gr.Row(equal_height=False):517 # ---------------- Inputs ----------------518 with gr.Column(scale=5, min_width=360):519 image_in = gr.Image(520 type="pil",521 label="Image",522 sources=["upload", "clipboard"],523 height=420,524 elem_classes=["fade"],525 )526 prompt_tb = gr.Textbox(527 label="Prompt",528 placeholder="e.g. horse",529 info="Short noun phrases work best, e.g. \u201chorse\u201d or \u201csaddle\u201d.",530 autofocus=True,531 )532 example_dd = gr.Dropdown(533 choices=EXAMPLE_PROMPTS,534 label="Example prompts",535 value=None,536 interactive=True,537 )538 with gr.Row():539 run_btn = gr.Button("Run segmentation", variant="primary", scale=3)540 reset_btn = gr.Button("Reset", variant="secondary", scale=1)541 542 status = gr.HTML(IDLE_STATUS)543 544 with gr.Accordion("Advanced", open=False):545 threshold = gr.Slider(546 minimum=0.05, maximum=0.95, value=0.5, step=0.05,547 label="Confidence threshold",548 info="Lower to reveal more instances; higher to keep only strong matches.",549 )550 show_boxes = gr.Checkbox(551 value=False, label="Show boxes & confidence scores"552 )553 gr.Markdown(554 "**Multiple prompts** — one per line. They share a single vision "555 "pass, so adding prompts is fast."556 )557 multi_tb = gr.Textbox(558 label="Prompts (one per line)",559 placeholder="horse\nsaddle\nperson",560 lines=3,561 )562 run_many_btn = gr.Button("Run all prompts", variant="secondary")563 564 # ---------------- Outputs (the hero) ----------------565 with gr.Column(scale=7, min_width=420):566 with gr.Tabs():567 with gr.Tab("Overlay"):568 overlay_out = gr.Image(569 label=None, height=540, interactive=False,570 show_label=False, elem_classes=["fade"],571 )572 with gr.Tab("Mask only"):573 mask_out = gr.Image(574 label=None, height=540, interactive=False,575 show_label=False, elem_classes=["fade"],576 )577 with gr.Tab("Original"):578 original_out = gr.Image(579 label=None, height=540, interactive=False,580 show_label=False, elem_classes=["fade"],581 )582 583 with gr.Accordion("Prompt history", open=False):584 history_gallery = gr.Gallery(585 label=None, show_label=False, columns=4, height=240,586 object_fit="cover", preview=False,587 )588 clear_btn = gr.Button("Clear history", variant="secondary", size="sm")589 590 # Optional bundled examples (image + prompt). Lights up only if files exist,591 # so the Space runs fine without any image assets checked in.592 ex_dir = "examples"593 ex_pairs = []594 if os.path.isdir(ex_dir):595 for fn, pr in [("horse.jpg", "horse"), ("street.jpg", "person"),596 ("kitchen.jpg", "handle")]:597 p = os.path.join(ex_dir, fn)598 if os.path.exists(p):599 ex_pairs.append([p, pr])600 if ex_pairs:601 gr.Examples(examples=ex_pairs, inputs=[image_in, prompt_tb],602 label="Try an example")603 604 gr.Markdown(605 "<sub>Built on Meta's Segment Anything Model 3.1 (Promptable Concept "606 "Segmentation). SAM 3 / 3.1 weights are gated on Hugging Face. "607 "Very descriptive phrases are less reliable than short nouns — for the "608 "reins, \u201cbridle\u201d or \u201creins\u201d will usually beat "609 "\u201cobject used for riding control\u201d.</sub>"610 )611 612 # ----- wiring -----613 out_targets = [overlay_out, mask_out, original_out, status,614 history_gallery, history_state]615 616 run_btn.click(617 run_single,618 inputs=[image_in, prompt_tb, threshold, show_boxes, history_state],619 outputs=out_targets,620 )621 prompt_tb.submit(622 run_single,623 inputs=[image_in, prompt_tb, threshold, show_boxes, history_state],624 outputs=out_targets,625 )626 run_many_btn.click(627 run_many,628 inputs=[image_in, multi_tb, threshold, show_boxes, history_state],629 outputs=out_targets,630 )631 example_dd.change(fill_prompt, inputs=example_dd, outputs=prompt_tb)632 reset_btn.click(633 reset_all,634 outputs=[image_in, prompt_tb, example_dd, overlay_out, mask_out,635 original_out, status],636 )637 clear_btn.click(clear_history, outputs=[history_gallery, history_state])638 639 return demo640 641 642if __name__ == "__main__":643 demo = build_demo()644 demo.queue(max_size=20).launch()645