Update streamlit_app.py

streamlit_app.py

@@ -1,11 +1,13 @@
 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
@@ -15,7 +17,9 @@ import tempfile
 import os
 from ultralytics import YOLO
 import cv2 as cv
-import pandas as pd
+
+# --- NEW: Import your refactored video processing logic ---
+from video_processor import process_video_with_progress
 
 model_path="best.pt"
 
@@ -31,19 +35,17 @@ st.set_page_config(
 st.sidebar.title("🚗 Driver Distraction System")
 st.sidebar.write("Choose an option below:")
 
-# Sidebar navigation
+# --- Sidebar navigation ---
 page = st.sidebar.radio("Select Feature", [
-    "Distraction System", 
-    "Real-time Drowsiness Detection",
-    "Video Drowsiness Detection"
+    "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']
-
-# Sidebar Class Name Display
-st.sidebar.subheader("Class Names")
 for idx, class_name in enumerate(class_names):
     st.sidebar.write(f"{idx}: {class_name}")
 
@@ -86,233 +88,80 @@ if page == "Distraction System":
                 else:
                     st.warning("No distractions detected.")
 
-    else:  # Video processing
-        uploaded_video = st.file_uploader("Upload Video", type=["mp4", "avi", "mov", "mkv", "webm"])
-
-        if uploaded_video is not None:
-            tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
-            tfile.write(uploaded_video.read())
-            temp_input_path = tfile.name
-            temp_output_path = tempfile.mktemp(suffix="_distraction_detected.mp4")
-
-            st.subheader("Video Information")
-            cap = cv.VideoCapture(temp_input_path)
-            fps = cap.get(cv.CAP_PROP_FPS)
-            width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
-            height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
-            total_frames = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
-            duration = total_frames / fps if fps > 0 else 0
-            cap.release()
-
-            col1, col2 = st.columns(2)
-            with col1:
-                st.metric("Duration", f"{duration:.2f} seconds")
-                st.metric("Original FPS", f"{fps:.2f}")
-            with col2:
-                st.metric("Resolution", f"{width}x{height}")
-                st.metric("Total Frames", total_frames)
-
-            st.subheader("Original Video Preview")
-            st.video(uploaded_video)
-
-            if st.button("Process Video for Distraction Detection"):
-                TARGET_PROCESSING_FPS = 10
-                # --- NEW: Hyperparameter for the temporal smoothing logic ---
-                PERSISTENCE_CONFIDENCE_THRESHOLD = 0.40  # Stick with old class if found with >= 40% confidence
-
-                st.info(f"🚀 For faster results, video will be processed at ~{TARGET_PROCESSING_FPS} FPS.")
-                st.info(f"🧠 Applying temporal smoothing to reduce status flickering (Persistence Threshold: {PERSISTENCE_CONFIDENCE_THRESHOLD*100:.0f}%).")
-                
-                progress_bar = st.progress(0, text="Starting video processing...")
-                
-                with st.spinner(f"Processing video... This may take a while."):
-                    model = YOLO(model_path)
-                    cap = cv.VideoCapture(temp_input_path)
-                    
-                    fourcc = cv.VideoWriter_fourcc(*'mp4v')
-                    out = cv.VideoWriter(temp_output_path, fourcc, fps, (width, height))
-                    
-                    frame_skip_interval = max(1, round(fps / TARGET_PROCESSING_FPS))
-                    
-                    frame_count = 0
-                    last_best_box_coords = None
-                    last_best_box_label = ""
-                    last_status_text = "Status: Initializing..."
-                    last_status_color = (128, 128, 128)
-                    # --- NEW: State variable to store the last confirmed class ---
-                    last_confirmed_class_name = 'safe driving'
-
-                    while cap.isOpened():
-                        ret, frame = cap.read()
-                        if not ret:
-                            break
-                        
-                        frame_count += 1
-                        progress = int((frame_count / total_frames) * 100) if total_frames > 0 else 0
-                        progress_bar.progress(progress, text=f"Analyzing frame {frame_count}/{total_frames}")
-                        
-                        annotated_frame = frame.copy()
-
-                        if frame_count % frame_skip_interval == 0:
-                            results = model(annotated_frame)
-                            result = results[0]
-                            
-                            last_best_box_coords = None # Reset box for this processing cycle
-
-                            if len(result.boxes) > 0:
-                                boxes = result.boxes
-                                class_names_dict = result.names
-                                confidences = boxes.conf.cpu().numpy()
-                                classes = boxes.cls.cpu().numpy()
-
-                                # --- NEW STABILITY LOGIC ---
-                                final_box_to_use = None
-                                
-                                # 1. Check if the last known class exists with reasonable confidence
-                                for i in range(len(boxes)):
-                                    current_class_name = class_names_dict[int(classes[i])]
-                                    if current_class_name == last_confirmed_class_name and confidences[i] >= PERSISTENCE_CONFIDENCE_THRESHOLD:
-                                        final_box_to_use = boxes[i]
-                                        break
-                                
-                                # 2. If not, fall back to the highest confidence detection in the current frame
-                                if final_box_to_use is None:
-                                    max_conf_idx = confidences.argmax()
-                                    final_box_to_use = boxes[max_conf_idx]
-                                # --- END OF NEW LOGIC ---
-
-                                # Now, process the determined "final_box_to_use"
-                                x1, y1, x2, y2 = final_box_to_use.xyxy[0].cpu().numpy()
-                                confidence = final_box_to_use.conf[0].cpu().numpy()
-                                class_id = int(final_box_to_use.cls[0].cpu().numpy())
-                                class_name = class_names_dict[class_id]
-
-                                # Update the state for the next frames
-                                last_confirmed_class_name = class_name 
-                                last_best_box_coords = (int(x1), int(y1), int(x2), int(y2))
-                                last_best_box_label = f"{class_name}: {confidence:.2f}"
-                                
-                                if class_name != 'safe driving':
-                                    last_status_text = f"Status: {class_name.replace('_', ' ').title()}"
-                                    last_status_color = (0, 0, 255)
-                                else:
-                                    last_status_text = "Status: Safe Driving"
-                                    last_status_color = (0, 128, 0)
-                            else:
-                                # No detections, reset to safe driving
-                                last_confirmed_class_name = 'safe driving'
-                                last_status_text = "Status: Safe Driving"
-                                last_status_color = (0, 128, 0)
-                        
-                        # Draw annotations on EVERY frame using the last known data
-                        if last_best_box_coords:
-                            cv.rectangle(annotated_frame, (last_best_box_coords[0], last_best_box_coords[1]), 
-                                         (last_best_box_coords[2], last_best_box_coords[3]), (0, 255, 0), 2)
-                            cv.putText(annotated_frame, last_best_box_label, 
-                                       (last_best_box_coords[0], last_best_box_coords[1] - 10), 
-                                       cv.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)
-
-                        # Draw status text
-                        font_scale, font_thickness = 1.0, 2
-                        (text_w, text_h), _ = cv.getTextSize(last_status_text, cv.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
-                        padding = 10
-                        rect_start = (padding, padding)
-                        rect_end = (padding + text_w + padding, padding + text_h + padding)
-                        cv.rectangle(annotated_frame, rect_start, rect_end, last_status_color, -1)
-                        text_pos = (padding + 5, padding + text_h + 5)
-                        cv.putText(annotated_frame, last_status_text, text_pos, cv.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness)
-                        
-                        out.write(annotated_frame)
-                    
-                    cap.release()
-                    out.release()
-                    progress_bar.progress(100, text="Video processing completed!")
-                    
-                    st.success("Video processed successfully!")
-                    
-                    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"distraction_detected_{uploaded_video.name}",
-                            mime="video/mp4",
-                            key="download_distraction_video"
-                        )
-                        
-                        st.subheader("Sample Frame from Processed Video")
-                        cap_out = cv.VideoCapture(temp_output_path)
-                        ret, frame = cap_out.read()
-                        if ret:
-                            frame_rgb = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
-                            st.image(frame_rgb, caption="Sample frame with distraction detection", use_container_width=True)
-                        cap_out.release()
-                
-                try:
-                    os.unlink(temp_input_path)
-                    if os.path.exists(temp_output_path): os.unlink(temp_output_path)
-                except Exception as e:
-                    st.warning(f"Failed to clean up temporary files: {e}")
-
 # --- Feature: Real-time Drowsiness Detection ---
 elif page == "Real-time Drowsiness Detection":
     st.title("🧠 Real-time Drowsiness Detection")
-    st.write("This will open your webcam and run the detection script.")
+    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"):
-        with st.spinner("Launching webcam..."):
+        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("Drowsiness detection started in a separate window. Press 'q' in that window to quit.")
+            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 download the processed video.")
+    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...")
-            with st.spinner("Processing video... This may take a while."):
-                process = subprocess.Popen([
-                    "python3", "drowsiness_detection.py", 
-                    "--mode", "video",
-                    "--input", temp_input_path,
-                    "--output", temp_output_path
-                ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
-                stdout, stderr = process.communicate()
-                if process.returncode == 0:
-                    progress_bar.progress(100, text="Video processing completed!")
-                    if os.path.exists(temp_output_path):
-                        st.success("Video processed successfully!")
-                        if stdout: st.code(stdout)
-                        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",
-                            key="download_processed_video"
-                        )
-                        st.subheader("Sample Frame from Processed Video")
-                        cap = cv.VideoCapture(temp_output_path)
-                        ret, frame = cap.read()
-                        if ret: st.image(cv.cvtColor(frame, cv.COLOR_BGR2RGB), caption="Sample frame with drowsiness detection", use_container_width=True)
-                        cap.release()
-                    else:
-                        st.error("Error: Processed video file not found.")
-                        if stderr: st.code(stderr)
-                else:
-                    st.error("An error occurred during video processing.")
-                    if stderr: st.code(stderr)
+            
+            # --- 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:
-                if os.path.exists(temp_input_path): os.unlink(temp_input_path)
-                if os.path.exists(temp_output_path): os.unlink(temp_output_path)
+                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.warning(f"Failed to clean up temporary files: {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}")