Update app.py

app.py

@@ -1,21 +1,9 @@
 import gradio as gr
 from transformers import pipeline
 
-# ViT-Modell (dein trainiertes Modell)
-vit_classifier = pipeline(
-    "image-classification",
-    model="dewiri/vit-base-oxford-iiit-pets",
-    top_k=3,
-    device=0
-)
-
-# CLIP-Modell für Zero-Shot Klassifikation
-clip_detector = pipeline(
-    model="openai/clip-vit-large-patch14",
-    task="zero-shot-image-classification",
-    top_k=3,
-    device=0
-)
+# Load models
+vit_classifier = pipeline("image-classification", model="dewiri/vit-base-oxford-iiit-pets")
+clip_detector = pipeline(model="openai/clip-vit-large-patch14", task="zero-shot-image-classification")
 
 labels_oxford_pets = [
     'Siamese', 'Birman', 'shiba inu', 'staffordshire bull terrier', 'basset hound', 'Bombay', 'japanese chin',
@@ -27,15 +15,12 @@ labels_oxford_pets = [
 
 def classify_pet(image):
     vit_results = vit_classifier(image)
-    vit_output = {res['label']: round(res['score'], 3) for res in vit_results}
+    vit_output = {result['label']: result['score'] for result in vit_results}
     
     clip_results = clip_detector(image, candidate_labels=labels_oxford_pets)
-    clip_output = {res['label']: round(res['score'], 3) for res in clip_results}
+    clip_output = {result['label']: result['score'] for result in clip_results}
     
-    return {
-        "ViT Classification (dewiri)": vit_output,
-        "CLIP Zero-Shot Classification": clip_output
-    }
+    return {"ViT Classification": vit_output, "CLIP Zero-Shot Classification": clip_output}
 
 example_images = [
     ["example_images/dog1.jpeg"],
@@ -50,8 +35,8 @@ iface = gr.Interface(
     inputs=gr.Image(type="filepath"),
     outputs=gr.JSON(),
     title="Pet Classification Comparison",
-    description="Compare a fine-tuned ViT model (dewiri/vit-base-oxford-iiit-pets) with OpenAI's CLIP zero-shot classifier.",
+    description="Upload an image of a pet, and compare results from a trained ViT model and a zero-shot CLIP model.",
     examples=example_images
 )
 
-iface.launch()
+iface.launch()