app.py

from __future__ import annotations

import base64
import html
import json
import os
from uuid import uuid4
from io import BytesIO

import gradio as gr
from PIL import Image

from grounding_demo.catalog import ROOT, load_cases
from grounding_demo.client import GroundingClient
from grounding_demo.console import TASK_HELP, compile_query, metrics_html, query_from_prompt, targets_html, task_help_html
from grounding_demo.protocol import parse_gam, parse_reference_boxes
from grounding_demo.render import draw_primitives

ROUTE_META = {
    "grounding": ("Grounding", "Open-vocabulary object localization"),
    "dense": ("Dense", "Dense instance localization"),
    "referring": ("Referring", "Language-guided referring"),
    "dense_point": ("Dense Point", "Dense point prediction"),
    "referring_point_in_mask": ("Referring Point in Mask", "Point inside a referred target"),
    "dense_point_in_mask": ("Dense Point in Mask", "Point inside multiple targets"),
    "visual_prompt": ("Visual Prompt", "Exemplar-guided localization"),
    "ocr": ("OCR", "Text recognition and localization"),
    "gui": ("GUI", "Interface element grounding"),
}
CASES = load_cases()
BY_ID = {row["id"]: row for row in CASES}
BY_ROUTE = {route: [row for row in CASES if row["route"] == route] for route in ROUTE_META}


def route_choices(route):
    return [(row["title"], row["id"]) for row in BY_ROUTE[route]]


def gallery_items(route):
    return [
        (str(ROOT / row["image"]), case_caption(row, i))
        for i, row in enumerate(BY_ROUTE[route], 1)
    ]


def case_caption(row, index):
    return f"{ROUTE_META[row['route']][0]} {index:02d}"


def input_view(path, canonical, route):
    """Show only user-supplied reference boxes; the client still receives raw path."""
    references = parse_reference_boxes(canonical) if route == "visual_prompt" else ()
    if not references or not path:
        return path
    with Image.open(path) as source:
        return draw_primitives(source, (), references, "INPUT · example boxes")


def query_update(canonical, route):
    return gr.update(value=query_from_prompt(canonical, route), placeholder=TASK_HELP[route][3])


def case_meta(row):
    return (
        "<div class='source-meta'><span>TRAINING SOURCE</span>"
        f"<b>{html.escape(row['source'])}</b><span>IMAGE SHA256</span>"
        f"<code>{row.get('image_sha256', '')[:16]}</code>"
        "<small>Curated research examples for qualitative exploration, not a held-out benchmark. Inference receives only the image and task prompt.</small></div>"
    )


def status_html(text, kind="ready"):
    return (
        f"<div class='stage-status {kind}' data-phase='{kind}' role='status' aria-live='polite'>"
        f"<i aria-hidden='true'></i><span>{html.escape(text)}</span>"
        + ("<time class='wait-clock' aria-hidden='true'>0.0 s</time>" if kind == "busy" else "")
        + "</div>"
    )


def reset_image(path, case_id, route, canonical, meta):
    return (
        path,
        case_id,
        input_view(path, canonical, route),
        query_update(canonical, route),
        canonical,
        meta,
        metrics_html(ROUTE_META[route][0]),
        targets_html([]),
        "[]",
        "",
        status_html("Image ready · Run grounding to extract visual evidence"),
        task_help_html(route),
    )


def load_case(case_id):
    row = BY_ID[case_id]
    return reset_image(
        str(ROOT / row["image"]), case_id, row["route"], row["prompt"], case_meta(row)
    )


def change_route(route):
    row = BY_ROUTE[route][0]
    return (
        gr.update(choices=route_choices(route), value=row["id"]),
        gr.update(value=gallery_items(route), selected_index=0),
        *load_case(row["id"]),
    )


def choose_case(case_id):
    row = BY_ID[case_id]
    return gr.update(selected_index=BY_ROUTE[row["route"]].index(row)), *load_case(case_id)


def choose_gallery(route, evt: gr.SelectData):
    index = evt.index[0] if isinstance(evt.index, tuple | list) else evt.index
    row = BY_ROUTE[route][index]
    return gr.update(value=row["id"]), *load_case(row["id"])


def upload_image(path, route, canonical):
    if not path:
        raise gr.Error("Select an image to continue.")
    with Image.open(path) as image:
        image.verify()
    return reset_image(
        path,
        None,
        route,
        canonical,
        "<div class='source-meta'><span>YOUR IMAGE</span><b>Uploaded image</b>"
        "<small>Select a task and describe the visual evidence to extract.</small></div>",
    )


def choose_upload(path, route, canonical):
    return (
        gr.update(value=None),
        gr.update(selected_index=None),
        *upload_image(path, route, canonical),
    )


def edit_query(text, route, path):
    return (
        compile_query(text, route),
        input_view(path, compile_query(text, route), route),
        metrics_html(ROUTE_META[route][0]),
        targets_html([]),
        "[]",
        "",
        status_html("Prompt updated · Run grounding to refresh the overlay"),
    )


def result_bundle(path, canonical, route, raw, elapsed=None, usage=None):
    parsed = parse_gam(raw, route)
    items = [dict(label=p.label, coordinates=list(p.coordinates)) for p in parsed.primitives]
    valid_items = items if parsed.valid else []
    image = Image.open(path).convert("RGB")
    source_bytes = BytesIO()
    image.save(source_bytes, format="PNG")
    image_source = "data:image/png;base64," + base64.b64encode(source_bytes.getvalue()).decode()
    rendered = (
        draw_primitives(
            image,
            parsed.primitives,
            parse_reference_boxes(canonical),
        )
        if parsed.valid
        else image
    )
    state = "Completed"
    if not parsed.valid:
        state = "Protocol error"
    message = (
        f"Complete · {len(valid_items)} visual targets"
        if parsed.valid
        else "Output validation failed · " + "; ".join(parsed.errors)
    )
    return dict(
        image=rendered,
        metrics=metrics_html(ROUTE_META[route][0], state, len(valid_items), elapsed, usage),
        targets=targets_html(valid_items, image_source=image_source),
        normalized=json.dumps(valid_items, ensure_ascii=False, indent=2),
        raw=raw,
        status=status_html(
            message, "error" if not parsed.valid else "ready"
        ),
    )


def display_bundle(bundle):
    return tuple(
        bundle[key] for key in ("image", "metrics", "targets", "normalized", "raw", "status")
    )


def inference_error(exc):
    timed_out = "timeout" in type(exc).__name__.lower()
    return (
        "Request timed out. Check the connection and retry; your image and prompt are preserved."
        if timed_out else
        "Inference unavailable. Check the connection and retry; your image and prompt are preserved."
    )


def run_inference(path, query_text, route):
    if not path:
        raise gr.Error("Choose an example or upload an image first.")
    canonical = compile_query(query_text, route)
    try:
        raw, elapsed, usage = GroundingClient().infer(path, canonical)
        bundle = result_bundle(path, canonical, route, raw, elapsed, usage)
    except Exception as exc:
        message = inference_error(exc)
        return (
            path,
            metrics_html(ROUTE_META[route][0], "API error"),
            targets_html([]),
            "[]",
            f"{type(exc).__name__}: {exc}",
            status_html(message, "error"),
        )
    return display_bundle(bundle)


RUN_LABEL = "Run Grounding"


def run_inference_stream(path, query_text, route):
    """Expose real request state without an opaque progress mask or synthetic steps."""
    if not path:
        raise gr.Error("Choose an example or upload an image first.")
    yield (
        path,
        metrics_html(ROUTE_META[route][0], "Awaiting inference"),
        targets_html([], pending=True),
        "[]",
        "",
        status_html("Awaiting inference", "busy"),
        gr.update(value="Grounding in progress…", interactive=False),
    )
    try:
        result = run_inference(path, query_text, route)
    except Exception as exc:
        result = (
            path,
            metrics_html(ROUTE_META[route][0], "API error"),
            targets_html([]),
            "[]",
            f"{type(exc).__name__}: {exc}",
            status_html(inference_error(exc), "error"),
        )
    yield (*result, gr.update(value=RUN_LABEL, interactive=True))


def run_console(path, query_text, route):
    # Identifies this request even when a proxy coalesces busy + final SSE events.
    run_id = uuid4().hex
    for state in run_inference_stream(path, query_text, route):
        state = list(state)
        state[5] = state[5].replace("data-phase=", f"data-run-id='{run_id}' data-phase=", 1)
        yield (*state, state[-1])


def service_status():
    ok, detail = GroundingClient().health()
    return (
        f"<div class='api-status {'online' if ok else 'offline'}' role='status' aria-live='polite'>"
        f"<i aria-hidden='true'></i><b>{'Inference ready' if ok else 'Inference offline'}</b>"
        f"<small>{'Model loaded. Ready to process your image.' if ok else 'Explore examples while the inference service is unavailable.'}</small></div>"
        f"<details class='connection-detail'><summary>Runtime details</summary><code>{html.escape(detail)}</code></details>"
    )


CSS = (ROOT / "assets" / "paper.css").read_text(encoding="utf-8")
CSS += "\n" + (ROOT / "assets" / "original.css").read_text(encoding="utf-8")
CSS += "\n" + (ROOT / "assets" / "studio.css").read_text(encoding="utf-8")
CSS = CSS.replace(
    "__MANROPE_DATA__",
    base64.b64encode((ROOT / "assets" / "Manrope-variable.ttf").read_bytes()).decode(),
)
# Gradio 5 scopes css= rules inside .contain, including media-query selectors.
# Host sizing must live in head; workbench styling remains scoped.
HOST_HEAD = """<style id="grounding-host-layout">
gradio-app:has(#console) .gradio-container { width:100%!important; max-width:1440px!important; min-width:0!important; box-sizing:border-box; overflow:clip!important; padding:28px 32px 32px!important; margin:0 auto!important; }
gradio-app:has(#console) .gradio-container>main { width:100%!important; max-width:none!important; min-width:0!important; padding:0!important; }
gradio-app:has(#console) .gradio-container .contain { max-width:none!important; }
gradio-app:has(#console) footer:not(.studio-footer) { display:none!important; }
@media(max-width:760px) { gradio-app:has(#console) .gradio-container { padding:16px 14px!important; } }
@media(max-width:430px) { gradio-app:has(#console) .gradio-container { padding:12px 10px!important; } }
</style>"""
JS = (ROOT / "assets" / "console.js").read_text(encoding="utf-8")
first_route = "grounding"
first_case = BY_ROUTE[first_route][2]

with gr.Blocks(
    css=CSS,
    js=JS,
    head=HOST_HEAD,
    title="GroundingPI · Perception Studio",
    theme=gr.themes.Base(font=["Manrope", "sans-serif"], font_mono=["ui-monospace", "monospace"]),
    analytics_enabled=False,
) as demo:
    image_state = gr.State(str(ROOT / first_case["image"]))
    case_state = gr.State(first_case["id"])
    with gr.Column(elem_id="console"):
        with gr.Row(elem_id="workspace-top", equal_height=True):
            with gr.Column(scale=3, min_width=280, elem_id="controls"):
                gr.HTML(
                    """<div class="control-brand">
                      <h2>Grounding<span>PI</span></h2>
                      <p>Visual primitives for physical intelligence.</p>
                    </div>"""
                )
                with gr.Column(elem_id="control-settings"):
                    with gr.Row(elem_id="control-fields"):
                        gr.HTML(
                            """<div class="media-field"><span>Media Type</span><div>
                              <span class="media-option is-active">Image</span>
                              <button class="media-option" type="button" disabled
                                title="Video is not available">Video</button></div></div>""",
                            elem_id="media-type",
                        )
                        route = gr.Dropdown(
                            [(name, key) for key, (name, _) in ROUTE_META.items()],
                            value=first_route,
                            label="Task Type",
                            filterable=False,
                            elem_id="route-select",
                        )
                    case = gr.Dropdown(
                        route_choices(first_route),
                        value=first_case["id"],
                        label="Selected example",
                        filterable=False,
                        visible=False,
                        elem_id="case-select",
                    )
                    upload = gr.UploadButton(
                        "Upload Image", file_types=["image"], type="filepath", elem_id="upload-button"
                    )
                    with gr.Accordion("Advanced Parameters", open=False, elem_id="advanced-options"):
                        gr.HTML(
                            "<div class='runtime-options'><span>Model</span><b>"
                            + html.escape(os.getenv("GROUNDING_MODEL", "grounding-pi-vlm"))
                            + "</b><span>Max output tokens</span><b>"
                            + html.escape(os.getenv("GROUNDING_MAX_TOKENS", "1024"))
                            + "</b><small>Read-only runtime configuration. Only the current image and task prompt enter inference.</small></div>"
                        )
                run_btn = gr.Button(RUN_LABEL, variant="primary", elem_id="run-button")
                with gr.Column(elem_id="service-panel"):
                    health = gr.HTML("<div class='api-status'>Checking inference service…</div>")
                    health_btn = gr.Button("Refresh connection", size="sm", elem_id="health-button")
            with gr.Column(scale=7, min_width=400, elem_id="visual-stage"):
                task_help = gr.HTML(task_help_html(first_route), elem_id="task-guide")
                with gr.Row(elem_id="stage-toolbar"):
                    status = gr.HTML(status_html("Image ready"), elem_id="stage-status")
                with gr.Column(elem_id="canvas-viewport"):
                    stage = gr.Image(
                        value=input_view(str(ROOT / first_case["image"]), first_case["prompt"], first_route),
                        type="filepath",
                        interactive=False,
                        show_label=False,
                        height=400,
                        elem_id="stage-image",
                        show_download_button=False,
                        show_share_button=False,
                        show_fullscreen_button=False,
                    )
                    gr.HTML(
                        """<div class="canvas-tools" aria-label="Canvas tools">
                      <span class="image-dimensions">Image canvas</span>
                      <div class="zoom-tools">
                        <button type="button" data-canvas="out" aria-label="Zoom out" title="Zoom out">
                          <svg viewBox="0 0 20 20" aria-hidden="true"><path d="M5 10h10"/></svg></button>
                        <output class="zoom-value" aria-label="Canvas zoom">100%</output>
                        <button type="button" data-canvas="in" aria-label="Zoom in" title="Zoom in">
                          <svg viewBox="0 0 20 20" aria-hidden="true"><path d="M5 10h10M10 5v10"/></svg></button>
                        <span class="tool-divider"></span>
                        <button type="button" data-canvas="fit" aria-label="Fit to canvas" title="Fit to canvas">
                          <svg viewBox="0 0 20 20" aria-hidden="true"><path d="M7 3H3v4m10-4h4v4M3 13v4h4m10-4v4h-4"/></svg></button>
                        <a class="canvas-download" download="grounding.png" aria-label="Download image" title="Download image">
                          <svg viewBox="0 0 20 20" aria-hidden="true"><path d="M10 3v10m-4-4 4 4 4-4M4 14v3h12v-3"/></svg></a>
                      </div>
                    </div>""",
                        elem_id="canvas-chrome",
                    )
                with gr.Row(elem_id="query-row"):
                    query = gr.Textbox(
                        value=query_from_prompt(first_case["prompt"], first_route),
                        label="What would you like to locate?",
                        show_label=False,
                        lines=2,
                        max_lines=4,
                        placeholder=TASK_HELP[first_route][3],
                        elem_id="query-dock",
                    )
                    gr.HTML("<p class='query-hint'>Enter to run · Shift + Enter for a new line</p>")
                    mobile_run = gr.Button(RUN_LABEL, variant="primary", elem_id="mobile-run-button")
                with gr.Column(elem_id="targets-panel"):
                    targets = gr.HTML(targets_html([]), elem_id="targets-content")
        with gr.Row(elem_id="review-row", equal_height=True):
            with gr.Column(scale=4, min_width=320, elem_id="sandbox-panel"):
                gr.HTML(
                    "<div class='film-heading'><h2>Explore Examples</h2><small>Curated visual scenarios</small></div>"
                )
                gallery = gr.Gallery(
                    value=gallery_items(first_route),
                    columns=3,
                    rows=1,
                    height=170,
                    show_label=False,
                    allow_preview=False,
                    object_fit="cover",
                    selected_index=2,
                    interactive=False,
                    elem_id="case-gallery",
                )
                with gr.Accordion("Data provenance", open=False, elem_id="source-details"):
                    meta = gr.HTML(case_meta(first_case))
            with gr.Column(scale=6, min_width=420, elem_id="review-summary"):
                with gr.Row(elem_id="result-columns", equal_height=True):
                    with gr.Column(scale=4, min_width=200, elem_id="metrics-panel"):
                        gr.HTML("<h2 class='inspector-subheading'>Inference Metrics</h2>")
                        metrics = gr.HTML(metrics_html(ROUTE_META[first_route][0]))
                    with gr.Column(scale=6, min_width=240, elem_id="sidebar"):
                        gr.HTML(
                            "<div class='inspector-heading'><h2>Target Detail</h2><span class='live-mark'>Select a category below the canvas to magnify its original pixels.</span></div>"
                        )
                        gr.HTML(
                            "<div class='evidence-detail' hidden><div class='detail-heading'><span>Detail View</span><small>Original pixels · magnified</small></div><div class='detail-image'></div><div class='detail-coordinates'></div></div>",
                            elem_id="detail-panel",
                        )
    with gr.Accordion("Prompt Console", open=False, elem_id="prompt-console"):
        prompt = gr.Textbox(
            value=first_case["prompt"], label="COMPILED PROMPT", lines=3, interactive=False
        )
        gr.HTML(
            "<div class='console-note'>Transparent inputs: only this image and compiled prompt are sent to the model. No annotations, training answers, or expected counts are included.</div>"
        )
    with gr.Accordion("Structured Output", open=False, elem_id="output-console"):
        normalized = gr.Code(
            value="[]", language="json", label="PARSED COORDINATES · 0–999", interactive=False
        )
        with gr.Accordion("Native Token Sequence", open=False):
            raw = gr.Textbox(label="RAW GAM SPECIAL TOKENS", lines=4, interactive=False)
    gr.HTML(
        "<footer class='studio-footer'><span>GroundingPI · Towards Physical Intelligence</span><span>Research demonstration <i>Qualitative examples, not benchmark results.</i></span></footer>"
    )

    reset_outputs = [
        image_state,
        case_state,
        stage,
        query,
        prompt,
        meta,
        metrics,
        targets,
        normalized,
        raw,
        status,
        task_help,
    ]
    event_options = dict(concurrency_id="console", concurrency_limit=1, show_progress="hidden")
    case.input(choose_case, case, [gallery, *reset_outputs], **event_options)
    route.input(change_route, route, [case, gallery, *reset_outputs], **event_options)
    gallery.select(choose_gallery, route, [case, *reset_outputs], **event_options)
    upload.upload(
        choose_upload, [upload, route, prompt], [case, gallery, *reset_outputs], **event_options
    )
    query.input(
        edit_query,
        [query, route, image_state],
        [prompt, stage, metrics, targets, normalized, raw, status],
        trigger_mode="always_last",
        **event_options,
    )
    run_btn.click(
        run_console,
        [image_state, query, route],
        [stage, metrics, targets, normalized, raw, status, run_btn, mobile_run],
        **event_options,
    )
    query.submit(
        run_console,
        [image_state, query, route],
        [stage, metrics, targets, normalized, raw, status, run_btn, mobile_run],
        **event_options,
    )
    mobile_run.click(
        run_console, [image_state, query, route],
        [stage, metrics, targets, normalized, raw, status, run_btn, mobile_run],
        **event_options,
    )
    health_btn.click(service_status, outputs=health, show_progress="hidden")
    demo.load(service_status, outputs=health, show_progress="hidden")


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch(
        server_name="0.0.0.0",
        server_port=int(os.getenv("DEMO_PORT", "7860")),
        root_path=os.getenv("DEMO_ROOT_PATH") or None,
        show_error=True,
    )