diff --git a/requirements-hw.txt b/requirements-hw.txt index 0aa037d8c5..544ae24234 100644 --- a/requirements-hw.txt +++ b/requirements-hw.txt @@ -8,3 +8,4 @@ diffimg==0.2.3 tensorflow>=1.13.1,<1.14 matplotlib>=3.0.2 seaborn>=0.9.0 +xeger==0.3.5 diff --git a/requirements.txt b/requirements.txt index aa94c00f1e..400fc78223 100644 --- a/requirements.txt +++ b/requirements.txt @@ -4,3 +4,4 @@ opencv-python>=4.2.0.32 tqdm>=4.23.0 beautifulsoup4>=4.6.0 diffimg==0.2.3 +xeger==0.3.5 diff --git a/trdg/data_generator.py b/trdg/data_generator.py index 46a60df78b..fe13827851 100644 --- a/trdg/data_generator.py +++ b/trdg/data_generator.py @@ -1,7 +1,7 @@ import os import random as rnd -from PIL import Image, ImageFilter, ImageStat +from PIL import Image, ImageFilter from trdg import computer_text_generator, background_generator, distorsion_generator @@ -49,9 +49,8 @@ def generate( output_mask, word_split, image_dir, - stroke_width=0, - stroke_fill="#282828", - image_mode="RGB", + save_bbox, + save_bbox_dir ): image = None @@ -77,8 +76,6 @@ def generate( character_spacing, fit, word_split, - stroke_width, - stroke_fill, ) random_angle = rnd.randint(0 - skewing_angle, skewing_angle) @@ -131,7 +128,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), Image.NEAREST) + 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 @@ -144,7 +141,7 @@ def generate( (size - horizontal_margin, new_height), Image.ANTIALIAS ) resized_mask = distorted_mask.resize( - (size - horizontal_margin, new_height), Image.NEAREST + (size - horizontal_margin, new_height), Image.ANTIALIAS ) background_width = size background_height = new_height + vertical_margin @@ -174,26 +171,6 @@ def generate( "RGB", (background_width, background_height), (0, 0, 0) ) - ############################################################## - # Comparing average pixel value of text and background image # - ############################################################## - try: - resized_img_st = ImageStat.Stat(resized_img, resized_mask.split()[2]) - background_img_st = ImageStat.Stat(background_img) - - resized_img_px_mean = sum(resized_img_st.mean[:2]) / 3 - background_img_px_mean = sum(background_img_st.mean) / 3 - - if abs(resized_img_px_mean - background_img_px_mean) < 15: - print("value of mean pixel is too similar. Ignore this image") - - print("resized_img_st \n {}".format(resized_img_st.mean)) - print("background_img_st \n {}".format(background_img_st.mean)) - - return - except Exception as err: - return - ############################# # Place text with alignment # ############################# @@ -224,29 +201,19 @@ def generate( (background_width - new_text_width - margin_right, margin_top), ) - ####################### + ################################## # Apply gaussian 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) - - ############################################ - # Change image mode (RGB, grayscale, etc.) # - ############################################ - - final_image = final_image.convert(image_mode) - final_mask = final_mask.convert(image_mode) ##################################### # Generate name for resulting image # ##################################### - # We remove spaces if space_width == 0 - if space_width == 0: - text = text.replace(" ", "") if name_format == 0: image_name = "{}_{}.{}".format(text, str(index), extension) mask_name = "{}_{}_mask.png".format(text, str(index)) @@ -261,12 +228,23 @@ def generate( image_name = "{}_{}.{}".format(text, str(index), extension) mask_name = "{}_{}_mask.png".format(text, str(index)) + image_name = image_name.replace(" ", "_") + image_name = image_name.replace("/", "-") + + mask_name = mask_name.replace(" ", "_") + mask_name = mask_name.replace("/", "-") + # Save the image if out_dir is not None: - final_image.save(os.path.join(out_dir, image_name)) + final_image.convert("RGB").save(os.path.join(out_dir, image_name)) if output_mask == 1: - final_mask.save(os.path.join(out_dir, mask_name)) + final_mask.convert("RGB").save(os.path.join(out_dir, mask_name)) else: if output_mask == 1: - return final_image, final_mask - return final_image + return final_image.convert("RGB"), final_mask.convert("RGB") + return final_image.convert("RGB") + + if save_bbox: + with open(os.path.join(save_bbox_dir, image_name.rsplit('.', 1)[0]) + ".txt", "w") as f: + bbox = final_mask.getbbox() + f.write(f"{bbox[0]}, {bbox[1]}, {bbox[2]}, {bbox[3]}") #x1, y1, x2, y2 \ No newline at end of file diff --git a/trdg/run.py b/trdg/run.py index 33ba250441..163f70ea6e 100755 --- a/trdg/run.py +++ b/trdg/run.py @@ -1,6 +1,5 @@ import argparse -import errno -import os +import os, errno import sys sys.path.append(os.path.join(os.path.dirname(__file__), "..")) @@ -8,16 +7,18 @@ import random as rnd import string import sys -from multiprocessing import Pool from tqdm import tqdm - -from trdg.data_generator import FakeTextDataGenerator -from trdg.string_generator import (create_strings_from_dict, - create_strings_from_file, - create_strings_from_wikipedia, - create_strings_randomly) +from trdg.string_generator import ( + create_strings_from_dict, + create_strings_from_file, + create_strings_from_wikipedia, + create_strings_randomly, + create_strings_from_regex +) from trdg.utils import load_dict, load_fonts +from trdg.data_generator import FakeTextDataGenerator +from multiprocessing import Pool def margins(margin): @@ -46,12 +47,20 @@ def parse_arguments(): help="When set, this argument uses a specified text file as source for the text", default="", ) + parser.add_argument( + "-re", + "--regex", + type=str, + nargs="?", + help="Generate words with some regex. For example emails/account numbers/....", + default="", + ) parser.add_argument( "-l", "--language", type=str, nargs="?", - help="The language to use, should be fr (French), en (English), es (Spanish), de (German), ar (Arabic), cn (Chinese), ja (Japanese) or hi (Hindi)", + help="The language to use, should be fr (French), en (English), es (Spanish), de (German), ar (Arabic), cn (Chinese), or hi (Hindi)", default="en", ) parser.add_argument( @@ -264,7 +273,7 @@ def parse_arguments(): type=margins, nargs="?", help="Define the margins around the text when rendered. In pixels", - default=(5, 5, 5, 5), + default=(64, 64, 64, 64), ) parser.add_argument( "-fi", @@ -309,28 +318,18 @@ def parse_arguments(): default=False, ) parser.add_argument( - "-stw", - "--stroke_width", - type=int, - nargs="?", - help="Define the width of the strokes", - default=0, - ) - parser.add_argument( - "-stf", - "--stroke_fill", - type=str, - nargs="?", - help="Define the color of the contour of the strokes, if stroke_width is bigger than 0", - default="#282828", + "-bb", + "--save_bbox", + action="store_true", + help="generates and saves bounding boxes of texts", + default=False ) parser.add_argument( - "-im", - "--image_mode", + "-bbdir", + "--save_bbox_dir", type=str, - nargs="?", - help="Define the image mode to be used. RGB is default, L means 8-bit grayscale images, 1 means 1-bit binary images stored with one pixel per byte, etc.", - default="RGB", + help="Directory path to save bounding boxes", + default="./bboxes" ) return parser.parse_args() @@ -350,6 +349,13 @@ def main(): if e.errno != errno.EEXIST: raise + if args.save_bbox: + try: + os.makedirs(args.save_bbox_dir) + except OSError as e: + if e.errno != errno.EEXIST: + raise + # Creating word list if args.dict: lang_dict = [] @@ -383,6 +389,8 @@ def main(): strings = create_strings_from_wikipedia(args.length, args.count, args.language) elif args.input_file != "": strings = create_strings_from_file(args.input_file, args.count) + elif args.regex: + strings = create_strings_from_regex(args.length, args.regex, args.count) elif args.random_sequences: strings = create_strings_randomly( args.length, @@ -407,11 +415,10 @@ def main(): if args.language == "ar": from arabic_reshaper import ArabicReshaper - from bidi.algorithm import get_display arabic_reshaper = ArabicReshaper() strings = [ - " ".join([get_display(arabic_reshaper.reshape(w)) for w in s.split(" ")[::-1]]) + " ".join([arabic_reshaper.reshape(w) for w in s.split(" ")[::-1]]) for s in strings ] if args.case == "upper": @@ -452,9 +459,8 @@ def main(): [args.output_mask] * string_count, [args.word_split] * string_count, [args.image_dir] * string_count, - [args.stroke_width] * string_count, - [args.stroke_fill] * string_count, - [args.image_mode] * string_count, + [args.save_bbox] * string_count, + [args.save_bbox_dir] * string_count, ), ), total=args.count, @@ -469,10 +475,7 @@ def main(): ) as f: for i in range(string_count): file_name = str(i) + "." + args.extension - label = strings[i] - if args.space_width == 0: - label = label.replace(" ", "") - f.write("{} {}\n".format(file_name, label)) + f.write("{} {}\n".format(file_name, strings[i])) if __name__ == "__main__": diff --git a/trdg/string_generator.py b/trdg/string_generator.py index 6853a419c8..f00d37b975 100644 --- a/trdg/string_generator.py +++ b/trdg/string_generator.py @@ -4,6 +4,7 @@ import requests from bs4 import BeautifulSoup +from xeger import Xeger def create_strings_from_file(filename, count): @@ -139,3 +140,13 @@ def create_strings_randomly(length, allow_variable, count, let, num, sym, lang): current_string += " " strings.append(current_string[:-1]) return strings + + +def create_strings_from_regex(length, regex, count): + re = Xeger(limit=length) + strings = [] + + for _ in range(0, count): + current_string = re.xeger(regex) + strings.append(current_string) + return strings