app.py

"""Drosophila behavior runnable example.

Upload a top-down video of flies in an arena; track them and quantify behavior
(locomotion, sleep, social interactions). Two engines behind one contract:
  * fast  — background subtraction + identity linking (always runnable)
  * pose  — deep pose estimation (SLEAP/DeepLabCut), wired; provide a model

Contract surfaces (do not edit): /process via gr.api (also the MCP tool),
?file_url=... loader, temp-file cleanup daemon.
"""
from __future__ import annotations

import json
import os
import tempfile
import urllib.request

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

from core import cleanup, viz
from core.io import APP_TMP_DIR, new_tmp_path, register_protected, resolve_to_video_path
from core.processing import ENGINES, process, simulate_full

APP_TITLE = "Drosophila behavior"
APP_DESCRIPTION = (
    "Upload a top-down video of fruit flies in an arena (TIFF stack or mp4). The "
    "**fast** engine tracks each fly (background subtraction + identity linking) and "
    "quantifies locomotion, sleep, and social interactions; **pose** is a wired slot "
    "for a SLEAP/DeepLabCut model. A synthetic example video is baked in."
)
INPUT_FILE_TYPES = [".tif", ".tiff", ".mp4", ".avi", ".mov"]

cleanup.start()


def _baked_example() -> str | None:
    p = APP_TMP_DIR / "example.tif"
    return register_protected(p) if p.exists() else None


EXAMPLE_PATH = _baked_example()


def _save_png(arr: np.ndarray, basename: str) -> str:
    out_path = new_tmp_path(basename)
    Image.fromarray(arr.astype(np.uint8)).save(out_path)
    return out_path


# ----------------- Public API (also the MCP tool) -----------------

def process_api(
    file_url: str,
    engine: str = "fast",
    fps: float = 15.0,
    px_per_mm: float = 4.0,
) -> gr.FileData:
    """Track flies and return a trajectory + ethogram summary PNG.

    Parameters
    ----------
    file_url : fly video (TIFF stack or mp4) as a URL or local path
    engine : 'fast' (centroid tracking) or 'pose' (deep model)
    fps : video frame rate (Hz)
    px_per_mm : spatial scale, for speed/distance in mm
    """
    video_path = resolve_to_video_path(file_url)
    img, report = process(video_path, engine=engine, fps=float(fps), px_per_mm=float(px_per_mm))
    out_path = _save_png(np.asarray(img), basename="behavior.png")
    print("[process_api] " + json.dumps({k: v for k, v in report.items() if k != "per_fly"}))
    return gr.FileData(path=out_path, orig_name=os.path.basename(out_path), mime_type="image/png")


def analyze_api(
    file_url: str,
    engine: str = "fast",
    fps: float = 15.0,
    px_per_mm: float = 4.0,
) -> dict:
    """Return the behavior report (per-fly locomotion/sleep + interactions) as JSON."""
    video_path = resolve_to_video_path(file_url)
    _img, report = process(video_path, engine=engine, fps=float(fps), px_per_mm=float(px_per_mm))
    return report


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

def _run_ui(file_obj, engine, fps, px_per_mm, make_video, video_secs,
            progress=gr.Progress(track_tqdm=True)):
    if file_obj is None:
        return None, "Upload a fly video first.", None
    progress(0.1, desc="Loading video…")
    video_path = resolve_to_video_path(file_obj)
    progress(0.4, desc=f"Tracking ({engine})…")
    r = simulate_full(video_path, engine=engine, fps=float(fps), px_per_mm=float(px_per_mm))
    video_path_out = None
    if make_video:
        progress(0.75, desc="Rendering annotated video…")
        video_path_out = viz.activity_video(r["movie"], r["res"], seconds=float(video_secs))
    progress(1.0, desc="Done")
    return r["summary"], json.dumps(r["report"], indent=2), video_path_out


with gr.Blocks(title=APP_TITLE, delete_cache=(1800, 21600)) as demo:
    gr.api(
        process_api,
        api_name="process",
        api_description=(
            "Input: a top-down VIDEO of flies in an arena (TIFF stack or mp4). "
            "Track flies in a video (URL, file, or bytes) and return a trajectory + "
            "ethogram summary PNG. engine: 'fast' or 'pose'."
        ),
    )
    gr.api(
        analyze_api,
        api_name="analyze",
        api_description="Same input as `process` (a top-down multi-fly arena video). Return the behavior report (per-fly locomotion/sleep + interactions) as JSON.",
    )

    gr.Markdown(f"# {APP_TITLE}")
    gr.Markdown(APP_DESCRIPTION)

    last_url_state = gr.State("")

    with gr.Row():
        file_input = gr.File(file_types=INPUT_FILE_TYPES, file_count="single",
                             label="Upload a fly video (TIFF stack / mp4)")
        with gr.Column():
            engine_input = gr.Dropdown(choices=ENGINES, value="fast", label="Engine")
            fps_input = gr.Slider(1.0, 60.0, value=15.0, step=1.0, label="Frame rate [Hz]")
            scale_input = gr.Slider(1.0, 20.0, value=4.0, step=0.5, label="Pixels per mm")
            video_input = gr.Checkbox(value=True, label="Render annotated video")
            video_secs_input = gr.Slider(2.0, 10.0, value=5.0, step=0.5, label="Video length [s]")

    if EXAMPLE_PATH:
        gr.Examples(
            examples=[[EXAMPLE_PATH, "fast", 15.0, 4.0]],
            inputs=[file_input, engine_input, fps_input, scale_input],
            label="Try the baked-in synthetic arena!",
        )

    run_btn = gr.Button("Run", variant="primary")
    summary_out = gr.Image(label="Trajectories + ethogram", type="numpy")
    with gr.Row():
        video_out = gr.Video(label="Annotated tracking (with timeline)", autoplay=True)
        report_box = gr.Code(label="Behavior report", language="json")

    run_btn.click(
        fn=_run_ui,
        inputs=[file_input, engine_input, fps_input, scale_input, video_input, video_secs_input],
        outputs=[summary_out, report_box, video_out],
        show_api=False,
    )

    @demo.load(inputs=[last_url_state], outputs=[last_url_state, file_input], show_api=False)
    def _load_from_query(prev_url, request: gr.Request):
        url = (request.query_params or {}).get("file_url") or ""
        if not url or url == prev_url:
            return [gr.skip(), gr.skip()]
        suffix = os.path.splitext(url)[1] or ".mp4"
        fd, tmp_path = tempfile.mkstemp(suffix=suffix, dir=str(APP_TMP_DIR))
        os.close(fd)
        try:
            urllib.request.urlretrieve(url, tmp_path)
        except Exception as e:  # noqa: BLE001
            try:
                os.remove(tmp_path)
            except Exception:  # noqa: BLE001
                pass
            raise gr.Error(f"Failed to download file_url: {e}")
        return [url, gr.update(value=tmp_path)]


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=16).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
        mcp_server=True,
    )