Upload gradcam.py

gradcam.py

@@ -0,0 +1,77 @@
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
+from torchvision import transforms
+from torch.nn.functional import interpolate
+
+class GradCAM:
+    def __init__(self, model, target_layer):
+        self.model = model
+        self.target_layer = target_layer
+        self.activations = []
+        self.gradients = []
+
+        # Register hooks
+        target_layer.register_forward_hook(self.save_activations)
+        target_layer.register_backward_hook(self.save_gradients)
+
+    def save_activations(self, module, input, output):
+        self.activations.append(output.detach())
+
+    def save_gradients(self, module, grad_input, grad_output):
+        self.gradients.append(grad_output[0].detach())
+
+    def forward(self, input_tensor):
+        return self.model(input_tensor)
+
+    def generate(self, input_tensor, target_class):
+        # Forward pass
+        output = self.forward(input_tensor)
+
+        # Backward pass for specific class
+        self.model.zero_grad()
+        loss = output[:, target_class].mean()
+        loss.backward(retain_graph=True)
+
+        # Get activations and gradients
+        activations = self.activations[0].cpu().data.numpy()[0]
+        gradients = self.gradients[0].cpu().data.numpy()[0]
+
+        # Compute weights
+        weights = np.mean(gradients, axis=(1, 2))
+
+        # Create CAM
+        cam = np.zeros(activations.shape[1:], dtype=np.float32)
+        for i, w in enumerate(weights):
+            cam += w * activations[i, :, :]
+
+        # Post-process CAM
+        cam = np.maximum(cam, 0)
+        cam = interpolate(torch.from_numpy(cam[None, None]),
+                        size=(224, 224), mode='bilinear').numpy()
+        cam = cam.squeeze()
+        if cam.max() != 0:
+            cam /= cam.max()
+
+        return cam
+
+def generate_gradcam(image, target_class, model, target_layer):
+    # Preprocess image
+    preprocess = transforms.Compose([
+        transforms.ToTensor(),
+    ])
+
+    if not isinstance(image, torch.Tensor):
+        image = preprocess(image)
+
+    image_preprocessed = image.unsqueeze(0).requires_grad_(True).to(device)
+
+    # Initialize Grad-CAM
+    gradcam = GradCAM(model, target_layer)
+
+    # Generate CAM
+    image = image.to(device) 
+
+    cam = gradcam.generate(image_preprocessed, target_class)
+    return cam
+