Added alpha layer support in augmentations.

augraphy/augmentations/inkmottling.py

@@ -69,6 +69,7 @@ def __call__(self, image, layer=None, force=False):
                 255,
                 size=(int(ysize / ink_mottling_noise_scale), int(xsize / ink_mottling_noise_scale)),
             ).astype("uint8")
+
             image_random = cv2.cvtColor(image_random, cv2.COLOR_GRAY2BGR)
 
             # apply gaussian blur to the mask of noise
@@ -89,6 +90,10 @@ def __call__(self, image, layer=None, force=False):
                     interpolation=cv2.INTER_AREA,
                 )
 
+            # add alpha layer to random image
+            if image.shape[2] == 4:
+                image_random = np.dstack((image_random, image[:, :, 3]))
+
             # blend noise mask with image ink based on the input alpha
             ink_mottling_alpha = random.uniform(self.ink_mottling_alpha_range[0], self.ink_mottling_alpha_range[1])
             image_blend = cv2.addWeighted(image, (1 - ink_mottling_alpha), image_random, ink_mottling_alpha, 0)

augraphy/augmentations/inkshifter.py

@@ -172,7 +172,9 @@ def put_fading(self, img, fade, f=0.5):
 
     def __call__(self, image, layer=None, force=None):
         if force or self.should_run():
-            h, w, _ = image.shape
+
+            h, w = image.shape[:2]
+
             text_shift_scale = random.randint(self.text_shift_scale_range[0], self.text_shift_scale_range[1])
             text_shift_factor = random.randint(self.text_shift_factor_range[0], self.text_shift_factor_range[1])
 

augraphy/augmentations/jpeg.py

@@ -1,6 +1,7 @@
 import random
 
 import cv2
+import numpy as np
 
 from augraphy.base.augmentation import Augmentation
 
@@ -32,10 +33,20 @@ def __repr__(self):
     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]
+
             encode_param = [
                 int(cv2.IMWRITE_JPEG_QUALITY),
                 random.randint(self.quality_range[0], self.quality_range[1]),
             ]
             result, encimg = cv2.imencode(".jpg", image, encode_param)
             image = cv2.imdecode(encimg, 1)
+
+            if has_alpha:
+                image = np.dstack((image, image_alpha))
+
             return image

augraphy/augmentations/letterpress.py

@@ -94,7 +94,6 @@ def __call__(self, image, layer=None, force=False):
 
             # initialize empty noise mask and noise mask with random values
             noise_mask = np.copy(image)
-            noise_mask2 = (np.random.random((image.shape[0], image.shape[1])) * 255).astype("uint8")
             noise_mask2 = np.random.randint(
                 self.value_range[0],
                 self.value_range[1],
@@ -104,8 +103,8 @@ def __call__(self, image, layer=None, force=False):
 
             # insert noise value according to generate points
             if len(image.shape) > 2:
-
-                for i in range(image.shape[2]):
+                # skip alpha layer
+                for i in range(3):
                     noise_mask[generated_points_y, generated_points_x, i] = noise_mask2[
                         generated_points_y,
                         generated_points_x,

augraphy/augmentations/lightinggradient.py

@@ -201,15 +201,18 @@ def apply_decay_value_linear(mask, canvas_y, max_value, padding_center, decay_ra
     # Applies the Augmentation to input data.
     def __call__(self, image, layer=None, force=False):
         if force or self.should_run():
-            image = image.copy()
+            frame = image.copy()
+
+            has_alpha = 0
+            if len(frame.shape) > 2 and frame.shape[2] == 4:
+                has_alpha = 1
+                frame, image_alpha = frame[:, :, :3], frame[:, :, 3]
 
             if self.transparency is None:
                 transparency = random.uniform(0.5, 0.85)
             else:
                 transparency = self.transparency
 
-            frame = image
-
             height, width = frame.shape[:2]
             if len(frame.shape) > 2:
                 hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
@@ -230,4 +233,8 @@ def __call__(self, image, layer=None, force=False):
             frame = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
             frame[frame > 255] = 255
             frame = np.asarray(frame, dtype=np.uint8)
+
+            if has_alpha:
+                frame = np.dstack((frame, image_alpha))
+
             return frame

augraphy/augmentations/linesdegradation.py

@@ -196,7 +196,8 @@ def __call__(self, image, layer=None, force=False):
 
             # replace detected lines with line value
             if len(image_output.shape) > 2:
-                for i in range(image_output.shape[2]):
+                # skip alpha layer
+                for i in range(3):
                     image_output[ystart:yend, xstart:xend, i][mask > 0] = replacement_mask[mask > 0]
             else:
                 image_output[ystart:yend, xstart:xend][mask > 0] = replacement_mask[mask > 0]

augraphy/augmentations/lowinkline.py

@@ -67,13 +67,14 @@ def add_transparency_line(self, mask, y, alpha=None):
                 low_ink_line_top[indices] = alpha
 
                 if len(mask.shape) > 2:
-                    low_ink_line_top = np.dstack(
-                        [
-                            low_ink_line_top[:, 0],
-                            low_ink_line_top[:, 0],
-                            low_ink_line_top[:, 0],
-                        ],
-                    )[0]
+                    stacked_line = [
+                        low_ink_line_top[:, 0],
+                        low_ink_line_top[:, 0],
+                        low_ink_line_top[:, 0],
+                    ]
+                    if mask.shape[2] == 4:
+                        stacked_line += [low_ink_line_top[:, 3]]
+                    low_ink_line_top = np.dstack(stacked_line)[0]
 
             if y + 1 < mask.shape[0]:
                 if len(mask.shape) > 2:
@@ -84,13 +85,14 @@ def add_transparency_line(self, mask, y, alpha=None):
                 low_ink_line_bottom[indices] = alpha
 
                 if len(mask.shape) > 2:
-                    low_ink_line_bottom = np.dstack(
-                        [
-                            low_ink_line_bottom[:, 0],
-                            low_ink_line_bottom[:, 0],
-                            low_ink_line_bottom[:, 0],
-                        ],
-                    )[0]
+                    stacked_line = [
+                        low_ink_line_bottom[:, 0],
+                        low_ink_line_bottom[:, 0],
+                        low_ink_line_bottom[:, 0],
+                    ]
+                    if mask.shape[2] == 4:
+                        stacked_line += [low_ink_line_bottom[:, 3]]
+                    low_ink_line_bottom = np.dstack(stacked_line)[0]
 
         else:
             low_ink_line = (np.random.random((xsize)) * 255).astype("uint8")
@@ -127,21 +129,21 @@ def add_transparency_line(self, mask, y, alpha=None):
 
         indices = mask[y, :] < low_ink_line
         if len(mask.shape) > 2:
-            mask[y, :][indices.reshape(xsize, mask.shape[2])] = low_ink_line[indices]
+            mask[y, :, :3][indices[:, :3]] = low_ink_line[:, :3][indices[:, :3]]
         else:
             mask[y, :][indices] = low_ink_line[indices]
 
         if y - 1 >= 0:
             indices = mask[y - 1, :] < low_ink_line_top
             if len(mask.shape) > 2:
-                mask[y - 1, :][indices.reshape(xsize, mask.shape[2])] = low_ink_line_top[indices]
+                mask[y - 1, :, :3][indices[:, :3]] = low_ink_line_top[:, :3][indices[:, :3]]
             else:
                 mask[y - 1, :][indices] = low_ink_line_top[indices]
 
         if y + 1 < mask.shape[0]:
             indices = mask[y - 1, :] < low_ink_line_bottom
             if len(mask.shape) > 2:
-                mask[y + 1, :][indices.reshape(xsize, mask.shape[2])] = low_ink_line_bottom[indices]
+                mask[y + 1, :, :3][indices[:, :3]] = low_ink_line_bottom[:, :3][indices[:, :3]]
             else:
                 mask[y + 1, :][indices] = low_ink_line_bottom[indices]
 

augraphy/augmentations/lowlightnoise.py

@@ -76,9 +76,15 @@ def _add_bias(self, image, value):
         bias_im = np.zeros_like(image) + value
         # add random column noise to the image
         col_pattern = np.random.uniform(0, int(value * 0.1) // 2, size=shape[0])
-        for channel in range(shape[2]):
+        if len(shape) > 2:
+            # skip alpha channel
+            for channel in range(3):
+                for column in columns:
+                    bias_im[:, column, channel] = value + col_pattern
+        else:
             for column in columns:
-                bias_im[:, column, channel] = value + col_pattern
+                bias_im[:, column] = value + col_pattern
+
         return bias_im
 
     def _apply_filter(self, image):
@@ -106,12 +112,15 @@ def _sensor_noise(self, image, current, exposure_time, value, gain=0.1):
         """
         base_current = current * exposure_time / gain  # noise due to thermal heat of the sensor
         dark_im = np.random.poisson(base_current, size=image.shape)
-        y_max, x_max, channels = dark_im.shape
+        y_max, x_max = dark_im.shape[:2]
         n_hot = int(0.00001 * x_max * y_max)
         hot_x = np.random.randint(0, x_max, size=n_hot)
         hot_y = np.random.randint(0, y_max, size=n_hot)
-        for channel in range(channels):
-            dark_im[hot_y, hot_x, channel] = current * exposure_time / gain
+        if len(dark_im.shape) > 2:
+            for channel in range(3):
+                dark_im[hot_y, hot_x, channel] = current * exposure_time / gain
+        else:
+            dark_im[hot_y, hot_x] = current * exposure_time / gain
         bias_im = self._add_bias(image.copy(), value)  # noise due to accumulation of photon on the screen
         noise = 0.1 * bias_im + 0.1 * dark_im
         return noise
@@ -152,6 +161,13 @@ def _simulate_low_light_image(self, image, alpha, beta, gamma, bias, photons):
     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:
+                if 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])
             beta = random.uniform(self.beta_range[0], self.beta_range[1])
@@ -165,4 +181,8 @@ def __call__(self, image, layer=None, force=False):
                 bias,
                 photons,
             )
+
+        if has_alpha:
+            result = np.dstack((result, image_alpha))
+
         return result