Fix grad heatmap image

app.py

@@ -1,11 +1,10 @@
 import gradio as gr
 import numpy as np
-import torch
 from PIL import Image
 from matplotlib import cm
 
 from guided_backprop import GuidedBackprop
-from utils import range_norm, denormalize
+from utils import range_norm, grad2heatmapped
 
 
 class GradioApp:
@@ -48,7 +47,7 @@ class GradioApp:
                 gr.Markdown("## Visualize activation maps")
                 with gr.Row():
                     with gr.Column(scale=1):
-                        chosen_layer = gr.Dropdown(label="Layer", info="Choose from the layers", interactive=True)
+                        chosen_layer = gr.Dropdown(label="Layer", value="conv_layer0", info="Choose from the layers", interactive=True)
                         chosen_filter = gr.Slider(label="Filter", info="Choose from the filters", interactive=True)
                         color = gr.Radio(["heatmap", "gray"], value="heatmap", label="Color", info="Choose the color of the activation map")
                     
@@ -127,15 +126,16 @@ class GradioApp:
         input_grad = self.input_grad[0].permute(1,2,0).detach()
         input_grad = range_norm(input_grad)
         input_grad = Image.fromarray(np.uint8(input_grad*255))
+        self.input_grad_img = input_grad
         return input_grad
 
     def apply_input_grad(self, input_image, grad_ratio):
-        grad_applied_img = self.gp.transforms(input_image)
-        mean, std = torch.FloatTensor(self.gp.transforms.mean) , torch.FloatTensor(self.gp.transforms.std)
-        grad_applied_img = grad_applied_img - grad_ratio * self.input_grad[0].detach()
-        grad_applied_img = range_norm(denormalize(grad_applied_img, mean, std))
-        grad_applied_img = Image.fromarray(np.uint8(grad_applied_img.permute(1,2,0)*255))
-        return grad_applied_img
+        heatmapped = grad2heatmapped(
+            input_image, 
+            self.input_grad_img, 
+            grad_ratio)
+        
+        return heatmapped
 
     def launch(self):
         self.app.launch()
@@ -143,4 +143,4 @@ class GradioApp:
 
 if __name__ == "__main__":
     gradio_app = GradioApp()
-    gradio_app.launch()
+    gradio_app.launch()

utils.py

@@ -1,8 +1,33 @@
+import numpy as np
+import cv2
+
+
 def range_norm(img):
     min = img.min()
     max = img.max()
     eps = 1e-6
     return (img-min)/(max-min+eps)
 
-def denormalize(img, mean, std):
-    return img*std.view(-1,1,1) + mean.view(-1,1,1)
+
+def grad2heatmapped(input_image, grad_image, ratio):
+    input_image = np.array(input_image)
+    grad_image = np.array(grad_image)
+    
+    # Invert negative pixels
+    grad_image[grad_image<100] += 128
+
+    # Apply thresholding and blur to obtain heatmap
+    th = cv2.threshold(grad_image, 140, 255, cv2.THRESH_BINARY)[1]
+    blur = cv2.GaussianBlur(th, (11,11), 11)
+    heatmap = cv2.applyColorMap(blur, cv2.COLORMAP_JET)
+
+    # Apply edge padding to heatmap to have 256x256 size
+    heatmap = np.pad(heatmap, ((16,16),(16,16),(0,0)), 'edge')
+
+    # Upsample heatmap to input_image size
+    heatmap = cv2.resize(heatmap, (input_image.shape[1], input_image.shape[0]))
+
+    # Superimpose heatmap on input_image
+    heatmapped = cv2.addWeighted(input_image, 1-ratio, heatmap, ratio, 0)
+
+    return heatmapped