from transformers import pipeline
import gradio as gr
from transformers import AutoConfig



label_map = {
    "0": "Necrosis/rust ",
    "1": "Muerto",
    "2": "Antracnosis/Soy Disease Classifier",
    "3": "Healthy",
    "4": "Marchitez bacteriana"}



pipe = pipeline(
    "image-classification", 
    model="sbaner24/vit-base-patch16-224-Soybean_11-46"
)


def classify_image(image):
    results = pipe(image)
    return {label_map.get(r['label'], r['label']): float(r['score']) for r in results}

    
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 soja")
    gr.Markdown("Sube una imagen de una hoja de soja para detectar posibles enfermedades.")
    gr.HTML(""" <a href="http://localhost:8081/enfSoja" 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=classify_image, inputs=inp, outputs=out)

demo.launch(debug=True)