Merge pull request #351 from kwcckw/dev

Added alpha layer support in augmentations.

augraphy/augmentations/lowlightnoise.py

@@ -163,10 +163,9 @@ def __call__(self, image, layer=None, force=False):
             result = image.copy()
 
             has_alpha = 0
-            if len(result.shape) > 2:
-                if result.shape[2] == 4:
-                    has_alpha = 1
-                    result, image_alpha = result[:, :, :3], result[:, :, 3]
+            if len(result.shape) > 2 and result.shape[2] == 4:
+                has_alpha = 1
+                result, image_alpha = result[:, :, :3], result[:, :, 3]
 
             photons = random.randint(self.num_photons_range[0], self.num_photons_range[1])
             alpha = random.uniform(self.alpha_range[0], self.alpha_range[1])

augraphy/augmentations/markup.py

@@ -186,8 +186,13 @@ def draw_line(self, p1, p2, markup_mask, markup_thickness, markup_color, reverse
     def __call__(self, image, layer=None, force=False):
 
         # change to 3 channels BGR format
+        has_alpha = 0
         if len(image.shape) < 3:
             image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)
+        elif image.shape[2] == 4:
+            has_alpha = 1
+            image, image_alpha = image[:, :, :3], image[:, :, 3]
+
         markup_image = image.copy()
 
         if self.markup_color == "random":
@@ -400,4 +405,7 @@ def __call__(self, image, layer=None, force=False):
 
             markup_image = ink_generator.generate_ink()
 
+        if has_alpha:
+            markup_image = np.dstack((markup_image, image_alpha))
+
         return markup_image

augraphy/augmentations/noisetexturize.py

@@ -88,6 +88,11 @@ def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
             image = image.copy()
 
+            has_alpha = 0
+            if len(image.shape) > 2 and image.shape[2] == 4:
+                has_alpha = 1
+                image, image_alpha = image[:, :, :3], image[:, :, 3]
+
             sigma = random.randint(self.sigma_range[0], self.sigma_range[1])
             turbulence = random.randint(
                 self.turbulence_range[0],
@@ -106,6 +111,9 @@ def __call__(self, image, layer=None, force=False):
                 result += self.noise(cols, rows, channel, ratio, sigma=sigma)
                 ratio = (ratio // turbulence) or 1
             cut = np.clip(result, 0, 255)
-
             cut = cut.astype(np.uint8)
+
+            if has_alpha:
+                cut = np.dstack((cut, image_alpha))
+
             return cut

augraphy/augmentations/noisylines.py

@@ -158,8 +158,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)
@@ -187,5 +191,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/pageborder.py

@@ -588,8 +588,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)
@@ -629,5 +633,15 @@ 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:
+                oysize, oxsize = image_output.shape[:2]
+                # check and make sure image alpha has a same size with image output
+                if oysize != height or oxsize != width:
+                    image_alpha = cv2.resize(
+                        image_alpha,
+                        (oxsize, oysize),
+                        interpolation=cv2.INTER_AREA,
+                    )
+                image_output = np.dstack((image_output, image_alpha))
 
             return image_output

augraphy/augmentations/quasicrystal.py

@@ -115,8 +115,12 @@ def apply_augmentation(ndim, pattern_image, frequency, phase, n_rotation):
     def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
             result = image.copy()
-            h, w = result.shape[:2]
+            has_alpha = 0
+            if len(result.shape) > 2 and result.shape[2] == 4:
+                has_alpha = 1
+                result, image_alpha = result[:, :, :3], result[:, :, 3]
 
+            h, w = result.shape[:2]
             self.n_rotation = random.randint(self.n_rotation_range[0], self.n_rotation_range[1])
             pattern_image = np.zeros((self.imgy, self.imgx, 3), dtype=np.uint8)
             frequency = random.random() * 100 + 18  # determines the frequency of pattern
@@ -144,4 +148,8 @@ def __call__(self, image, layer=None, force=False):
 
             # overlay pattern into image
             result = sw.superimpose(result, invert)
+
+            if has_alpha:
+                result = np.dstack((result, image_alpha))
+
             return result

augraphy/augmentations/reflectedlight.py

@@ -68,8 +68,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)
@@ -275,5 +279,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/scribbles.py

@@ -106,6 +106,11 @@ def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
             image = image.copy()
 
+            has_alpha = 0
+            if len(image.shape) > 2 and image.shape[2] == 4:
+                has_alpha = 1
+                image, image_alpha = image[:, :, :3], image[:, :, 3]
+
             if self.scribbles_type == "text" or self.scribbles_type == "random":
                 # create fonts directory
                 os.makedirs(self.fonts_directory, exist_ok=True)
@@ -245,4 +250,7 @@ def __call__(self, image, layer=None, force=False):
 
             image_output = ink_generator.generate_ink()
 
+            if has_alpha:
+                image_output = np.dstack((image_output, image_alpha))
+
             return image_output

augraphy/augmentations/sectionshift.py

@@ -86,9 +86,13 @@ def apply_shift(self, image, shift_box, section_shift_x, section_shift_y):
         # fill the shifted area with value
         if self.section_shift_fill_value != -1:
             if self.section_shift_fill_value == "random":
-                image[y0:yn, x0:xn] = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
+                fill_value = (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
             else:
-                image[y0:yn, x0:xn] = self.section_shift_fill_value
+                fill_value = self.section_shift_fill_value
+            # add alpha value
+            if image.shape[2] == 4:
+                fill_value = (fill_value[0], fill_value[1], fill_value[2], 255)
+            image[y0:yn, x0:xn] = fill_value
 
         # shift the section of image
         image[y0 + section_shift_y : yn + section_shift_y, x0 + section_shift_x : xn + section_shift_x] = image_section

augraphy/augmentations/shadowcast.py

@@ -68,8 +68,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)
@@ -196,5 +200,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/subtlenoise.py

@@ -53,7 +53,8 @@ def __call__(self, image, layer=None, force=False):
             if len(image.shape) > 2:
                 # convert to int to enable negative
                 image = image.astype("int")
-                for i in range(image.shape[2]):
+                # skip alpha layer
+                for i in range(3):
                     image[:, :, i] = self.add_subtle_noise(image[:, :, i])
             # single channel image
             else:

augraphy/augmentations/voronoi.py

@@ -170,6 +170,12 @@ def apply_augmentation(self):
     def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
             result = image.copy()
+
+            has_alpha = 0
+            if len(result.shape) > 2 and result.shape[2] == 4:
+                has_alpha = 1
+                result, image_alpha = result[:, :, :3], result[:, :, 3]
+
             h, w = result.shape[:2]
             if self.noise_type == "random":
                 self.perlin = random.choice([True, False])
@@ -220,4 +226,8 @@ def __call__(self, image, layer=None, force=False):
             # original image is padded and voronoi_mesh passes through it like a sliding window
             result = sw.make_patterns(result, voronoi_mesh, self.ws)
             result = result[self.ws : h + self.ws, self.ws : w + self.ws]
+
+            if has_alpha:
+                result = np.dstack((result, image_alpha))
+
             return result

augraphy/augmentations/watermark.py

@@ -238,7 +238,15 @@ def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
             image = image.copy()
 
+            has_alpha = 0
+            if len(image.shape) > 2 and image.shape[2] == 4:
+                has_alpha = 1
+                image, image_alpha = image[:, :, :3], image[:, :, 3]
+
             watermark_foreground = self.create_watermark()
             watermark_image = self.apply_watermark(watermark_foreground, image)
 
+            if has_alpha:
+                watermark_image = np.dstack((watermark_image, image_alpha))
+
             return watermark_image