File size: 4,251 Bytes
fbbdeab
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
"""
Grad-CAM implementation for explainable AI
Shows which parts of the image influenced the model's decision
"""

import torch
import torch.nn.functional as F
import numpy as np
import cv2
from PIL import Image

class GradCAM:
    """
    Grad-CAM for VGG16 to visualize model decisions
    """
    
    def __init__(self, model, target_layer):
        """
        Initialize Grad-CAM
        
        Args:
            model: Trained VGG16 model
            target_layer: Layer to visualize (e.g., model.features[-1])
        """
        self.model = model
        self.target_layer = target_layer
        self.gradients = None
        self.activations = None
        
        # Register hooks
        self.target_layer.register_forward_hook(self.save_activation)
        self.target_layer.register_backward_hook(self.save_gradient)
    
    def save_activation(self, module, input, output):
        """Save forward activation"""
        self.activations = output.detach()
    
    def save_gradient(self, module, grad_input, grad_output):
        """Save backward gradient"""
        self.gradients = grad_output[0].detach()
    
    def generate_cam(self, input_tensor, class_idx):
        """
        Generate Class Activation Map
        
        Args:
            input_tensor: Preprocessed input image tensor
            class_idx: Target class index
            
        Returns:
            numpy array: Heatmap (0-255)
        """
        # Forward pass
        output = self.model(input_tensor)
        
        # Zero gradients
        self.model.zero_grad()
        
        # Backward pass for target class
        class_score = output[0, class_idx]
        class_score.backward()
        
        # Get gradients and activations
        gradients = self.gradients[0]  # [C, H, W]
        activations = self.activations[0]  # [C, H, W]
        
        # Calculate weights (global average pooling of gradients)
        weights = gradients.mean(dim=(1, 2), keepdim=True)  # [C, 1, 1]
        
        # Weighted combination of activation maps
        cam = (weights * activations).sum(dim=0)  # [H, W]
        
        # Apply ReLU (only positive influence)
        cam = F.relu(cam)
        
        # Normalize to 0-1
        cam = cam - cam.min()
        cam = cam / cam.max()
        
        # Convert to numpy and resize to 224x224
        cam = cam.cpu().numpy()
        cam = cv2.resize(cam, (224, 224))
        
        # Convert to 0-255 range
        cam = np.uint8(255 * cam)
        
        return cam
    
    def overlay_heatmap(self, image, heatmap, alpha=0.5):
        """
        Overlay heatmap on original image
        
        Args:
            image: Original PIL Image
            heatmap: Heatmap numpy array (0-255)
            alpha: Transparency (0-1)
            
        Returns:
            PIL Image with heatmap overlay
        """
        # Resize image to 224x224
        image = image.resize((224, 224))
        image_np = np.array(image)
        
        # Apply colormap to heatmap (red = high activation)
        heatmap_colored = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
        heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
        
        # Overlay
        overlayed = cv2.addWeighted(image_np, 1-alpha, heatmap_colored, alpha, 0)
        
        return Image.fromarray(overlayed)


def generate_gradcam_visualization(model, image, predicted_class_idx, transform):
    """
    Generate Grad-CAM visualization for a prediction
    
    Args:
        model: Trained VGG16 model
        image: PIL Image
        predicted_class_idx: Index of predicted class
        transform: Image transformation pipeline
        
    Returns:
        PIL Image: Image with heatmap overlay
    """
    # Set model to eval mode
    model.eval()
    
    # Preprocess image
    input_tensor = transform(image).unsqueeze(0)
    
    # Create Grad-CAM instance (target last convolutional layer of VGG16)
    gradcam = GradCAM(model, model.features[-1])
    
    # Generate heatmap
    heatmap = gradcam.generate_cam(input_tensor, predicted_class_idx)
    
    # Overlay on original image
    visualization = gradcam.overlay_heatmap(image, heatmap, alpha=0.4)
    
    return visualization