Merge pull request #29 from Belval/text-orientation-and-variable-spacing

Text orientation and variable spacing

Author: Belval

README.md

@@ -22,11 +22,11 @@ tqdm
  You can simply use `pip install -r requirements.txt` too.
 
 ## New
+- 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`
-- Generate random strings with letters, numbers and symbols (Thank you @FHainzl)
-- Save the labels in a file instead of in the file name (Thank you @FHainzl)
 - Add support for Simplified and Traditional Chinese
 
 ## How does it work?

TextRecognitionDataGenerator/computer_text_generator.py

@@ -4,9 +4,51 @@
 
 class ComputerTextGenerator(object):
     @classmethod
-    def generate(cls, text, font, text_color, font_size):
+    def generate(cls, text, font, text_color, font_size, orientation, space_width):
+        if orientation == 0:
+            return cls.__generate_horizontal_text(text, font, text_color, font_size, space_width)
+        elif orientation == 1:
+            return cls.__generate_vertical_text(text, font, text_color, font_size, space_width)
+        else:
+            raise ValueError("Unknown orientation " + str(orientation))
+    
+    @classmethod
+    def __generate_horizontal_text(cls, text, font, text_color, font_size, space_width):
+        image_font = ImageFont.truetype(font=font, size=font_size)
+        words = text.split(' ')
+        space_width = image_font.getsize(' ')[0] * space_width
+
+        words_width = [image_font.getsize(w)[0] for w in words]
+        text_width =  sum(words_width) + int(space_width) * (len(words) - 1)
+        text_height = max([image_font.getsize(w)[1] for w in words])
+
+        txt_img = Image.new('RGBA', (text_width, text_height), (0, 0, 0, 0))
+
+        txt_draw = ImageDraw.Draw(txt_img)
+
+        colors = [ImageColor.getrgb(c) for c in text_color.split(',')]
+        c1, c2 = colors[0], colors[-1]
+
+        fill = (
+            random.randint(c1[0], c2[0]),
+            random.randint(c1[1], c2[1]),
+            random.randint(c1[2], c2[2])
+        )
+
+        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
+
+    @classmethod
+    def __generate_vertical_text(cls, text, font, text_color, font_size, space_width):
         image_font = ImageFont.truetype(font=font, size=font_size)
-        text_width, text_height = image_font.getsize(text)
+        
+        space_height = int(image_font.getsize(' ')[1] * space_width)
+
+        char_heights = [image_font.getsize(c)[1] if c != ' ' else space_height for c in text]
+        text_width = max([image_font.getsize(c)[0] for c in text])
+        text_height = sum(char_heights)
 
         txt_img = Image.new('RGBA', (text_width, text_height), (0, 0, 0, 0))
 
@@ -21,6 +63,7 @@ def generate(cls, text, font, text_color, font_size):
             random.randint(c1[2], c2[2])
         )
 
-        txt_draw.text((0, 0), text, fill=fill, font=image_font)
+        for i, c in enumerate(text):
+            txt_draw.text((0, sum(char_heights[0:i])), c, fill=fill, font=image_font)
 
         return txt_img

TextRecognitionDataGenerator/data_generator.py

@@ -21,16 +21,18 @@ def generate_from_tuple(cls, t):
         cls.generate(*t)
 
     @classmethod
-    def generate(cls, index, text, font, out_dir, height, extension, skewing_angle, random_skew, blur, random_blur, background_type, distorsion_type, distorsion_orientation, is_handwritten, name_format, width, alignment, text_color):
+    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):
         image = None
 
         ##########################
         # Create picture of text #
         ##########################
         if is_handwritten:
+            if orientation == 1:
+                raise ValueError("Vertical handwritten text is unavailable")
             image = HandwrittenTextGenerator.generate(text)
         else:
-            image = ComputerTextGenerator.generate(text, font, text_color, height)
+            image = ComputerTextGenerator.generate(text, font, text_color, size, orientation, space_width)
 
         random_angle = random.randint(0-skewing_angle, skewing_angle)
 
@@ -64,23 +66,32 @@ def generate(cls, index, text, font, out_dir, height, extension, skewing_angle,
         # Resize image to desired format #
         ##################################
 
-        new_width = int(float(distorted_img.size[0] + 10) * (float(height) / float(distorted_img.size[1] + 10)))
-
-        resized_img = distorted_img.resize((new_width, height - 10), Image.ANTIALIAS)
-
-        background_width = width if width > 0 else new_width + 10
+        # 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
+            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)
+            background_width = size
+            background_height = new_height + 10
+        else:
+            raise ValueError("Invalid orientation")
 
         #############################
         # Generate background image #
         #############################
         if background_type == 0:
-            background = BackgroundGenerator.gaussian_noise(height, background_width)
+            background = BackgroundGenerator.gaussian_noise(background_height, background_width)
         elif background_type == 1:
-            background = BackgroundGenerator.plain_white(height, background_width)
+            background = BackgroundGenerator.plain_white(background_height, background_width)
         elif background_type == 2:
-            background = BackgroundGenerator.quasicrystal(height, background_width)
+            background = BackgroundGenerator.quasicrystal(background_height, background_width)
         else:
-            background = BackgroundGenerator.picture(height, background_width)
+            background = BackgroundGenerator.picture(background_height, background_width)
 
         #############################
         # Place text with alignment #

TextRecognitionDataGenerator/run.py

@@ -103,7 +103,7 @@ def parse_arguments():
         "--format",
         type=int,
         nargs="?",
-        help="Define the height of the produced images",
+        help="Define the height of the produced images if horizontal, else the width",
         default=32,
     )
     parser.add_argument(
@@ -212,6 +212,14 @@ def parse_arguments():
         help="Define the alignment of the text in the image. Only used if the width parameter is set. 0: left, 1: center, 2: right",
         default=1
     )
+    parser.add_argument(
+        "-or",
+        "--orientation",
+        type=int,
+        nargs="?",
+        help="Define the orientation of the text. 0: Horizontal, 1: Vertical",
+        default=0
+    )
     parser.add_argument(
         "-tc",
         "--text_color",
@@ -220,6 +228,14 @@ def parse_arguments():
         help="Define the text's color, should be either a single hex color or a range in the ?,? format.",
         default='#282828'
     )
+    parser.add_argument(
+        "-sw",
+        "--space_width",
+        type=float,
+        nargs="?",
+        help="Define the width of the spaces between words. 2.0 means twice the normal space width",
+        default=1.0
+    )
 
     return parser.parse_args()
 
@@ -304,7 +320,9 @@ def main():
             [args.name_format] * string_count,
             [args.width] * string_count,
             [args.alignment] * string_count,
-            [args.text_color] * string_count
+            [args.text_color] * string_count,
+            [args.orientation] * string_count,
+            [args.space_width] * string_count
         )
     ), total=args.count):
         pass

TextRecognitionDataGenerator/string_generator.py

@@ -17,7 +17,7 @@ def create_strings_from_file(filename, count):
         if len(lines) == 0:
             raise Exception("No lines could be read in file")
         while len(strings) < count:
-            if len(lines) > count - len(strings):
+            if len(lines) >= count - len(strings):
                 strings.extend(lines[0:count - len(strings)])
             else:
                 strings.extend(lines)