Spaces:
Sleeping
Sleeping
from transformers import AutoImageProcessor, AutoModelForImageClassification | |
from PIL import Image | |
import torch | |
import gradio as gr | |
image_processor = AutoImageProcessor.from_pretrained("wesleyacheng/dog-breeds-multiclass-image-classification-with-vit") | |
model = AutoModelForImageClassification.from_pretrained("wesleyacheng/dog-breeds-multiclass-image-classification-with-vit") | |
def classify_dog(image): | |
inputs = image_processor(images=image, return_tensors="pt") | |
with torch.no_grad(): | |
outputs = model(**inputs) | |
logits = outputs.logits | |
predicted_class_idx = logits.argmax(-1).item() | |
predicted_breed = model.config.id2label[predicted_class_idx] | |
return f"Predicted Dog Breed: {predicted_breed}" | |
demo = gr.Interface( | |
fn=classify_dog, | |
inputs=gr.Image(type="pil"), | |
outputs="text", | |
title="Dog Breed Classifier", | |
description="Upload an image of a dog and the model will classify its breed (120 breeds supported)." | |
) | |
demo.launch() | |