Update app.py

added grad-cam feature

app.py

@@ -4,9 +4,11 @@ from PIL import Image
 import gradio as gr
 import matplotlib.pyplot as plt
 import pandas as pd
+import numpy as np
 import os
 import hashlib
 from huggingface_hub import hf_hub_download
+import cv2
 
 # Modell laden vom Hugging Face Model Hub
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -39,18 +41,71 @@ def get_image_hash(image):
     img_bytes = image.tobytes()
     return hashlib.md5(img_bytes).hexdigest()
 
-# Plot-Funktion
+# Plot-Funktion für Wahrscheinlichkeiten
 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.invert_yaxis()
     ax.set_xlabel('Confidence')
     ax.set_title('Emotion Probabilities')
     plt.tight_layout()
     return fig
 
+# Grad-CAM Hilfsfunktion
+def generate_gradcam(image, model, class_idx):
+    model.eval()
+
+    # Hook für Features und Gradients
+    gradients = []
+    activations = []
+
+    def save_gradient(grad):
+        gradients.append(grad)
+
+    def forward_hook(module, input, output):
+        activations.append(output)
+        output.register_hook(save_gradient)
+
+    # Letztes Convolutional Layer
+    target_layer = model.layer4[1].conv2
+    handle = target_layer.register_forward_hook(forward_hook)
+
+    image_tensor = transform(image).unsqueeze(0).to(device)
+    output = model(image_tensor)
+
+    # Softmax -> Klasse auswählen
+    pred_class = output.argmax(dim=1).item()
+
+    model.zero_grad()
+    class_score = output[0, class_idx]
+    class_score.backward()
+
+    # Gradients und Activations holen
+    gradients = gradients[0].cpu().data.numpy()[0]
+    activations = activations[0].cpu().data.numpy()[0]
+
+    weights = np.mean(gradients, axis=(1, 2))
+    gradcam = np.zeros(activations.shape[1:], dtype=np.float32)
+
+    for i, w in enumerate(weights):
+        gradcam += w * activations[i, :, :]
+
+    gradcam = np.maximum(gradcam, 0)
+    gradcam = cv2.resize(gradcam, (224, 224))
+    gradcam = gradcam - np.min(gradcam)
+    gradcam = gradcam / np.max(gradcam)
+
+    # Bild zurückkonvertieren
+    heatmap = cv2.applyColorMap(np.uint8(255 * gradcam), cv2.COLORMAP_JET)
+    image_np = np.array(image.resize((224, 224)).convert("RGB"))
+    overlay = cv2.addWeighted(image_np, 0.6, heatmap, 0.4, 0)
+
+    handle.remove()  # Hook entfernen
+
+    return Image.fromarray(overlay)
+
 # Prediction-Funktion
 def predict_emotion(image):
     image = image.convert("RGB")
@@ -80,7 +135,10 @@ def predict_emotion(image):
     # Bild-Hash für spätere Zuordnung
     img_hash = get_image_hash(image)
 
-    return prediction, f"Confidence: {confidence.item()*100:.2f}%\n{prediction_status}", top3, fig, img_hash
+    # Grad-CAM Overlay
+    gradcam_img = generate_gradcam(image, model, predicted.item())
+
+    return prediction, f"Confidence: {confidence.item()*100:.2f}%\n{prediction_status}", top3, fig, gradcam_img, img_hash
 
 # Funktion um Feedback zu speichern
 def save_feedback(img_hash, model_prediction, user_feedback, confidence):
@@ -108,9 +166,9 @@ def download_feedback():
 
 # Kombinierte Funktion
 def full_pipeline(image, user_feedback):
-    prediction, confidence_text, top3, fig, img_hash = predict_emotion(image)
+    prediction, confidence_text, top3, fig, gradcam_img, img_hash = predict_emotion(image)
     feedback_message = save_feedback(img_hash, prediction, user_feedback, confidence_text.split("\n")[0])
-    return prediction, confidence_text, top3, fig, feedback_message
+    return prediction, confidence_text, top3, fig, gradcam_img, feedback_message
 
 # Gradio Interface
 with gr.Blocks() as interface:
@@ -125,12 +183,13 @@ with gr.Blocks() as interface:
             confidence_output = gr.Textbox(label="Confidence + Einschätzung")
             top3_output = gr.Dataframe(headers=["Emotion", "Wahrscheinlichkeit (%)"], label="Top 3 Emotionen")
             plot_output = gr.Plot(label="Verteilung der Emotionen")
+            gradcam_output = gr.Image(label="Grad-CAM Visualisierung")
             feedback_message_output = gr.Textbox(label="Feedback-Status")
 
     submit_btn.click(
         fn=full_pipeline,
         inputs=[image_input, feedback_input],
-        outputs=[prediction_output, confidence_output, top3_output, plot_output, feedback_message_output]
+        outputs=[prediction_output, confidence_output, top3_output, plot_output, gradcam_output, feedback_message_output]
     )
 
     download_btn.click(