guided_gradcam.py

import os, sys
import torch
import numpy as np
import cv2
from torchvision import transforms
from PIL import Image

# Add model import path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from models.efficientnet_b0 import EfficientNetB0Classifier

class GuidedBackprop:
    def __init__(self, model):
        self.model = model.eval()
        self.gradients = None
        self.forward_relu_outputs = []
        
        # Replace ReLU with GuidedReLU
        self.relu_targets = []
        for module in self.model.modules():
            if isinstance(module, torch.nn.ReLU):
                module.register_forward_hook(self.forward_hook)
                module.register_full_backward_hook(self.backward_hook)
                self.relu_targets.append(module)
    
    def forward_hook(self, module, input, output):
        self.forward_relu_outputs.append(output)
        
    def backward_hook(self, module, grad_in, grad_out):
        corresponding_forward_output = self.forward_relu_outputs[-len(self.relu_targets):]
        corresponding_forward_output = corresponding_forward_output[self.relu_targets.index(module)]
        positive_mask = torch.gt(corresponding_forward_output, 0.).float()
        positive_grad = torch.gt(grad_in[0], 0.).float()
        grad = positive_mask * positive_grad * grad_in[0]
        return (grad,)
        
    def generate_gradients(self, input_image, target_class=None):
        self.forward_relu_outputs = []
        
        # Ensure input requires gradients
        input_image.requires_grad_(True)
        
        # Forward pass
        model_output = self.model(input_image)
        output_score = model_output.squeeze()
        
        # Backward pass
        self.model.zero_grad()
        output_score.backward(retain_graph=True)
        
        gradients_as_arr = input_image.grad.data.cpu().numpy()[0]
        return gradients_as_arr

class GuidedGradCAM:
    def __init__(self, model, target_layer):
        self.model = model.eval()
        self.target_layer = target_layer
        self.activations = None
        self.gradients = None
        self.guided_backprop = GuidedBackprop(model)
        
        self.hooks = [
            target_layer.register_forward_hook(self.save_activation),
            target_layer.register_full_backward_hook(self.save_gradient)
        ]
    
    def save_activation(self, module, input, output):
        self.activations = output.detach()
        
    def save_gradient(self, module, grad_input, grad_output):
        self.gradients = grad_output[0]
    
    def generate_cam(self, input_tensor, target_class=None):
        # Get gradients and activations
        model_output = self.model(input_tensor)
        
        self.model.zero_grad()
        output_score = model_output.squeeze()  # For binary classification
        output_score.backward(retain_graph=True)
        
        if self.gradients is None:
            raise RuntimeError("No gradients found. The model architecture might not be compatible.")
        
        # Generate GradCAM
        gradients = self.gradients.squeeze()
        activations = self.activations.squeeze()
        
        weights = torch.mean(gradients, dim=(1, 2))
        cam = torch.zeros(activations.shape[1:], dtype=torch.float32, device=activations.device)
        
        for i, w in enumerate(weights):
            cam += w * activations[i]
        
        cam = torch.relu(cam)
        cam = cam - torch.min(cam)
        cam = cam / (torch.max(cam) + 1e-7)
        
        # Generate Guided Backprop
        guided_grads = self.guided_backprop.generate_gradients(input_tensor)
        
        # Resize CAM to match input size
        cam_np = cam.cpu().numpy()
        cam_np = cv2.resize(cam_np, (380, 380))
        
        # Combine them
        cam_gb = np.multiply(cam_np[..., np.newaxis], guided_grads.transpose(1, 2, 0))
        
        return cam_np, guided_grads, cam_gb

    def __call__(self, input_tensor):
        return self.generate_cam(input_tensor)
    
    def remove_hooks(self):
        for hook in self.hooks:
            hook.remove()

def preprocess_and_show_results(img_path, cam, guided_grads, cam_gb, save_path):
    # Read and preprocess original image
    raw = cv2.imread(img_path)
    raw = cv2.resize(raw, (380, 380))
    
    # Process GradCAM
    heatmap = np.uint8(255 * cam)
    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
    gradcam_overlay = cv2.addWeighted(raw, 0.6, heatmap, 0.4, 0)
    
    # Process Guided Backprop
    guided_grads = guided_grads.transpose(1, 2, 0)
    guided_grads = np.maximum(guided_grads, 0)
    guided_grads = (guided_grads - guided_grads.min()) / (guided_grads.max() - guided_grads.min() + 1e-8)
    guided_grads = np.uint8(255 * guided_grads)
    
    # Process Guided GradCAM
    cam_gb = np.maximum(cam_gb, 0)
    cam_gb = (cam_gb - cam_gb.min()) / (cam_gb.max() - cam_gb.min() + 1e-8)
    cam_gb = np.uint8(255 * cam_gb)
    
    # Save all visualizations
    cv2.imwrite(save_path.replace('.png', '_gradcam.png'), gradcam_overlay)
    cv2.imwrite(save_path.replace('.png', '_guided_bp.png'), guided_grads)
    cv2.imwrite(save_path.replace('.png', '_guided_gradcam.png'), cam_gb)
    cv2.imwrite(save_path.replace('.png', '_original.png'), raw)

if __name__ == "__main__":
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
    # Load model
    model = EfficientNetB0Classifier()
    model.load_state_dict(torch.load("results_efficientnet_b0/efficientnet_best9912.pth", map_location=device))
    model = model.eval()
    
    # Enable gradients for feature visualization
    for param in model.parameters():
        param.requires_grad = True
    
    model.to(device)
    
    # Image transform
    transform = transforms.Compose([
        transforms.Resize((380, 380)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    
    # Use the final convolutional layer
    target_layer = model.base_model.features[-1]
    guided_gradcam = GuidedGradCAM(model, target_layer)
    
    # Process test images
    image_paths = np.load("test_paths.npy", allow_pickle=True).astype(str)
    np.random.seed(42)
    selected_indices = np.random.choice(len(image_paths), 5, replace=False)
    
    os.makedirs("results/guided_gradcam", exist_ok=True)
    
    for i in selected_indices:
        input_path = image_paths[i]
        output_path = f"results/guided_gradcam/visualization_{i}.png"
        
        # Process image
        img = Image.open(input_path).convert("RGB")
        input_tensor = transform(img).unsqueeze(0).to(device)
        
        # Generate visualizations
        cam, guided_grads, cam_gb = guided_gradcam(input_tensor)
        
        # Save results
        preprocess_and_show_results(input_path, cam, guided_grads, cam_gb, output_path)
        print(f"✅ Saved visualizations for image {i}")
    
    guided_gradcam.remove_hooks()