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

# 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 classificação
def classify_image(inp):
    # Redimensione a imagem para o formato esperado pelo modelo (192x256)
    img = Image.fromarray(inp).resize((256, 192))

    # Converta a imagem para um array numpy e normalize-a (escala de 0 a 1)
    img = np.array(img) / 255.0

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

    # Obtém a data e hora atual
    data_hora = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # Formate a saída em HTML
    result_html = f"""
    <h2>Resultado da Classificação</h2>
    <p><strong>Categoria Predita:</strong> {predicted_class}</p>
    <p><strong>Data e Hora:</strong> {data_hora}</p>
    """

    return result_html

# Crie uma interface Gradio
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(shape=(192, 256)),
    outputs=gr.outputs.HTML(),  # Saída formatada com HTML
    live=True,
    title="Classificador de Catarata",
    description="Este é um classificador de imagens médicas para detectar catarata. Faça o upload de uma imagem e clique em 'Classificar' para obter a categoria predita e a data/hora da classificação.",
    allow_screenshot=False,
    allow_flagging=False,
    layout="vertical",
    capture_session=True,
    theme="default",
)

# Inicie a interface Gradio
iface.launch()