Update app.py

app.py

@@ -1,27 +1,42 @@
 import gradio as gr
 import numpy as np
 import tensorflow as tf
+from tensorflow.keras.models import load_model
 
-# Load your model here
-model = tf.keras.models.load_model("pain_analysis.h5")  # Ensure you replace this with your actual model path
+# Load your trained pain classification model
+model = load_model("path/to/your/pain_analysis.h5")  # Adjust the path as necessary
 
 def predict(image):
-    image = np.array(image) / 255.0  # Normalize the image
-    # Ensure the image is in the shape your model expects
+    # 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
-    # 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)
+    # Perform prediction
+    result = model.predict(image)
+    predicted_class = np.argmax(result, axis=1)  # Get the predicted class
 
-    return predicted_class  # Return the predicted class or whatever output format you need
+    # Map the predicted class index to pain levels
+    pain_levels = {
+        0: "No Pain",
+        1: "Mild Pain",
+        2: "Moderate Pain",
+        3: "Severe Pain",
+    }
+    
+    return pain_levels[predicted_class[0]]  # Return the corresponding pain level
 
-# Use gr.Image directly for inputs
+# Define the Gradio interface
 iface = gr.Interface(
     fn=predict,
-    inputs=gr.Image(type="numpy"),  # Use gr.Image directly
-    outputs="label"  # Specify the output type as needed
+    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."
 )
 
-iface.launch(share=True)  # Set share=True to create a public link
+# Launch the app
+iface.launch()