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

# Load your model here
model = tf.keras.models.load_model("pain_analysis.h5")  # Ensure you replace this with your actual model path

def predict(image):
    image = np.array(image) / 255.0  # Normalize the image
    # Ensure the image is in the shape your model expects
    image = np.expand_dims(image, axis=0)  # Add batch dimension
    # Make predictions using your model
    result = model.predict(image)  # Perform prediction

    # If your model outputs a probability distribution, you might want to take the argmax
    predicted_class = np.argmax(result, axis=1)

    return predicted_class  # Return the predicted class or whatever output format you need

# Use gr.Image directly for inputs
iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),  # Use gr.Image directly
    outputs="label"  # Specify the output type as needed
)

iface.launch(share=True)  # Set share=True to create a public link