katkarrohit203's picture
Update app.py
010ee15 verified
import streamlit as st
import numpy as np
import cv2 as cv
from tensorflow.keras.models import load_model
from tensorflow.keras.preprocessing.image import img_to_array
import pygame
import os
import time
# Load the pre-trained CNN model for drowsiness detection
#model = load_model('drowsiness.h5')
# Define labels
labels = ['Closed', 'Open']
# Load face and eye cascade classifiers
face_cascade = cv.CascadeClassifier(cv.data.haarcascades + "haarcascade_frontalface_default.xml")
eye_cascade = cv.CascadeClassifier(cv.data.haarcascades + "haarcascade_eye.xml")
# Function to initialize audio
def initialize_audio():
try:
if not os.environ.get("DISPLAY"):
raise RuntimeError("No display found. Skipping audio initialization.")
pygame.mixer.init()
pygame.mixer.music.load("alarm.wav")
except Exception as e:
print(f"Audio initialization failed: {e}")
initialize_audio()
# Initialize webcam
cam = cv.VideoCapture(0)
FRAME_WINDOW = st.image([])
run = st.button("Start")
stop = st.button("Stop")
# Function to preprocess the frame for model input
def preprocess_frame(frame):
frame = cv.resize(frame, (24, 24)) # Resize to match model input shape
frame = frame.astype("float") / 255.0 # Normalize
frame = img_to_array(frame)
frame = np.expand_dims(frame, axis=0)
return frame
running = False
closed_start_time = None
alarm_triggered = False
if run:
running = True
while running:
ret, frame = cam.read()
if not ret:
st.warning("Failed to access webcam.")
break
frame = cv.flip(frame, 1)
gray_frame = cv.cvtColor(frame, cv.COLOR_BGR2GRAY) # Convert to grayscale
# Detect faces
faces = face_cascade.detectMultiScale(gray_frame, scaleFactor=1.1, minNeighbors=5)
eye_status = 'Closed'
for (x, y, w, h) in faces:
face_roi = frame[y:y+h, x:x+w]
roi_gray = gray_frame[y:y+h, x:x+w]
eyes = eye_cascade.detectMultiScale(roi_gray, scaleFactor=1.1, minNeighbors=5)
# Draw rectangle around face
cv.rectangle(frame, (x, y), (x+w, y+h), (255, 0, 0), 2)
# Draw ellipses around detected eyes
if len(eyes) > 0:
eye_status = 'Open'
for (ex, ey, ew, eh) in eyes:
center = (x + ex + ew // 2, y + ey + eh // 2)
axes = (ew // 2, eh // 2)
cv.ellipse(frame, center, axes, 0, 0, 360, (0, 255, 0), 2)
closed_start_time = None
if alarm_triggered:
pygame.mixer.music.stop()
alarm_triggered = False
else:
if closed_start_time is None:
closed_start_time = time.time()
elif time.time() - closed_start_time >= 3:
if not alarm_triggered:
try:
pygame.mixer.music.play()
except Exception as e:
print("Alarm triggered, but no audio device found.")
alarm_triggered = True
# Display eye status text
cv.putText(frame, f'Eyes: {eye_status}', (30, 30), cv.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
frame = cv.cvtColor(frame, cv.COLOR_BGR2RGB)
FRAME_WINDOW.image(frame)
# Stop condition
if stop:
running = False
cam.release()
cv.destroyAllWindows()
break
cam.release()
cv.destroyAllWindows()