Spaces:
Sleeping
Sleeping
File size: 5,114 Bytes
86b2bd4 |
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 |
import streamlit as st
import cv2
from ultralytics import YOLO
import cvzone
import math
import pygame
import numpy as np
# Initialize pygame mixer
pygame.mixer.init()
# Load sound
alert_sound = pygame.mixer.Sound('alarm.mp3')
# Load the model
model = YOLO('best.pt')
# Reading the classes
classnames = ['Drowsy', 'Awake']
# Streamlit UI
st.set_page_config(layout="wide") # Set wide layout
# Add the logo to the sidebar
logo_path = "logo.png" # Use the uploaded file path
st.sidebar.empty() # Add empty space
st.sidebar.image(logo_path, use_column_width=True)
# Create a sidebar for navigation
st.sidebar.title("Options")
page = st.sidebar.selectbox("Choose a page", ["Webcam Detection", "Image Upload"])
st.title("Drowsiness Detection")
if page == "Webcam Detection":
st.header("Real-Time Drowsiness Detection")
# Layout
col1, col2 = st.columns(2)
with col1:
start_button = st.button('Start Webcam')
with col2:
stop_button = st.button('Stop Webcam')
alert_placeholder = st.empty() # Placeholder for alerts
stframe = st.empty()
status_text = st.empty()
message_text = st.empty()
if start_button:
cap = cv2.VideoCapture(0)
drowsy_count = 0 # Counter for consecutive "Drowsy" detections
while cap.isOpened():
ret, frame = cap.read()
if not ret:
status_text.write("Failed to grab frame")
break
frame = cv2.resize(frame, (640, 480))
# Run the model on the frame
result = model(frame, stream=True)
# Flag to track if "Drowsy" is detected in this frame
drowsy_detected = False
# Getting bbox, confidence, and class name information to work with
for info in result:
boxes = info.boxes
for box in boxes:
confidence = box.conf[0]
confidence = math.ceil(confidence * 100)
Class = int(box.cls[0])
if confidence > 50:
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 5)
cvzone.putTextRect(frame, f'{classnames[Class]} {confidence}%', [x1 + 8, y1 + 100],
scale=1.5, thickness=2)
if classnames[Class] == 'Drowsy':
drowsy_detected = True
# Increment the counter if "Drowsy" is detected, otherwise reset the counter
if drowsy_detected:
drowsy_count += 1
status_text.write("Drowsiness detected!")
else:
drowsy_count = 0
status_text.write("Monitoring...")
# Play alert sound and send message if "Drowsy" is detected 3 or more times
if drowsy_count >= 3:
pygame.mixer.Sound.play(alert_sound)
alert_placeholder.markdown(
f'<div style="color: red; font-size: 24px; border: 2px solid red; padding: 10px;">**Be careful! Drowsiness detected!**</div>',
unsafe_allow_html=True,
)
drowsy_count = 0 # Reset the counter after playing the sound
# Convert image back to RGB for Streamlit
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Display the image
stframe.image(frame, channels="RGB")
# Check if stop button is pressed
if stop_button:
break
cap.release()
status_text.write("Webcam stopped.")
message_text.write("")
alert_placeholder.empty()
elif page == "Image Upload":
st.header("Drowsiness Detection on Image")
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
# Read the image
file_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
frame = cv2.imdecode(file_bytes, 1)
# Perform prediction
results = model(frame, stream=True)
# Process the results
for result in results:
boxes = result.boxes
for box in boxes:
confidence = box.conf[0]
confidence = math.ceil(confidence * 100)
Class = int(box.cls[0])
if confidence > 50:
x1, y1, x2, y2 = box.xyxy[0]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 0, 255), 5)
cv2.putText(frame, f'{classnames[Class]} {confidence}%', (x1 + 8, y1 + 100),
cv2.FONT_HERSHEY_SIMPLEX, 1.5, (255, 255, 255), 2, cv2.LINE_AA)
# Convert image back to RGB for Streamlit
frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
# Display the image
st.image(frame, channels="RGB") |