Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import cv2 | |
| from ultralytics import YOLO | |
| import cvzone | |
| import math | |
| import os | |
| os.environ["SDL_AUDIODRIVER"] = "dummy" | |
| 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("") | |