app.py

import atexit
import os
import sys
from pathlib import Path

import gradio as gr

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "shared"))
sys.path.insert(0, str(Path(__file__).resolve().parent / "shared"))

from env_utils import load_dotenv_if_present, require_secrets
from eula_tab import build_eula_tab
from eve_app_tabs import build_offline_inference_tab, build_live_inference_tab
from eve_inference_handlers import EveAppHandlers, patch_video_for_external_urls
from eve_worker_pool import EveWorkerPool
from live_inference import (
    TAB_SWITCH_AUTO_STOP_JS,
    RtcConfigProvider,
    patch_aioice_stun_transaction,
    patch_aiortc_h264_nvenc,
    patch_fastrtc_frame_queue,
    patch_fastrtc_yuv420p_output,
)
from live_stream_manager import LiveStreamManager
from log_utils import log_cpu_info, setup_logger
from mod_models import DEFAULT_MOD_MODEL, MOD_MODEL_REGISTRY, MOD_MODELS
from session_tracker import SessionTracker
from usage_analytics import UsageTracker
from video_processing import VideoLimits

_VIDEO_EXTS = (".mp4", ".avi", ".mov", ".mkv", ".webm")


def _build_feature_radio(hint: str = "") -> gr.Radio:
    """Build the model-selection radio — exactly one MOD model active at a time."""
    label = "Object Detection Model"
    if hint:
        label += f" — {hint}"
    return gr.Radio(
        choices=MOD_MODELS,
        value=DEFAULT_MOD_MODEL,
        label=label,
        info="Pick one object detector: GMOD Base Model (generic 80-class), Automotive Object Detector (8-class)"
        ", or Office Object Detector (8-class).",
        interactive=True,
    )


def _scan_examples(folder: Path, exts: tuple[str, ...]) -> list[list[str]]:
    """Scan ``folder`` for files with any extension in ``exts``."""
    if not folder.is_dir():
        return []
    return [[str(p)] for p in sorted(folder.iterdir()) if p.suffix.lower() in exts]


if __name__ == "__main__":
    patch_fastrtc_frame_queue()
    patch_fastrtc_yuv420p_output()
    patch_aioice_stun_transaction()
    patch_aiortc_h264_nvenc()

    load_dotenv_if_present()
    logger = setup_logger(name="app")
    log_cpu_info(logger)
    require_secrets("MODEL_ACCESS_TOKEN")

    tracker = UsageTracker(
        repo_id=os.environ.get("ANALYTICS_REPO_ID", "LatticeSemi/STAGING-Demo-Analytics-v1.0"),
    )
    tracker.log("server", "server_start")

    max_workers = int(os.environ.get("MAX_WORKERS", os.cpu_count()))
    max_ram_gb = float(os.environ.get("MAX_RAM_GB", 32))
    pool = EveWorkerPool(max_workers=max_workers, max_ram_gb=max_ram_gb, ram_headroom_gb=2.0)

    max_fps_raw = float(os.environ.get("MAX_TARGET_FPS", "24"))
    min_fps_raw = float(os.environ.get("MIN_TARGET_FPS", "15"))
    max_fps: float | None = max_fps_raw if max_fps_raw > 0 else None
    min_fps: float | None = min_fps_raw if min_fps_raw > 0 else None

    camera_width = int(os.environ.get("CAMERA_WIDTH", 640))
    camera_height = int(os.environ.get("CAMERA_HEIGHT", 360))

    stream_manager = LiveStreamManager(
        pool,
        session_lifetime_seconds=60 * 4,
        max_fps=max_fps,
        min_fps=min_fps,
        tracker=tracker,
    )

    session_tracker = SessionTracker(pool=pool, tracker=tracker, logger=logger)

    handlers = EveAppHandlers(
        pool=pool,
        stream_manager=stream_manager,
        sessions=session_tracker,
        logger=logger,
        max_fps=max_fps,
        min_fps=min_fps,
        mod_model_registry=MOD_MODEL_REGISTRY,
    )

    def _shutdown() -> None:
        session_tracker.shutdown()
        stream_manager.shutdown()
        tracker.log("server", "server_stop")
        tracker.shutdown()

    atexit.register(_shutdown)

    examples_dir = Path(__file__).resolve().parent / "examples"
    video_examples = _scan_examples(examples_dir, _VIDEO_EXTS)
    video_limits = VideoLimits()
    rtc_config_provider = RtcConfigProvider()

    with gr.Blocks(
        title="Object Detection Demo",
        theme=gr.themes.Default(
            text_size=gr.themes.sizes.text_lg,
            primary_hue=gr.themes.colors.yellow,
        ),
        css=(
            f"#webrtc-stream-col {{ max-width: {camera_width}px !important; margin: 0 auto; }}"
            " .gradio-container h1, .gradio-container .md h1 { font-size: 2.25rem !important; }"
            " .gradio-container h2, .gradio-container .md h2 { font-size: 1.75rem !important; }"
            " .gradio-container h3, .gradio-container .md h3 { font-size: 1.4rem !important; }"
            " .gradio-container button[role='tab'],"
            " .gradio-container button[role='tab'] *"
            " { text-decoration: underline !important; }"
            " .gradio-container .tab-container {"
            " height: auto !important;"
            " overflow: visible !important;"
            " gap: 4px !important;"
            " border-bottom: 2px solid var(--border-color-primary) !important; }"
            " .gradio-container .tab-container::after { display: none !important; }"
            " .gradio-container button[role='tab'] {"
            " height: auto !important;"
            " padding: 10px 20px !important;"
            " border: 1px solid var(--border-color-primary) !important;"
            " border-bottom: none !important;"
            " border-radius: 8px 8px 0 0 !important;"
            " background: var(--background-fill-secondary) !important;"
            " margin-bottom: -2px !important; }"
            " .gradio-container button[role='tab'].selected {"
            " background: var(--primary-500) !important;"
            " color: var(--neutral-950) !important;"
            " border-color: var(--primary-500) !important;"
            " font-weight: 600 !important; }"
            " .gradio-container button[role='tab'].selected::after { display: none !important; }"
        ),
        head=TAB_SWITCH_AUTO_STOP_JS,
    ) as demo:
        gr.Markdown("# Lattice sensAI — Generic Object Detection")
        gr.Markdown(
            "Run object detection on videos or a live camera feed using "
            "one of three EVE SDK models:\n\n"
            "- **GMOD Base Model** — 80-class generic multi-object detector.\n\n"
            "   - Classes, grouped by category:\n\n"
            "     - **People and animals:** person, bird, cat, dog, horse, sheep, cow, elephant, bear, zebra and giraffe.\n"
            "     - **Vehicles:** bicycle, car, motorcycle, airplane, bus, train, truck and boat.\n"
            "     - **Outdoor:** traffic light, fire hydrant, stop sign, parking meter and bench.\n"
            "     - **Accessories:** backpack, umbrella, handbag, tie and suitcase.\n"
            "     - **Sports:** frisbee, skis, snowboard, sports ball, kite, baseball bat, baseball glove, skateboard, surfboard and tennis racket.\n"
            "     - **Kitchen:** bottle, wine glass, cup, fork, knife, spoon and bowl.\n"
            "     - **Food:** banana, apple, sandwich, orange, broccoli, carrot, hot dog, pizza, donut and cake.\n"
            "     - **Furniture:** chair, couch, potted plant, bed, dining table and toilet.\n"
            "     - **Electronics:** tv, laptop, mouse, remote, keyboard and cell phone.\n"
            "     - **Appliances:** microwave, oven, toaster, sink and refrigerator.\n"
            "     - **Indoor:** book, clock, vase, scissors, teddy bear, hair drier and toothbrush.\n\n"
            "- **Automotive Object Detector** — 8-class automotive object detector finetuned from GMOD Base Model.\n\n"
            "   - Classes: person, bicycle, car, motorcycle, bus, truck, traffic light and stop sign.\n"
            "- **Office Object Detector** — 8-class office object detector finetuned from GMOD Base Model.\n\n"
            "   - Classes: bottle, cup, potted plant, laptop, mouse, keyboard, cell phone and book.\n\n"
            "Only one model is active at a time. Pick the model with the radio button below."
        )
        gr.Markdown(
            "A guide on how to finetune the GMOD Base Model with your own data will be available soon, along with the release 7.3 of the EVE SDK.\n\n"
            "For any questions or support, please reach out to us at evehelp@latticesemi.com."
        )

        # Required by EveAppHandlers signatures (Face ID gallery state). Empty
        # for this demo — Face ID is not exposed.
        session_registry = gr.State(value={})
        session_hash_state = gr.State(value="unknown")
        # Constant False states for the four Face/Person/FaceID/Hand flags
        # that EveAppHandlers expects but this demo does not expose.
        false_state = gr.State(value=False)

        with gr.Tabs() as tabs:
            live_tab, webrtc_stream, live_radio = build_live_inference_tab(
                rtc_configuration=lambda: rtc_config_provider.get(),
                feature_checkbox_builder=_build_feature_radio,
                max_fps=int(max_fps) if max_fps else 30,
                width=camera_width,
                height=camera_height,
            )

            _, video_input, video_output, offline_radio, process_btn, video_example_dataset = build_offline_inference_tab(
                feature_checkbox_builder=_build_feature_radio,
                example_videos=video_examples,
                video_limits=video_limits,
            )
            patch_video_for_external_urls(video_output)

            build_eula_tab()

        gr.Markdown(
            "<p style='font-size:1.0em;text-align:center'>"
            "<strong>Note:</strong> For demo purposes, execution of the AI pipeline and image draw "
            "operations are all performed on a Hugging Face CPU server. Performance may vary based on "
            "the number of concurrent users.</p>"
        )

        # --- Examples wire-up: clicking a sample loads it into the matching input ---
        video_example_dataset.click(
            fn=lambda sample: sample[0],
            inputs=[video_example_dataset],
            outputs=[video_input],
        )

        # --- Process button gating ---
        video_input.change(
            fn=lambda video_path: gr.update(interactive=bool(video_path)),
            inputs=[video_input],
            outputs=[process_btn],
        )

        # --- Process dispatcher ---
        def _process_dispatch(
            video_path: str | None,
            mod_model: str,
            registry: dict,
            request: gr.Request,
            progress: gr.Progress = gr.Progress(),
        ):
            if not video_path:
                raise gr.Error("Please upload a video.")
            output_path, registry = handlers.run_eve_inference(
                video_path,
                False,
                False,
                False,
                False,
                registry,
                mod_model,
                request,
                progress,
            )
            return output_path, registry

        process_btn.click(
            fn=_process_dispatch,
            inputs=[
                video_input,
                offline_radio,
                session_registry,
            ],
            outputs=[
                video_output,
                session_registry,
            ],
            concurrency_limit=pool.worker_count,
        )

        # --- Live feature change tracking ---
        live_radio.change(
            fn=handlers.on_live_feature_change,
            inputs=[false_state, false_state, false_state, false_state, live_radio],
            outputs=[],
        )

        # --- Live inference wiring ---
        webrtc_stream.stream(
            fn=handlers.process_live_frame,
            inputs=[
                webrtc_stream,
                false_state,
                false_state,
                false_state,
                false_state,
                session_registry,
                session_hash_state,
                live_radio,
            ],
            outputs=[webrtc_stream],
            concurrency_limit=pool.worker_count + 8,
        )

        # --- Tab switching analytics ---
        tabs.select(fn=session_tracker.on_tab_switch, inputs=[], outputs=[])

        # --- Session lifecycle ---
        demo.load(session_tracker.on_load, outputs=[session_hash_state])
        demo.unload(handlers.cleanup_session)

        if tracker.enabled:
            gr.HTML(
                "<p style='font-size:0.8em;color:#888;text-align:center;margin-top:1em'>"
                "This demo collects anonymous usage data (session activity, feature usage) "
                "to improve the experience. No personal information is stored.</p>"
            )

    demo.queue()
    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)