import gradio as gr
from PIL import Image
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification


# Load model and processor
MODEL_NAME = "nexusbert/resnet50-cassava-finetuned"

processor = AutoImageProcessor.from_pretrained(MODEL_NAME)
model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
id2label = model.config.id2label


def predict(image: Image.Image):
    """
    Run inference on a single image and return class probabilities.
    Gradio's Label component expects a dict: {label: probability}.
    """
    inputs = processor(images=image, return_tensors="pt")

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)[0]

    # Map class indices to human-readable labels and probabilities (as floats)
    result = {
        id2label[i]: float(probs[i])
        for i in range(len(probs))
    }
    return result


demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Upload cassava leaf image"),
    outputs=gr.Label(num_top_classes=3, label="Predicted classes"),
    title="Cassava Leaf Disease Classifier",
    description=(
        "Upload an image of a cassava leaf and this app will predict the most likely disease class "
        "using the `nexusbert/resnet50-cassava-finetuned` model from Hugging Face."
    ),
    examples=None,
)


if __name__ == "__main__":
    # Set share=True if you want a public shareable link
    demo.launch()