Spaces:
Sleeping
Sleeping
File size: 2,894 Bytes
fd0a8b8 |
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 |
import streamlit as st
import cv2
from ultralytics import YOLO
import cvzone
import math
import pygame
# Initialize pygame mixer
pygame.mixer.init()
# Load sound
alert_sound = pygame.mixer.Sound('siren-alert-96052.mp3') # Using raw string
# Load the model
model = YOLO('best.pt')
# Reading the classes
classnames = ['Drowsy', 'Awake']
# Streamlit UI
st.title("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')
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)
message_text.write("**Be careful!** Drowsiness detected multiple times!")
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("")
|