Update app.py

app.py

@@ -1,44 +1,34 @@
 import gradio as gr
 import tensorflow as tf
-import cv2
+from tensorflow.keras.models import load_model
+from PIL import Image
 import numpy as np
 
-# Define the custom layer 'FixedDropout'
-class FixedDropout(tf.keras.layers.Layer):
-    def __init__(self, rate, **kwargs):
-        super(FixedDropout, self).__init__(**kwargs)
-        self.rate = rate
-
-    def call(self, inputs, training=None):
-        return tf.keras.layers.Dropout(self.rate)(inputs, training=training)
-
-# Load the TensorFlow model with custom layer handling
-def load_model_with_custom_objects(model_path):
-    with tf.keras.utils.custom_object_scope({'FixedDropout': FixedDropout}):
-        model = tf.keras.models.load_model(model_path)
-    return model
-
+# Load the TensorFlow model
 tf_model_path = 'modelo_treinado.h5'  # Update with the path to your TensorFlow model
-tf_model = load_model_with_custom_objects(tf_model_path)
+tf_model = load_model(tf_model_path)
 
+# Class labels for the model
 class_labels = ["Normal", "Cataract"]
 
-# Define a Gradio interface
-def classify_image(input_image):
-    # Preprocess the input image
-    input_image = cv2.resize(input_image, (224, 224))  # Resize the image to match the model's input size
-    input_image = np.expand_dims(input_image, axis=0)  # Add batch dimension
-    input_image = input_image / 255.0  # Normalize pixel values (assuming input range [0, 255])
+# Define a function for prediction
+def predict(image):
+    # Preprocess the input image (resize and normalize)
+    image = image.resize((224, 224))  # Adjust the size as needed
+    image = np.array(image) / 255.0  # Normalize pixel values
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
 
-    # Make predictions using the loaded model
-    predictions = tf_model.predict(input_image)
-    class_index = np.argmax(predictions, axis=1)[0]
-    predicted_class = class_labels[class_index]
+    # Make a prediction using the loaded TensorFlow model
+    predictions = tf_model.predict(image)
 
-    return predicted_class
+    # Get the predicted class label
+    predicted_label = class_labels[np.argmax(predictions)]
 
-# Create a Gradio interface
-input_image = gr.inputs.Image(shape=(224, 224, 3))  # Define the input image shape
-output_label = gr.outputs.Label()  # Define the output label
+    return predicted_label
 
-gr.Interface(fn=classify_image, inputs=input_image, outputs=output_label).launch()
+# Create the Gradio interface
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=gr.Label(num_top_classes=2),
+).launch()