app.py

"""
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()