streamlit_app.py

import asyncio
import sys

# --- Boilerplate for compatibility ---
if sys.platform.startswith('linux') and sys.version_info >= (3, 8):
    try:
        asyncio.set_event_loop_policy(asyncio.DefaultEventLoopPolicy())
    except Exception:
        pass

import streamlit as st
from PIL import Image
import numpy as np
import subprocess
import time
import tempfile
import os
from ultralytics import YOLO
import cv2 as cv

# --- NEW: Import your refactored video processing logic ---
from video_processor import process_video_with_progress

model_path="best.pt"

# --- Page Configuration ---
st.set_page_config(
    page_title="Driver Distraction System",
    page_icon="🚗",
    layout="wide",
    initial_sidebar_state="expanded",
)

# --- Sidebar ---
st.sidebar.title("🚗 Driver Distraction System")
st.sidebar.write("Choose an option below:")

# --- Sidebar navigation ---
page = st.sidebar.radio("Select Feature", [
    "Distraction System",
    "Video Drowsiness Detection",
    "Real-time Drowsiness Detection"
])

# --- Class Labels (for YOLO model) ---
st.sidebar.subheader("Class Names")
class_names = ['drinking', 'hair and makeup', 'operating the radio', 'reaching behind',
               'safe driving', 'talking on the phone', 'talking to passenger', 'texting']
for idx, class_name in enumerate(class_names):
    st.sidebar.write(f"{idx}: {class_name}")

# --- Feature: YOLO Distraction Detection ---
if page == "Distraction System":
    st.title("Driver Distraction System")
    st.write("Upload an image or video to detect distractions using YOLO model.")

    # File type selection
    file_type = st.radio("Select file type:", ["Image", "Video"])

    if file_type == "Image":
        uploaded_file = st.file_uploader("Upload Image", type=["jpg", "jpeg", "png"])
        if uploaded_file is not None:
            image = Image.open(uploaded_file).convert('RGB')
            image_np = np.array(image)
            col1, col2 = st.columns([1, 1])
            with col1:
                st.subheader("Uploaded Image")
                st.image(image, caption="Original Image", use_container_width=True)
            with col2:
                st.subheader("Detection Results")
                model = YOLO(model_path)
                start_time = time.time()
                results = model(image_np)
                end_time = time.time()
                prediction_time = end_time - start_time
                result = results[0]
                if len(result.boxes) > 0:
                    boxes = result.boxes
                    confidences = boxes.conf.cpu().numpy()
                    classes = boxes.cls.cpu().numpy()
                    class_names_dict = result.names
                    max_conf_idx = confidences.argmax()
                    predicted_class = class_names_dict[int(classes[max_conf_idx])]
                    confidence_score = confidences[max_conf_idx]
                    st.markdown(f"### Predicted Class: **{predicted_class}**")
                    st.markdown(f"### Confidence Score: **{confidence_score:.4f}**  ({confidence_score*100:.1f}%)")
                    st.markdown(f"Inference Time: {prediction_time:.2f} seconds")
                else:
                    st.warning("No distractions detected.")

# --- Feature: Real-time Drowsiness Detection ---
elif page == "Real-time Drowsiness Detection":
    st.title("🧠 Real-time Drowsiness Detection")
    st.info("This feature requires a local webcam and will open a new window.")
    st.warning("This feature is intended for local use and will not function in the cloud deployment.")
    if st.button("Start Drowsiness Detection"):
        try:
            # This call is fine, as your new drowsiness_detection.py is set up to handle it.
            subprocess.Popen(["python3", "drowsiness_detection.py", "--mode", "webcam"])
            st.success("Attempted to launch detection window. Please check your desktop.")
        except Exception as e:
            st.error(f"Failed to start process: {e}")

# --- Feature: Video Drowsiness Detection ---
elif page == "Video Drowsiness Detection":
    st.title("📹 Video Drowsiness Detection")
    st.write("Upload a video file to detect drowsiness and generate a report.")
    uploaded_video = st.file_uploader("Upload Video", type=["mp4", "avi", "mov", "mkv", "webm"])

    if uploaded_video is not None:
        # Create a temporary file to hold the uploaded video
        tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
        tfile.write(uploaded_video.read())
        temp_input_path = tfile.name
        temp_output_path = tempfile.mktemp(suffix="_processed.mp4")

        st.subheader("Original Video Preview")
        st.video(uploaded_video)

        if st.button("Process Video for Drowsiness Detection"):
            progress_bar = st.progress(0, text="Preparing to process video...")
            
            # --- NEW: Define a callback function for the progress bar ---
            def streamlit_progress_callback(current, total):
                if total > 0:
                    percent_complete = int((current / total) * 100)
                    progress_bar.progress(percent_complete, text=f"Analyzing frame {current}/{total}...")

            try:
                with st.spinner("Processing video... This may take a while."):
                    # --- NEW: Directly call your robust video processing function ---
                    # No more complex subprocess logic needed!
                    stats = process_video_with_progress(
                        input_path=temp_input_path,
                        output_path=temp_output_path,
                        progress_callback=streamlit_progress_callback
                    )
                
                progress_bar.progress(100, text="Video processing completed!")
                st.success("Video processed successfully!")

                # --- NEW: Display the returned statistics ---
                st.subheader("Detection Results")
                col1, col2, col3 = st.columns(3)
                col1.metric("Drowsy Events", stats.get('drowsy_events', 0))
                col2.metric("Yawn Events", stats.get('yawn_events', 0))
                col3.metric("Head Down Events", stats.get('head_down_events', 0))

                # Offer the processed video for download
                if os.path.exists(temp_output_path):
                    with open(temp_output_path, "rb") as file:
                        video_bytes = file.read()
                    st.download_button(
                        label="📥 Download Processed Video",
                        data=video_bytes,
                        file_name=f"drowsiness_detected_{uploaded_video.name}",
                        mime="video/mp4"
                    )
            except Exception as e:
                st.error(f"An error occurred during video processing: {e}")
            finally:
                # Cleanup temporary files
                try:
                    if os.path.exists(temp_input_path): os.unlink(temp_input_path)
                    if os.path.exists(temp_output_path): os.unlink(temp_output_path)
                except Exception as e_clean:
                    st.warning(f"Failed to clean up temporary files: {e_clean}")