from PIL import Image
from torchvision import transforms
from transformers import AutoModelForImageClassification
import gradio as gr
import torch

def predict(inp):
    inputs = data_transforms(inp)[None]
    model.eval()
    with torch.no_grad():
        logits = model(inputs)['logits']
        probs = torch.softmax(logits,dim=1)
        print(logits)
        confidences = {labels[i]: probs[0][i] for i in range(10)}
    return confidences


data_transforms = transforms.Compose([
	    transforms.Resize((32,32)),  # Resize the images to a specific size
	    transforms.ToTensor(),           # Convert images to tensors
	    transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))  # Normalize the image data
	])


model = AutoModelForImageClassification.from_pretrained("Manu8/vit_cifar", trust_remote_code=True)# Load saved weights
labels = [
    'airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck'
]
gr.Interface(fn=predict,
             inputs=gr.Image(type="pil"),
             outputs=gr.Label(num_top_classes=3)).launch()