import tensorflow as tf
from keras.utils import custom_object_scope
import gradio as gr
import numpy as np

# Defina a camada personalizada FixedDropout
class FixedDropout(tf.keras.layers.Dropout):
    def __init__(self, rate, **kwargs):
        super().__init__(rate, **kwargs)
        self._rate = rate

    def call(self, inputs):
        return tf.nn.dropout(inputs, self._rate)

# Registre a camada personalizada FixedDropout com o TensorFlow
with custom_object_scope({'FixedDropout': FixedDropout}):
    # Carregue o modelo
    loaded_model = tf.keras.models.load_model('modelo_treinado.h5')

# Crie uma lista de classes
class_names = ["Normal", "Cataract"]

# Defina a função de pré-processamento para redimensionar a imagem
def preprocess_image(inp):
    img = tf.image.resize(inp, (192, 256))  # Redimensiona a imagem para (192, 256)
    return img

# Defina a função de classificação
def classify_image(inp):
    # Faça uma previsão usando o modelo treinado
    prediction = loaded_model.predict(inp).flatten()

    # Retorna a classe prevista
    return class_names[np.argmax(prediction)]

# Crie uma interface Gradio
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(shape=(224, 224)),
    outputs=gr.outputs.Label(),
    capture_session=True
)

# Inicie a interface Gradio
iface.launch(debug=True)