From bbb285b15a8348b09275385f2686c39f69df2ea2 Mon Sep 17 00:00:00 2001 From: Kok Wei <52521979+kwcckw@users.noreply.github.com> Date: Fri, 18 Aug 2023 17:11:30 +0800 Subject: [PATCH] Added alpha layer support in augmentations. --- augraphy/augmentations/brightness.py | 2 +- augraphy/augmentations/colorshift.py | 1 - augraphy/augmentations/dirtyrollers.py | 6 ------ augraphy/augmentations/dithering.py | 24 ++++++++---------------- augraphy/augmentations/dotmatrix.py | 6 ------ augraphy/augmentations/geometric.py | 13 ++++++++----- augraphy/augmentations/glitcheffect.py | 10 ++++++++++ augraphy/augmentations/hollow.py | 3 ++- augraphy/augmentations/inkbleed.py | 6 ++++++ augraphy/augmentations/inkcolorswap.py | 7 +++++++ 10 files changed, 42 insertions(+), 36 deletions(-) diff --git a/augraphy/augmentations/brightness.py b/augraphy/augmentations/brightness.py index 99270363..c5573971 100644 --- a/augraphy/augmentations/brightness.py +++ b/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 diff --git a/augraphy/augmentations/colorshift.py b/augraphy/augmentations/colorshift.py index 69600564..90ec4306 100644 --- a/augraphy/augmentations/colorshift.py +++ b/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) diff --git a/augraphy/augmentations/dirtyrollers.py b/augraphy/augmentations/dirtyrollers.py index be79d7eb..eb04d566 100644 --- a/augraphy/augmentations/dirtyrollers.py +++ b/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 diff --git a/augraphy/augmentations/dithering.py b/augraphy/augmentations/dithering.py index a5d2f232..81d8bc47 100644 --- a/augraphy/augmentations/dithering.py +++ b/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 diff --git a/augraphy/augmentations/dotmatrix.py b/augraphy/augmentations/dotmatrix.py index 2335244c..2704ae0e 100644 --- a/augraphy/augmentations/dotmatrix.py +++ b/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 diff --git a/augraphy/augmentations/geometric.py b/augraphy/augmentations/geometric.py index ac6f08c8..7db656ab 100644 --- a/augraphy/augmentations/geometric.py +++ b/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) diff --git a/augraphy/augmentations/glitcheffect.py b/augraphy/augmentations/glitcheffect.py index f9f8aaab..fe19b307 100644 --- a/augraphy/augmentations/glitcheffect.py +++ b/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 diff --git a/augraphy/augmentations/hollow.py b/augraphy/augmentations/hollow.py index 8075f697..bf9d1202 100644 --- a/augraphy/augmentations/hollow.py +++ b/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") diff --git a/augraphy/augmentations/inkbleed.py b/augraphy/augmentations/inkbleed.py index a87e21d1..995c544f 100644 --- a/augraphy/augmentations/inkbleed.py +++ b/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 diff --git a/augraphy/augmentations/inkcolorswap.py b/augraphy/augmentations/inkcolorswap.py index b7f926da..c10ef506 100644 --- a/augraphy/augmentations/inkcolorswap.py +++ b/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)