Update model_utils.py

model_utils.py

@@ -54,5 +54,46 @@ def predict(model, img_tensor, device):
     return dict(sorted_probs)
 
 
+import cv2
+import numpy as np
+import torch
+
+def generate_gradcam(model, input_tensor, target_class, device):
+    features = []
+    gradients = []
+
+    def forward_hook(module, input, output):
+        features.append(output.detach())
+
+    def backward_hook(module, grad_input, grad_output):
+        gradients.append(grad_output[0].detach())
+
+    last_conv_layer = model.features[-1]
+    forward_handle = last_conv_layer.register_forward_hook(forward_hook)
+    backward_handle = last_conv_layer.register_backward_hook(backward_hook)
+
+    model.zero_grad()
+    output = model(input_tensor.unsqueeze(0).to(device))
+    class_score = output[0][target_class]
+    class_score.backward()
+
+    grads = gradients[0]
+    fmap = features[0]
+
+    weights = grads.mean(dim=[2, 3], keepdim=True)
+    cam = (weights * fmap).sum(dim=1).squeeze()
+    cam = torch.relu(cam).cpu().numpy()
+
+    cam = cam - cam.min()
+    cam = cam / cam.max()
+    cam = cv2.resize(cam, (224, 224))
+
+    forward_handle.remove()
+    backward_handle.remove()
+
+    return cam
+
+
+