[FEAT] add GradCAM

gradio-inference.py

@@ -4,15 +4,20 @@
 
 import os
 import numpy as np
+import cv2
 import traceback
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 
 import gradio as gr
 from PIL import Image
 import logging
 
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
 from main import get_transform
 
 logging.basicConfig(level=logging.INFO)
@@ -76,54 +81,166 @@ def load_model(model_type: str = "efficientvit") -> nn.Module:
         logging.warning(
             f"Default model path '{model_path}' not found. Using untrained model."
         )
-    # Set model to evaluation mode
+    # Move model to device and set to evaluation mode
+    model.to(device)
     model.eval()
     return model
 
 
-def predict_image(image: np.ndarray, model_type: str) -> dict:
+def get_target_layers(model, model_type):
+    """
+    Get the target layers for GradCAM based on model type.
+
+    Args:
+        model: The model
+        model_type: Type of model
+
+    Returns:
+        target_layers: List of layers to use for GradCAM
+    """
+    try:
+        if model_type == "mobilenetv4":
+            # For MobileNetV4, use the last convolutional layer in features
+            return [model.features[-1]]
+        elif model_type == "levit":
+            # For LeViT (transformer), use the last block
+            return [model.blocks[-1]]
+        elif model_type == "efficientvit":
+            # For EfficientViT, use the last stage
+            return [model.stages[-1]]
+        elif model_type == "gernet":
+            # For GENet, use the last stage
+            return [model.stages[-1]]
+        elif model_type == "regnetx":
+            # For RegNetX, use the last trunk layer
+            return [model.trunk[-1]]
+        else:
+            # Default: try to get the last feature layer
+            if hasattr(model, "features"):
+                return [model.features[-1]]
+            elif hasattr(model, "stages"):
+                return [model.stages[-1]]
+            elif hasattr(model, "blocks"):
+                return [model.blocks[-1]]
+            else:
+                raise ValueError(
+                    f"Cannot determine target layer for model type: {model_type}"
+                )
+    except Exception as e:
+        logging.warning(f"Error getting target layer: {e}. Using fallback.")
+        # Fallback: try to get any reasonable last conv layer
+        for module in reversed(list(model.modules())):
+            if isinstance(module, nn.Conv2d):
+                return [module]
+        raise ValueError("Could not find suitable target layer for GradCAM")
+
+
+def apply_heatmap_on_image(img, cam, alpha=0.4):
+    """
+    Apply CAM heatmap overlay on the original image.
+
+    Args:
+        img: Original image (PIL Image or numpy array)
+        cam: Class activation map (grayscale, values 0-1)
+        alpha: Overlay transparency (not used with show_cam_on_image, kept for compatibility)
+
+    Returns:
+        Heatmap overlay image as numpy array
+    """
+    # Convert PIL to numpy if needed
+    if isinstance(img, Image.Image):
+        img = np.array(img)
+
+    # Normalize image to 0-1 range for show_cam_on_image
+    img_float = img.astype(np.float32) / 255.0
+
+    # Resize CAM to match image size
+    h, w = img.shape[:2]
+    cam_resized = cv2.resize(cam, (w, h))
+
+    # Use pytorch_grad_cam utility to overlay
+    # This function expects img in 0-1 range and cam in 0-1 range
+    overlay = show_cam_on_image(img_float, cam_resized, use_rgb=True)
+
+    return overlay
+
+
+def predict_image(image: np.ndarray, model_type: str) -> tuple[dict, np.ndarray]:
     """
-    Predict eye disease from an uploaded image.
+    Predict eye disease from an uploaded image and generate attention heatmap.
 
     Args:
         image: Input image from Gradio
-        model_path: Path to the model state dict
         model_type: Type of model architecture
 
     Returns:
-        Dictionary of class probabilities
+        Tuple of (Dictionary of class probabilities, Heatmap overlay image)
     """
     try:
-
         logging.info("Starting prediction...")
+
+        # Handle None image
+        if image is None:
+            logging.warning("No image provided.")
+            return {cls: 0.0 for cls in CLASSES}, None
+
         # Load model
         model = load_model(model_type)
 
         # Preprocess image
         logging.info("Preprocessing image...")
-        if image is None:
-            logging.warning("No image provided.")
-            return {cls: 0.0 for cls in CLASSES}
         transform = get_transform()
-        if image is None:
-            return {cls: 0.0 for cls in CLASSES}
 
-        # Convert numpy array to PIL Image
-        img = Image.fromarray(image).convert("RGB")
-        img_tensor = transform(img).unsqueeze(0).to(device)
+        # Convert numpy array to PIL Image and keep original for heatmap
+        img_pil = Image.fromarray(image).convert("RGB")
+        img_tensor = transform(img_pil).unsqueeze(0).to(device)
         logging.info("Image preprocessed successfully.")
 
-        # Make prediction
-        with torch.no_grad():
-            outputs = model(img_tensor)
-            probabilities = torch.nn.functional.softmax(outputs, dim=1)[0].cpu().numpy()
+        # Get target layers for GradCAM
+        try:
+            target_layers = get_target_layers(model, model_type)
+            logging.info(f"Using target layers: {target_layers}")
+
+            # Initialize GradCAM from pytorch_grad_cam library
+            cam_extractor = GradCAM(model=model, target_layers=target_layers)
 
-        # Return probabilities for each class
-        return {cls: float(prob) for cls, prob in zip(CLASSES, probabilities)}
+            # Generate CAM - the library handles forward and backward passes
+            grayscale_cam = cam_extractor(input_tensor=img_tensor, targets=None)
+
+            # Get the CAM for the first image in batch
+            cam = grayscale_cam[0, :]
+
+            # Get model prediction
+            with torch.no_grad():
+                outputs = model(img_tensor)
+
+            # Generate heatmap overlay
+            heatmap_overlay = apply_heatmap_on_image(img_pil, cam)
+
+            # Clean up
+            del cam_extractor
+
+        except Exception as e:
+            logging.error(f"Error generating heatmap: {e}")
+            traceback.print_exc()
+            # Fallback: just do prediction without heatmap
+            with torch.no_grad():
+                outputs = model(img_tensor)
+            heatmap_overlay = np.array(img_pil)  # Return original image
+
+        # Get probabilities
+        probabilities = F.softmax(outputs, dim=1)[0].cpu().detach().numpy()
+
+        # Return probabilities and heatmap
+        result_dict = {cls: float(prob) for cls, prob in zip(CLASSES, probabilities)}
+
+        logging.info("Prediction completed successfully.")
+        return result_dict, heatmap_overlay
 
     except Exception as e:
+        logging.error(f"Error during prediction: {e}")
         traceback.print_exc()
-        return {cls: 0.0 for cls in CLASSES}
+        return {cls: 0.0 for cls in CLASSES}, None
 
 
 def main():
@@ -159,12 +276,13 @@ def main():
 
             with gr.Column():
                 output_chart = gr.Label(label="Prediction")
+                output_heatmap = gr.Image(label="Attention Heatmap")
 
         # Process the image when the button is clicked
         submit_btn.click(
             fn=predict_image,
             inputs=[input_image, model_type],
-            outputs=output_chart,
+            outputs=[output_chart, output_heatmap],
         )
 
         # Examples section
@@ -172,7 +290,7 @@ def main():
         gr.Examples(
             examples=[],  # Add example paths here
             inputs=input_image,
-            outputs=[output_chart],
+            outputs=[output_chart, output_heatmap],
             fn=predict_image,
             cache_examples=True,
         )
@@ -188,7 +306,15 @@ def main():
                 - Enter the path to your trained model file (.pth)
                 - Select the model architecture that was used for training
             3. **Analyze**: Click the "Analyze Image" button to get results
-            4. **Interpret results**: The system will show the detected condition and probability distribution
+            4. **Interpret results**: The system will show the detected condition, probability distribution, and an attention heatmap
+            
+            ## Attention Heatmap:
+            
+            The attention heatmap visualizes which regions of the fundus image the model is focusing on when making its prediction.
+            - **Red/Yellow areas**: Regions the model considers most important for the diagnosis
+            - **Blue/Green areas**: Regions with less influence on the prediction
+            
+            This helps in understanding and validating the model's decision-making process.
             
             ## Model Information:
             

pyproject.toml

@@ -6,8 +6,11 @@ readme = "README.md"
 requires-python = ">=3.12.9"
 dependencies = [
   "gradio>=5.29.0",
+  "grad-cam>=1.5.0",
   "matplotlib>=3.10.3",
+  "opencv-python>=4.8.0",
   "pandas>=2.2.3",
+  "pillow>=10.0.0",
   "scikit-learn>=1.6.1",
   "seaborn>=0.13.2",
   "timm>=1.0.15",