import gradio as gr
import tensorflow as tf
from tensorflow.keras.models import load_model
from tensorflow.keras.layers import Layer
import numpy as np
from PIL import Image

# Define the custom 'FixedDropout' layer
class FixedDropout(Layer):
    def __init__(self, rate, **kwargs):
        super(FixedDropout, self).__init__(**kwargs)
        self.rate = rate

    def call(self, inputs, training=None):
        if training is None:
            training = K.learning_phase()

        if training == 1:
            return K.dropout(inputs, self.rate)
        return inputs

# Register the custom layer in a custom object scope
custom_objects = {"FixedDropout": FixedDropout}

# Load the TensorFlow model with the custom object scope
tf_model_path = 'modelo_treinado.h5'  # Update with the path to your model
tf_model = load_model(tf_model_path, custom_objects=custom_objects)

# Class labels for the model
class_labels = ["Normal", "Cataract"]

# Define a function for prediction
def predict(image):
    # Preprocess the input image
    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 a prediction using the loaded TensorFlow model
    predictions = tf_model.predict(image)

    # Get the predicted class label
    predicted_label = class_labels[np.argmax(predictions)]

    return predicted_label

# Create the Gradio interface
gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=2)
).launch()