Update app.py

app.py

@@ -17,27 +17,21 @@ model.eval()
 print("Model and processor loaded successfully.")
 
 # --- 2. MODIFIED Define the Explainability (Grad-CAM) Function ---
-# This function generates the heatmap showing which parts of the image the model focused on.
 def generate_heatmap(image_tensor, original_image, target_class_index):
-    # --- THIS IS THE FIX ---
-    # The original code assumed a ConvNeXT model. This model is a Swin Transformer.
-    # We now target the final layer normalization of the Swin Transformer's main body,
-    # which is a standard and effective layer for Grad-CAM on this architecture.
+    # This part is correct from our last fix.
     target_layer = model.swin.layernorm
 
     # Initialize LayerGradCam
     lgc = LayerGradCam(model, target_layer)
 
-    # Generate attributions (the "importance" of each pixel)
-    # The baselines are a reference point, typically a black image.
-    baselines = torch.zeros_like(image_tensor)
-    attributions = lgc.attribute(image_tensor, target=target_class_index, baselines=baselines, relu_attributions=True)
+    # --- 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.
+    attributions = lgc.attribute(image_tensor, target=target_class_index, relu_attributions=True)
 
-    # The output of LayerGradCam is a heatmap. We process it for visualization.
-    # We take the mean across the color channels and format it correctly.
+    # The rest of the function remains the same.
     heatmap = np.transpose(attributions.squeeze(0).cpu().detach().numpy(), (1, 2, 0))
 
-    # Use Captum's visualization tool to overlay the heatmap on the original image.
     visualized_image, _ = viz.visualize_image_attr(
         heatmap,
         np.array(original_image),