Add output mask output format

Author: Belval

tests.py

@@ -150,6 +150,7 @@ def test_generate_data_with_format(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -184,6 +185,7 @@ def test_generate_data_with_extension(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -218,6 +220,7 @@ def test_generate_data_with_skew_angle(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -252,6 +255,7 @@ def test_generate_data_with_blur(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -286,6 +290,7 @@ def test_generate_data_with_sine_distorsion(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -320,6 +325,7 @@ def test_generate_data_with_cosine_distorsion(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -354,6 +360,7 @@ def test_generate_data_with_left_alignment(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -388,6 +395,7 @@ def test_generate_data_with_center_alignment(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -422,6 +430,7 @@ def test_generate_data_with_right_alignment(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -457,6 +466,7 @@ def test_raise_if_handwritten_and_vertical(self):
                 0,
                 (5, 5, 5, 5),
                 0,
+                0,
             )
             raise Exception("Vertical handwritten did not throw")
         except ValueError:
@@ -487,6 +497,7 @@ def test_generate_vertical_text(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -521,6 +532,7 @@ def test_generate_horizontal_text_with_variable_space(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -555,6 +567,7 @@ def test_generate_vertical_text_with_variable_space(self):
             0,
             (5, 5, 5, 5),
             0,
+            0,
         )
 
         self.assertTrue(
@@ -590,6 +603,7 @@ def test_generate_text_with_unknown_orientation(self):
                 0,
                 (5, 5, 5, 5),
                 0,
+                0,
             )
             raise Exception("Unknown orientation did not throw")
         except ValueError:
@@ -620,6 +634,7 @@ def test_generate_data_with_fit(self):
             0,
             (0, 0, 0, 0),
             1,
+            0,
         )
 
         self.assertTrue(
@@ -663,18 +678,6 @@ def test_generate_data_with_white_background(self):
 
         os.remove("tests/out/white_background.jpg")
 
-    def test_generate_data_with_gaussian_background(self):
-        background_generator.gaussian_noise(64, 128).convert("RGB").save(
-            "tests/out/gaussian_background.jpg"
-        )
-
-        self.assertTrue(
-            md5("tests/out/gaussian_background.jpg")
-            == md5("tests/expected_results/gaussian_background.jpg")
-        )
-
-        os.remove("tests/out/gaussian_background.jpg")
-
     def test_generate_data_with_quasicrystal_background(self):
         bkgd = background_generator.quasicrystal(64, 128)
 

tests/expected_results/TEST TEST TEST_10.jpg

@@ -30,8 +30,11 @@ def _generate_horizontal_text(
     text_height = max([image_font.getsize(c)[1] for c in text])
 
     txt_img = Image.new("RGBA", (text_width, text_height), (0, 0, 0, 0))
+    txt_mask = Image.new("RGB", (text_width, text_height), (0, 0, 0))
 
-    txt_draw = ImageDraw.Draw(txt_img)
+    txt_img_draw = ImageDraw.Draw(txt_img)
+    txt_mask_draw = ImageDraw.Draw(txt_mask, mode="RGB")
+    txt_mask_draw.fontmode = "1"
 
     colors = [ImageColor.getrgb(c) for c in text_color.split(",")]
     c1, c2 = colors[0], colors[-1]
@@ -43,17 +46,23 @@ def _generate_horizontal_text(
     )
 
     for i, c in enumerate(text):
-        txt_draw.text(
+        txt_img_draw.text(
             (sum(char_widths[0:i]) + i * character_spacing, 0),
             c,
             fill=fill,
             font=image_font,
         )
+        txt_mask_draw.text(
+            (sum(char_widths[0:i]) + i * character_spacing, 0),
+            c,
+            fill=((i + 1) // (255 * 255), (i + 1) // 255, (i + 1) % 255),
+            font=image_font,
+        )
 
     if fit:
-        return txt_img.crop(txt_img.getbbox())
+        return txt_img.crop(txt_img.getbbox()), txt_mask.crop(txt_img.getbbox())
     else:
-        return txt_img
+        return txt_img, txt_mask
 
 
 def _generate_vertical_text(
@@ -70,8 +79,10 @@ def _generate_vertical_text(
     text_height = sum(char_heights) + character_spacing * len(text)
 
     txt_img = Image.new("RGBA", (text_width, text_height), (0, 0, 0, 0))
+    txt_mask = Image.new("RGBA", (text_width, text_height), (0, 0, 0, 0))
 
-    txt_draw = ImageDraw.Draw(txt_img)
+    txt_img_draw = ImageDraw.Draw(txt_img)
+    txt_mask_draw = ImageDraw.Draw(txt_img)
 
     colors = [ImageColor.getrgb(c) for c in text_color.split(",")]
     c1, c2 = colors[0], colors[-1]
@@ -83,14 +94,20 @@ def _generate_vertical_text(
     )
 
     for i, c in enumerate(text):
-        txt_draw.text(
+        txt_img_draw.text(
             (0, sum(char_heights[0:i]) + i * character_spacing),
             c,
             fill=fill,
             font=image_font,
         )
+        txt_mask_draw.text(
+            (0, sum(char_heights[0:i]) + i * character_spacing),
+            c,
+            fill=(i // (255 * 255), i // 255, i % 255),
+            font=image_font,
+        )
 
     if fit:
-        return txt_img.crop(txt_img.getbbox())
+        return txt_img.crop(txt_img.getbbox()), txt_mask.crop(txt_img.getbbox())
     else:
-        return txt_img
+        return txt_img, txt_mask

tests/expected_results/TEST TEST TEST_12.jpg

@@ -46,6 +46,7 @@ def generate(
         character_spacing,
         margins,
         fit,
+        output_mask,
     ):
         image = None
 
@@ -61,7 +62,7 @@ def generate(
                 raise ValueError("Vertical handwritten text is unavailable")
             image = handwritten_text_generator.generate(text, text_color)
         else:
-            image = computer_text_generator.generate(
+            image, mask = computer_text_generator.generate(
                 text,
                 font,
                 text_color,
@@ -71,33 +72,40 @@ def generate(
                 character_spacing,
                 fit,
             )
-
         random_angle = rnd.randint(0 - skewing_angle, skewing_angle)
 
         rotated_img = image.rotate(
             skewing_angle if not random_skew else random_angle, expand=1
         )
 
+        rotated_mask = mask.rotate(
+            skewing_angle if not random_skew else random_angle, expand=1
+        )
+
         #############################
         # Apply distorsion to image #
         #############################
         if distorsion_type == 0:
             distorted_img = rotated_img  # Mind = blown
+            distorted_mask = rotated_mask
         elif distorsion_type == 1:
-            distorted_img = distorsion_generator.sin(
+            distorted_img, distorted_mask = distorsion_generator.sin(
                 rotated_img,
+                rotated_mask,
                 vertical=(distorsion_orientation == 0 or distorsion_orientation == 2),
                 horizontal=(distorsion_orientation == 1 or distorsion_orientation == 2),
             )
         elif distorsion_type == 2:
-            distorted_img = distorsion_generator.cos(
+            distorted_img, distorted_mask = distorsion_generator.cos(
                 rotated_img,
+                rotated_mask,
                 vertical=(distorsion_orientation == 0 or distorsion_orientation == 2),
                 horizontal=(distorsion_orientation == 1 or distorsion_orientation == 2),
             )
         else:
-            distorted_img = distorsion_generator.random(
+            distorted_img, distorted_mask = distorsion_generator.random(
                 rotated_img,
+                rotated_mask,
                 vertical=(distorsion_orientation == 0 or distorsion_orientation == 2),
                 horizontal=(distorsion_orientation == 1 or distorsion_orientation == 2),
             )
@@ -115,6 +123,7 @@ def generate(
             resized_img = distorted_img.resize(
                 (new_width, size - vertical_margin), Image.ANTIALIAS
             )
+            resized_mask = distorted_mask.resize((new_width, size - vertical_margin))
             background_width = width if width > 0 else new_width + horizontal_margin
             background_height = size
         # Vertical text
@@ -126,6 +135,9 @@ def generate(
             resized_img = distorted_img.resize(
                 (size - horizontal_margin, new_height), Image.ANTIALIAS
             )
+            resized_mask = distorted_mask.resize(
+                (size - horizontal_margin, new_height), Image.ANTIALIAS
+            )
             background_width = size
             background_height = new_height + vertical_margin
         else:
@@ -135,21 +147,22 @@ def generate(
         # Generate background image #
         #############################
         if background_type == 0:
-            background = background_generator.gaussian_noise(
+            background_img = background_generator.gaussian_noise(
                 background_height, background_width
             )
         elif background_type == 1:
-            background = background_generator.plain_white(
+            background_img = background_generator.plain_white(
                 background_height, background_width
             )
         elif background_type == 2:
-            background = background_generator.quasicrystal(
+            background_img = background_generator.quasicrystal(
                 background_height, background_width
             )
         else:
-            background = background_generator.picture(
+            background_img = background_generator.picture(
                 background_height, background_width
             )
+        background_mask = Image.new("RGB", (background_width, background_height), (0, 0, 0))
 
         #############################
         # Place text with alignment #
@@ -158,45 +171,62 @@ def generate(
         new_text_width, _ = resized_img.size
 
         if alignment == 0 or width == -1:
-            background.paste(resized_img, (margin_left, margin_top), resized_img)
+            background_img.paste(resized_img, (margin_left, margin_top), resized_img)
+            background_mask.paste(resized_mask, (margin_left, margin_top))
         elif alignment == 1:
-            background.paste(
+            background_img.paste(
                 resized_img,
                 (int(background_width / 2 - new_text_width / 2), margin_top),
                 resized_img,
             )
+            background_mask.paste(
+                resized_mask,
+                (int(background_width / 2 - new_text_width / 2), margin_top),
+            )
         else:
-            background.paste(
+            background_img.paste(
                 resized_img,
                 (background_width - new_text_width - margin_right, margin_top),
                 resized_img,
             )
+            background_mask.paste(
+                resized_mask,
+                (background_width - new_text_width - margin_right, margin_top),
+            )
 
         ##################################
         # Apply gaussian blur #
         ##################################
 
-        final_image = background.filter(
-            ImageFilter.GaussianBlur(
-                radius=(blur if not random_blur else rnd.randint(0, blur))
-            )
+        gaussian_filter = ImageFilter.GaussianBlur(
+            radius=blur if not random_blur else rnd.randint(0, blur)
         )
+        final_image = background_img.filter(gaussian_filter)
+        final_mask = background_mask.filter(gaussian_filter)
 
         #####################################
         # Generate name for resulting image #
         #####################################
         if name_format == 0:
             image_name = "{}_{}.{}".format(text, str(index), extension)
+            mask_name = "{}_{}_mask.png".format(text, str(index))
         elif name_format == 1:
             image_name = "{}_{}.{}".format(str(index), text, extension)
+            mask_name = "{}_{}_mask.png".format(str(index), text)
         elif name_format == 2:
             image_name = "{}.{}".format(str(index), extension)
+            mask_name = "{}_mask.png".format(str(index))
         else:
             print("{} is not a valid name format. Using default.".format(name_format))
             image_name = "{}_{}.{}".format(text, str(index), extension)
+            mask_name = "{}_{}_mask.png".format(text, str(index))
 
         # Save the image
         if out_dir is not None:
             final_image.convert("RGB").save(os.path.join(out_dir, image_name))
+            if output_mask == 1:
+                final_mask.convert("RGB").save(os.path.join(out_dir, mask_name))
         else:
+            if output_mask == 1:
+                return final_image.convert("RGB"), final_mask.convert("RGB")
             return final_image.convert("RGB")