Thegranger/Retinal-XAI-Classifier

1import torch2import torch.nn.functional as F3from torchvision import models4from captum.attr import IntegratedGradients, Occlusion, GradientAttribution5from captum.attr import visualization as viz    6import matplotlib.pyplot as plt7import numpy as np8import cv29 10# ✅ GradCAM Class11class GradCAM(GradientAttribution):12    def __init__(self, model, target_layer):13        super().__init__(model)14        self.model = model15        self.target_layer = target_layer16        self.gradients = None17        self.activations = None18        self._register_hooks()19 20    def _register_hooks(self):21        def forward_hook(module, input, output):22            self.activations = output23 24        def backward_hook(module, grad_in, grad_out):25            self.gradients = grad_out[0]26 27        self.target_layer.register_forward_hook(forward_hook)28        self.target_layer.register_full_backward_hook(backward_hook)29 30    def attribute(self, inputs, target=None):31        outputs = self.model(inputs)32        if target is None:33            target = torch.argmax(outputs, dim=1)34 35        self.model.zero_grad()36        loss = outputs[:, target]37        loss.backward(retain_graph=True)38 39        pooled_gradients = torch.mean(self.gradients, dim=[2, 3], keepdim=True)40        activations = self.activations41        for i in range(activations.shape[1]):42            activations[:, i, :, :] *= pooled_gradients[:, i, :, :]43 44        heatmap = torch.mean(activations, dim=1).squeeze()45        heatmap = F.relu(heatmap)46        heatmap /= torch.max(heatmap)47        return heatmap.detach().cpu().numpy()48 49# ✅ Model Loaders50def load_resnet_model():51    model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)52    model.eval()53    return model, model.layer4[2].conv354 55def load_convnext_model(path="model.pth"):56    model = models.convnext_base(weights=None)57    model.classifier[2] = torch.nn.Linear(model.classifier[2].in_features, 5)58    model.load_state_dict(torch.load(path, map_location=torch.device("cpu")))59    model.eval()60    return model61 62# ✅ Visual Methods63def apply_gradcam(model, layer, input_tensor):64    cam = GradCAM(model, layer)65    heatmap = cam.attribute(input_tensor)66    return heatmap67 68def integrated_gradients(model, input_tensor, target_class):69    ig = IntegratedGradients(model)70    attr = ig.attribute(input_tensor, target=target_class, n_steps=50)71    attr = attr.squeeze().cpu().detach().numpy().transpose(1, 2, 0)72    73    # Sum the absolute values across all channels74    integrated_attr = np.sum(np.abs(attr), axis=2)75    76    # Create a white background (3 channels for RGB)77    white_bg = np.ones_like(integrated_attr) * 255  # Create a white background78    white_bg = np.uint8(white_bg)79    80    # Convert the white background to 3 channels (RGB)81    white_bg_rgb = np.stack([white_bg] * 3, axis=-1)  # Convert to RGB82    83    # Mask the white background with the speckles (heatmap)84    speckles = np.uint8(integrated_attr * 255)  # Convert to a scale for display85    mask = np.clip(speckles, 0, 255)  # Ensure the values are within valid range86    87    # Apply color map to create a heatmap88    result = cv2.applyColorMap(mask, cv2.COLORMAP_JET)89    90    # Now both `white_bg_rgb` and `result` have the same shape (3 channels, H, W)91    final_output = cv2.addWeighted(white_bg_rgb, 0.5, result, 0.5, 0)92 93    # Display the final output94    plt.figure(figsize=(8, 8))95    plt.imshow(final_output)96    plt.axis('off')97    plt.show()98 99    return integrated_attr100 101 102 103def occlusion_map(model, input_tensor, target_class):104    occlusion = Occlusion(model)105    attr = occlusion.attribute(106        input_tensor,107        target=target_class,108        strides=(3, 8, 8),109        sliding_window_shapes=(3, 15, 15),110        baselines=0111    )112    attr = attr.squeeze().cpu().detach().numpy().transpose(1, 2, 0)113    return np.sum(np.abs(attr), axis=2)114 115# ✅ Extract and Convert Feature Maps to Images116def extract_feature_maps(model, layer, input_tensor, visualize=False):117    activations = []118 119    def hook_fn(module, input, output):120        activations.append(output.detach())121 122    handle = layer.register_forward_hook(hook_fn)123 124    with torch.no_grad():125        _ = model(input_tensor)126 127    handle.remove()128 129    feature_map = activations[0].squeeze(0).cpu()  # shape: [C, H, W]130    selected_maps = feature_map[:8]  # First 8 feature maps131 132    feature_images = []133    for fmap in selected_maps:134        fmap_np = fmap.numpy()135 136        # Normalize to 0–1137        fmap_norm = (fmap_np - np.min(fmap_np)) / (np.max(fmap_np) - np.min(fmap_np) + 1e-6)138 139        # Convert to grayscale image (0–255)140        fmap_gray = np.uint8(fmap_norm * 255)141 142        # Convert grayscale to RGB (structure is still grayscale)143        fmap_rgb = np.stack([fmap_gray] * 3, axis=-1)144 145        # Create light blue tint (e.g., soft blue: R=200, G=220, B=255)146        tint_color = np.array([200, 220, 255], dtype=np.uint8)147 148        # Blend the grayscale image with the tint (light tint)149        blended = (0.9 * fmap_rgb + 0.1 * tint_color).astype(np.uint8)150 151        feature_images.append(blended)152 153    return feature_images154 155 156 157 158 159# ✅ Overlay for GradCAM160def create_colormap_overlay(image, mask):161    mask_resized = cv2.resize(mask, (image.shape[1], image.shape[0]))162    heatmap = cv2.applyColorMap(np.uint8(255 * mask_resized), cv2.COLORMAP_JET)163    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)164    overlay = 0.4 * heatmap + 0.6 * image165    return np.uint8(overlay)166