app.py

"""Free SAM3 text-to-mask demo, powered by SegmentationAPI (segmentationapi.com).

Flow per request: presign upload -> PUT image to S3 -> create job with text
prompts -> poll until done -> render masks as labeled overlays.
"""

import io
import json
import os
import threading
import time
import zipfile
from collections import deque
from datetime import date
from urllib.parse import urlparse

import gradio as gr
import numpy as np
import requests
from PIL import Image

API_BASE = "https://api.segmentationapi.com/v1"
API_KEY = os.environ.get("SEGMENTATIONAPI_KEY", "")
HEADERS = {"x-api-key": API_KEY}

WEBSITE_URL = "https://www.segmentationapi.com"
DOCS_URL = "https://www.segmentationapi.com/docs"

# Demo limits - these are the upgrade path, keep them visible in the UI copy.
MAX_SIDE = int(os.environ.get("DEMO_MAX_SIDE", "1024"))
MAX_PROMPTS = int(os.environ.get("DEMO_MAX_PROMPTS", "5"))
REQUESTS_PER_HOUR = int(os.environ.get("DEMO_REQUESTS_PER_HOUR", "10"))
GLOBAL_DAILY_LIMIT = int(os.environ.get("DEMO_GLOBAL_DAILY_LIMIT", "1000"))
POLL_TIMEOUT_S = int(os.environ.get("DEMO_POLL_TIMEOUT_S", "120"))

UPGRADE_CTA = (
    f"Get your own free API key at [segmentationapi.com]({WEBSITE_URL}) "
    "for full resolution, video segmentation, and batch processing."
)


# --------------------------------------------------------------------------- #
# Rate limiting (in-memory; resets on Space restart, which is fine for a demo)
# --------------------------------------------------------------------------- #

_rl_lock = threading.Lock()
_requests_by_ip: dict[str, deque] = {}
_global_count = {"day": date.today().isoformat(), "count": 0}


def _client_ip(request: gr.Request | None) -> str:
    if request is None:
        return "unknown"
    forwarded = request.headers.get("x-forwarded-for", "")
    if forwarded:
        return forwarded.split(",")[0].strip()
    return request.client.host if request.client else "unknown"


def _check_rate_limit(ip: str) -> None:
    now = time.time()
    today = date.today().isoformat()
    with _rl_lock:
        if _global_count["day"] != today:
            _global_count["day"] = today
            _global_count["count"] = 0
        if _global_count["count"] >= GLOBAL_DAILY_LIMIT:
            raise gr.Error(
                "The free demo has hit today's global limit. "
                f"Come back tomorrow - or {UPGRADE_CTA}"
            )
        history = _requests_by_ip.setdefault(ip, deque())
        while history and now - history[0] > 3600:
            history.popleft()
        if len(history) >= REQUESTS_PER_HOUR:
            raise gr.Error(
                f"Demo limit reached ({REQUESTS_PER_HOUR} requests/hour). "
                f"{UPGRADE_CTA}"
            )
        history.append(now)
        _global_count["count"] += 1


# --------------------------------------------------------------------------- #
# SegmentationAPI client
# --------------------------------------------------------------------------- #

def _api_error(resp: requests.Response, step: str) -> gr.Error:
    try:
        detail = resp.json().get("message", resp.text[:200])
    except Exception:
        detail = resp.text[:200]
    return gr.Error(f"SegmentationAPI {step} failed ({resp.status_code}): {detail}")


def _upload_image(img: Image.Image) -> str:
    """Presign, upload, return the task ID."""
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    data = buf.getvalue()

    resp = requests.post(
        f"{API_BASE}/uploads/presign",
        headers={**HEADERS, "Content-Type": "application/json"},
        json={"contentType": "image/png"},
        timeout=30,
    )
    if resp.status_code == 401:
        raise gr.Error(
            "The demo backend is not configured yet (invalid or missing API key). "
            "If you run this Space, set the SEGMENTATIONAPI_KEY secret."
        )
    if not resp.ok:
        raise _api_error(resp, "presign")
    presign = resp.json()

    put = requests.put(
        presign["uploadUrl"],
        data=data,
        headers={"Content-Type": "image/png"},
        timeout=60,
    )
    if not put.ok:
        raise gr.Error(f"Image upload failed ({put.status_code}). Please try again.")
    return presign["taskId"]


def _create_job(task_id: str, prompts: list[str], threshold: float, mask_threshold: float) -> str:
    resp = requests.post(
        f"{API_BASE}/jobs",
        headers={**HEADERS, "Content-Type": "application/json"},
        json={
            "type": "image",
            "tasks": [task_id],
            "prompts": prompts,
            "threshold": threshold,
            "maskThreshold": mask_threshold,
            "generatePreview": True,
        },
        timeout=30,
    )
    if not resp.ok:
        raise _api_error(resp, "job creation")
    return resp.json()["jobId"]


def _poll_job(job_id: str) -> dict:
    """Poll until every task is terminal: ~2.5s start, backoff capped at 10s."""
    deadline = time.time() + POLL_TIMEOUT_S
    delay = 2.5
    while time.time() < deadline:
        time.sleep(delay)
        resp = requests.get(f"{API_BASE}/jobs/{job_id}", headers=HEADERS, timeout=30)
        if not resp.ok:
            raise _api_error(resp, "job polling")
        job = resp.json()
        tasks = job.get("tasks", [])
        if tasks and all(t.get("status") in ("success", "failed") for t in tasks):
            return job
        delay = min(delay * 1.6, 10)
    raise gr.Error("Segmentation timed out. Please try again.")


# --------------------------------------------------------------------------- #
# Results - the status endpoint returns no mask data. Results are a zip
# archive (POST /v1/jobs/{jobId}/download, then poll GET until "ready")
# containing output_manifest.json plus the mask PNGs:
#   items: [{"taskId": ..., "previewUrl": "...",
#            "masks": [{"maskIndex": 0, "url": "...", "confidence": 0.97}]}]
# --------------------------------------------------------------------------- #

def _fetch_results(job_id: str) -> tuple[list[dict], zipfile.ZipFile]:
    """Request the results archive, wait for it, return (manifest items, zip)."""
    resp = requests.post(f"{API_BASE}/jobs/{job_id}/download", headers=HEADERS, timeout=30)
    if not resp.ok:
        raise _api_error(resp, "results request")
    info = resp.json()
    deadline = time.time() + POLL_TIMEOUT_S
    while info.get("status") in ("pending", "processing") and time.time() < deadline:
        time.sleep(info.get("retryAfterSeconds") or 3)
        poll = requests.get(f"{API_BASE}/jobs/{job_id}/download", headers=HEADERS, timeout=30)
        if not poll.ok:
            raise _api_error(poll, "results polling")
        info = poll.json()
    if info.get("status") != "ready" or not info.get("downloadUrl"):
        raise gr.Error(
            f"Preparing results failed: {info.get('error') or 'timed out'}. Please try again."
        )

    archive_resp = requests.get(info["downloadUrl"], timeout=120)
    archive_resp.raise_for_status()
    archive = zipfile.ZipFile(io.BytesIO(archive_resp.content))
    manifest_name = next(
        (n for n in archive.namelist() if n.endswith("output_manifest.json")), None
    )
    if manifest_name is None:
        raise gr.Error("Results archive is missing its manifest. Please try again.")
    manifest = json.loads(archive.read(manifest_name))
    return manifest.get("items", []), archive


def _read_artifact(ref: str, archive: zipfile.ZipFile) -> bytes:
    """Resolve a manifest file reference: an archive member matched by the
    tail of the reference's path, falling back to fetching it as a URL."""
    parts = [p for p in urlparse(ref).path.split("/") if p]
    names = archive.namelist()
    for tail_len in range(min(3, len(parts)), 0, -1):
        tail = "/".join(parts[-tail_len:])
        member = next((n for n in names if n.endswith(tail)), None)
        if member:
            return archive.read(member)
    resp = requests.get(ref, timeout=60)
    resp.raise_for_status()
    return resp.content


def _load_mask(ref: str, archive: zipfile.ZipFile, size: tuple[int, int]) -> np.ndarray:
    mask_img = Image.open(io.BytesIO(_read_artifact(ref, archive))).convert("L")
    if mask_img.size != size:
        mask_img = mask_img.resize(size, Image.NEAREST)
    arr = np.asarray(mask_img)
    binary = (arr > 127).astype(np.uint8)
    if not binary.any():
        binary = (arr > 0).astype(np.uint8)
    return binary


# --------------------------------------------------------------------------- #
# Main handler
# --------------------------------------------------------------------------- #

def segment(image: Image.Image | None, prompts_text: str, threshold: float,
            mask_threshold: float, request: gr.Request):
    if image is None:
        raise gr.Error("Please upload an image (or pick an example below).")
    prompts = [p.strip() for chunk in prompts_text.split("\n") for p in chunk.split(",")]
    prompts = list(dict.fromkeys(p for p in prompts if p))
    if not prompts:
        raise gr.Error('Type what to segment, e.g. "person" or "car, traffic light".')
    if len(prompts) > MAX_PROMPTS:
        raise gr.Error(
            f"The free demo supports up to {MAX_PROMPTS} prompts per request. {UPGRADE_CTA}"
        )
    if not API_KEY:
        raise gr.Error(
            "The demo backend is not configured yet (SEGMENTATIONAPI_KEY secret is missing)."
        )

    _check_rate_limit(_client_ip(request))

    image = image.convert("RGB")
    if max(image.size) > MAX_SIDE:
        image.thumbnail((MAX_SIDE, MAX_SIDE), Image.LANCZOS)

    started = time.time()
    task_id = _upload_image(image)
    job_id = _create_job(task_id, prompts, threshold, mask_threshold)
    job = _poll_job(job_id)

    failed = next((t for t in job.get("tasks", []) if t.get("status") == "failed"), None)
    if job.get("error") or failed:
        reason = job.get("error") or failed.get("error")
        raise gr.Error(f"Segmentation failed{': ' + str(reason) if reason else ''}. "
                       "Try a different image or prompt.")

    items, archive = _fetch_results(job_id)
    item = next((i for i in items if i.get("masks")), None)
    if item is None:
        raise gr.Error(
            f'No matches found for "{", ".join(prompts)}" - try a different phrase.'
        )
    preview_url = item.get("previewUrl")

    # Masks are detected instances; the manifest doesn't map them back to
    # individual prompts, so only a single prompt gets a meaningful label.
    annotations = []
    masks = sorted(item["masks"], key=lambda m: m.get("maskIndex", 0))
    for i, entry in enumerate(masks):
        try:
            mask = _load_mask(entry["url"], archive, image.size)
        except Exception:
            continue
        if not mask.any():
            continue
        label = prompts[0] if len(prompts) == 1 else f"object {entry.get('maskIndex', i) + 1}"
        confidence = entry.get("confidence")
        if isinstance(confidence, (int, float)):
            label = f"{label} ({confidence:.2f})"
        annotations.append((mask, label))

    elapsed = time.time() - started

    if annotations:
        status = (
            f"✅ Found **{len(annotations)}** mask{'s' if len(annotations) != 1 else ''} "
            f"in {elapsed:.1f}s. ⚡ {UPGRADE_CTA}"
        )
        return (np.asarray(image), annotations), status

    if preview_url:
        try:
            preview = Image.open(io.BytesIO(_read_artifact(preview_url, archive))).convert("RGB")
        except Exception:
            preview = None
        if preview is not None:
            return (np.asarray(preview), []), (
                f"✅ Done in {elapsed:.1f}s (showing server-rendered preview). ⚡ {UPGRADE_CTA}"
            )

    raise gr.Error(
        f'No matches found for "{", ".join(prompts)}" - try a different phrase.'
    )


# --------------------------------------------------------------------------- #
# UI
# --------------------------------------------------------------------------- #

EXAMPLE_SPECS = [
    ("examples/people.jpg", "person"),
    ("examples/aerial.jpg", "building, road"),
    ("examples/shelf.jpg", "bottle"),
    ("examples/street.jpg", "car, traffic light"),
]
EXAMPLES = [[path, prompt] for path, prompt in EXAMPLE_SPECS if os.path.exists(path)]

# Brand: segmentationapi.com - dark UI, warm orange accent, mono detail text.
THEME = gr.themes.Soft(
    primary_hue="orange",
    secondary_hue="amber",
    neutral_hue="stone",
    font=[gr.themes.GoogleFont("Outfit"), "ui-sans-serif", "system-ui", "sans-serif"],
    font_mono=[gr.themes.GoogleFont("JetBrains Mono"), "ui-monospace", "monospace"],
)

CSS = """
.app-header {text-align: center; padding: 8px 0 0 0;}
.app-header h1 {margin-bottom: 4px;}
.powered-by {
    text-align: center; padding: 10px 14px; border-radius: 12px;
    background: linear-gradient(90deg, rgba(249,115,22,.12), rgba(20,184,166,.08));
    border: 1px solid rgba(249,115,22,.35); margin-bottom: 8px;
}
.footer-cta {
    text-align: center; padding: 14px; border-radius: 12px; margin-top: 12px;
    background: linear-gradient(90deg, rgba(249,115,22,.15), rgba(20,184,166,.10));
    border: 1px solid rgba(249,115,22,.4); font-size: 1.02em;
}
"""

# Match the site's dark look regardless of the visitor's system preference.
FORCE_DARK_JS = """
function refresh() {
    const url = new URL(window.location);
    if (url.searchParams.get('__theme') !== 'dark') {
        url.searchParams.set('__theme', 'dark');
        window.location.href = url.href;
    }
}
"""

with gr.Blocks(title="Free SAM3 Text-to-Mask") as demo:
    gr.Markdown(
        """
        # ✂️ SAM 3 Text → Mask - Free Demo
        Type what you want to segment in plain English and get pixel-perfect masks, powered by Meta's **SAM 3**.
        """,
        elem_classes=["app-header"],
    )
    gr.Markdown(
        f"⚡ Powered by **[SegmentationAPI]({WEBSITE_URL})** - this free demo is "
        f"rate-limited ({REQUESTS_PER_HOUR} requests/hour, images downscaled to "
        f"{MAX_SIDE}px, up to {MAX_PROMPTS} prompts). "
        f"[Get an API key]({WEBSITE_URL}) for full resolution, **video segmentation**, "
        f"and batch processing · [API docs]({DOCS_URL})",
        elem_classes=["powered-by"],
    )

    with gr.Row():
        with gr.Column(scale=1):
            image_in = gr.Image(type="pil", label="Image", sources=["upload", "clipboard"])
            prompts_in = gr.Textbox(
                label="What should we segment?",
                placeholder='e.g. "person" - separate multiple targets with commas',
                lines=1,
            )
            with gr.Accordion("Advanced settings", open=False):
                threshold_in = gr.Slider(
                    0.0, 1.0, value=0.5, step=0.05, label="Detection threshold",
                    info="Higher = fewer, more confident detections",
                )
                mask_threshold_in = gr.Slider(
                    0.0, 1.0, value=0.5, step=0.05, label="Mask threshold",
                    info="Controls mask binarization",
                )
            submit_btn = gr.Button("✂️ Segment", variant="primary")

        with gr.Column(scale=1):
            result_out = gr.AnnotatedImage(label="Masks")
            status_out = gr.Markdown("")

    if EXAMPLES:
        gr.Examples(examples=EXAMPLES, inputs=[image_in, prompts_in], label="Try one of these")

    gr.Markdown(
        f"🚀 **Like what you see?** This demo runs on the exact same API you'd ship with. "
        f"**[Get your API key at segmentationapi.com]({WEBSITE_URL})** - full resolution, "
        f"video & batch segmentation, no rate limits.",
        elem_classes=["footer-cta"],
    )

    submit_btn.click(
        segment,
        inputs=[image_in, prompts_in, threshold_in, mask_threshold_in],
        outputs=[result_out, status_out],
        api_name=False,  # no programmatic access via the Gradio client
    )
    prompts_in.submit(
        segment,
        inputs=[image_in, prompts_in, threshold_in, mask_threshold_in],
        outputs=[result_out, status_out],
        api_name=False,
    )

demo.queue(max_size=20, default_concurrency_limit=4)

if __name__ == "__main__":
    demo.launch(theme=THEME, css=CSS, js=FORCE_DARK_JS)