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

model = AutoModelForImageClassification.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")
processor = ViTImageProcessor.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

id2label = model.config.id2label


def predict(image):
  """Predicts the class of the input image using the fine-tuned student model."""

  # Convert the Gradio image to a PIL Image
  image = Image.fromarray(image)

  # Preprocess the image
  inputs = processor(image, return_tensors="pt").to(device)

  # Make prediction
  with torch.no_grad():
    outputs = model(**inputs)
    predicted_class_idx = torch.argmax(outputs.logits, dim=1).item()

  # Get predicted class label
  predicted_class = id2label[predicted_class_idx]

  return predicted_class

iface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy"),
    outputs="text",
    title="Pets Image Classification",
    description="Upload an image of a cat or dog to get its breed prediction.",
).launch()