app.py

import gradio as gr
import spaces, torch, os, json
from PIL import Image
from typing import Union
import numpy as np
from samv2_handler import (
    load_sam_image_model,
    run_sam_im_inference,
    load_sam_video_model,
    run_sam_video_inference,
    logger,
)
from toolbox.mask_encoding import b64_mask_decode

torch.autocast(device_type="cuda", dtype=torch.bfloat16).__enter__()
if torch.cuda.get_device_properties(0).major >= 8:
    torch.backends.cuda.matmul.allow_tf32 = True
    torch.backends.cudnn.allow_tf32 = True


@spaces.GPU
def load_im_model(variant):
    return load_sam_image_model(variant=variant, device="cuda")


@spaces.GPU
def load_vid_model(variant):
    return load_sam_video_model(variant=variant, device="cuda")


@spaces.GPU
@torch.inference_mode()
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def process_image(
    im: Image.Image,
    variant: str,
    bboxes: Union[list, str] = None,
    points: Union[list, str] = None,
    point_labels: Union[list, str] = None,
) -> list[str]:
    """SAM2 image segmentation: segment objects in an image from box and/or point prompts and return their masks.

    Provide prompts as bounding boxes and/or points (points require point_labels). Returns a JSON list of segmentation masks, one mask per prompt -- one mask per bounding box, or a single mask for a points-only prompt -- in the same order as the input prompts. Each mask is a base64-encoded 1-bit (binary) PNG string; when decoded the PNG has the same width and height as the original input image, with pixel value 1 inside the object and 0 elsewhere.

    Args:
        im: The RGB image to segment, as a Pillow Image.
        variant: SAM2 model variant to load; one of "tiny", "small", "base_plus", or "large".
        bboxes: Object bounding boxes to segment, as a JSON list of dicts [{"x0":..., "y0":..., "x1":..., "y1":...}, ...] in ABSOLUTE PIXEL coordinates of the input image (top-left corner x0/y0, bottom-right corner x1/y1) -- the albumentations "pascal_voc" format; either bboxes or points must be provided.
        points: Object points to segment, as a JSON list of dicts [{"x":..., "y":...}, ...] in ABSOLUTE PIXEL coordinates of the input image; requires point_labels.
        point_labels: JSON list of integers, one per point, where 1 marks a foreground (include) point and 0 a background (exclude) point.
    """
    # input validation
    has_bboxes = type(bboxes) != type(None) and bboxes != ""
    has_points = type(points) != type(None) and points != ""
    has_point_labels = type(point_labels) != type(None) and point_labels != ""
    assert has_bboxes or has_points, f"either bboxes or points must be provided."
    if has_points:
        assert has_point_labels, f"point_labels is required if points are provided."

    bboxes = json.loads(bboxes) if isinstance(bboxes, str) and has_bboxes else bboxes
    points = json.loads(points) if isinstance(points, str) and has_points else points
    point_labels = (
        json.loads(point_labels)
        if isinstance(point_labels, str) and has_point_labels
        else point_labels
    )
    if has_points:
        assert len(points) == len(
            point_labels
        ), f"{len(points)} points provided but there are {len(point_labels)} labels."

    model = load_im_model(variant=variant)
    return run_sam_im_inference(
        model,
        image=im,
        bboxes=bboxes,
        points=points,
        point_labels=point_labels,
        b64_encode_mask=True,
    )


@spaces.GPU(
    duration=120
)  # user must have 2-minute of inference time left at the time of calling
@torch.inference_mode()
@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
def process_video(
    video_path: str,
    variant: str,
    masks: Union[list, str],
    drop_masks: bool = False,
    ref_frame_idx: int = 0,
    async_frame_load: bool = True,
) -> list[dict]:
    """SAM2 video segmentation: track objects through a video given their segmentation masks on a reference frame.

    Seed the objects to track by supplying their masks on the reference frame (ref_frame_idx); SAM2 then propagates each object forward -- and, when ref_frame_idx is nonzero, also backward -- through every frame. Returns a JSON list of per-frame detections, one entry per tracked object per frame in which it appears (frames where an object is absent are skipped). Each detection is a dict with: "frame" (integer frame index); "track_id" (integer object id, stable across frames); "x", "y", "w", "h" (the object's bounding box as top-left-x, top-left-y, width, height, each NORMALIZED to 0.0-1.0 by dividing by the frame width or height -- this is the albumentations "coco" layout [x_min, y_min, width, height] but normalized to 0-1 rather than absolute pixels, and it is NOT [x_min, y_min, x_max, y_max]; box formats documented at https://albumentations.ai/docs/3-basic-usage/bounding-boxes-augmentations/#bounding-box-formats ); "conf" (always 1); and, unless drop_masks is true, "mask_b64" (a base64-encoded 1-bit PNG string the same width and height as the video frame, 1 inside the object and 0 elsewhere).

    Args:
        video_path: Filesystem path to the input video file.
        variant: SAM2 model variant to load; one of "tiny", "small", "base_plus", or "large".
        masks: JSON list of base64-encoded 1-bit PNG masks for the reference frame, one per object to track, e.g. ["b'iVBORw0KGgo...'", ...]; the b'...' literal wrapper is accepted and stripped.
        drop_masks: When true, omit the "mask_b64" field from every detection so only bounding-box information is returned.
        ref_frame_idx: Frame index the provided masks correspond to; a nonzero value triggers bidirectional tracking (forward and backward from this frame).
        async_frame_load: When true, load video frames in parallel with propagation to reduce inference time.
    """
    model = load_vid_model(variant=variant)
    masks = json.loads(masks) if isinstance(masks, str) else masks
    logger.debug(f"masks---\n{masks}")
    masks = [
        m[2:-1].encode() if m.startswith("b'") and m.endswith("'") else m for m in masks
    ]  # expect the b'' literal to be included
    masks = np.array([b64_mask_decode(m).astype(np.uint8) for m in masks])
    logger.debug(f"masks---\n{masks}")
    return run_sam_video_inference(
        model,
        video_path=video_path,
        masks=masks,
        device="cuda",
        do_tidy_up=True,
        drop_mask=drop_masks,
        async_frame_load=async_frame_load,
        ref_frame_idx=ref_frame_idx,
    )


with gr.Blocks() as demo:
    with gr.Tab("Images"):
        gr.Interface(
            fn=process_image,
            inputs=[
                gr.Image(label="Input Image", type="pil"),
                gr.Dropdown(
                    label="Model Variant",
                    choices=["tiny", "small", "base_plus", "large"],
                ),
                gr.Textbox(
                    label="Bounding Boxes",
                    value=None,
                    lines=5,
                    placeholder='JSON list of dicts: [{"x0":..., "y0":..., "x1":..., "y1":...}, ...]',
                ),
                gr.Textbox(
                    label="Points",
                    lines=3,
                    placeholder='JSON list of dicts: [{"x":..., "y":...}, ...]',
                ),
                gr.Textbox(label="Points' Labels", placeholder="JSON List of Integars"),
            ],
            outputs=gr.JSON(label="Output JSON"),
            title="SAM2 for Images",
            api_name="process_image",
        )
    with gr.Tab("Videos"):
        gr.Interface(
            fn=process_video,
            inputs=[
                gr.Video(label="Input Video"),
                gr.Dropdown(
                    label="Model Variant",
                    choices=["tiny", "small", "base_plus", "large"],
                ),
                gr.Textbox(
                    label="Masks for Objects of Interest in the First Frame",
                    value=None,
                    lines=5,
                    placeholder="""
                    JSON list of base64 encoded masks, e.g.: ["b'iVBORw0KGgoAAAANSUhEUgAABDgAAAeAAQAAAAADGtqnAAAXz...'",...]
                    """,
                ),
                gr.Checkbox(
                    label="Drop Masks",
                    info="remove base64 encoded masks from result JSON",
                    value=True,
                ),
                gr.Number(
                    label="Reference Frame Index",
                    info="frame index for the provided object masks",
                    value=0,
                    precision=0,
                ),
                gr.Checkbox(
                    label="async frame load",
                    info="start inference in parallel to frame loading",
                ),
            ],
            outputs=gr.JSON(label="Output JSON"),
            title="SAM2 for Videos",
            api_name="process_video",
        )

# Model weights are auto-downloaded from the HuggingFace Hub on first use.
demo.launch(
    mcp_server=True, app_kwargs={"docs_url": "/docs"}  # add FastAPI Swagger API Docs
)