app.py

import atexit
import os
import sys
from pathlib import Path

import gradio as gr

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

from env_utils import load_dotenv_if_present, require_secrets
from eula_tab import build_eula_tab
from eve_app_tabs import build_live_inference_tab, build_offline_inference_tab
from eve_inference_handlers import EveAppHandlers, patch_video_for_external_urls
from eve_worker_pool import EveWorkerPool
from face_id_tab import FaceIdTab
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 session_tracker import SessionTracker
from usage_analytics import UsageTracker
from video_file_server import VideoFileServer
from video_processing import VideoLimits, get_example_videos, wire_video_upload


def _build_feature_checkboxes(
    hint: str = "",
) -> tuple[gr.Checkbox, gr.Checkbox, gr.Checkbox, gr.Checkbox]:
    """Build the EVE feature checkbox group used by both tabs.

    Must be called inside a Gradio layout context.
    """
    with gr.Group():
        label = "<p><b>Features</b>"
        if hint:
            label += f" — <em>{hint}</em>"
        label += "</p>"
        gr.HTML(label)
        cb_face = gr.Checkbox(label="Face Detection", value=True)
        cb_person = gr.Checkbox(label="Person Detection", value=True)
        cb_face_id = gr.Checkbox(label="Face Identification", value=False)
        cb_hand_gesture = gr.Checkbox(label="Hand Detection", value=False)
    return cb_face, cb_person, cb_face_id, cb_hand_gesture


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)

    # FPS cap: each live stream runs between MAX_TARGET_FPS (idle) and
    # MIN_TARGET_FPS (full load).  Set either to 0 to disable the cap.
    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)

    # Separate HTTP server for video output — bypasses Chrome's per-origin
    # connection limit that blocks /file= requests while SSE connections are open.
    video_server: VideoFileServer | None = None
    # TODO: Put that under an env var, for debug purposes, locally, for more than 2 tabs.
    """if not os.environ.get("SPACE_ID"):
        gradio_cache = os.path.join(tempfile.gettempdir(), "gradio")
        os.makedirs(gradio_cache, exist_ok=True)
        video_server = VideoFileServer(root_dir=gradio_cache, logger=logger)
        video_server.start()"""

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

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

    atexit.register(_shutdown)

    EXAMPLE_VIDEOS = get_example_videos(str(APP_DIR / "examples"))
    VIDEO_LIMITS = VideoLimits()
    rtc_config_provider = RtcConfigProvider()

    face_id_tab = FaceIdTab(
        pool,
        max_users=2,
        examples_dir=str(APP_DIR / "examples_fid"),
        accept_video=True,
        video_limits=VideoLimits(max_file_size_mb=150, max_duration_seconds=60),
        tracker=tracker,
    )

    with gr.Blocks(
        title="Eve HMI 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 Edge Vision Engine SDK")
        gr.Markdown(
            "Our SDK solves the human sensing challenges by outputting ready-to-use data."
            " Our models have a low computation footprint and are ideal for 0.5-2 TOPS devices"
            " like FPGAs, SOCs and small NPUs.\n\n"
            # TODO: Insert performance summary table
            "To access the EVE SDK, fill the form on this [page](https://huggingface.co/LatticeSemi/sensAI-Edge-Vision-Engine-SDK-Packages) and follow the instructions for download."
            " For any questions or support, please reach out to us at evehelp@latticesemi.com.\n\n"
            "You can also preview the EVE SDK with the following tabs:\n\n"
            "- <u>Live Inference</u> to run it live from your webcam\n"
            "- <u>Offline Inference</u> to test it with videos you can upload\n"
            "- Use the <u>Face ID Registration</u> tab to register face(s) you can use in either Live or Offline Inference to test the Face ID model.\n\n"
            "> Note: 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.\n"
            "\n\n"
        )

        session_registry = gr.State(value={})
        session_hash_state = gr.State(value="unknown")

        with gr.Tabs() as tabs:
            (
                _live_tab,
                webrtc_stream,
                (live_cb_face, live_cb_person, live_cb_face_id, live_cb_hand_gesture),
            ) = build_live_inference_tab(
                rtc_configuration=lambda: rtc_config_provider.get(),
                feature_checkbox_builder=_build_feature_checkboxes,
                extras_builder=lambda: face_id_tab.build_summary(height=80, scale=1),
                max_fps=int(max_fps) if max_fps else 30,
                width=camera_width,
                height=camera_height,
            )

            (
                video_tab,
                input_video,
                output_video,
                (cb_face, cb_person, cb_face_id, cb_hand_gesture),
                process_btn,
                example_dataset,
            ) = build_offline_inference_tab(
                feature_checkbox_builder=_build_feature_checkboxes,
                example_videos=EXAMPLE_VIDEOS,
                video_limits=VIDEO_LIMITS,
                extras_builder=lambda: face_id_tab.build_summary(height=80, scale=1),
            )
            patch_video_for_external_urls(output_video)

            face_id_tab.build()

            build_eula_tab()

            if os.environ.get("ENABLE_PROFILER", "").strip() not in ("", "0", "false"):
                from profiler_tab import build_profiler_tab

                build_profiler_tab(pool)

        # --- Offline Inference wiring ---
        wire_video_upload(input_video, output_video, process_btn, example_dataset, VIDEO_LIMITS)

        process_btn.click(
            fn=handlers.run_eve_inference,
            inputs=[
                input_video,
                cb_face,
                cb_person,
                cb_face_id,
                cb_hand_gesture,
                session_registry,
            ],
            outputs=[output_video, session_registry],
            concurrency_limit=pool.worker_count,
        ).then(
            fn=face_id_tab.refresh_all,
            inputs=[session_registry],
            outputs=face_id_tab.all_slot_components,
        )

        # --- Face ID wiring (self-contained in FaceIdTab) ---
        face_id_tab.wire(session_registry)

        # Refresh summary thumbnails when switching to tabs that show them
        for tab in (video_tab, _live_tab):
            tab.select(
                fn=face_id_tab.refresh_summary,
                inputs=[session_registry],
                outputs=face_id_tab.summary_components,
            )

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

        live_feature_inputs = [live_cb_face, live_cb_person, live_cb_face_id, live_cb_hand_gesture]
        for cb in live_feature_inputs:
            cb.change(fn=handlers.on_live_feature_change, inputs=live_feature_inputs, outputs=[])

        # --- Live Inference wiring ---
        # FastRTC defaults concurrency_limit to 1 which causes Gradio's queue
        # to reject new WebRTC connections with "Too many concurrent connections".
        # Since process_live_frame is non-blocking (shows an overlay while
        # waiting for a worker), we allow MORE streams than workers so that
        # queued users see a "waiting" overlay instead of a connection error.
        webrtc_stream.stream(
            fn=handlers.process_live_frame,
            inputs=[
                webrtc_stream,
                live_cb_face,
                live_cb_person,
                live_cb_face_id,
                live_cb_hand_gesture,
                session_registry,
                session_hash_state,
            ],
            outputs=[webrtc_stream],
            concurrency_limit=pool.worker_count + 8,
        )

        # --- Session tracking ---
        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)