Update app.py

app.py

@@ -16,17 +16,25 @@ model = AutoModelForImageClassification.from_pretrained(model_id, torch_dtype=to
 model.eval()
 print("Model and processor loaded successfully.")
 
-# --- 2. MODIFIED Define the Explainability (Grad-CAM) Function ---
+# --- 2. Define the Explainability (Grad-CAM) Function ---
 def generate_heatmap(image_tensor, original_image, target_class_index):
+
+    # --- THIS IS THE FIX ---
+    # We define a wrapper function that ensures our model returns a simple tensor,
+    # which is what Captum expects. It takes the model's output object and
+    # extracts the 'logits' tensor from it.
+    def model_forward_wrapper(input_tensor):
+        outputs = model(pixel_values=input_tensor)
+        return outputs.logits
+
     # This part is correct from our last fix.
     target_layer = model.swin.layernorm
 
-    # Initialize LayerGradCam
-    lgc = LayerGradCam(model, target_layer)
+    # Initialize LayerGradCam, but pass our new wrapper function instead of the raw model.
+    # Captum will now use this wrapper to get the model's output.
+    lgc = LayerGradCam(model_forward_wrapper, target_layer)
 
-    # --- THIS IS THE FIX ---
-    # The 'baselines' argument is not used by LayerGradCam, so we remove it.
-    # The call is now simpler and correct for this specific method.
+    # This call now works because `lgc` gets a proper tensor from our wrapper.
     attributions = lgc.attribute(image_tensor, target=target_class_index, relu_attributions=True)
 
     # The rest of the function remains the same.