from transformers import AutoModel

model = AutoModel.from_pretrained("KFrimps/oxford-pets-vit-from-scratch")

import gradio as gr

from PIL import Image

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()