import cv2
import spaces
import numpy as np
from tensorflow.keras.models import load_model
import gradio as gr
import tempfile
import os
os.environ['CUDA_VISIBLE_DEVICES'] = "0"
# Load your pre-trained model
model = load_model('cnn_lstm1.h5')

# Function to preprocess each frame
def preprocess_frame(frame):
    resized_frame = cv2.resize(frame, (224, 224))  # Adjust size based on your model's input shape
    normalized_frame = resized_frame / 255.0
    return np.expand_dims(normalized_frame, axis=0)  # Add batch dimension

@spaces.GPU(duration=120)
def predict_drowsiness(video_path):
    # Open the video file
    cap = cv2.VideoCapture(video_path)
    frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(cap.get(cv2.CAP_PROP_FPS))
    
    # Create a temporary file for the output video
    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=False) as temp_output:
        temp_output_path = temp_output.name
    
    # Output video settings
    out = cv2.VideoWriter(temp_output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (frame_width, frame_height))
    
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Preprocess frame
        preprocessed_frame = preprocess_frame(frame)
        
        # Use the model to predict drowsiness
        prediction = model.predict(preprocessed_frame)
        drowsiness = np.argmax(prediction)
        
        # Add label to frame
        label = 'Drowsy' if drowsiness == 0 else 'Alert'
        cv2.putText(frame, label, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # Write the frame with label to the output video
        out.write(frame)
    
    # Release resources
    cap.release()
    out.release()
    
    return temp_output_path  # Return the path to the temporary output video

# Gradio interface
interface = gr.Interface(
    fn=predict_drowsiness,
    inputs=gr.Video(),  # Video input from webcam or upload
    outputs="video",  # Return a playable video with predictions
    title="Drowsiness Detection in Video",
    description="Upload a video or record one, and this tool will detect if the person is drowsy.",
)

# Launch the app
if __name__ == "__main__":
    interface.launch()