app.py

"""
app.py
------
Traffic Scene Interpretation System — Streamlit front end.

Scope for this build (agreed scope):
  - Image upload -> detection
  - Video upload -> detection
  - Vehicle counting + basic scene interpretation (congestion label)
  - Download annotated result

Webcam mode is intentionally left out of this version — image + video
upload was confirmed to be sufficient for the supervisor's requirements.
A `webcam_stub.py` stub is included separately with notes on how to add it
later if needed, so it doesn't block delivery of the core system.

Run with:
    streamlit run app.py
"""

import os
import tempfile

import cv2
import numpy as np
import streamlit as st
from PIL import Image

import theme
from detector import TrafficDetector

st.set_page_config(page_title="Traffic Scene Interpretation System", page_icon="🚦", layout="wide")
theme.inject()


@st.cache_resource(show_spinner="Loading YOLO model (first run only)...")
def load_detector() -> TrafficDetector:
    return TrafficDetector()


def render_stats(stats, cam_tag: str):
    avg_counts = stats.per_frame_average()
    theme.chip_row(avg_counts)

    col1, col2 = st.columns([1, 1])
    with col1:
        theme.congestion_badge(stats.congestion_label())
    with col2:
        if stats.fps:
            st.caption(f"PROCESSING SPEED · {stats.fps} FPS")


def image_tab(detector: TrafficDetector):
    uploaded = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"], key="img")

    if uploaded is not None:
        pil_image = Image.open(uploaded).convert("RGB")
        bgr_image = cv2.cvtColor(np.array(pil_image), cv2.COLOR_RGB2BGR)

        with st.spinner("Running detection..."):
            annotated, stats = detector.detect_image(bgr_image)

        annotated_rgb = cv2.cvtColor(annotated, cv2.COLOR_BGR2RGB)

        col1, col2 = st.columns(2)
        with col1:
            with st.container(border=True):
                theme.feed_caption("CAM 01", "RAW FEED")
                st.image(pil_image, use_container_width=True)
        with col2:
            with st.container(border=True):
                theme.feed_caption("CAM 01", "DETECTION OVERLAY")
                st.image(annotated_rgb, use_container_width=True)

        st.markdown('<hr class="tsis-hr">', unsafe_allow_html=True)

        with st.container(border=True):
            theme.feed_caption("SUMMARY", "DETECTION READOUT")
            render_stats(stats, "CAM 01")

            result_pil = Image.fromarray(annotated_rgb)
            buf_path = os.path.join(tempfile.gettempdir(), "annotated_result.png")
            result_pil.save(buf_path)
            with open(buf_path, "rb") as f:
                st.download_button("Download result image", f, file_name="detection_result.png")


def video_tab(detector: TrafficDetector):
    uploaded = st.file_uploader("Upload a video", type=["mp4", "avi", "mov", "mkv"], key="vid")

    if uploaded is not None:
        # Save upload to a temp file since OpenCV needs a real file path
        in_path = os.path.join(tempfile.gettempdir(), f"input_{uploaded.name}")
        out_path = os.path.join(tempfile.gettempdir(), "annotated_output.mp4")
        with open(in_path, "wb") as f:
            f.write(uploaded.read())

        progress_bar = st.progress(0, text="Starting...")

        def update_progress(current, total):
            if total:
                progress_bar.progress(min(current / total, 1.0), text=f"Processing frame {current}/{total}")
            else:
                progress_bar.progress(0, text=f"Processing frame {current}")

        with st.spinner("Running detection on video... this can take a while for longer clips."):
            stats = detector.detect_video(in_path, out_path, progress_callback=update_progress)

        progress_bar.empty()

        with st.container(border=True):
            theme.feed_caption("CAM 02", "DETECTION OVERLAY")
            st.video(out_path)

        st.markdown('<hr class="tsis-hr">', unsafe_allow_html=True)

        with st.container(border=True):
            theme.feed_caption("SUMMARY", "DETECTION READOUT")
            render_stats(stats, "CAM 02")

            with open(out_path, "rb") as f:
                st.download_button("Download result video", f, file_name="detection_result.mp4")


def main():
    theme.masthead()
    theme.hero(
        eyebrow="VEHICLE DETECTION · SCENE ANALYSIS",
        title="Traffic Scene Interpretation System",
        subtitle=(
            "Upload footage from a traffic camera to detect vehicles and pedestrians, "
            "count them by class, and read the overall state of the scene."
        ),
    )

    detector = load_detector()

    tab1, tab2 = st.tabs(["Image feed", "Video feed"])
    with tab1:
        image_tab(detector)
    with tab2:
        video_tab(detector)


if __name__ == "__main__":
    main()