Update app.py

app.py

@@ -45,48 +45,108 @@ def generate_heatmap(image_tensor, original_image, target_class_index):
         
         # Define wrapper function for model forward pass
         def model_forward_wrapper(input_tensor):
-            with torch.no_grad():  # Save memory during attribution
-                outputs = model(pixel_values=input_tensor)
-                return outputs.logits
+            outputs = model(pixel_values=input_tensor)
+            return outputs.logits
 
         # Get the target layer for Grad-CAM
-        # For SWIN transformer, use the layer normalization layer
-        target_layer = model.swin.layernorm
+        # For SWIN transformer, try different layers for better visualization
+        try:
+            # Try the encoder's last layer first
+            target_layer = model.swin.encoder.layers[-1].blocks[-1].layernorm_after
+        except:
+            try:
+                # Fallback to the main layernorm
+                target_layer = model.swin.layernorm
+            except:
+                # Final fallback to pooler if available
+                target_layer = model.swin.pooler.layernorm if hasattr(model.swin, 'pooler') else model.swin.layernorm
 
         # Initialize LayerGradCam with the wrapper function
         lgc = LayerGradCam(model_forward_wrapper, target_layer)
 
-        # Generate attributions
-        with torch.no_grad():
-            attributions = lgc.attribute(
-                image_tensor, 
-                target=target_class_index, 
-                relu_attributions=True
-            )
+        # Generate attributions - remove torch.no_grad() to allow gradients
+        attributions = lgc.attribute(
+            image_tensor, 
+            target=target_class_index, 
+            relu_attributions=False  # Changed to False to see both positive and negative attributions
+        )
 
         # Convert attributions to numpy for visualization
-        heatmap = np.transpose(
-            attributions.squeeze(0).cpu().detach().numpy(), 
-            (1, 2, 0)
-        )
+        attr_np = attributions.squeeze(0).cpu().detach().numpy()
+        
+        # Normalize attributions to [0, 1] range for better visualization
+        attr_min = attr_np.min()
+        attr_max = attr_np.max()
+        if attr_max > attr_min:
+            attr_np = (attr_np - attr_min) / (attr_max - attr_min)
+        
+        # Transpose for visualization (channels last)
+        if len(attr_np.shape) == 3:
+            heatmap = np.transpose(attr_np, (1, 2, 0))
+        else:
+            # If single channel, expand to 3 channels
+            heatmap = np.expand_dims(attr_np, axis=-1)
+            heatmap = np.repeat(heatmap, 3, axis=-1)
 
-        # Create visualization
+        # Create visualization with enhanced parameters
         visualized_image, _ = viz.visualize_image_attr(
             heatmap,
             np.array(original_image),
             method="blended_heat_map",
-            sign="all",
+            sign="all",  # Show both positive and negative attributions
             show_colorbar=True,
             title="AI Detection Heatmap",
-            alpha_overlay=0.6
+            alpha_overlay=0.5,  # Reduced alpha for better visibility
+            cmap="RdYlBu_r",  # Red-Yellow-Blue colormap (reversed)
+            outlier_perc=2  # Remove outliers for better contrast
         )
         
         return visualized_image
         
     except Exception as e:
         print(f"Error generating heatmap: {e}")
-        # Return original image if heatmap generation fails
-        return np.array(original_image)
+        print(f"Attribution shape: {attributions.shape if 'attributions' in locals() else 'Not generated'}")
+        
+        # Create a simple fallback heatmap using GradCAM on a different layer
+        try:
+            from captum.attr import GradCam
+            
+            # Use GradCAM instead of LayerGradCAM as fallback
+            gc = GradCam(model_forward_wrapper, target_layer)
+            attributions = gc.attribute(image_tensor, target=target_class_index)
+            
+            # Process the attributions
+            attr_np = attributions.squeeze().cpu().detach().numpy()
+            
+            # Normalize
+            attr_min = attr_np.min()
+            attr_max = attr_np.max()
+            if attr_max > attr_min:
+                attr_np = (attr_np - attr_min) / (attr_max - attr_min)
+            
+            # Create a simple overlay
+            import matplotlib.pyplot as plt
+            import matplotlib.cm as cm
+            
+            # Resize attribution to match image size
+            from PIL import Image as PILImage
+            attr_resized = PILImage.fromarray((attr_np * 255).astype(np.uint8)).resize(original_image.size)
+            attr_resized = np.array(attr_resized) / 255.0
+            
+            # Apply colormap
+            colored_attr = cm.jet(attr_resized)[:, :, :3]  # Remove alpha channel
+            
+            # Blend with original image
+            original_np = np.array(original_image) / 255.0
+            blended = 0.6 * original_np + 0.4 * colored_attr
+            blended = (blended * 255).astype(np.uint8)
+            
+            return blended
+            
+        except Exception as e2:
+            print(f"Fallback heatmap also failed: {e2}")
+            # Return original image if all heatmap generation fails
+            return np.array(original_image)
 
 # --- 3. Main Prediction Function ---
 def predict(image_upload: Image.Image, image_url: str):