app.py

"""
segment-everything-api — one call in, every mask out.

SAM 2.1 automatic mask generation, returned as a **label map** rather than a
picture of one. The reply is a lossless PNG where each segment is a distinct
index encoded in the red and green channels (`id = R + G*256`, 0 = unsegmented),
plus a JSON manifest of each segment's area and bounding box.

That shape is the point: a client can decode the PNG straight into a per-pixel
label array and do hit-testing, hover highlighting and selection locally, with no
further round trips. Overlay screenshots — what most segmentation demos return —
can't be decoded back into masks, so they're offered only as a human preview.

API:  /segment(image, points_per_crop, min_area_frac, pred_iou_thresh,
               stability_score_thresh) -> (labels.png, manifest, preview)
"""

import os
import tempfile

import spaces  # must precede torch — it patches torch.cuda before CUDA init
import gradio as gr
import numpy as np
import torch
from PIL import Image
from scipy import ndimage
from transformers import pipeline

# SAM's own defaults (32 points, IoU .88, stability .95) are tuned for photos and
# leave illustration badly under-segmented — measured on a busy anime piece they
# covered 13% of the frame, and merged unrelated things into single giant blobs.
# The defaults on `segment` below are what fixed that, measured on the same piece:
#
#   setting                     parts   biggest region
#   SAM defaults                  ~40   (13% of frame even covered)
#   64 pts / IoU .5 / stab .75     117   24%   <- one blob held a quarter of the art
#   96 pts / IoU .35 / stab .5     284   10%
#   96 pts / IoU .3  / stab .3     443    9%   <- current, finest that stays sane
#
# `stability_score_thresh` is the one that lets the finer, "less stable" masks
# through, and the DENSE GRID is what actually splits a big region: more sample
# points inside it means more competing fine masks. At 64 points the giants
# survive whatever else is loosened. The cost is ~45 s instead of ~15 s, paid once
# per version because the client caches the label map.
MODEL_ID = os.environ.get("SAM_MODEL", "facebook/sam2.1-hiera-large")
MAX_EDGE = 2048  # bounds VRAM and the returned payload
MAX_SEGMENTS = 65535  # what the two-channel label encoding holds

generator = pipeline("mask-generation", model=MODEL_ID, device=0)


def _fit(image: Image.Image) -> Image.Image:
    """Downscale so the long edge is at most MAX_EDGE; leave smaller art alone."""
    w, h = image.size
    scale = MAX_EDGE / max(w, h)
    if scale >= 1:
        return image
    return image.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.LANCZOS)


def _close_gaps(labels: np.ndarray, reach: int) -> np.ndarray:
    """
    Absorb the hairline gutters between neighbouring masks.

    SAM's masks don't tile the plane — adjacent ones leave a pixel or two of
    unlabelled space between them. Left alone those gutters shatter into
    thousands of thread-thin components downstream (measured: 199 segments became
    6072 clickable parts), which makes hovering useless.

    So each unlabelled pixel within `reach` of a segment joins its NEAREST
    segment. The distance cap is the point: gutters get closed, while a genuinely
    unsegmented expanse stays unsegmented instead of being annexed by whatever
    happened to border it.
    """
    empty = labels == 0
    if reach <= 0 or not empty.any() or empty.all():
        return labels
    dist, (iy, ix) = ndimage.distance_transform_edt(empty, return_indices=True)
    return np.where(empty & (dist <= reach), labels[iy, ix], labels)


def _paint_labels(masks, height: int, width: int, min_area: int, gap_reach: int = 4):
    """
    Flatten possibly-overlapping masks into one label map.

    Painted largest-first, so where masks overlap the SMALLER one wins — a big
    background mask never buries the little parts a user actually wants to click.
    Areas and boxes are measured on the finished map, not on the input masks, so
    they describe what's really there after the overwrites and the gap closing.
    """
    order = sorted(range(len(masks)), key=lambda i: -int(masks[i].sum()))
    raw = np.zeros((height, width), dtype=np.int32)
    for slot, i in enumerate(order[:MAX_SEGMENTS], start=1):
        raw[np.asarray(masks[i], dtype=bool)] = slot

    # Drop whatever ended up too small (or fully covered), then renumber densely.
    areas = np.bincount(raw.ravel(), minlength=int(raw.max()) + 1)
    keep = [i for i in range(1, len(areas)) if areas[i] >= min_area]
    remap = np.zeros(len(areas), dtype=np.int32)
    for new_id, old_id in enumerate(keep, start=1):
        remap[old_id] = new_id
    # Close gaps AFTER culling, not before: culling hands its rejects back to the
    # unlabelled class, and those rejects are exactly the debris the fill exists
    # to absorb. Filling first left thousands of specks behind.
    labels = _close_gaps(remap[raw], gap_reach)

    boxes = ndimage.find_objects(labels)
    final_areas = np.bincount(labels.ravel(), minlength=len(keep) + 1)
    regions = []
    for new_id in range(1, len(keep) + 1):
        box = boxes[new_id - 1]
        if box is None:
            continue
        ys, xs = box
        regions.append(
            {
                "id": new_id,
                "area": int(final_areas[new_id]),
                "bbox": [int(xs.start), int(ys.start), int(xs.stop), int(ys.stop)],
            }
        )
    return labels, regions


def _encode(labels: np.ndarray) -> str:
    """Label ids into R + G*256 of a lossless, untagged PNG."""
    height, width = labels.shape
    rgb = np.zeros((height, width, 3), dtype=np.uint8)
    rgb[..., 0] = (labels & 0xFF).astype(np.uint8)
    rgb[..., 1] = ((labels >> 8) & 0xFF).astype(np.uint8)
    path = os.path.join(tempfile.mkdtemp(), "labels.png")
    Image.fromarray(rgb, mode="RGB").save(path, format="PNG", optimize=True)
    return path


def _preview(image: np.ndarray, labels: np.ndarray) -> np.ndarray:
    """A human-readable tint of the label map. Never parse this — use the PNG."""
    rng = np.random.default_rng(7)
    palette = rng.integers(40, 235, size=(int(labels.max()) + 1, 3), dtype=np.uint8)
    palette[0] = (0, 0, 0)
    tint = palette[labels]
    return ((image.astype(np.uint16) + tint.astype(np.uint16) * 2) // 3).astype(np.uint8)


# NOTE: `crops_n_layers` (SAM's multi-scale crop pass, the other way to find
# finer parts) is unusable on transformers 5.15 — its SAM2 processor tries to
# torch.stack crops of different sizes and raises. Re-test before reaching for it.
@spaces.GPU(duration=300)
def _run(pixels: np.ndarray, points_per_crop: int, pred_iou_thresh: float, stability: float):
    return generator(
        Image.fromarray(pixels),
        points_per_crop=int(points_per_crop),
        points_per_batch=64,
        pred_iou_thresh=float(pred_iou_thresh),
        stability_score_thresh=float(stability),
    )


def segment(
    image: Image.Image,
    points_per_crop: int = 96,
    min_area_frac: float = 0.0001,
    pred_iou_thresh: float = 0.3,
    stability_score_thresh: float = 0.3,
    gap_reach: int = 32,
):
    """
    Segment every object in an image and return the result as a label map.

    Args:
        image: The image to segment.
        points_per_crop: Grid density of the automatic prompt sampler — higher
            finds smaller parts and costs more time (8 coarse … 96 fine).
        min_area_frac: Drop segments smaller than this fraction of the frame.
        pred_iou_thresh: Minimum predicted mask quality to keep a segment.
        stability_score_thresh: Minimum mask stability to keep a segment.
        gap_reach: How far (px) an unlabelled pixel may reach to join its nearest
            segment. The default closes the debris between masks completely (a
            measured 2559 stray blobs down to 0) while leaving a genuinely large
            unsegmented expanse alone, since nothing is within reach of it. 0 off.

    Returns:
        labels.png — a lossless PNG whose pixels carry the segment id as
            `R + G*256`, with 0 meaning unsegmented;
        manifest — {width, height, count, model, regions:[{id, area, bbox}]};
        preview — a tinted image for eyeballing only, not for parsing.
    """
    if image is None:
        raise gr.Error("No image supplied.")
    image = _fit(image.convert("RGB"))
    pixels = np.asarray(image)
    height, width = pixels.shape[:2]

    outputs = _run(pixels, points_per_crop, pred_iou_thresh, stability_score_thresh)
    masks = outputs.get("masks", []) if isinstance(outputs, dict) else []
    min_area = max(1, int(min_area_frac * width * height))
    labels, regions = _paint_labels(masks, height, width, min_area, int(gap_reach))

    manifest = {
        "width": width,
        "height": height,
        "count": len(regions),
        "model": MODEL_ID,
        "regions": regions,
    }
    return _encode(labels), manifest, _preview(pixels, labels)


with gr.Blocks(title="Segment Everything API") as demo:
    gr.Markdown(
        "# Segment Everything API\n"
        "SAM 2.1 automatic mask generation returned as a **label map PNG** "
        "(`id = R + G*256`, 0 = unsegmented) plus a JSON manifest — so a client can "
        "decode the masks and do its own hit-testing offline. Built for API use; "
        "the preview is only for eyeballing."
    )
    with gr.Row():
        with gr.Column():
            image_in = gr.Image(label="Image", type="pil")
            points = gr.Slider(8, 128, value=96, step=8, label="Points per crop (detail)")
            min_area = gr.Slider(
                0.0, 0.01, value=0.0001, step=0.0001, label="Min segment area (fraction)"
            )
            iou = gr.Slider(0.2, 0.99, value=0.3, step=0.01, label="Predicted IoU threshold")
            stability = gr.Slider(0.2, 0.99, value=0.3, step=0.01, label="Stability threshold")
            gap = gr.Slider(0, 64, value=32, step=1, label="Gap close reach (px)")
            run = gr.Button("Segment", variant="primary")
        with gr.Column():
            preview_out = gr.Image(label="Preview (not machine-readable)")
            labels_out = gr.File(label="labels.png")
            manifest_out = gr.JSON(label="manifest")

    run.click(
        fn=segment,
        inputs=[image_in, points, min_area, iou, stability, gap],
        outputs=[labels_out, manifest_out, preview_out],
        api_name="segment",
    )

demo.launch(mcp_server=True)