Update app.py

app.py

@@ -3,24 +3,27 @@ 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")
+# 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")
 
-def classify_dog(image):
-    inputs = image_processor(images=image, return_tensors="pt")
+# 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():
-        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}"
+        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}!"
 
-demo = gr.Interface(
-    fn=classify_dog,
+# Set up the Gradio web interface
+app = gr.Interface(
+    fn=detect_breed,
     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)."
+    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."
 )
 
-demo.launch()
+app.launch()