import tensorflow as tf
import efficientnet.tfkeras as efn
import gradio as gr
import numpy as np

# Dimensões da imagem
IMG_HEIGHT = 512
IMG_WIDTH = 512

# Função para construir o modelo
def build_original_model():
    inp = tf.keras.layers.Input(shape=(IMG_HEIGHT, IMG_WIDTH, 3))
    efnet = efn.EfficientNetB3(  # Usando EfficientNetB3
        input_shape=(IMG_HEIGHT, IMG_WIDTH, 3),
        weights='imagenet',
        include_top=False
    )
    x = efnet(inp)
    x = tf.keras.layers.GlobalAveragePooling2D()(x)
    x = tf.keras.layers.Dense(2, activation='softmax')(x)
    model = tf.keras.Model(inputs=inp, outputs=x)
    opt = tf.keras.optimizers.Adam(learning_rate=0.000003)
    loss = tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.01)
    model.compile(optimizer=opt, loss=loss, metrics=['accuracy'])
    return model

# Carregue o modelo treinado
original_model = build_original_model()
original_model.load_weights('modelo_treinado.h5')

# Função para realizar o pré-processamento da imagem de entrada
def preprocess_image(input_image):
    # Redimensione a imagem para as dimensões esperadas pelo modelo
    input_image = tf.image.resize(input_image, (IMG_HEIGHT, IMG_WIDTH))
    
    # Normalização dos valores de pixel para o intervalo [0, 1]
    input_image = input_image / 255.0

    return input_image

# Função para fazer previsões usando o modelo treinado
def predict(input_image):
    # Realize o pré-processamento na imagem de entrada
    input_image = preprocess_image(input_image)

    # Faça uma previsão usando o modelo carregado
    input_image = tf.expand_dims(input_image, axis=0)
    prediction = original_model.predict(input_image)

    # A saída será uma matriz de previsões
    class_names = ["Normal", "Cataract"]
    probabilities = prediction[0]

    # Determine a classe mais provável
    predicted_class = class_names[np.argmax(probabilities)]

    # Adicione informações adicionais à saída
    confidence = probabilities[np.argmax(probabilities)]

    return {
        "Imagem de Entrada": input_image[0].numpy(),
        "Classificação": predicted_class,
        "Confiança": confidence,
        "Probabilidades": {class_names[i]: probabilities[i] for i in range(len(class_names))}
    }

# Crie uma interface Gradio para fazer previsões
iface = gr.Interface(
    predict,
    inputs=gr.inputs.Image(label="Carregue uma imagem da região ocular", type="pil"),
    outputs=[
        gr.outputs.Image(label="Imagem de Entrada"),
        gr.outputs.Textbox(label="Classificação"),
        gr.outputs.Textbox(label="Confiança"),
        gr.outputs.Table(label="Probabilidades")
    ],
    theme="auto",  # Usar o tema padrão do Gradio
)

# Execute a interface Gradio
iface.launch()