import tensorflow as tf
from keras.utils import custom_object_scope
import gradio as gr
import numpy as np
from PIL import Image, ImageDraw

# 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
def preprocess_image(inp):
    # Redimensione a imagem para o formato esperado pelo modelo (192x256)
    img = Image.fromarray(inp)
    img = img.resize((256, 192))
    
    # Converta a imagem para um array numpy
    img = np.array(img)
    
    # Normalize a imagem (escala de 0 a 1)
    img = img / 255.0

    return img

# Defina a função de classificação
def classify_image(inp):
    # Pré-processar a imagem
    img = preprocess_image(inp)

    # Faça uma previsão usando o modelo treinado
    prediction = loaded_model.predict(np.expand_dims(img, axis=0)).flatten()

    # Obtém a classe prevista
    predicted_class = class_names[np.argmax(prediction)]

    # Crie uma imagem de saída com a classe prevista
    output_image = Image.new("RGB", (256, 192))
    output_image.paste((255, 255, 255), (0, 0, 256, 192))
    output_image_draw = ImageDraw.Draw(output_image)
    output_image_draw.text((10, 10), f"Classe Prevista: {predicted_class}", fill=(0, 0, 0))

    # Converte a imagem de saída de volta para numpy array
    output_image = np.array(output_image)

    # Retorna uma lista com a imagem de entrada e o rótulo
    return [inp, output_image], predicted_class

# Crie uma interface Gradio
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(shape=(192, 256)),
    outputs=[gr.outputs.Multi([gr.outputs.Image(type="numpy"), gr.outputs.Image(type="numpy")]), gr.outputs.Label()],
    capture_session=True,
    title="Detecção de Catarata",
    description="Esta interface permite a detecção de catarata em imagens de olhos.",
    instructions="Carregue uma imagem de um olho e clique em 'Classificar' para obter a previsão.",
    theme="default",  # Você pode escolher um tema predefinido ou personalizar as cores manualmente.
    button_text="Classificar Imagem",  # Personalize o texto do botão.
)

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