import gradio as gr
import tensorflow as tf
import numpy as np

# Dicionário para mapear os nomes amigáveis para os arquivos reais
MODELOS = {
    "EfficientNetB4 (Alta Precisão)": "EfficientNetB4_adam.keras",
    "MobileNet (Alta Velocidade)": "MobileNet_adam.keras"
}

# Cache para não precisar recarregar o arquivo toda vez que clicar no botão
cache_modelos = {}

def carregar_modelo(nome_selecionado):
    arquivo = MODELOS[nome_selecionado]
    if arquivo not in cache_modelos:
        print(f"Carregando {arquivo} para o cache...")
        cache_modelos[arquivo] = tf.keras.models.load_model(arquivo)
    return cache_modelos[arquivo]

def realizar_predicao(imagem, nome_modelo):
    # 1. Carregar o modelo selecionado no dropdown
    model = carregar_modelo(nome_modelo)
    
    # 2. Definir o tamanho baseado no modelo
    # EfficientNetB4 usa 380x380, MobileNet usa 224x224
    tamanho = (380, 380) if "EfficientNetB4" in nome_modelo else (224, 224)
    
    # 3. Pré-processar a imagem
    img_redimensionada = tf.image.resize(imagem, tamanho)
    img_array = np.expand_dims(img_redimensionada, axis=0)
    
    # 4. Fazer a predição
    predicoes = model.predict(img_array)
    classe_id = np.argmax(predicoes, axis=1)[0]
    confianca = float(np.max(predicoes))
    
    return f"Modelo: {nome_modelo}\nClasse Detectada: {classe_id}\nConfiança: {confianca:.2%}"

# Criar a Interface Visual
with gr.Blocks() as demo:
    gr.Markdown("Teste de modelo Tensorflow")
    gr.Markdown("Selecione qual modelo a IA deve usar para analisar a foto.")
    
    with gr.Row():
        with gr.Column():
            seletor = gr.Dropdown(
                choices=list(MODELOS.keys()), 
                value=list(MODELOS.keys())[0], 
                label="Escolha o Modelo"
            )
            input_img = gr.Image(label="Upload da Imagem")
            btn = gr.Button("Analisar")
        
        with gr.Column():
            output_text = gr.Textbox(label="Resultado da Predição")

    # Ação do botão
    btn.click(fn=realizar_predicao, inputs=[input_img, seletor], outputs=output_text)

if __name__ == "__main__":
    demo.launch()