Merge pull request #349 from kwcckw/dev

Added alpha layer support in augmentations.

augraphy/augmentations/brightness.py

@@ -116,4 +116,4 @@ def __call__(self, image, layer=None, force=False):
             if has_alpha:
                 image_output = np.dstack((image_output, image_alpha))
 
-            return image
+            return image_output

augraphy/augmentations/colorshift.py

@@ -174,7 +174,6 @@ def __call__(self, image, layer=None, force=False):
                 # increase kernel value in each iteration to create a betetr effect
                 kernel_value += 2
 
-            # return image follows the input image color channel
             # return image follows the input image color channel
             if is_gray:
                 image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2GRAY)

augraphy/augmentations/dirtyrollers.py

@@ -191,12 +191,8 @@ def __call__(self, image, layer=None, force=False):
             image = image.copy()
 
             # check and convert image into BGR format
-            has_alpha = 0
             if len(image.shape) > 2:
                 is_gray = 0
-                if image.shape[2] == 4:
-                    has_alpha = 1
-                    image, image_alpha = image[:, :, :3], image[:, :, 3]
             else:
                 is_gray = 1
                 image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
@@ -225,7 +221,5 @@ def __call__(self, image, layer=None, force=False):
 
             if is_gray:
                 image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2GRAY)
-            if has_alpha:
-                image_output = np.dstack((image_output, image_alpha))
 
             return image_output

augraphy/augmentations/dithering.py

@@ -48,18 +48,17 @@ def dither_Floyd_Steinberg(self, image):
         :type image: numpy.array (numpy.uint8)
         """
 
+        ysize, xsize = image.shape[:2]
+        img_dither_fs = image.copy().astype("float")
         if len(image.shape) > 2:  # coloured image
-            ysize, xsize, dim = image.shape
-            img_dither_fs = image.copy().astype("float")
-            for channel_num in range(dim):
+            # skip alpha channel
+            for channel_num in range(3):
                 self.apply_Floyd_Steinberg(
                     img_dither_fs[:, :, channel_num],
                     ysize,
                     xsize,
                 )
         else:  # grayscale or binary
-            ysize, xsize = image.shape
-            img_dither_fs = image.copy().astype("float")
             self.apply_Floyd_Steinberg(img_dither_fs, ysize, xsize)
 
         return img_dither_fs.astype("uint8")
@@ -132,10 +131,11 @@ def dither_Ordered(self, image, order=5):
                 ordered_matrix[y][x] = np.floor((value / total_number) * 255)
         ordered_matrix = np.array(ordered_matrix, dtype="float64")
 
+        ysize, xsize = image.shape[:2]
+        img_dither_ordered = image.copy().astype("float")
         if len(image.shape) > 2:  # coloured image
-            ysize, xsize, dim = image.shape
-            img_dither_ordered = image.copy().astype("float")
-            for channel_num in range(dim):
+            # skip alpha channel
+            for channel_num in range(3):
                 self.apply_Ordered(
                     img_dither_ordered[:, :, channel_num],
                     ysize,
@@ -144,8 +144,6 @@ def dither_Ordered(self, image, order=5):
                     ordered_matrix,
                 )
         else:  # grayscale or binary
-            ysize, xsize = image.shape
-            img_dither_ordered = image.copy().astype("float")
             self.apply_Ordered(
                 img_dither_ordered,
                 ysize,
@@ -202,12 +200,8 @@ def __call__(self, image, layer=None, force=False):
             image = image.copy()
 
             # check and convert image into BGR format
-            has_alpha = 0
             if len(image.shape) > 2:
                 is_gray = 0
-                if image.shape[2] == 4:
-                    has_alpha = 1
-                    image, image_alpha = image[:, :, :3], image[:, :, 3]
             else:
                 is_gray = 1
                 image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
@@ -224,7 +218,5 @@ def __call__(self, image, layer=None, force=False):
 
             if is_gray:
                 image_dither = cv2.cvtColor(image_dither, cv2.COLOR_BGR2GRAY)
-            if has_alpha:
-                image_dither = np.dstack((image_dither, image_alpha))
 
             return image_dither

augraphy/augmentations/dotmatrix.py

@@ -279,12 +279,8 @@ def __call__(self, image, layer=None, force=False):
             ysize, xsize = image.shape[:2]
 
             # check and convert image into BGR format
-            has_alpha = 0
             if len(image.shape) > 2:
                 is_gray = 0
-                if image.shape[2] == 4:
-                    has_alpha = 1
-                    image, image_alpha = image[:, :, :3], image[:, :, 3]
             else:
                 is_gray = 1
                 image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
@@ -571,7 +567,5 @@ def __call__(self, image, layer=None, force=False):
             # return image follows the input image color channel
             if is_gray:
                 image_dot_matrix = cv2.cvtColor(image_dot_matrix, cv2.COLOR_BGR2GRAY)
-            if has_alpha:
-                image_dot_matrix = np.dstack((image_dot_matrix, image_alpha))
 
             return image_dot_matrix

augraphy/augmentations/geometric.py

@@ -149,7 +149,10 @@ def __call__(self, image, layer=None, force=False):
 
                 # convert from rgb to grayscale using their average
                 if len(image.shape) < 3:
-                    self.padding_value = np.mean(self.padding_value)
+                    padding_value = np.mean(self.padding_value)
+                elif image.shape[2] == 4:
+                    # add alpha value
+                    padding_value = (self.padding_value[0], self.padding_value[1], self.padding_value[2], 255)
 
                 # padding on left side
                 if self.padding[0] > 0:
@@ -170,7 +173,7 @@ def __call__(self, image, layer=None, force=False):
                     elif self.padding_type == "mirror":
                         image_padding = np.fliplr(image[:, : self.padding[0]].copy())
                     else:
-                        image_padding = np.full(padding_shape, fill_value=self.padding_value, dtype="uint8")
+                        image_padding = np.full(padding_shape, fill_value=padding_value, dtype="uint8")
                     # combine padding image and original image
                     image = np.concatenate([image_padding, image], axis=1)
 
@@ -193,7 +196,7 @@ def __call__(self, image, layer=None, force=False):
                     elif self.padding_type == "mirror":
                         image_padding = np.fliplr(image[:, -self.padding[1] :].copy())
                     else:
-                        image_padding = np.full(padding_shape, fill_value=self.padding_value, dtype="uint8")
+                        image_padding = np.full(padding_shape, fill_value=padding_value, dtype="uint8")
                     # combine padding image and original image
                     image = np.concatenate([image, image_padding], axis=1)
 
@@ -216,7 +219,7 @@ def __call__(self, image, layer=None, force=False):
                     elif self.padding_type == "mirror":
                         image_padding = np.flipud(image[: self.padding[2], :].copy())
                     else:
-                        image_padding = np.full(padding_shape, fill_value=self.padding_value, dtype="uint8")
+                        image_padding = np.full(padding_shape, fill_value=padding_value, dtype="uint8")
                     # combine padding image and original image
                     image = np.concatenate([image_padding, image], axis=0)
 
@@ -239,7 +242,7 @@ def __call__(self, image, layer=None, force=False):
                     elif self.padding_type == "mirror":
                         image_padding = np.flipud(image[-self.padding[3] :, :].copy())
                     else:
-                        image_padding = np.full(padding_shape, fill_value=self.padding_value, dtype="uint8")
+                        image_padding = np.full(padding_shape, fill_value=padding_value, dtype="uint8")
                     # combine padding image and original image
                     image = np.concatenate([image, image_padding], axis=0)
 

augraphy/augmentations/glitcheffect.py

@@ -115,6 +115,13 @@ def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
             image = image.copy()
 
+            # check and convert image into BGR format
+            if len(image.shape) > 2:
+                is_gray = 0
+            else:
+                is_gray = 1
+                image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGRA)
+
             # apply color shift before the glitch effect
             color_shift = ColorShift(
                 color_shift_offset_x_range=(3, 5),
@@ -146,4 +153,7 @@ def __call__(self, image, layer=None, force=False):
                     image_output = np.rot90(self.apply_glitch(np.rot90(image_output, 1)), 3)
                     image_output = self.apply_glitch(image_output)
 
+            if is_gray:
+                image_output = cv2.cvtColor(image_output, cv2.COLOR_BGRA2GRAY)
+
             return image_output

augraphy/augmentations/hollow.py

@@ -235,7 +235,8 @@ def __call__(self, image, layer=None, force=False):
 
             # get background by removing the detected contours
             image_output = image.copy()
-            image_output[image_mask > 0] = image_median[image_mask > 0]
+            for i in range(3):
+                image_output[:, :, i][image_mask > 0] = image_median[:, :, i][image_mask > 0]
 
             # create a rando mask
             image_random = np.random.randint(0, 255, size=image_mask.shape, dtype="uint8")

augraphy/augmentations/inkbleed.py

@@ -47,8 +47,12 @@ def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
 
             # convert and make sure image is color image
+            has_alpha = 0
             if len(image.shape) > 2:
                 is_gray = 0
+                if image.shape[2] == 4:
+                    has_alpha = 1
+                    image, image_alpha = image[:, :, :3], image[:, :, 3]
             else:
                 is_gray = 1
                 image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
@@ -82,5 +86,7 @@ def __call__(self, image, layer=None, force=False):
             # return image follows the input image color channel
             if is_gray:
                 image_output = cv2.cvtColor(image_output, cv2.COLOR_BGR2GRAY)
+            if has_alpha:
+                image_output = np.dstack((image_output, image_alpha))
 
             return image_output

augraphy/augmentations/inkcolorswap.py

@@ -193,8 +193,15 @@ def __call__(self, image, layer=None, force=False):
             else:
                 ink_swap_color = self.ink_swap_color
 
+            # add alpha value
+            if image.shape[2] == 4:
+                ink_swap_color = (ink_swap_color[0], ink_swap_color[1], ink_swap_color[2], 255)
+
             # create a mask of swap color
             image_color = np.full_like(image, fill_value=ink_swap_color, dtype="uint8")
+            # update alpha
+            if image.shape[2] == 4:
+                image_color[:, :, 3] = image[:, :, 3].copy()
 
             # blend image with swap color
             image_color = cv2.addWeighted(image, 1.0, image_color, 1.0, 0)