app.py

import gradio as gr
import torch
import torchvision
from torchvision import transforms
import matplotlib.pyplot as plt
from vit_class import ViT

from typing import List, Tuple

from PIL import Image

class_names = ['jerry', 'tom']
image_size = (224, 224)
model_save_path = "vit_model.pth"
device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

vit = ViT(num_classes=len(class_names))
checkpoint = torch.load(model_save_path, map_location=device)
vit.load_state_dict(checkpoint["model_state_dict"])
vit.eval()
vit.to(device)


def predict(image):
    # return "Hello", "Hello"
    if image is None:
        return "No image captured", "0.0"

    image_transform = transforms.Compose(
            [
                transforms.Resize(image_size),
                transforms.ToTensor(),
                transforms.Normalize(
                    mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
                ),
            ]
        )

    with torch.inference_mode():
        # Transform and add an extra dimension to image (model requires samples in [batch_size, color_channels, height, width])
        transformed_image = image_transform(image).unsqueeze(dim=0)

        # Make a prediction on image with an extra dimension and send it to the target device
        target_image_pred = vit(transformed_image.to(device))

    # Convert logits -> prediction probabilities (using torch.softmax() for multi-class classification)
    target_image_pred_probs = torch.softmax(target_image_pred, dim=1)

    # Convert prediction probabilities -> prediction labels
    target_image_pred_label = torch.argmax(target_image_pred_probs, dim=1)

    return class_names[target_image_pred_label], (target_image_pred_probs.max().item())

demo = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="pil", label="Input Image"),
    outputs=[gr.Textbox(label="Character", lines=1), gr.Textbox(label="Probability", lines=1)],
)
demo.launch(share=True)