"""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, )