Update app.py

app.py

@@ -1,13 +1,20 @@
 import gradio as gr
 from transformers import pipeline
 
-# ViT-Modell für klassische Klassifikation
-vit_classifier = pipeline("image-classification", model="dewiri/vit-base-oxford-iiit-pets")
+# ViT-Modell (dein trainiertes Modell von Hugging Face)
+vit_classifier = pipeline(
+    "image-classification",
+    model="dewiri/vit-base-oxford-iiit-pets",
+    top_k=3,
+    device=0  # GPU, falls verfügbar
+)
 
-# SIGLIP-Modell für Zero-Shot Klassifikation (anstelle von CLIP)
+# SIGLIP für Zero-Shot Klassifikation
 siglip_detector = pipeline(
     model="google/siglip-so400m-patch14-384",
-    task="zero-shot-image-classification"
+    task="zero-shot-image-classification",
+    top_k=3,
+    device=0
 )
 
 labels_oxford_pets = [
@@ -20,13 +27,13 @@ labels_oxford_pets = [
 
 def classify_pet(image):
     vit_results = vit_classifier(image)
-    vit_output = {result['label']: result['score'] for result in vit_results}
+    vit_output = {res['label']: round(res['score'], 3) for res in vit_results}
     
     siglip_results = siglip_detector(image, candidate_labels=labels_oxford_pets)
-    siglip_output = {result['label']: result['score'] for result in siglip_results}
+    siglip_output = {res['label']: round(res['score'], 3) for res in siglip_results}
     
     return {
-        "ViT Classification": vit_output,
+        "ViT Classification (dewiri)": vit_output,
         "SIGLIP Zero-Shot Classification": siglip_output
     }
 
@@ -43,7 +50,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="filepath"),
     outputs=gr.JSON(),
     title="Pet Classification Comparison",
-    description="Upload an image of a pet, and compare results from a trained ViT model and a zero-shot SIGLIP model.",
+    description="Compare a fine-tuned ViT model (dewiri/vit-base-oxford-iiit-pets) with a SIGLIP zero-shot classifier.",
     examples=example_images
 )