import tensorflow as tf
import numpy as np
import cv2
import gradio as gr

# Load model once at startup
try:
    model = tf.keras.models.load_model('eyeStateModel.h5', compile=False)
    print("Model loaded successfully!")
except Exception as e:
    print(f"Error loading model: {str(e)}")
    model = None

def preprocess_image(image):
    # Convert to RGB if needed
    if len(image.shape) == 2:
        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
    elif len(image.shape) == 3 and image.shape[2] == 4:
        image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
    
    # Resize image to 224x224
    image = cv2.resize(image, (224, 224))
    
    # Normalize pixel values
    image = image / 255.0
    
    # Add batch dimension
    image = np.expand_dims(image, axis=0)
    return image

def predict_eye_state(image):
    if model is None:
        return {"Error": 1.0}
        
    try:
        if image is None:
            return {"Error": 1.0}
        
        # Preprocess the image
        processed_image = preprocess_image(image)
        
        # Make prediction
        prediction = model.predict(processed_image)
        
        # Get confidence scores
        closed_conf = float(prediction[0][0])
        open_conf = float(prediction[0][1])
        
        # Return dictionary with numeric values
        return {
            "Eyes Closed": closed_conf,
            "Eyes Open": open_conf
        }
        
    except Exception as e:
        print(f"Error in prediction: {str(e)}")
        return {"Error": 1.0}

# Create Gradio interface
demo = gr.Interface(
    fn=predict_eye_state,
    inputs=gr.Image(),
    outputs=gr.Label(num_top_classes=2),
    title="Eye State Detection",
    description="Upload an image of eyes to detect if they are open or closed."
)

if __name__ == "__main__":
    demo.launch()