Pain-Analysis / app.py
szili2011's picture
Update app.py
57dda0e verified
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()