Add --fit and --margins parameters

Fixes: #49

Author: Belval

README.md

@@ -22,11 +22,11 @@ tqdm
  You can simply use `pip install -r requirements.txt` too.
 
 ## New
+- Add `--fit` and `--margins` for finer layout control
 - Change the text orientation using the `-or` parameter
 - Change the space width using the `-sw` parameter
 - Specify text color range using `-tc '#000000,#FFFFFF'`, please note that the quotes are **necessary**  
 - Explicit alignement when using `-al` with fixed width (0: Left, 1: Center, 2: Right)
-- Fixed width using `-wd`
 - Add support for Simplified and Traditional Chinese
 
 ## How does it work?

TextRecognitionDataGenerator/computer_text_generator.py

@@ -4,16 +4,16 @@
 
 class ComputerTextGenerator(object):
     @classmethod
-    def generate(cls, text, font, text_color, font_size, orientation, space_width):
+    def generate(cls, text, font, text_color, font_size, orientation, space_width, fit):
         if orientation == 0:
-            return cls.__generate_horizontal_text(text, font, text_color, font_size, space_width)
+            return cls.__generate_horizontal_text(text, font, text_color, font_size, space_width, fit)
         elif orientation == 1:
-            return cls.__generate_vertical_text(text, font, text_color, font_size, space_width)
+            return cls.__generate_vertical_text(text, font, text_color, font_size, space_width, fit)
         else:
             raise ValueError("Unknown orientation " + str(orientation))
 
     @classmethod
-    def __generate_horizontal_text(cls, text, font, text_color, font_size, space_width):
+    def __generate_horizontal_text(cls, text, font, text_color, font_size, space_width, fit):
         image_font = ImageFont.truetype(font=font, size=font_size)
         words = text.split(' ')
         space_width = image_font.getsize(' ')[0] * space_width
@@ -38,10 +38,13 @@ def __generate_horizontal_text(cls, text, font, text_color, font_size, space_wid
         for i, w in enumerate(words):
             txt_draw.text((sum(words_width[0:i]) + i * int(space_width), 0), w, fill=fill, font=image_font)
 
-        return txt_img
+        if fit:
+            return txt_img.crop(txt_img.getbbox())
+        else:
+            return txt_img
 
     @classmethod
-    def __generate_vertical_text(cls, text, font, text_color, font_size, space_width):
+    def __generate_vertical_text(cls, text, font, text_color, font_size, space_width, fit):
         image_font = ImageFont.truetype(font=font, size=font_size)
 
         space_height = int(image_font.getsize(' ')[1] * space_width)
@@ -66,4 +69,7 @@ def __generate_vertical_text(cls, text, font, text_color, font_size, space_width
         for i, c in enumerate(text):
             txt_draw.text((0, sum(char_heights[0:i])), c, fill=fill, font=image_font)
 
-        return txt_img
+        if fit:
+            return txt_img.crop(txt_img.getbbox())
+        else:
+            return txt_img

TextRecognitionDataGenerator/data_generator.py

@@ -21,18 +21,22 @@ def generate_from_tuple(cls, t):
         cls.generate(*t)
 
     @classmethod
-    def generate(cls, index, text, font, out_dir, size, extension, skewing_angle, random_skew, blur, random_blur, background_type, distorsion_type, distorsion_orientation, is_handwritten, name_format, width, alignment, text_color, orientation, space_width):
+    def generate(cls, index, text, font, out_dir, size, extension, skewing_angle, random_skew, blur, random_blur, background_type, distorsion_type, distorsion_orientation, is_handwritten, name_format, width, alignment, text_color, orientation, space_width, margins, fit):
         image = None
 
+        margin_top, margin_left, margin_bottom, margin_right = margins
+        horizontal_margin = margin_left + margin_right
+        vertical_margin = margin_top + margin_bottom
+
         ##########################
         # Create picture of text #
         ##########################
         if is_handwritten:
             if orientation == 1:
                 raise ValueError("Vertical handwritten text is unavailable")
-            image = HandwrittenTextGenerator.generate(text, text_color)
+            image = HandwrittenTextGenerator.generate(text, text_color, fit)
         else:
-            image = ComputerTextGenerator.generate(text, font, text_color, size, orientation, space_width)
+            image = ComputerTextGenerator.generate(text, font, text_color, size, orientation, space_width, fit)
 
         random_angle = random.randint(0-skewing_angle, skewing_angle)
 
@@ -68,16 +72,16 @@ def generate(cls, index, text, font, out_dir, size, extension, skewing_angle, ra
 
         # Horizontal text
         if orientation == 0:
-            new_width = int(float(distorted_img.size[0] + 10) * (float(size) / float(distorted_img.size[1] + 10)))
-            resized_img = distorted_img.resize((new_width, size - 10), Image.ANTIALIAS)
-            background_width = width if width > 0 else new_width + 10
+            new_width = int(distorted_img.size[0] * (float(size - vertical_margin) / float(distorted_img.size[1])))
+            resized_img = distorted_img.resize((new_width, size - vertical_margin), Image.ANTIALIAS)
+            background_width = width if width > 0 else new_width + horizontal_margin
             background_height = size
         # Vertical text
         elif orientation == 1:
-            new_height = int(float(distorted_img.size[1] + 10) * (float(size) / float(distorted_img.size[0] + 10)))
-            resized_img = distorted_img.resize((size - 10, new_height), Image.ANTIALIAS)
+            new_height = int(float(distorted_img.size[1]) * (float(size - horizontal_margin) / float(distorted_img.size[0])))
+            resized_img = distorted_img.resize((size - horizontal_margin, new_height), Image.ANTIALIAS)
             background_width = size
-            background_height = new_height + 10
+            background_height = new_height + vertical_margin
         else:
             raise ValueError("Invalid orientation")
 
@@ -99,12 +103,12 @@ def generate(cls, index, text, font, out_dir, size, extension, skewing_angle, ra
 
         new_text_width, _ = resized_img.size
 
-        if alignment == 0:
-            background.paste(resized_img, (5, 5), resized_img)
+        if alignment == 0 or width == -1:
+            background.paste(resized_img, (margin_left, margin_top), resized_img)
         elif alignment == 1:
-            background.paste(resized_img, (int(background_width / 2 - new_text_width / 2), 5), resized_img)
+            background.paste(resized_img, (int(background_width / 2 - new_text_width / 2), margin_top), resized_img)
         else:
-            background.paste(resized_img, (background_width - new_text_width - 5, 5), resized_img)
+            background.paste(resized_img, (background_width - new_text_width - margin_right, margin_top), resized_img)
 
         ##################################
         # Apply gaussian blur #

TextRecognitionDataGenerator/handwritten_text_generator.py

@@ -116,7 +116,7 @@ def __join_images(cls, images):
         return compound_image
 
     @classmethod
-    def generate(cls, text, text_color):
+    def generate(cls, text, text_color, fit):
         with open(os.path.join('handwritten_model', 'translation.pkl'), 'rb') as file:
             translation = pickle.load(file)
 

TextRecognitionDataGenerator/run.py

@@ -13,10 +13,11 @@
 from data_generator import FakeTextDataGenerator
 from multiprocessing import Pool
 
-def valid_range(s):
-    if len(s.split(',')) > 2:
-        raise argparse.ArgumentError("The given range is invalid, please use ?,? format.")
-    return tuple([int(i) for i in s.split(',')])
+def margins(margin):
+    margins = margin.split(',')
+    if len(margins) == 1:
+        return [margins[0]] * 4
+    return [int(m) for m in margins]
 
 def parse_arguments():
     """
@@ -236,6 +237,22 @@ def parse_arguments():
         help="Define the width of the spaces between words. 2.0 means twice the normal space width",
         default=1.0
     )
+    parser.add_argument(
+        "-m",
+        "--margins",
+        type=margins,
+        nargs="?",
+        help="Define the margins around the text when rendered. In pixels",
+        default=(5, 5, 5, 5)
+    )
+    parser.add_argument(
+        "-fi",
+        "--fit",
+        action="store_true",
+        help="Apply a tight crop around the rendered text",
+        default=False
+    )
+
 
     return parser.parse_args()
 
@@ -322,7 +339,9 @@ def main():
             [args.alignment] * string_count,
             [args.text_color] * string_count,
             [args.orientation] * string_count,
-            [args.space_width] * string_count
+            [args.space_width] * string_count,
+            [args.margins] * string_count,
+            [args.fit] * string_count
         )
     ), total=args.count):
         pass

tests.py

@@ -86,7 +86,9 @@ def test_generate_data_with_format(self):
             0,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -116,7 +118,9 @@ def test_generate_data_with_extension(self):
             0,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -146,7 +150,9 @@ def test_generate_data_with_skew_angle(self):
             0,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -176,7 +182,9 @@ def test_generate_data_with_blur(self):
             0,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -206,7 +214,9 @@ def test_generate_data_with_sine_distorsion(self):
             0,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -236,7 +246,9 @@ def test_generate_data_with_cosine_distorsion(self):
             0,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -266,7 +278,9 @@ def test_generate_data_with_left_alignment(self):
             0,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -296,7 +310,9 @@ def test_generate_data_with_center_alignment(self):
             1,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -326,7 +342,9 @@ def test_generate_data_with_right_alignment(self):
             2,
             '#010101',
             0,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -357,7 +375,9 @@ def test_raise_if_handwritten_and_vertical(self):
                 2,
                 '#010101',
                 1,
-                1
+                1,
+                (5,5,5,5),
+                0
             )
             raise Exception("Vertical handwritten did not throw")
         except ValueError:
@@ -384,7 +404,9 @@ def test_generate_vertical_text(self):
             0,
             '#010101',
             1,
-            1
+            1,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -414,7 +436,9 @@ def test_generate_horizontal_text_with_variable_space(self):
             0,
             '#010101',
             0,
-            4
+            4,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -444,7 +468,9 @@ def test_generate_vertical_text_with_variable_space(self):
             0,
             '#010101',
             1,
-            2
+            2,
+            (5,5,5,5),
+            0
         )
 
         self.assertTrue(
@@ -475,12 +501,46 @@ def test_generate_text_with_unknown_orientation(self):
                 0,
                 '#010101',
                 100,
-                2
+                2,
+                (5,5,5,5),
+                0
             )
             raise Exception("Unknown orientation did not throw")
         except ValueError:
             pass
 
+    def test_generate_data_with_fit(self):
+        FakeTextDataGenerator.generate(
+            13,
+            'TEST TEST TEST',
+            'tests/font.ttf',
+            'tests/out/',
+            64,
+            'jpg',
+            0,
+            False,
+            0,
+            False,
+            1,
+            0,
+            0,
+            False,
+            0,
+            -1,
+            0,
+            '#010101',
+            0,
+            1,
+            (0,0,0,0),
+            1
+        )
+
+        self.assertTrue(
+            md5('tests/out/TEST TEST TEST_13.jpg') == md5('tests/expected_results/TEST TEST TEST_13.jpg')
+        )
+
+        os.remove('tests/out/TEST TEST TEST_13.jpg')
+
     def test_generate_string_with_letters(self):
         s = create_strings_randomly(1, False, 1, True, False, False, 'en')[0]