added features

Top 3 Emotionen
Confidence Einschätzung
Balkendiagramm (Bar Chart)
Schöner Gradio Output

app.py

@@ -2,6 +2,7 @@ import torch
 from torchvision import models, transforms
 from PIL import Image
 import gradio as gr
+import matplotlib.pyplot as plt
 from huggingface_hub import hf_hub_download
 
 # Modell laden vom Hugging Face Model Hub
@@ -12,7 +13,6 @@ model_path = hf_hub_download(
     filename="emotion_model.pt"
 )
 
-
 model = models.resnet18()
 model.fc = torch.nn.Linear(model.fc.in_features, 9)
 model.load_state_dict(torch.load(model_path, map_location=device))
@@ -28,6 +28,17 @@ transform = transforms.Compose([
     transforms.ToTensor()
 ])
 
+def plot_probabilities(probabilities, labels):
+    probs = probabilities.cpu().numpy().flatten()
+    fig, ax = plt.subplots(figsize=(8, 4))
+    ax.barh(labels, probs)
+    ax.set_xlim(0, 1)
+    ax.invert_yaxis()  # Highest probability on top
+    ax.set_xlabel('Confidence')
+    ax.set_title('Emotion Probabilities')
+    plt.tight_layout()
+    return fig
+
 def predict_emotion(image):
     image = image.convert("RGB")
     image = transform(image).unsqueeze(0).to(device)
@@ -35,20 +46,39 @@ def predict_emotion(image):
     with torch.no_grad():
         outputs = model(image)
         probs = torch.softmax(outputs, dim=1)
-        confidence, predicted = torch.max(probs, 1)
 
+    # Top 3 Predictions
+    top3_prob, top3_idx = torch.topk(probs, 3)
+    top3 = [(labels[i], f"{p.item()*100:.2f}%") for i, p in zip(top3_idx[0], top3_prob[0])]
+
+    # Overall Prediction
+    confidence, predicted = torch.max(probs, 1)
+    prediction = labels[predicted.item()]
+
+    # Unsicherheitswarnung
     if confidence.item() < 0.7:
-        return "Unbekannte Emotion", f"{confidence.item()*100:.2f}%"
+        prediction_status = "⚠️ Unsichere Vorhersage"
     else:
-        return labels[predicted.item()], f"{confidence.item()*100:.2f}%"
+        prediction_status = "✅ Sichere Vorhersage"
+
+    # Bar Chart
+    fig = plot_probabilities(probs, labels)
+
+    # Ausgabe
+    return prediction, f"Confidence: {confidence.item()*100:.2f}%\n{prediction_status}", top3, fig
 
 # Gradio Interface
 interface = gr.Interface(
     fn=predict_emotion,
     inputs=gr.Image(type="pil"),
-    outputs=["text", "text"],
+    outputs=[
+        gr.Textbox(label="Vorhergesagte Emotion"),
+        gr.Textbox(label="Confidence + Einschätzung"),
+        gr.Dataframe(headers=["Emotion", "Wahrscheinlichkeit (%)"], label="Top 3 Emotionen"),
+        gr.Plot(label="Verteilung der Emotionen")
+    ],
     title="Emotion Recognition App",
-    description="Lade ein Bild hoch und erkenne die Emotion."
+    description="Lade ein Bild hoch und erkenne die Emotion. Zeigt auch die Top 3 Emotionen und alle Wahrscheinlichkeiten als Balkendiagramm."
 )
 
 interface.launch()