import streamlit as st
import torch
from pygame import mixer
import numpy as np
import cv2
cv2.__version__
# Load YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'custom', path=r"C:\Users\gokul\PycharmProjects\driver_drowsiness\last.pt", force_reload=True)

# Initialize the mixer
mixer.init()
sound = mixer.Sound(r"C:\Users\gokul\PycharmProjects\driver_drowsiness\alram.wav")

# Initialize the score variable

# Function to start the camera and process frames
def start_camera():
    global score
    score = 0

    cap = cv2.VideoCapture(0)

    while cap.isOpened():
        ret, frame = cap.read()
        # Make detections
        results = model(frame)

        height, width = frame.shape[0:2]  # Extract height and width
        class_id = 16
        if class_id in results.pred[0][:, 5]:
            # Increment the score variable
            score += 1
            if score > 15:
                try:
                    sound.play()
                except:
                    pass

        else:
            print((results.pred[0][:, 5]))
            score -= 1
            if score < 0:
                score = 0

                # Display the score on the frame
        cv2.putText(frame, 'Score: ' + str(score), (100, height - 20),
                    fontFace=cv2.FONT_HERSHEY_COMPLEX_SMALL,
                    fontScale=1, color=(255, 255, 255),
                    thickness=1, lineType=cv2.LINE_AA)

        rendered = np.squeeze(results.render())
        frame[0:rendered.shape[0], 0:rendered.shape[1]] = rendered

        cv2.imshow('YOLO', frame)
        if cv2.waitKey(10) & 0xFF == ord('q'):
            break
    cap.release()
    cv2.destroyAllWindows()
# Streamlit App
st.title("Drowsiness Detection with YOLOv5")

# Button to start the camera
if st.button("Start Camera"):
    start_camera()