File size: 7,272 Bytes
530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 7317aac |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 |
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}") |