Improvements on image saving

Author: utkuozbulak

src/gradcam.py

@@ -86,12 +86,13 @@ def generate_cam(self, input_image, target_class=None):
         cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam))  # Normalize between 0-1
         cam = np.uint8(cam * 255)  # Scale between 0-255 to visualize
         cam = np.uint8(Image.fromarray(cam).resize((input_image.shape[2],
-                       input_image.shape[3]), Image.ANTIALIAS))
+                       input_image.shape[3]), Image.ANTIALIAS))/255
         # ^ I am extremely unhappy with this line. Originally resizing was done in cv2 which
-        # supports resizing numpy matrices, however, when I moved the repository to PIL, this
-        # option is out of the window. So, in order to use resizing with ANTIALIAS feature of PIL,
+        # supports resizing numpy matrices with antialiasing, however,
+        # when I moved the repository to PIL, this option was out of the window.
+        # So, in order to use resizing with ANTIALIAS feature of PIL,
         # I briefly convert matrix to PIL image and then back.
-        # If there is a more beautiful way, send a PR.
+        # If there is a more beautiful way, do not hesitate to send a PR.
         return cam
 
 

src/misc_functions.py

@@ -98,30 +98,43 @@ def apply_colormap_on_image(org_im, activation, colormap_name):
     return no_trans_heatmap, heatmap_on_image
 
 
+def format_np_output(np_arr):
+    """
+        This is a (kind of) bandaid fix to streamline saving procedure.
+        It converts all the outputs to the same format which is 3xWxH
+        with using sucecssive if clauses.
+    Args:
+        im_as_arr (Numpy array): Matrix of shape 1xWxH or WxH or 3xWxH
+    """
+    # Phase/Case 1: The np arr only has 2 dimensions
+    # Result: Add a dimension at the beginning
+    if len(np_arr.shape) == 2:
+        np_arr = np.expand_dims(np_arr, axis=0)
+    # Phase/Case 2: Np arr has only 1 channel (assuming first dim is channel)
+    # Result: Repeat first channel and convert 1xWxH to 3xWxH
+    if np_arr.shape[0] == 1:
+        np_arr = np.repeat(np_arr, 3, axis=0)
+    # Phase/Case 3: Np arr is of shape 3xWxH
+    # Result: Convert it to WxHx3 in order to make it saveable by PIL
+    if np_arr.shape[0] == 3:
+        np_arr = np_arr.transpose(1, 2, 0)
+    # Phase/Case 4: NP arr is normalized between 0-1
+    # Result: Multiply with 255 and change type to make it saveable by PIL
+    if np.max(np_arr) <= 1:
+        np_arr = (np_arr*255).astype(np.uint8)
+    return np_arr
+
+
 def save_image(im, path):
     """
-        Saves a numpy matrix of shape D(1 or 3) x W x H as an image
+        Saves a numpy matrix or PIL image as an image
     Args:
         im_as_arr (Numpy array): Matrix of shape DxWxH
         path (str): Path to the image
-
-    TODO: Streamline image saving, it is ugly.
-    """
-    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
-        # A bandaid fix to an issue with gradcam
-        if im.shape[0] == 3 and np.max(im) == 1:
-            im = im.transpose(1, 2, 0) * 255
-        elif im.shape[0] == 3 and np.max(im) > 1:
-            im = im.transpose(1, 2, 0)
-        im = Image.fromarray(im.astype(np.uint8))
+    """
+    if isinstance(im, (np.ndarray, np.generic)):
+        im = format_np_output(im)
+        im = Image.fromarray(im)
     im.save(path)