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

# 1. Load a modern SOTA model (DINOv2)
# This model is faster and more accurate than the original ViT
model_name = "facebook/dinov2-base-imagenet1k-1-layer"
processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

def classify_image(img):
    # 2. Pre-process image
    inputs = processor(images=img, return_tensors="pt")
    
    # 3. Inference
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 4. Convert logits to probabilities (0% to 100%)
    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
    
    # 5. Extract top 5 results for a better UI
    top5_prob, top5_indices = torch.topk(probabilities, 5)
    
    # Create a dictionary of {Label: Probability} for Gradio's Label component
    confidences = {
        model.config.id2label[idx.item()]: float(prob) 
        for prob, idx in zip(top5_prob, top5_indices)
    }
    
    return confidences

# 6. Build a modern UI
demo = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil", label="Upload Image"),
    # gr.Label automatically creates a beautiful bar chart for probabilities
    outputs=gr.Label(num_top_classes=5, label="Predictions"),
    title="Next-Gen Image Classification",
    description="Running on **Meta's DINOv2** foundation model. Upload any image to see the top 5 predicted categories.",
    theme="soft" 
)

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