""" SAM 3 / SAM 3.1 — Image / Video / Panorama / Multi-Class Segmentation Hugging Face Spaces app (ZeroGPU). Requires: - ZeroGPU hardware selected under Settings -> Hardware - An HF_TOKEN secret (Settings -> Variables and secrets) with access to the gated facebook/sam3 and facebook/sam3.1 checkpoints (request access at https://huggingface.co/facebook/sam3 and https://huggingface.co/facebook/sam3.1) """ import glob import hashlib import subprocess from datetime import datetime import os import cv2 import numpy as np import torch import gradio as gr import spaces from PIL import Image, ImageDraw from huggingface_hub import login, list_repo_files, hf_hub_download # ---------------------------------------------------------------------- # Auth (Space secret) -- gated checkpoint download needs this at startup. # ---------------------------------------------------------------------- HF_TOKEN = os.environ.get("hf_token") if HF_TOKEN: login(token=HF_TOKEN) print("Logged in to Hugging Face using HF_TOKEN secret.") else: print("WARNING: HF_TOKEN secret not set. Gated checkpoint download will fail.") print("Add it in this Space's Settings -> Variables and secrets.") os.makedirs("/tmp/output", exist_ok=True) IMAGE_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp", ".tif", ".tiff"} VIDEO_EXTS = {".mp4", ".mov", ".avi", ".mkv", ".webm"} MAX_CLASSES = 5 VERSION_REPOS = { "SAM 3.1": "facebook/sam3.1", "SAM 3": "facebook/sam3", } DEFAULT_VERSION = "SAM 3.1" def get_device(): """ZeroGPU only attaches a GPU inside a @spaces.GPU-decorated call, so this must be evaluated at call time, never at module import time.""" return "cuda" if torch.cuda.is_available() else "cpu" # ---------------------------------------------------------------------- # Lazy, version-aware model loaders # ---------------------------------------------------------------------- _model_cache = {} # version -> {"processor": ..., "predictor": ...} _current_version = {"value": None} # which version is currently loaded on GPU def _find_bpe_path(): """Locate the bundled BPE tokenizer file regardless of how sam3 was installed.""" import sam3 pkg_dir = os.path.dirname(sam3.__file__) if getattr(sam3, "__file__", None) else None search_roots = [pkg_dir] if pkg_dir else [] search_roots += ["/usr/local/lib", "/usr/lib", os.getcwd()] for root in search_roots: if not root: continue matches = glob.glob(os.path.join(root, "**", "bpe_simple_vocab_16e6.txt.gz"), recursive=True) if matches: return matches[0] raise FileNotFoundError("Could not locate bpe_simple_vocab_16e6.txt.gz -- check sam3 install.") BPE_PATH = _find_bpe_path() print(f"Using bpe_path: {BPE_PATH}") def _resolve_checkpoint_path(repo_id): """Lists the repo's actual files and picks the right weight file automatically, instead of assuming a filename that may differ between repos/versions.""" files = list_repo_files(repo_id) weight_files = [f for f in files if f.endswith(".pt") or f.endswith(".safetensors")] assert weight_files, f"No .pt/.safetensors checkpoint file found in {repo_id}" pt_files = [f for f in weight_files if f.endswith(".pt")] chosen = pt_files[0] if pt_files else weight_files[0] print(f"Resolved checkpoint file in {repo_id}: {chosen}") return hf_hub_download(repo_id=repo_id, filename=chosen) def _unload_current_model(): """Frees GPU memory from whichever version is currently loaded before switching.""" if _current_version["value"] is not None: old = _current_version["value"] _model_cache.pop(old, None) if torch.cuda.is_available(): torch.cuda.empty_cache() print(f"Unloaded {old} to free GPU memory.") def get_image_processor(version=DEFAULT_VERSION): if version not in VERSION_REPOS: raise ValueError(f"Unknown version '{version}', expected one of {list(VERSION_REPOS)}") if version not in _model_cache or "processor" not in _model_cache.get(version, {}): if _current_version["value"] is not None and _current_version["value"] != version: _unload_current_model() from sam3.model_builder import build_sam3_image_model from sam3.model.sam3_image_processor import Sam3Processor repo_id = VERSION_REPOS[version] print(f"Loading {version} image model from {repo_id} (downloads/caches on first call)...") ckpt_path = _resolve_checkpoint_path(repo_id) model = build_sam3_image_model(bpe_path=BPE_PATH, checkpoint_path=ckpt_path, device=get_device()) model = model.to(get_device()) _model_cache.setdefault(version, {})["processor"] = Sam3Processor(model) _current_version["value"] = version print(f"{version} image model ready.") return _model_cache[version]["processor"] def get_video_predictor(version=DEFAULT_VERSION): if version not in VERSION_REPOS: raise ValueError(f"Unknown version '{version}', expected one of {list(VERSION_REPOS)}") if version not in _model_cache or "predictor" not in _model_cache.get(version, {}): if _current_version["value"] is not None and _current_version["value"] != version: _unload_current_model() from sam3.model_builder import build_sam3_video_predictor repo_id = VERSION_REPOS[version] print(f"Loading {version} video predictor from {repo_id} (downloads/caches on first call)...") ckpt_path = _resolve_checkpoint_path(repo_id) predictor = build_sam3_video_predictor(bpe_path=BPE_PATH, checkpoint_path=ckpt_path, async_loading_frames=True) _model_cache.setdefault(version, {})["predictor"] = predictor _current_version["value"] = version print(f"{version} video predictor ready.") return _model_cache[version]["predictor"] # ---------------------------------------------------------------------- # Video repair (fixes "moov atom not found" / truncated uploads) # ---------------------------------------------------------------------- def ensure_readable_video(file_path): """cv2/SAM3 can fail on MP4s with a missing/misplaced moov atom (common with large or truncated uploads). Remux through ffmpeg first -- this fixes most container-level issues without re-encoding. Raises a clear gr.Error if the file is genuinely corrupted.""" fixed_path = file_path + "_fixed.mp4" try: result = subprocess.run( ["ffmpeg", "-y", "-v", "error", "-i", file_path, "-c", "copy", "-movflags", "+faststart", fixed_path], capture_output=True, text=True, timeout=180, ) except FileNotFoundError: raise gr.Error("ffmpeg is not installed on this Space -- check packages.txt includes 'ffmpeg'.") except subprocess.TimeoutExpired: raise gr.Error("Video repair step timed out -- the file may be too large or genuinely corrupted.") if result.returncode != 0 or not os.path.exists(fixed_path) or os.path.getsize(fixed_path) == 0: stderr_tail = result.stderr[-500:] if result and result.stderr else "unknown error" raise gr.Error( "This video file appears to be corrupted or incompletely uploaded " "(missing video index / moov atom). Try re-exporting or re-uploading it. " f"ffmpeg said: {stderr_tail}" ) return fixed_path # ---------------------------------------------------------------------- # Helpers: masks, panorama tiling, multi-class, video propagation # ---------------------------------------------------------------------- def _to_numpy_masks(masks): return masks.detach().cpu().numpy() if hasattr(masks, "detach") else np.array(masks) def is_panorama(image, ratio_threshold=2.5): w, h = image.size return (w / h) >= ratio_threshold def segment_panorama(processor, image, prompt, tile_overlap=0.25, min_tile_w=768): w, h = image.size tile_w = max(min_tile_w, h) tile_w = min(tile_w, w) stride = max(int(tile_w * (1 - tile_overlap)), 1) x_starts = list(range(0, max(w - tile_w, 0) + 1, stride)) if not x_starts: x_starts = [0] if x_starts[-1] + tile_w < w: x_starts.append(w - tile_w) combined_mask = np.zeros((h, w), dtype=bool) all_boxes, all_scores = [], [] for x0 in x_starts: x1 = min(x0 + tile_w, w) tile = image.crop((x0, 0, x1, h)) with torch.inference_mode(), torch.autocast(device_type=get_device(), dtype=torch.bfloat16): state = processor.set_image(tile) out = processor.set_text_prompt(state=state, prompt=prompt) masks_np = _to_numpy_masks(out["masks"]) boxes, scores = out["boxes"], out["scores"] for i in range(masks_np.shape[0]): m = masks_np[i] if m.ndim == 3: m = m[0] m = m.astype(bool) combined_mask[0:h, x0:x1] |= m[:, : (x1 - x0)] for b in boxes: bx = b.tolist() if hasattr(b, "tolist") else list(b) bx[0] += x0 bx[2] += x0 all_boxes.append(bx) for s in scores: all_scores.append(float(s)) return combined_mask, all_boxes, all_scores def parse_class_prompt(prompt, max_classes=MAX_CLASSES): class_names = [p.strip() for p in prompt.split(",") if p.strip()] class_names = list(dict.fromkeys(class_names)) if not class_names: raise ValueError("Please enter at least one class.") if len(class_names) > max_classes: raise ValueError(f"Please enter at most {max_classes} classes (you entered {len(class_names)}).") return class_names def get_class_colors(class_names): colors = {} for name in class_names: h = hashlib.md5(name.encode()).hexdigest() r, g, b = int(h[0:2], 16), int(h[2:4], 16), int(h[4:6], 16) colors[name] = tuple(max(80, c) for c in (r, g, b)) return colors def draw_legend(image, class_colors): img = image.copy() draw = ImageDraw.Draw(img) pad, box, line_h = 8, 14, 22 w = max(120, max(len(c) for c in class_colors) * 8 + 40) h = line_h * len(class_colors) + pad * 2 draw.rectangle([4, 4, 4 + w, 4 + h], fill=(0, 0, 0)) for i, (name, color) in enumerate(class_colors.items()): y = 4 + pad + i * line_h draw.rectangle([12, y, 12 + box, y + box], fill=color) draw.text((12 + box + 8, y - 2), name, fill=(255, 255, 255)) return img def segment_multiclass_image(processor, image, class_names, panorama=False): class_colors = get_class_colors(class_names) img_np = np.array(image.convert("RGB")).astype(np.uint8) overlay = img_np.copy() alpha = 0.55 per_class_counts = {} if not panorama: with torch.inference_mode(), torch.autocast(device_type=get_device(), dtype=torch.bfloat16): state = processor.set_image(image) for name in class_names: color = np.array(class_colors[name], dtype=np.float32) if panorama: combined_mask, boxes, scores = segment_panorama(processor, image, name) per_class_counts[name] = len(boxes) overlay[combined_mask] = ( overlay[combined_mask].astype(np.float32) * (1 - alpha) + color * alpha ).astype(np.uint8) else: with torch.inference_mode(), torch.autocast(device_type=get_device(), dtype=torch.bfloat16): out = processor.set_text_prompt(state=state, prompt=name) masks_np = _to_numpy_masks(out["masks"]) per_class_counts[name] = masks_np.shape[0] for i in range(masks_np.shape[0]): m = masks_np[i] if m.ndim == 3: m = m[0] m = m.astype(bool) overlay[m] = (overlay[m].astype(np.float32) * (1 - alpha) + color * alpha).astype(np.uint8) result = Image.fromarray(overlay) result = draw_legend(result, class_colors) return result, per_class_counts def propagate_in_video(video_predictor, session_id): outputs_per_frame = {} for frame_output in video_predictor.handle_stream_request( request=dict(type="propagate_in_video", session_id=session_id, propagation_direction="both") ): idx = frame_output["frame_index"] outputs_per_frame[idx] = frame_output["outputs"] return outputs_per_frame # ---------------------------------------------------------------------- # Main inference entrypoint # ---------------------------------------------------------------------- # ZeroGPU only grants a GPU for the duration of a decorated call. duration is the max # seconds this call is allowed to hold the GPU -- raise it if longer videos time out # (check your account's max allowed duration in the ZeroGPU docs if this errors). @spaces.GPU(duration=120) def run_sam3(file_path, prompt, model_version, progress=gr.Progress()): if file_path is None: raise gr.Error("Please upload an image or video first.") if not prompt or not prompt.strip(): raise gr.Error("Please type at least one class (e.g. 'person, dog, car').") try: class_names = parse_class_prompt(prompt, max_classes=MAX_CLASSES) except ValueError as e: raise gr.Error(str(e)) ext = os.path.splitext(file_path)[1].lower() timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") if ext in IMAGE_EXTS: progress(0.1, desc=f"Loading {model_version} image model...") processor = get_image_processor(model_version) image = Image.open(file_path).convert("RGB") panorama = is_panorama(image) progress(0.3, desc=f"Segmenting {len(class_names)} class(es)" + (" (panorama tiling)..." if panorama else "...")) result_image, per_class_counts = segment_multiclass_image(processor, image, class_names, panorama=panorama) status = f"[{model_version}] Found: " + ", ".join(f"{k}={v}" for k, v in per_class_counts.items()) progress(0.9, desc="Saving output...") out_path = f"/tmp/output/sam3_overlay_{timestamp}.png" result_image.save(out_path) progress(1.0, desc="Done!") return ( gr.update(value=result_image, visible=True), gr.update(visible=False), gr.update(value=out_path, visible=True), status, ) elif ext in VIDEO_EXTS: progress(0.05, desc="Checking video file...") file_path = ensure_readable_video(file_path) progress(0.15, desc=f"Loading {model_version} video predictor...") from sam3.visualization_utils import save_masklet_video video_predictor = get_video_predictor(model_version) video_prompt = class_names[0] note = "" if len(class_names) > 1: note = f" (only first class '{video_prompt}' is tracked for video)" progress(0.25, desc="Reading video frames...") cap = cv2.VideoCapture(file_path) video_fps = cap.get(cv2.CAP_PROP_FPS) or 24.0 video_frames = [] while True: ret, frame = cap.read() if not ret: break video_frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)) cap.release() progress(0.4, desc="Tracking object through video...") with torch.inference_mode(), torch.autocast(device_type=get_device(), dtype=torch.bfloat16): response = video_predictor.handle_request( request=dict(type="start_session", resource_path=file_path) ) session_id = response["session_id"] video_predictor.handle_request( request=dict(type="add_prompt", session_id=session_id, frame_index=0, text=video_prompt) ) outputs_per_frame = propagate_in_video(video_predictor, session_id) progress(0.85, desc="Rendering overlay video...") out_path = f"/tmp/output/sam3_overlay_{timestamp}.mp4" save_masklet_video(video_frames, outputs_per_frame, out_path=out_path, alpha=0.5, fps=video_fps) video_predictor.handle_request(request=dict(type="close_session", session_id=session_id)) progress(1.0, desc="Done!") status = f"[{model_version}] Tracked {len(outputs_per_frame)} frames for '{video_prompt}'.{note}" return ( gr.update(visible=False), gr.update(value=out_path, visible=True), gr.update(value=out_path, visible=True), status, ) else: raise gr.Error(f"Unsupported file type: {ext}") # ---------------------------------------------------------------------- # UI # ---------------------------------------------------------------------- with gr.Blocks(title="SAM 3 Segmentation") as demo: gr.Markdown( "## 🟦 SAM 3 / SAM 3.1 — Image / Video / Panorama / Multi-Class Segmentation\n" "Upload an image or video, pick a model version, type up to 5 comma-separated classes, hit Run. " "Panoramas (very wide images) are auto-tiled. Multiple classes are merged into one overlay with a legend." ) with gr.Row(): with gr.Column(scale=1): file_input = gr.File(label="Upload image or video", file_types=["image", "video"], type="filepath") model_toggle = gr.Radio( choices=["SAM 3.1", "SAM 3"], value=DEFAULT_VERSION, label="Model version", info="SAM 3.1 = faster multi-object video tracking (Object Multiplex). Toggle to SAM 3 for the original.", ) prompt_input = gr.Textbox( label="Text prompt(s) — comma-separated, up to 5 classes", placeholder="e.g. person, dog, red car", ) run_btn = gr.Button("▶ Run Segmentation", variant="primary") status_box = gr.Textbox(label="Status", interactive=False) with gr.Column(scale=1): image_output = gr.Image(label="Result (image)", visible=False) video_output = gr.Video(label="Result (video)", visible=False) download_output = gr.File(label="Download result", visible=False) run_btn.click( fn=run_sam3, inputs=[file_input, prompt_input, model_toggle], outputs=[image_output, video_output, download_output, status_box], ) if __name__ == "__main__": demo.launch()