import torch
from torchvision import transforms
import gradio as gr
import timm
from huggingface_hub import hf_hub_download

# Descargar y cargar el checkpoint desde Hugging Face
checkpoint_path = hf_hub_download(repo_id="monabarreiro/ModeloMaiz", filename="best_Model_vit.pt")
checkpoint = torch.load(checkpoint_path, map_location='cpu')

# Crear el modelo con el número correcto de clases
model = timm.create_model('vit_base_patch16_224', num_classes=len(checkpoint['class_names']))

# Ajustar el state_dict para cargar correctamente los pesos
state_dict = checkpoint['model_state_dict']
new_state_dict = {}
for k, v in state_dict.items():
    if k.startswith('model.'):
        new_state_dict[k[6:]] = v
    else:
        new_state_dict[k] = v
model.load_state_dict(new_state_dict)
model.eval()

# Traducción de clases al español
class_names_en = checkpoint['class_names']
class_names_es = {
    "Blight": "Tizón",
    "Common_Rust": "Roya común",
    "Gray_Leaf_Spot": "Mancha gris de la hoja",
    "Healthy": "Sano"
}
class_names = [class_names_es[name] for name in class_names_en]

# Preprocesamiento de imágenes
preprocess = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# Función de predicción
def predict(image):
    img_tensor = preprocess(image).unsqueeze(0)
    with torch.no_grad():
        outputs = model(img_tensor)
        probs = torch.nn.functional.softmax(outputs, dim=1)
    return {class_names[i]: float(probs[0][i]) for i in range(len(class_names))}

with gr.Blocks() as demo:
    gr.Markdown("Clasificador de enfermedades de maiz")
    gr.Markdown("Sube una imagen de una hoja de maíz para detectar posibles enfermedades.")
    gr.HTML(""" <a href="http://localhost:8081/enfMaiz" target="_blank">Mostrar Enfermedades</a>""")

    with gr.Row():
        inp = gr.Image(type="pil", label="Subí una imagen")
        out = gr.Textbox(label="Predicción")

    btn = gr.Button("Predecir")
    btn.click(fn=predict, inputs=inp, outputs=out)

demo.launch(debug=True)