import gradio as gr
import numpy as np
import tensorflow as tf
from tensorflow.keras.models import load_model

# Load your trained pain classification model
model = load_model("pain_analysis.h5")  # Adjust the path as necessary

def predict(image):
    # Resize the image to the expected input size of (148, 148)
    target_size = (148, 148)
    image = tf.image.resize(image, target_size)

    # Normalize the image to [0, 1]
    image = np.array(image) / 255.0
    image = np.expand_dims(image, axis=0)  # Add batch dimension

    # Perform prediction
    result = model.predict(image)
    predicted_class = np.argmax(result, axis=1)  # Get the predicted class

    # Map the predicted class index to pain levels
    pain_levels = {
        0: "No Pain",
        1: "Low Pain",
        2: "Medium Pain",
        3: "High Pain",
    }
    
    return pain_levels[predicted_class[0]]  # Return the corresponding pain level

# Define the Gradio interface
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),  # Expecting image input as numpy array
    outputs="text",  # Return the predicted pain level as text
    title="Pain Level Classification Model",
    description="Upload an image to classify the pain level using the trained model."
)

# Launch the app
iface.launch()