h3_sam3.py

"""Say what should change, and get a mask clip that tracks it.

The brush is the wrong tool for video: it paints one still, and a shape painted on the first frame is wrong by the
last one. SAM 3 takes a phrase — "the fox", "the person on the left" — segments it, and propagates the track through
every frame, which is the shape the rest of this Space already consumes: a mask clip, one frame per source frame.

The model is loaded lazily and only when a request asks for a mask, so a Space that is only ever driven through the
API with an uploaded mask never pays for it.
"""

from __future__ import annotations

import os
import tempfile

import numpy as np
import torch
from PIL import Image

SAM3_REPO = os.environ.get("H3_SAM3_REPO", "facebook/sam3")
# Masks are consumed as *hard* by default: a feathered boundary puts rows at intermediate timesteps and blends two
# images that differ in tone, which is what shows up as a soft band along the mask. Dilation is the knob that
# matters, not softness.
_model = None
_processor = None
_tracker = None
_tracker_processor = None


def available() -> tuple[bool, str]:
    """Whether SAM 3 can be used here. The checkpoint is gated, so this can fail on access rather than on code."""
    try:
        from transformers import Sam3VideoModel, Sam3VideoProcessor  # noqa: F401
    except ImportError:
        return False, "this `transformers` has no `Sam3VideoModel`"
    if not (os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")):
        return False, f"`{SAM3_REPO}` is gated and no `HF_TOKEN` is set on this Space"
    return True, ""


def load():
    """Fetch the model once. Called at startup so the download is not on GPU time."""
    global _model, _processor
    if _model is not None:
        return _model, _processor
    from transformers import Sam3VideoModel, Sam3VideoProcessor

    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    _model = Sam3VideoModel.from_pretrained(SAM3_REPO, token=token, dtype=torch.bfloat16)
    _processor = Sam3VideoProcessor.from_pretrained(SAM3_REPO, token=token)
    return _model, _processor


def load_tracker():
    """SAM 3's geometry-prompted tracker, which is what takes a click.

    The text-grounded video model and the tracker are two different heads of the same release: the first finds what a
    phrase names, the second follows what you point at. Only the tracker accepts points, and it is the faster of the
    two by a wide margin — a single click propagated over 124 frames in 15 s against 73 s for the phrase.
    """
    global _tracker, _tracker_processor
    if _tracker is not None:
        return _tracker, _tracker_processor
    from transformers import Sam3TrackerVideoModel, Sam3TrackerVideoProcessor

    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    _tracker = Sam3TrackerVideoModel.from_pretrained(SAM3_REPO, token=token, dtype=torch.bfloat16)
    _tracker_processor = Sam3TrackerVideoProcessor.from_pretrained(SAM3_REPO, token=token)
    return _tracker, _tracker_processor


def segment_from_points(
    frames: np.ndarray,
    points: list[tuple[int, int]],
    labels: list[int],
    frame_idx: int = 0,
    progress=None,
) -> np.ndarray | None:
    r"""
    Track whatever the clicks point at, through every frame.

    Args:
        frames (`np.ndarray` of shape `(num_frames, height, width, 3)`): The clip, `uint8` RGB.
        points (`list[tuple[int, int]]`): The clicked `(x, y)` pixels, on `frame_idx`.
        labels (`list[int]`): `1` for a point that is part of the subject, `0` for one that is not.
        frame_idx (`int`, defaults to 0): Which frame the clicks were made on. Tracking runs outward from it.
        progress (`callable`, *optional*): Gradio progress callback, called per frame.

    Returns:
        `np.ndarray` of shape `(num_frames, height, width)` of `bool`, or None when the clicks selected nothing.
    """
    if not points:
        return None
    model, processor = load_tracker()
    model.to("cuda")
    num_frames, height, width = frames.shape[:3]
    frame_idx = max(0, min(int(frame_idx), num_frames - 1))

    session = processor.init_video_session(
        video=frames, inference_device="cuda", processing_device="cpu",
        video_storage_device="cpu", dtype=torch.bfloat16,
    )
    processor.add_inputs_to_inference_session(
        inference_session=session,
        frame_idx=frame_idx,
        obj_ids=1,
        input_points=[[[list(map(int, point)) for point in points]]],
        input_labels=[[list(map(int, labels))]],
    )

    masks = np.zeros((num_frames, height, width), dtype=bool)
    # Two passes from the clicked frame: forward to the end, then backward to the start. Tracking only ever runs
    # outward from where the prompt was made, so a click on a middle frame would otherwise leave the front unmasked.
    for reverse in (False, True):
        for output in model.propagate_in_video_iterator(
            inference_session=session, start_frame_idx=frame_idx,
            max_frame_num_to_track=num_frames, reverse=reverse, show_progress_bar=False,
        ):
            index = output.frame_idx
            if not 0 <= index < num_frames:
                continue
            found = processor.post_process_masks(
                [output.pred_masks], original_sizes=[[height, width]], binarize=True
            )[0]
            masks[index] = found.cpu().numpy().reshape(-1, height, width).any(0)
            if progress is not None:
                done = abs(index - frame_idx) + 1
                progress(min(1.0, done / num_frames), desc=f"Tracking — frame {index + 1}/{num_frames}")
    return masks if masks.any() else None


def segment(frames: np.ndarray, phrase: str, progress=None) -> np.ndarray | None:
    r"""
    Track `phrase` through `frames` and return a boolean mask per frame.

    Args:
        frames (`np.ndarray` of shape `(num_frames, height, width, 3)`): The clip, `uint8` RGB.
        phrase (`str`): What to segment, in plain words.
        progress (`callable`, *optional*): Gradio progress callback, called per frame.

    Returns:
        `np.ndarray` of shape `(num_frames, height, width)` of `bool`, or None when nothing matched.
    """
    model, processor = load()
    model.to("cuda")
    num_frames, height, width = frames.shape[:3]

    session = processor.init_video_session(
        video=frames,
        inference_device="cuda",
        processing_device="cpu",
        video_storage_device="cpu",
        dtype=torch.bfloat16,
    )
    session = processor.add_text_prompt(inference_session=session, text=phrase)

    masks = np.zeros((num_frames, height, width), dtype=bool)
    # Per-frame progress by hand rather than a tqdm wrapper: wrapping SAM 3's generator breaks
    # `postprocess_outputs` when gradio's `track_tqdm` is active.
    # Disabling SAM 3's own progress bar here is load-bearing, not tidiness. It wraps this loop in tqdm, and
    # gradio's `Progress(track_tqdm=True)` patches tqdm with one whose `__next__` reads `self.iterables[-1]` — empty
    # here, because the iterator is created inside a ZeroGPU worker rather than in the handler gradio set up, which
    # raises `IndexError: list index out of range` from deep inside transformers. Progress is reported by hand below.
    for output in model.propagate_in_video_iterator(
        inference_session=session, max_frame_num_to_track=num_frames, show_progress_bar=False
    ):
        found = processor.postprocess_outputs(session, output)
        instances = found.get("masks")
        index = output.frame_idx
        if instances is not None and len(instances) and 0 <= index < num_frames:
            array = instances.float().cpu().numpy()
            # Every instance the phrase matched is part of the same mask.
            array = (array.reshape(-1, array.shape[-2], array.shape[-1]) > 0.5).any(axis=0)
            if array.shape != (height, width):
                array = (
                    np.asarray(
                        Image.fromarray((array * 255).astype(np.uint8)).resize((width, height), Image.NEAREST)
                    )
                    > 127
                )
            masks[index] = array
        if progress is not None and 0 <= index < num_frames:
            progress((index + 1) / num_frames, desc=f"Segmenting — frame {index + 1}/{num_frames}")

    return masks if masks.any() else None


def dilate(masks: np.ndarray, pixels: int) -> np.ndarray:
    r"""
    Grow every mask by `pixels`, always from the undilated original so repeated adjustments do not compound.

    A segmentation hugs its subject, and the model needs the room around it: what replaces the subject rarely has the
    same silhouette, and anything of the original left outside the mask is preserved and shows.
    """
    if pixels <= 0:
        return masks
    import cv2

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (2 * int(pixels) + 1, 2 * int(pixels) + 1))
    return np.stack([cv2.dilate(frame.astype(np.uint8), kernel).astype(bool) for frame in masks])


def to_clip(masks: np.ndarray, fps: int = 24) -> str:
    r"""Write a boolean mask batch as a mask clip, near-losslessly — its pixel values are read back as the mask."""
    import av

    path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    height, width = masks.shape[1:3]
    container = av.open(path, mode="w")
    stream = container.add_stream("libx264", rate=fps)
    stream.width, stream.height, stream.pix_fmt = width, height, "yuv420p"
    stream.options = {"crf": "6", "preset": "veryfast", "tune": "stillimage"}
    for frame in masks:
        plane = np.repeat((frame.astype(np.uint8) * 255)[:, :, None], 3, axis=2)
        container.mux(stream.encode(av.VideoFrame.from_ndarray(plane, format="rgb24")))
    for packet in stream.encode():
        container.mux(packet)
    container.close()
    return path


def preview(frames: np.ndarray, masks: np.ndarray, fps: int = 24) -> str:
    r"""The mask over the clip, so what will be repainted is visible before a request is booked."""
    import av

    path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    height, width = frames.shape[1:3]
    container = av.open(path, mode="w")
    stream = container.add_stream("libx264", rate=fps)
    stream.width, stream.height, stream.pix_fmt = width, height, "yuv420p"
    stream.options = {"crf": "20", "preset": "veryfast"}
    for frame, mask in zip(frames, masks):
        tinted = frame.astype(np.float32)
        tinted[..., 0] = np.minimum(255, tinted[..., 0] + mask * 90)
        tinted[..., 2] = np.maximum(0, tinted[..., 2] - mask * 40)
        container.mux(stream.encode(av.VideoFrame.from_ndarray(tinted.astype(np.uint8), format="rgb24")))
    for packet in stream.encode():
        container.mux(packet)
    container.close()
    return path