import streamlit as st
import tensorflow as tf
import numpy as np
import cv2
import os
from tensorflow.keras.models import load_model

# Load the trained model
@st.cache_resource
def load_trained_model():
    model = load_model("model.h5")  # Updated model name
    return model

# Preprocess a video to extract frames and prepare for prediction
def preprocess_video(video_path, img_height=224, img_width=224):
    cap = cv2.VideoCapture(video_path)
    frames = []
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        # Resize frame to match model input size
        frame = cv2.resize(frame, (img_height, img_width))
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB
        frames.append(frame)
    cap.release()
    frames = np.array(frames, dtype=np.float32) / 255.0  # Normalize
    return frames

# Predict from video frames
def predict_video(model, frames):
    predictions = model.predict(frames, batch_size=32)  # Batch size for efficient inference
    avg_prediction = np.mean(predictions, axis=0)  # Average over frames
    class_idx = np.argmax(avg_prediction)  # Get the class index
    return class_idx, avg_prediction

# App UI
st.title("Driver Distraction Detection")
st.write("Team 18 Video Project: Sayandip Bhattacharyya, Purnendu Rudrapal, Sridatta Das, Sidhartha Karjee")

# File uploader
uploaded_file = st.file_uploader("Upload a video", type=["mp4", "avi", "mov"])

# Classes (Ensure these match the class names from training)
class_names = ["Microsleep", "Yawning"]

if uploaded_file:
    # Save uploaded video temporarily
    temp_video_path = "temp_video.mp4"
    with open(temp_video_path, "wb") as f:
        f.write(uploaded_file.read())
    
    st.video(temp_video_path)

    st.write("Processing video...")
    # Load model
    model = load_trained_model()
    
    # Preprocess video
    frames = preprocess_video(temp_video_path)
    st.write(f"Extracted {len(frames)} frames for prediction.")
    
    # Make prediction
    class_idx, avg_prediction = predict_video(model, frames)
    
    # Display result
    st.write(f"**Prediction:** {class_names[class_idx]}")
    st.write(f"Confidence: {avg_prediction[class_idx] * 100:.2f}%")
    
    # Cleanup temporary file
    os.remove(temp_video_path)