Update app.py

app.py

@@ -38,140 +38,52 @@ except Exception as e:
     raise
 
 # --- 2. Define the Explainability (Grad-CAM) Function ---
+### FIX ###: This function is now more robust. It returns `None` on failure
+### instead of returning the original image, allowing the main function to handle it.
 def generate_heatmap(image_tensor, original_image, target_class_index):
+    """
+    Generates a Grad-CAM heatmap.
+    Returns a numpy array of the blended image, or None if it fails.
+    """
     try:
-        # Ensure tensor is on CPU and requires gradients
-        image_tensor = image_tensor.to(device)
-        image_tensor.requires_grad_(True)
-        
-        # Define wrapper function for model forward pass
-        def model_forward_wrapper(input_tensor):
-            outputs = model(pixel_values=input_tensor)
-            return outputs.logits
+        # LayerGradCam is often a good choice for transformer-based models.
+        # The target layer is chosen as one of the last normalization layers in the SWIN transformer.
+        # This might need adjustment for different model architectures.
+        target_layer = model.swin.encoder.layers[-1].blocks[-1].norm1
+        lgc = LayerGradCam(model.forward, target_layer)
 
-        # Try different approaches for better heatmap generation
-        try:
-            # First try: Use GradCam directly (often more reliable than LayerGradCam)
-            from captum.attr import GradCam
-            
-            # For SWIN transformer, target the last convolutional-like layer
-            try:
-                # Try to find a suitable layer in the SWIN model
-                target_layer = model.swin.encoder.layers[-1].blocks[-1].norm1
-            except:
-                try:
-                    target_layer = model.swin.encoder.layers[-1].blocks[0].norm1
-                except:
-                    target_layer = model.swin.layernorm
-            
-            gc = GradCam(model_forward_wrapper, target_layer)
-            
-            # Generate attributions
-            attributions = gc.attribute(image_tensor, target=target_class_index)
-            
-            # Process attributions
-            attr_np = attributions.squeeze().cpu().detach().numpy()
-            
-            print(f"Attribution stats: min={attr_np.min():.4f}, max={attr_np.max():.4f}, mean={attr_np.mean():.4f}")
-            
-            # Normalize to [0, 1] range
-            if attr_np.max() > attr_np.min():
-                attr_np = (attr_np - attr_np.min()) / (attr_np.max() - attr_np.min())
-            
-            # Resize to match original image size
-            from PIL import Image as PILImage
-            import cv2
-            
-            # Resize attribution map to original image size
-            attr_resized = cv2.resize(attr_np, original_image.size, interpolation=cv2.INTER_LINEAR)
-            
-            # Create a more visible heatmap
-            import matplotlib.pyplot as plt
-            import matplotlib.cm as cm
-            
-            # Apply a strong colormap (jet gives good red visualization)
-            colored_attr = cm.jet(attr_resized)[:, :, :3]  # Remove alpha channel
-            
-            # Convert original image to numpy
-            original_np = np.array(original_image) / 255.0
-            
-            # Create a stronger blend to make heatmap more visible
-            alpha = 0.6  # Higher alpha for more heatmap visibility
-            blended = (1 - alpha) * original_np + alpha * colored_attr
-            blended = (blended * 255).astype(np.uint8)
-            
-            return blended
-            
-        except Exception as e1:
-            print(f"GradCam failed: {e1}")
-            
-            # Fallback: Try LayerGradCam
-            try:
-                lgc = LayerGradCam(model_forward_wrapper, target_layer)
-                attributions = lgc.attribute(
-                    image_tensor, 
-                    target=target_class_index, 
-                    relu_attributions=False
-                )
-                
-                # Process the attributions
-                attr_np = attributions.squeeze(0).cpu().detach().numpy()
-                
-                # Handle different attribution shapes
-                if len(attr_np.shape) == 3:
-                    # Take mean across channels if multi-channel
-                    attr_np = np.mean(attr_np, axis=0)
-                
-                # Normalize
-                if attr_np.max() > attr_np.min():
-                    attr_np = (attr_np - attr_np.min()) / (attr_np.max() - attr_np.min())
-                
-                # Create visualization using captum's viz
-                if len(attr_np.shape) == 2:
-                    # Expand to 3 channels for visualization
-                    heatmap = np.expand_dims(attr_np, axis=-1)
-                    heatmap = np.repeat(heatmap, 3, axis=-1)
-                else:
-                    heatmap = np.transpose(attr_np, (1, 2, 0))
-                
-                visualized_image, _ = viz.visualize_image_attr(
-                    heatmap,
-                    np.array(original_image),
-                    method="blended_heat_map",
-                    sign="all",
-                    show_colorbar=True,
-                    title="AI Detection Heatmap",
-                    alpha_overlay=0.4,
-                    cmap="jet",  # Use jet colormap for strong red visualization
-                    outlier_perc=1
-                )
-                
-                return visualized_image
-                
-            except Exception as e2:
-                print(f"LayerGradCam also failed: {e2}")
-                
-                # Final fallback: Create a simple random heatmap for demonstration
-                print("Creating demonstration heatmap...")
-                
-                # Create a simple demonstration heatmap
-                h, w = original_image.size[1], original_image.size[0]
-                demo_attr = np.random.rand(h, w) * 0.5 + 0.3  # Random values between 0.3 and 0.8
-                
-                # Apply jet colormap
-                colored_attr = cm.jet(demo_attr)[:, :, :3]
-                
-                # Blend with original
-                original_np = np.array(original_image) / 255.0
-                blended = 0.7 * original_np + 0.3 * colored_attr
-                blended = (blended * 255).astype(np.uint8)
-                
-                return blended
+        # Generate attributions
+        attributions = lgc.attribute(
+            image_tensor,
+            target=target_class_index,
+            relu_attributions=True  # Use relu_attributions to focus on positive contributions
+        )
         
+        # Squeeze the attributions to a 2D map
+        attribution_map = attributions.squeeze(0).squeeze(0).cpu().detach().numpy()
+
+        ### FIX ###: Check if the attribution map is uniform (all zeros or same value).
+        # This happens when the model has no strong evidence for its decision,
+        # which is common in misclassifications.
+        if (attribution_map.max() - attribution_map.min()) < 1e-6:
+            print("Warning: Heatmap generation failed due to uniform gradients. The model likely has low confidence or is misclassifying.")
+            return None
+
+        # Use Captum's visualization tool to create a blended heatmap
+        blended_image, _ = viz.visualize_image_attr(
+            attribution_map,
+            np.array(original_image),
+            method="blended_heat_map",
+            sign="positive",  # Focus on what positively contributed to the decision
+            alpha_overlay=0.5, # Make the overlay reasonably transparent
+            cmap="jet",        # 'jet' colormap shows hot areas in red
+            show_colorbar=False
+        )
+        return blended_image
+
     except Exception as e:
-        print(f"Complete heatmap generation failed: {e}")
-        # Return original image if everything fails
-        return np.array(original_image)
+        print(f"Error during heatmap generation: {e}")
+        return None
 
 # --- 3. Main Prediction Function ---
 def predict(image_upload: Image.Image, image_url: str):