Update app.py

app.py

@@ -1,20 +1,47 @@
 import gradio as gr
+import torch
+from PIL import Image
+from transformers import AutoImageProcessor, ViTForImageClassification
 from transformers import pipeline
 
-# Load models
-cifar10_classifier = pipeline("image-classification", model="Fadri/results")
-clip_detector = pipeline(model="openai/clip-vit-large-patch14", task="zero-shot-image-classification")
-
-# CIFAR-10 Klassen
+# CIFAR-10 Klassenlabels
 labels_cifar10 = [
     'airplane', 'automobile', 'bird', 'cat', 'deer', 
     'dog', 'frog', 'horse', 'ship', 'truck'
 ]
 
+# Lade Modell und Processor separat
+processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224-in21k")
+model = ViTForImageClassification.from_pretrained("Fadri/results")
+
+# CLIP für Zero-Shot bleibt wie vorher
+clip_detector = pipeline(model="openai/clip-vit-large-patch14", task="zero-shot-image-classification")
+
+def predict_cifar10(image_path):
+    # Bild laden und vorverarbeiten
+    image = Image.open(image_path).convert("RGB")
+    inputs = processor(images=image, return_tensors="pt")
+    
+    # Modellvorhersage
+    with torch.no_grad():
+        outputs = model(**inputs)
+        logits = outputs.logits
+        predicted_class_idx = logits.argmax(-1).item()
+    
+    # Top-3 Ergebnisse mit Wahrscheinlichkeiten
+    probabilities = torch.nn.functional.softmax(logits, dim=-1)[0]
+    top3_probs, top3_indices = torch.topk(probabilities, 3)
+    
+    results = {}
+    for idx, prob in zip(top3_indices, top3_probs):
+        label = model.config.id2label[idx.item()]
+        results[label] = round(prob.item(), 4)
+    
+    return results
+
 def classify_image(image):
-    # Klassifikation mit deinem trainierten CIFAR-10 Modell
-    cifar10_results = cifar10_classifier(image)
-    cifar10_output = {result['label']: result['score'] for result in cifar10_results}
+    # Klassifikation mit deinem Modell
+    cifar10_output = predict_cifar10(image)
     
     # Zero-Shot-Klassifikation mit CLIP
     clip_results = clip_detector(image, candidate_labels=labels_cifar10)
@@ -25,14 +52,12 @@ def classify_image(image):
         "CLIP Zero-Shot Klassifikation": clip_output
     }
 
-# Beispielbilder - Du kannst diese später anpassen
+# Beispielbilder (Pfade anpassen)
 example_images = [
     ["examples/airplane.jpg"],
     ["examples/car.jpg"],
     ["examples/dog.jpg"],
-    ["examples/cat.jpg"],
-    ["examples/ship.jpg"],
-    ["examples/truck.jpg"]
+    ["examples/cat.jpg"]
 ]
 
 # Gradio Interface
@@ -41,7 +66,7 @@ iface = gr.Interface(
     inputs=gr.Image(type="filepath"),
     outputs=gr.JSON(),
     title="CIFAR-10 Klassifikation",
-    description="Lade ein Bild hoch und vergleiche die Ergebnisse zwischen deinem trainierten ViT Modell und dem Zero-Shot CLIP Modell für CIFAR-10 Klassen.",
+    description="Lade ein Bild hoch und vergleiche die Ergebnisse zwischen deinem trainierten ViT Modell und CLIP.",
     examples=example_images
 )