import gradio as gr
import torch
from PIL import Image
from transformers import ViTImageProcessor, ViTForImageClassification


def test(image):
    processor = ViTImageProcessor.from_pretrained('google/vit-base-patch16-224')
    model = ViTForImageClassification.from_pretrained('google/vit-base-patch16-224')
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    probabilities_list = probabilities.tolist()[0]  

    class_probabilities = {
        model.config.id2label[class_idx]: probability
        for class_idx, probability in enumerate(probabilities_list)
    }
    
    top_4_probabilities = dict(sorted(class_probabilities.items(), key=lambda item: item[1], reverse=True)[:4])

    return top_4_probabilities

demo = gr.Interface(fn=test, inputs=gr.Image(type="pil"), outputs=gr.Label("Top 4 Scores and Classes"))
demo.launch()