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

# Load the image processor and model from Hugging Face
processor = AutoImageProcessor.from_pretrained("wesleyacheng/dog-breeds-multiclass-image-classification-with-vit")
breed_model = AutoModelForImageClassification.from_pretrained("wesleyacheng/dog-breeds-multiclass-image-classification-with-vit")

# This function takes an uploaded image and returns the predicted dog breed
def detect_breed(img):
    inputs = processor(images=img, return_tensors="pt")
    with torch.no_grad():
        result = breed_model(**inputs)
    predictions = result.logits
    top_prediction = predictions.argmax(dim=1).item()
    breed_name = breed_model.config.id2label[top_prediction]
    return f"This looks like a {breed_name}!"

# Set up the Gradio web interface
app = gr.Interface(
    fn=detect_breed,
    inputs=gr.Image(type="pil"),
    outputs="text",
    title="Dog Breed Identifier 🐶",
    description="Upload a photo of a dog and find out what breed it is! The model can recognize 120 different dog breeds."
)

app.launch()