diff --git a/trdg/computer_text_generator.py b/trdg/computer_text_generator.py index eb016e2ef..e35a1a680 100644 --- a/trdg/computer_text_generator.py +++ b/trdg/computer_text_generator.py @@ -15,7 +15,7 @@ "0xe4d", "0xe4e", ] -TH_UNDER_VOWELS = ["0xe38", "0xe39", "\0xe3A"] +TH_UNDER_VOWELS = ["0xe38", "0xe39", "0xe3A"] TH_UPPER_VOWELS = ["0xe31", "0xe34", "0xe35", "0xe36", "0xe37"] diff --git a/trdg/data_generator.py b/trdg/data_generator.py index 75a79467f..e32958d79 100644 --- a/trdg/data_generator.py +++ b/trdg/data_generator.py @@ -54,6 +54,7 @@ def generate( stroke_fill: str = "#282828", image_mode: str = "RGB", output_bboxes: int = 0, + debug = False ) -> Image: image = None @@ -189,10 +190,11 @@ def generate( 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") + if debug: + 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)) + 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: @@ -258,7 +260,8 @@ def generate( elif name_format == 2: name = str(index) else: - print("{} is not a valid name format. Using default.".format(name_format)) + if debug: + print("{} is not a valid name format. Using default.".format(name_format)) name = "{}_{}".format(text, str(index)) name = make_filename_valid(name, allow_unicode=True) diff --git a/trdg/generators/from_strings.py b/trdg/generators/from_strings.py index ef7d22efa..b2ff12cfa 100644 --- a/trdg/generators/from_strings.py +++ b/trdg/generators/from_strings.py @@ -45,6 +45,8 @@ def __init__( image_mode: str = "RGB", output_bboxes: int = 0, rtl: bool = False, + fail_retry_count: int = 100, + debug = False, ): self.count = count self.strings = strings @@ -89,7 +91,8 @@ def __init__( self.stroke_width = stroke_width self.stroke_fill = stroke_fill self.image_mode = image_mode - + self.fail_retry_count = fail_retry_count + self.debug = debug def __iter__(self): return self @@ -100,42 +103,61 @@ def next(self): if self.generated_count == self.count: raise StopIteration self.generated_count += 1 - return ( - FakeTextDataGenerator.generate( - self.generated_count, - self.strings[(self.generated_count - 1) % len(self.strings)], - self.fonts[(self.generated_count - 1) % len(self.fonts)], - None, - self.size, - None, - self.skewing_angle, - self.random_skew, - self.blur, - self.random_blur, - self.background_type, - self.distorsion_type, - self.distorsion_orientation, - self.is_handwritten, - 0, - self.width, - self.alignment, - self.text_color, - self.orientation, - self.space_width, - self.character_spacing, - self.margins, - self.fit, - self.output_mask, - self.word_split, - self.image_dir, - self.stroke_width, - self.stroke_fill, - self.image_mode, - self.output_bboxes, - ), + generated_image = None + current_string = self.strings[(self.generated_count - 1) % len(self.strings)] + current_font = self.fonts[(self.generated_count - 1) % len(self.fonts)] + tries = 0 + while tries < self.fail_retry_count: + generated_image = FakeTextDataGenerator.generate( + self.generated_count, + current_string, + current_font, + None, + self.size, + None, + self.skewing_angle, + self.random_skew, + self.blur, + self.random_blur, + self.background_type, + self.distorsion_type, + self.distorsion_orientation, + self.is_handwritten, + 0, + self.width, + self.alignment, + self.text_color, + self.orientation, + self.space_width, + self.character_spacing, + self.margins, + self.fit, + self.output_mask, + self.word_split, + self.image_dir, + self.stroke_width, + self.stroke_fill, + self.image_mode, + self.output_bboxes, + self.debug + ) + if generated_image is not None: + break + elif self.debug: + print (f"Try {tries + 1} to generate image") + tries += 1 + + if generated_image is None: + print (f"Tried {self.fail_retry_count} times to generate image of {current_string}") + print (f"{current_font} But failed") + + elif generated_image is not None and tries > 1: + print (f"Tries {tries} times Succeed of {current_string}, {current_font}") + + return (generated_image, self.orig_strings[(self.generated_count - 1) % len(self.orig_strings)] if self.rtl - else self.strings[(self.generated_count - 1) % len(self.strings)], + else current_string, ) def reshape_rtl(self, strings: list, rtl_shaper: ArabicReshaper):