diff --git a/augraphy/augmentations/badphotocopy.py b/augraphy/augmentations/badphotocopy.py index ed17be17..ac8c153b 100644 --- a/augraphy/augmentations/badphotocopy.py +++ b/augraphy/augmentations/badphotocopy.py @@ -213,9 +213,13 @@ def apply_augmentation(self, image): :type image: numpy.array (numpy.uint8) """ - # convert and make sure image is color image + # 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) @@ -336,6 +340,8 @@ def apply_augmentation(self, image): # return image follows the input image color channel if is_gray: result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY) + if has_alpha: + result = np.dstack((result, image_alpha)) return result diff --git a/augraphy/augmentations/bindingsandfasteners.py b/augraphy/augmentations/bindingsandfasteners.py index 8c8c544c..11883614 100644 --- a/augraphy/augmentations/bindingsandfasteners.py +++ b/augraphy/augmentations/bindingsandfasteners.py @@ -786,6 +786,13 @@ def __call__(self, image, layer=None, force=False): image = image.copy() ysize, xsize = image.shape[:2] + # check for alpha layer + has_alpha = 0 + if len(image.shape) > 2: + if image.shape[2] == 4: + has_alpha = 1 + image, image_alpha = image[:, :, :3], image[:, :, 3] + # generate randomized overlay types if self.overlay_types == "random": overlay_types = random.choice( @@ -888,4 +895,10 @@ def __call__(self, image, layer=None, force=False): image_output = ob.build_overlay() + if has_alpha: + ysize, xsize = image_output.shape[:2] + if ysize != image_alpha.shape[0] or xsize != image_alpha.shape[1]: + image_alpha = cv2.resize(image_alpha, (xsize, ysize), interpolation=cv2.INTER_AREA) + image_output = np.dstack((image_output, image_alpha)) + return image_output diff --git a/augraphy/augmentations/bleedthrough.py b/augraphy/augmentations/bleedthrough.py index 0c58876e..9fe44d3d 100644 --- a/augraphy/augmentations/bleedthrough.py +++ b/augraphy/augmentations/bleedthrough.py @@ -173,6 +173,13 @@ def __call__(self, image, layer=None, force=False): if force or self.should_run(): image = image.copy() + # check for alpha layer + has_alpha = 0 + if len(image.shape) > 2: + if image.shape[2] == 4: + has_alpha = 1 + image, image_alpha = image[:, :, :3], image[:, :, 3] + image_bleedthrough_foreground = self.create_bleedthrough_foreground(image) image_bleed = self.generate_bleeding_ink( @@ -185,4 +192,7 @@ def __call__(self, image, layer=None, force=False): image_bleed_offset = self.generate_offset(image_bleed, self.offsets) image_bleedthrough = self.blend(image, image_bleed_offset, self.alpha) + if has_alpha: + image_bleedthrough = np.dstack((image_bleedthrough, image_alpha)) + return image_bleedthrough diff --git a/augraphy/augmentations/bookbinding.py b/augraphy/augmentations/bookbinding.py index 5fbdc307..19105dec 100644 --- a/augraphy/augmentations/bookbinding.py +++ b/augraphy/augmentations/bookbinding.py @@ -502,8 +502,12 @@ def __call__(self, image, layer=None, force=False): image = image.copy() # 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) @@ -533,5 +537,10 @@ 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: + ysize, xsize = image_output.shape[:2] + if ysize != image_alpha.shape[0] or xsize != image_alpha.shape[1]: + image_alpha = np.full((ysize, xsize), fill_value=255, dtype="uint8") + image_output = np.dstack((image_output, image_alpha)) return image_output diff --git a/augraphy/augmentations/brightness.py b/augraphy/augmentations/brightness.py index f1411beb..99270363 100644 --- a/augraphy/augmentations/brightness.py +++ b/augraphy/augmentations/brightness.py @@ -77,11 +77,20 @@ def adjust_min_brightness(image, min_brightness_value): # Applies the Augmentation to input data. def __call__(self, image, layer=None, force=False): if force or self.should_run(): - image = image.copy() - value = random.uniform(self.brightness_range[0], self.brightness_range[1]) - if len(image.shape) < 3: + + 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) + + value = random.uniform(self.brightness_range[0], self.brightness_range[1]) + hsv = cv2.cvtColor(image.astype("uint8"), cv2.COLOR_BGR2HSV) hsv = np.array(hsv, dtype=np.float64) @@ -99,6 +108,12 @@ def __call__(self, image, layer=None, force=False): hsv[:, :, 2] = v hsv = np.array(hsv, dtype=np.uint8) - image = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + image_output = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) + + # 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 diff --git a/augraphy/augmentations/brightnesstexturize.py b/augraphy/augmentations/brightnesstexturize.py index 968b91aa..9485a5ba 100644 --- a/augraphy/augmentations/brightnesstexturize.py +++ b/augraphy/augmentations/brightnesstexturize.py @@ -19,7 +19,7 @@ class BrightnessTexturize(Augmentation): :type p: float, optional """ - def __init__(self, texturize_range=(0.9, 0.99), deviation=0.03, p=1): + def __init__(self, texturize_range=(0.8, 0.99), deviation=0.08, p=1): """Constructor method""" super().__init__(p=p) self.low = texturize_range[0] @@ -34,18 +34,19 @@ def __repr__(self): # Applies the Augmentation to input data. def __call__(self, image, layer=None, force=False): if force or self.should_run(): - image_output = image.copy() + image = image.copy() - # for colour image + has_alpha = 0 if len(image.shape) > 2: - hsv = cv2.cvtColor(image_output.astype("uint8"), cv2.COLOR_BGR2HSV) - # for gray image + is_gray = 0 + if image.shape[2] == 4: + has_alpha = 1 + image, image_alpha = image[:, :, :3], image[:, :, 3] else: - bgr = cv2.cvtColor( - image_output.astype("uint8"), - cv2.COLOR_GRAY2BGR, - ) - hsv = cv2.cvtColor(bgr, cv2.COLOR_BGR2HSV) + is_gray = 1 + image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) + + hsv = cv2.cvtColor(image, cv2.COLOR_BGR2HSV) # compute random value value = random.uniform(self.low, self.high) @@ -88,8 +89,10 @@ def __call__(self, image, layer=None, force=False): image_output = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) - # convert back to gray - if len(image.shape) < 3: - image_output = cv2.cvtColor(hsv, cv2.COLOR_BGR2GRAY) + # 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/colorpaper.py b/augraphy/augmentations/colorpaper.py index b6a72698..03145d45 100644 --- a/augraphy/augmentations/colorpaper.py +++ b/augraphy/augmentations/colorpaper.py @@ -37,8 +37,16 @@ def add_color(self, image): :type image: numpy.array (numpy.uint8) """ - if len(image.shape) < 3: + 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) + ysize, xsize = image.shape[:2] # convert to hsv colorspace @@ -52,9 +60,15 @@ def add_color(self, image): image_hsv[:, :, 1] = image_s # convert back to bgr - color_image = cv2.cvtColor(image_hsv, cv2.COLOR_HSV2BGR) + image_color = cv2.cvtColor(image_hsv, cv2.COLOR_HSV2BGR) + + # return image follows the input image color channel + if is_gray: + image_color = cv2.cvtColor(image_color, cv2.COLOR_BGR2GRAY) + if has_alpha: + image_color = np.dstack((image_color, image_alpha)) - return color_image + return image_color # Applies the Augmentation to input data. def __call__(self, image, layer=None, force=False): diff --git a/augraphy/augmentations/colorshift.py b/augraphy/augmentations/colorshift.py index de458503..69600564 100644 --- a/augraphy/augmentations/colorshift.py +++ b/augraphy/augmentations/colorshift.py @@ -142,14 +142,18 @@ def apply_color_shift(self, image, kernel_value): def __call__(self, image, layer=None, force=False): if force or self.should_run(): - image_output = image.copy() + image = image.copy() # convert and make sure image is color image - if len(image_output.shape) > 2: + 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_output = cv2.cvtColor(image_output, cv2.COLOR_GRAY2BGR) + image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR) # generate random color shift iterations color_shift_iterations = random.randint(self.color_shift_iterations[0], self.color_shift_iterations[1]) @@ -163,13 +167,18 @@ def __call__(self, image, layer=None, force=False): if not (kernel_value % 2): kernel_value += 1 + # first assignment for the iterations + image_output = image for i in range(color_shift_iterations): image_output = self.apply_color_shift(image_output, kernel_value) # 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) + if has_alpha: + image_output = np.dstack((image_output, image_alpha)) return image_output diff --git a/augraphy/augmentations/delaunay.py b/augraphy/augmentations/delaunay.py index bc911eec..729ff184 100644 --- a/augraphy/augmentations/delaunay.py +++ b/augraphy/augmentations/delaunay.py @@ -253,6 +253,18 @@ def apply_augmentation(self): # Applies the Augmentation to input data. def __call__(self, image, layer=None, force=False): if force or self.should_run(): + + # 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) + self.width = self.height = random.choice( [400, 480, 500, 600, 640, 720], ) # randomly selecting the width and the height of the background pattern @@ -292,11 +304,13 @@ def __call__(self, image, layer=None, force=False): threshold = self.ws // 20 delaunay_mesh = delaunay_mesh[threshold : h - threshold, threshold : w - threshold] delaunay_mesh = cv2.resize(delaunay_mesh, (self.ws, self.ws), interpolation=cv2.INTER_LINEAR) - if len(image.shape) < 3: - delaunay_mesh = cv2.cvtColor(delaunay_mesh, cv2.COLOR_RGB2GRAY) - elif len(image.shape) == 3 and image.shape[2] == 1: - delaunay_mesh = cv2.cvtColor(delaunay_mesh, cv2.COLOR_RGB2GRAY) sw = PatternMaker(alpha=0.49) result = sw.make_patterns(image=result, mesh_img=delaunay_mesh, window_size=self.ws) result = result[self.ws : h + self.ws, self.ws : w + self.ws] + + if is_gray: + result = cv2.cvtColor(result, cv2.COLOR_BGR2GRAY) + if has_alpha: + result = np.dstack((result, image_alpha)) + return result diff --git a/augraphy/augmentations/dirtydrum.py b/augraphy/augmentations/dirtydrum.py index 4989987e..e3f8ebf4 100644 --- a/augraphy/augmentations/dirtydrum.py +++ b/augraphy/augmentations/dirtydrum.py @@ -242,6 +242,17 @@ def __call__(self, image, layer=None, force=False): if force or self.should_run(): 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) + if self.direction == -1: # Select random direction direction = random.choice([0, 1, 2]) @@ -286,4 +297,10 @@ def __call__(self, image, layer=None, force=False): image_dirty_drum = self.blend(image, image_dirty) + # return image follows the input image color channel + if is_gray: + image_dirty_drum = cv2.cvtColor(image_dirty_drum, cv2.COLOR_BGR2GRAY) + if has_alpha: + image_dirty_drum = np.dstack((image_dirty_drum, image_alpha)) + return image_dirty_drum diff --git a/augraphy/augmentations/dirtyrollers.py b/augraphy/augmentations/dirtyrollers.py index 6a56aa59..be79d7eb 100644 --- a/augraphy/augmentations/dirtyrollers.py +++ b/augraphy/augmentations/dirtyrollers.py @@ -190,9 +190,13 @@ def __call__(self, image, layer=None, force=False): if force or self.should_run(): image = image.copy() - # convert and make sure image is color image + # 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) @@ -219,8 +223,9 @@ def __call__(self, image, layer=None, force=False): if rotate: image_output = cv2.rotate(image_output, cv2.ROTATE_90_COUNTERCLOCKWISE) - # 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/dithering.py b/augraphy/augmentations/dithering.py index cc9b283c..a5d2f232 100644 --- a/augraphy/augmentations/dithering.py +++ b/augraphy/augmentations/dithering.py @@ -1,5 +1,6 @@ import random +import cv2 import numba as nb import numpy as np from numba import config @@ -200,6 +201,17 @@ def __call__(self, image, layer=None, force=False): if force or self.should_run(): 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) + if self.dither == "random": dither_type = random.choice(["ordered", "Floyd Steinberg"]) else: @@ -210,4 +222,9 @@ def __call__(self, image, layer=None, force=False): else: image_dither = self.dither_Floyd_Steinberg(image) + 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 2895c3c9..2335244c 100644 --- a/augraphy/augmentations/dotmatrix.py +++ b/augraphy/augmentations/dotmatrix.py @@ -278,9 +278,13 @@ def __call__(self, image, layer=None, force=False): ysize, xsize = image.shape[:2] - # convert and make sure image is color image + # 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) @@ -567,5 +571,7 @@ 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/faxify.py b/augraphy/augmentations/faxify.py index b4706fc0..0f632735 100644 --- a/augraphy/augmentations/faxify.py +++ b/augraphy/augmentations/faxify.py @@ -227,6 +227,17 @@ def __call__(self, image, layer=None, force=False): if force or self.should_run(): 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) + if self.monochrome == -1: monochrome = random.choice([0, 1]) else: @@ -348,4 +359,12 @@ def __call__(self, image, layer=None, force=False): # upscale image image_faxify = cv2.resize(image_out, (image.shape[1], image.shape[0])) + if is_gray and len(image_faxify.shape) > 2: + image_faxify = cv2.cvtColor(image_faxify, cv2.COLOR_BGR2GRAY) + if has_alpha: + # convert to BGRA if input has alpha layer + if len(image_faxify.shape) < 3: + image_faxify = cv2.cvtColor(image_faxify, cv2.COLOR_GRAY2BGR) + image_faxify = np.dstack((image_faxify, image_alpha)) + return image_faxify diff --git a/augraphy/base/augmentationpipeline.py b/augraphy/base/augmentationpipeline.py index b8cc63b0..54acf964 100644 --- a/augraphy/base/augmentationpipeline.py +++ b/augraphy/base/augmentationpipeline.py @@ -162,9 +162,9 @@ def augment_single_image(self, image): """ # Check if image has correct channel - if len(image.shape) > 2 and (image.shape[2] != 3): + if len(image.shape) > 2 and (image.shape[2] != 3 and image.shape[2] != 4): raise Exception( - "Image should have channel number of 3 (BGR), but actual dimensions were {}.".format( + "Image should have channel number of 3 (BGR) or 4 (BGRA), but actual dimensions were {}.".format( image.shape, ), ) @@ -222,6 +222,11 @@ def augment_single_image(self, image): image, ) + # add alpha layer for color image + if len(image.shape) > 2: + if image.shape[2] == 3: + image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA) + data = dict() # Store performance metadata and other logs here. diff --git a/doc/source/installation.rst b/doc/source/installation.rst index 4611dd43..2529613a 100644 --- a/doc/source/installation.rst +++ b/doc/source/installation.rst @@ -2,7 +2,7 @@ Installation ============ -Augraphy is tested and supported only in Python 3.8, 3.9, 3.10, 3.11. +Augraphy is tested and supported only in Python 3.8, 3.9, 3.10, 3.11. ----------------------