Remove CV2 dependencies, update 1

Author: utkuozbulak

src/gradcam.py

@@ -7,7 +7,7 @@
 import numpy as np
 import torch
 
-from misc_functions import get_params, save_class_activation_on_image
+from misc_functions import get_example_params, save_class_activation_images
 
 
 class CamExtractor():
@@ -93,11 +93,11 @@ def generate_cam(self, input_image, target_class=None):
     # Get params
     target_example = 2  # Snake
     (original_image, prep_img, target_class, file_name_to_export, pretrained_model) =\
-        get_params(target_example)
+        get_example_params(target_example)
     # Grad cam
     grad_cam = GradCam(pretrained_model, target_layer=11)
     # Generate cam mask
     cam = grad_cam.generate_cam(prep_img, target_class)
     # Save mask
-    save_class_activation_on_image(original_image, cam, file_name_to_export)
+    save_class_activation_images(original_image, cam, file_name_to_export)
     print('Grad cam completed')

src/guided_backprop.py

@@ -6,7 +6,7 @@
 import torch
 from torch.nn import ReLU
 
-from misc_functions import (get_params,
+from misc_functions import (get_example_params,
                             convert_to_grayscale,
                             save_gradient_images,
                             get_positive_negative_saliency)
@@ -66,7 +66,7 @@ def generate_gradients(self, input_image, target_class):
 if __name__ == '__main__':
     target_example = 0  # Snake
     (original_image, prep_img, target_class, file_name_to_export, pretrained_model) =\
-        get_params(target_example)
+        get_example_params(target_example)
 
     # Guided backprop
     GBP = GuidedBackprop(pretrained_model)

src/guided_gradcam.py

@@ -5,7 +5,7 @@
 """
 import numpy as np
 
-from misc_functions import (get_params,
+from misc_functions import (get_example_params,
                             convert_to_grayscale,
                             save_gradient_images)
 from gradcam import GradCam
@@ -29,7 +29,7 @@ def guided_grad_cam(grad_cam_mask, guided_backprop_mask):
     # Get params
     target_example = 0  # Snake
     (original_image, prep_img, target_class, file_name_to_export, pretrained_model) =\
-        get_params(target_example)
+        get_example_params(target_example)
 
     # Grad cam
     gcv2 = GradCam(pretrained_model, target_layer=11)

src/inverted_representation.py

@@ -9,7 +9,7 @@
 from torch.optim import SGD
 import os
 
-from misc_functions import get_params, recreate_image
+from misc_functions import get_example_params, recreate_image
 
 
 class InvertedRepresentation():
@@ -117,7 +117,7 @@ def generate_inverted_image_specific_layer(self, input_image, img_size, target_l
     # Get params
     target_example = 0  # Snake
     (original_image, prep_img, target_class, file_name_to_export, pretrained_model) =\
-        get_params(target_example)
+        get_example_params(target_example)
 
     inverted_representation = InvertedRepresentation(pretrained_model)
     image_size = 224  # width & height

src/misc_functions.py

@@ -5,25 +5,26 @@
 """
 import os
 import copy
-import cv2
 import numpy as np
+from PIL import Image
+import matplotlib.cm as mpl_color_map
 
 import torch
 from torch.autograd import Variable
 from torchvision import models
 
 
-def convert_to_grayscale(cv2im):
+def convert_to_grayscale(im_as_arr):
     """
         Converts 3d image to grayscale
 
     Args:
-        cv2im (numpy arr): RGB image with shape (D,W,H)
+        im_as_arr (numpy arr): RGB image with shape (D,W,H)
 
     returns:
         grayscale_im (numpy_arr): Grayscale image with shape (1,W,D)
     """
-    grayscale_im = np.sum(np.abs(cv2im), axis=0)
+    grayscale_im = np.sum(np.abs(im_as_arr), axis=0)
     im_max = np.percentile(grayscale_im, 99)
     im_min = np.min(grayscale_im)
     grayscale_im = (np.clip((grayscale_im - im_min) / (im_max - im_min), 0, 1))
@@ -41,59 +42,100 @@ def save_gradient_images(gradient, file_name):
     """
     if not os.path.exists('../results'):
         os.makedirs('../results')
+    # Normalize
     gradient = gradient - gradient.min()
     gradient /= gradient.max()
-    gradient = np.uint8(gradient * 255).transpose(1, 2, 0)
+    # Save image
     path_to_file = os.path.join('../results', file_name + '.jpg')
-    # Convert RBG to GBR
-    gradient = gradient[..., ::-1]
-    cv2.imwrite(path_to_file, gradient)
+    save_image(gradient, path_to_file)
 
 
-def save_class_activation_on_image(org_img, activation_map, file_name):
+def save_class_activation_images(org_img, activation_map, file_name):
     """
         Saves cam activation map and activation map on the original image
 
     Args:
         org_img (PIL img): Original image
-        activation_map (numpy arr): activation map (grayscale) 0-255
+        activation_map (numpy arr): Activation map (grayscale) 0-255
         file_name (str): File name of the exported image
     """
     if not os.path.exists('../results'):
         os.makedirs('../results')
     # Grayscale activation map
-    path_to_file = os.path.join('../results', file_name+'_Cam_Grayscale.jpg')
-    cv2.imwrite(path_to_file, activation_map)
-    # Heatmap of activation map
-    activation_heatmap = cv2.applyColorMap(activation_map, cv2.COLORMAP_HSV)
-    path_to_file = os.path.join('../results', file_name+'_Cam_Heatmap.jpg')
-    cv2.imwrite(path_to_file, activation_heatmap)
-    # Heatmap on picture
-    org_img = cv2.resize(org_img, (224, 224))
-    img_with_heatmap = np.float32(activation_heatmap) + np.float32(org_img)
-    img_with_heatmap = img_with_heatmap / np.max(img_with_heatmap)
-    path_to_file = os.path.join('../results', file_name+'_Cam_On_Image.jpg')
-    cv2.imwrite(path_to_file, np.uint8(255 * img_with_heatmap))
-
-
-def preprocess_image(cv2im, resize_im=True):
+    heatmap, heatmap_on_image = apply_colormap_on_image(org_img, activation_map, 'hsv')
+    # Save colored heatmap
+    path_to_file = os.path.join('../results', file_name+'_Cam_Heatmap.png')
+    save_image(heatmap, path_to_file)
+    # Save heatmap on iamge
+    path_to_file = os.path.join('../results', file_name+'_Cam_On_Image.png')
+    save_image(heatmap_on_image, path_to_file)
+    # SAve grayscale heatmap
+    path_to_file = os.path.join('../results', file_name+'_Cam_Grayscale.png')
+    save_image(activation_map, path_to_file)
+
+
+def apply_colormap_on_image(org_im, activation, colormap_name):
+    """
+        Apply heatmap on image
+    Args:
+        org_img (PIL img): Original image
+        activation_map (numpy arr): Activation map (grayscale) 0-255
+        colormap_name (str): Name of the colormap
+    """
+    # Get colormap
+    color_map = mpl_color_map.get_cmap(colormap_name)
+    no_trans_heatmap = color_map(activation)
+    # Change alpha channel in colormap to make sure original image is displayed
+    heatmap = copy.copy(no_trans_heatmap)
+    heatmap[:, :, 3] = 0.4
+    heatmap = Image.fromarray((heatmap*255).astype(np.uint8))
+    no_trans_heatmap = Image.fromarray((no_trans_heatmap*255).astype(np.uint8))
+
+    # Apply heatmap on iamge
+    heatmap_on_image = Image.new("RGBA", org_im.size)
+    heatmap_on_image = Image.alpha_composite(heatmap_on_image, org_im.convert('RGBA'))
+    heatmap_on_image = Image.alpha_composite(heatmap_on_image, heatmap)
+    return no_trans_heatmap, heatmap_on_image
+
+
+def save_image(im, path):
+    """
+        Saves a numpy matrix of shape D(1 or 3) x W x H as an image
+    Args:
+        im_as_arr (Numpy array): Matrix of shape DxWxH
+        path (str): Path to the image
+    """
+    if isinstance(im, np.ndarray):
+        if len(im.shape) == 2:
+            im = np.expand_dims(im, axis=0)
+        if im.shape[0] == 1:
+            # Converting an image with depth = 1 to depth = 3, repeating the same values
+            # For some reason PIL complains when I want to save channel image as jpg without
+            # additional format in the .save()
+            im = np.repeat(im, 3, axis=0)
+            # Convert to values to range 1-255 and W,H, D
+        im = im.transpose(1, 2, 0) * 255
+        im = Image.fromarray(im.astype(np.uint8))
+    im.save(path)
+
+
+def preprocess_image(pil_im, resize_im=True):
     """
         Processes image for CNNs
 
     Args:
         PIL_img (PIL_img): Image to process
         resize_im (bool): Resize to 224 or not
     returns:
-        im_as_var (Pytorch variable): Variable that contains processed float tensor
+        im_as_var (torch variable): Variable that contains processed float tensor
     """
     # mean and std list for channels (Imagenet)
     mean = [0.485, 0.456, 0.406]
     std = [0.229, 0.224, 0.225]
     # Resize image
     if resize_im:
-        cv2im = cv2.resize(cv2im, (224, 224))
-    im_as_arr = np.float32(cv2im)
-    im_as_arr = np.ascontiguousarray(im_as_arr[..., ::-1])
+        pil_im.thumbnail((512, 512))
+    im_as_arr = np.float32(pil_im)
     im_as_arr = im_as_arr.transpose(2, 0, 1)  # Convert array to D,W,H
     # Normalize the channels
     for channel, _ in enumerate(im_as_arr):
@@ -127,11 +169,6 @@ def recreate_image(im_as_var):
         recreated_im[c] -= reverse_mean[c]
     recreated_im[recreated_im > 1] = 1
     recreated_im[recreated_im < 0] = 0
-    recreated_im = np.round(recreated_im * 255)
-
-    recreated_im = np.uint8(recreated_im).transpose(1, 2, 0)
-    # Convert RBG to GBR
-    recreated_im = recreated_im[..., ::-1]
     return recreated_im
 
 
@@ -149,7 +186,7 @@ def get_positive_negative_saliency(gradient):
     return pos_saliency, neg_saliency
 
 
-def get_params(example_index):
+def get_example_params(example_index):
     """
         Gets used variables for almost all visualizations, like the image, model etc.
 
@@ -164,15 +201,14 @@ def get_params(example_index):
         pretrained_model(Pytorch model): Model to use for the operations
     """
     # Pick one of the examples
-    example_list = [['../input_images/snake.jpg', 56],
-                    ['../input_images/cat_dog.png', 243],
-                    ['../input_images/spider.png', 72]]
-    selected_example = example_index
-    img_path = example_list[selected_example][0]
-    target_class = example_list[selected_example][1]
+    example_list = (('../input_images/snake.jpg', 56),
+                    ('../input_images/cat_dog.png', 243),
+                    ('../input_images/spider.png', 72))
+    img_path = example_list[example_index][0]
+    target_class = example_list[example_index][1]
     file_name_to_export = img_path[img_path.rfind('/')+1:img_path.rfind('.')]
     # Read image
-    original_image = cv2.imread(img_path, 1)
+    original_image = Image.open(img_path).convert('RGB')
     # Process image
     prep_img = preprocess_image(original_image)
     # Define model

src/smooth_grad.py

@@ -8,7 +8,7 @@
 from torch.autograd import Variable
 import torch
 
-from misc_functions import (get_params,
+from misc_functions import (get_example_params,
                             convert_to_grayscale,
                             save_gradient_images)
 from vanilla_backprop import VanillaBackprop
@@ -49,7 +49,7 @@ def generate_smooth_grad(Backprop, prep_img, target_class, param_n, param_sigma_
     # Get params
     target_example = 0  # Snake
     (original_image, prep_img, target_class, file_name_to_export, pretrained_model) =\
-        get_params(target_example)
+        get_example_params(target_example)
 
     VBP = VanillaBackprop(pretrained_model)
     # GBP = GuidedBackprop(pretrained_model)  # if you want to use GBP dont forget to

src/vanilla_backprop.py

@@ -5,7 +5,7 @@
 """
 import torch
 
-from misc_functions import get_params, convert_to_grayscale, save_gradient_images
+from misc_functions import get_example_params, convert_to_grayscale, save_gradient_images
 
 
 class VanillaBackprop():
@@ -46,9 +46,9 @@ def generate_gradients(self, input_image, target_class):
 
 if __name__ == '__main__':
     # Get params
-    target_example = 0  # Snake
+    target_example = 1  # Snake
     (original_image, prep_img, target_class, file_name_to_export, pretrained_model) =\
-        get_params(target_example)
+        get_example_params(target_example)
     # Vanilla backprop
     VBP = VanillaBackprop(pretrained_model)
     # Generate gradients