File size: 14,249 Bytes
530a851 7317aac 530a851 7317aac 530a851 7317aac 530a851 6736f11 530a851 6736f11 7317aac 6736f11 530a851 7317aac 530a851 6736f11 530a851 7317aac 530a851 7317aac 530a851 6736f11 7317aac 6736f11 7317aac 6736f11 7317aac 6736f11 |
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 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 285 286 287 288 289 290 291 292 293 294 295 296 297 298 299 300 301 302 303 304 305 306 307 308 |
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
# --- FIXED: Model path handling ---
model_path = "best.pt"
if not os.path.exists(model_path):
st.error(f"Model file '{model_path}' not found. Please ensure it's included in your deployment.")
st.stop()
# --- 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:")
# --- FIXED: Disable webcam feature for cloud deployment ---
if os.getenv("SPACE_ID"): # Running on Hugging Face Spaces
available_features = [
"Distraction System",
"Video Drowsiness Detection"
]
st.sidebar.info("💡 Note: Real-time webcam detection is not available in cloud deployment.")
else:
available_features = [
"Distraction System",
"Video Drowsiness Detection",
"Real-time Drowsiness Detection"
]
# --- Sidebar navigation ---
page = st.sidebar.radio("Select Feature", available_features)
# --- 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:
try:
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")
# Load model with error handling
try:
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.")
except Exception as e:
st.error(f"Error loading or running model: {str(e)}")
st.info("Please ensure the model file 'best.pt' is present and valid.")
except Exception as e:
st.error(f"Error processing image: {str(e)}")
elif file_type == "Video":
uploaded_file = st.file_uploader("Upload Video", type=["mp4", "avi", "mov", "mkv", "webm"])
if uploaded_file is not None:
try:
# Create a temporary file to hold the uploaded video
tfile = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4")
tfile.write(uploaded_file.read())
temp_input_path = tfile.name
temp_output_path = tempfile.mktemp(suffix="_processed.mp4")
st.subheader("Original Video Preview")
st.video(uploaded_file)
if st.button("Process Video for Distraction Detection"):
progress_bar = st.progress(0, text="Preparing to process video...")
try:
model = YOLO(model_path)
cap = cv.VideoCapture(temp_input_path)
total_frames = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
fps = cap.get(cv.CAP_PROP_FPS)
# Get video properties
width = int(cap.get(cv.CAP_PROP_FRAME_WIDTH))
height = int(cap.get(cv.CAP_PROP_FRAME_HEIGHT))
# Setup video writer
fourcc = cv.VideoWriter_fourcc(*'mp4v')
out = cv.VideoWriter(temp_output_path, fourcc, fps, (width, height))
frame_count = 0
detections = []
while True:
ret, frame = cap.read()
if not ret:
break
frame_count += 1
# Process frame with YOLO
results = model(frame)
result = results[0]
# Draw detections on frame
annotated_frame = result.plot()
out.write(annotated_frame)
# Store detection info
if len(result.boxes) > 0:
boxes = result.boxes
for i in range(len(boxes)):
conf = boxes.conf[i].cpu().numpy()
cls = int(boxes.cls[i].cpu().numpy())
class_name = result.names[cls]
detections.append({
'frame': frame_count,
'class': class_name,
'confidence': conf
})
# Update progress
progress = int((frame_count / total_frames) * 100)
progress_bar.progress(progress, text=f"Processing frame {frame_count}/{total_frames}")
cap.release()
out.release()
st.success("Video processed successfully!")
# Show results
st.subheader("Detection Results")
if detections:
# Count detections by class
class_counts = {}
for det in detections:
class_name = det['class']
if class_name not in class_counts:
class_counts[class_name] = 0
class_counts[class_name] += 1
# Display metrics
cols = st.columns(len(class_counts))
for i, (class_name, count) in enumerate(class_counts.items()):
cols[i].metric(class_name.title(), count)
else:
st.info("No distractions detected in the video.")
# Offer 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"distraction_detected_{uploaded_file.name}",
mime="video/mp4"
)
except Exception as e:
st.error(f"Error processing video: {str(e)}")
finally:
# Cleanup
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:
st.warning(f"Failed to clean up temporary files: {e}")
except Exception as e:
st.error(f"Error handling video upload: {str(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:
try:
# 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...")
# --- 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."):
# Call your robust video processing function
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!")
# 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}")
st.info("Please ensure all required model files are present and the video format is supported.")
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}")
except Exception as e:
st.error(f"Error handling video upload: {str(e)}")
# --- Footer ---
st.sidebar.markdown("---")
st.sidebar.markdown("### 📝 Notes")
st.sidebar.markdown("""
- **Image Detection**: Upload JPG, PNG images
- **Video Detection**: Upload MP4, AVI, MOV videos
- **Cloud Limitations**: Webcam access not available in cloud deployment
- **Model**: Uses YOLO for distraction detection
""") |