import gradio as gr
import tensorflow as tf
import cv2
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.backend.dropout(inputs, level=self.rate)

# 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

tf_model_path = 'modelo_treinado.h5'  # Update with the path to your TensorFlow model
tf_model = load_model_with_custom_objects(tf_model_path)

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])

    # 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]

    return predicted_class

# 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

gr.Interface(fn=classify_image, inputs=input_image, outputs=output_label).launch()