Add -obb parameter to output bounding boxes

PyaePhyoKhant · web-flow · commit e7217dbd561c · 2022-01-18T21:25:15.000-05:00
diff --git a/tests.py b/tests.py
@@ -1073,6 +1073,92 @@ def test_generate_data_with_hindi_text(self):
 
         os.remove("tests/out/22_परीक्षा परीक्षा परीक्षा.png")
 
+    def test_generate_data_with_output_bounding_box(self):
+        FakeTextDataGenerator.generate(
+            21,
+            "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,
+            (5, 5, 5, 5),
+            0,
+            0,
+            False,
+            os.path.join(os.path.split(os.path.realpath(__file__))[0], "trdg/images"),
+            output_bboxes=1,
+        )
+
+        self.assertLess(
+            diff(
+                "tests/out/TEST TEST TEST_21.jpg",
+                "tests/expected_results/TEST TEST TEST_21.jpg",
+                delete_diff_file=True,
+            ),
+            0.11
+        )
+
+        os.remove("tests/out/TEST TEST TEST_21.jpg")
+        os.remove("tests/out/TEST TEST TEST_21_boxes.txt")
+
+    def test_generate_data_with_tesseract_output_bounding_box(self):
+        FakeTextDataGenerator.generate(
+            22,
+            "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,
+            (5, 5, 5, 5),
+            0,
+            0,
+            False,
+            os.path.join(os.path.split(os.path.realpath(__file__))[0], "trdg/images"),
+            output_bboxes=2,
+        )
+
+        self.assertLess(
+            diff(
+                "tests/out/TEST TEST TEST_22.jpg",
+                "tests/expected_results/TEST TEST TEST_22.jpg",
+                delete_diff_file=True,
+            ),
+            0.11
+        )
+
+        os.remove("tests/out/TEST TEST TEST_22.jpg")
+        os.remove("tests/out/TEST TEST TEST_22.box")
+
     def test_generate_string_with_letters(self):
         s = create_strings_randomly(1, False, 1, True, False, False, "en")[0]
 
diff --git a/tests/expected_results/TEST TEST TEST_21.jpg b/tests/expected_results/TEST TEST TEST_21.jpg
diff --git a/tests/expected_results/TEST TEST TEST_21_boxes.txt b/tests/expected_results/TEST TEST TEST_21_boxes.txt
@@ -0,0 +1,14 @@
+5 22 31 58
+36 22 57 58
+62 22 85 59
+88 22 114 58
+115 22 128 59
+129 22 155 58
+160 22 181 58
+186 22 209 59
+212 22 238 58
+239 22 252 59
+253 22 279 58
+284 22 305 58
+310 22 333 59
+336 22 362 58
diff --git a/tests/expected_results/TEST TEST TEST_22.box b/tests/expected_results/TEST TEST TEST_22.box
@@ -0,0 +1,14 @@
+T 5 6 31 42 0
+E 36 6 57 42 0
+S 62 5 85 42 0
+T 88 6 114 42 0
+  115 5 128 43 0
+T 129 6 155 42 0
+E 160 6 181 42 0
+S 186 5 209 42 0
+T 212 6 238 42 0
+  239 5 252 43 0
+T 253 6 279 42 0
+E 284 6 305 42 0
+S 310 5 333 42 0
+T 336 6 362 42 0
diff --git a/tests/expected_results/TEST TEST TEST_22.jpg b/tests/expected_results/TEST TEST TEST_22.jpg
diff --git a/trdg/computer_text_generator.py b/trdg/computer_text_generator.py
@@ -132,6 +132,7 @@ def _generate_vertical_text(
 
     txt_img_draw = ImageDraw.Draw(txt_img)
     txt_mask_draw = ImageDraw.Draw(txt_mask)
+    txt_mask_draw.fontmode = "1"
 
     colors = [ImageColor.getrgb(c) for c in text_color.split(",")]
     c1, c2 = colors[0], colors[-1]
@@ -163,7 +164,7 @@ def _generate_vertical_text(
         txt_mask_draw.text(
             (0, sum(char_heights[0:i]) + i * character_spacing),
             c,
-            fill=(i // (255 * 255), i // 255, i % 255),
+            fill=((i + 1) // (255 * 255), (i + 1) // 255, (i + 1) % 255),
             font=image_font,
             stroke_width=stroke_width,
             stroke_fill=stroke_fill,
diff --git a/trdg/data_generator.py b/trdg/data_generator.py
@@ -4,6 +4,7 @@
 from PIL import Image, ImageFilter, ImageStat
 
 from trdg import computer_text_generator, background_generator, distorsion_generator
+from trdg.utils import mask_to_bboxes
 
 try:
     from trdg import handwritten_text_generator
@@ -52,6 +53,7 @@ def generate(
         stroke_width=0, 
         stroke_fill="#282828",
         image_mode="RGB", 
+        output_bboxes=0,
     ):
         image = None
 
@@ -248,24 +250,36 @@ def generate(
         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))
+            name = "{}_{}".format(text, str(index))
         elif name_format == 1:
-            image_name = "{}_{}.{}".format(str(index), text, extension)
-            mask_name = "{}_{}_mask.png".format(str(index), text)
+            name = "{}_{}".format(str(index), text)
         elif name_format == 2:
-            image_name = "{}.{}".format(str(index), extension)
-            mask_name = "{}_mask.png".format(str(index))
+            name = 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))
+            name = "{}_{}".format(text, str(index))
+
+        image_name = "{}.{}".format(name, extension)
+        mask_name = "{}_mask.png".format(name)
+        box_name = "{}_boxes.txt".format(name)
+        tess_box_name = "{}.box".format(name)
+
 
         # Save the image
         if out_dir is not None:
             final_image.save(os.path.join(out_dir, image_name))
             if output_mask == 1:
                 final_mask.save(os.path.join(out_dir, mask_name))
+            if output_bboxes == 1:
+                bboxes = mask_to_bboxes(final_mask)
+                with open(os.path.join(out_dir, box_name), "w") as f:
+                    for bbox in bboxes:
+                        f.write(" ".join([str(v) for v in bbox]) + "\n")
+            if output_bboxes == 2:
+                bboxes = mask_to_bboxes(final_mask, tess=True)
+                with open(os.path.join(out_dir, tess_box_name), "w") as f:
+                    for bbox, char in zip(bboxes, text):
+                        f.write(" ".join([char] + [str(v) for v in bbox] + ['0']) + "\n")
         else:
             if output_mask == 1:
                 return final_image, final_mask
diff --git a/trdg/generators/from_dict.py b/trdg/generators/from_dict.py
@@ -41,6 +41,7 @@ def __init__(
         stroke_width=0, 
         stroke_fill="#282828",
         image_mode="RGB",
+        output_bboxes=0,
     ):
         self.count = count
         self.length = length
@@ -74,6 +75,7 @@ def __init__(
             stroke_width,
             stroke_fill,
             image_mode,
+            output_bboxes,
         )
 
     def __iter__(self):
diff --git a/trdg/generators/from_random.py b/trdg/generators/from_random.py
@@ -44,6 +44,7 @@ def __init__(
             stroke_width=0,
             stroke_fill="#282828",
             image_mode="RGB",
+            output_bboxes=0,
     ):
         self.generated_count = 0
         self.count = count
@@ -89,6 +90,7 @@ def __init__(
             stroke_width,
             stroke_fill,
             image_mode,
+            output_bboxes,
         )
 
     def __iter__(self):
diff --git a/trdg/generators/from_strings.py b/trdg/generators/from_strings.py
@@ -42,6 +42,7 @@ def __init__(
         stroke_width=0, 
         stroke_fill="#282828",
         image_mode="RGB",
+        output_bboxes=0,
     ):
         self.count = count
         self.strings = strings
@@ -77,6 +78,7 @@ def __init__(
         self.output_mask = output_mask
         self.word_split = word_split
         self.image_dir = image_dir
+        self.output_bboxes = output_bboxes
         self.generated_count = 0
         self.stroke_width = stroke_width
         self.stroke_fill = stroke_fill
@@ -123,6 +125,7 @@ def next(self):
                 self.stroke_width,
                 self.stroke_fill,
                 self.image_mode, 
+                self.output_bboxes,
             ),
             self.orig_strings[(self.generated_count - 1) % len(self.orig_strings)] if self.rtl else self.strings[(self.generated_count - 1) % len(self.strings)],
         )
diff --git a/trdg/generators/from_wikipedia.py b/trdg/generators/from_wikipedia.py
@@ -40,6 +40,7 @@ def __init__(
             stroke_width=0,
             stroke_fill="#282828",
             image_mode="RGB",
+            output_bboxes=0,
     ):
         self.generated_count = 0
         self.count = count
@@ -73,6 +74,7 @@ def __init__(
             stroke_width,
             stroke_fill,
             image_mode,
+            output_bboxes,
         )
 
     def __iter__(self):
diff --git a/trdg/run.py b/trdg/run.py
@@ -194,6 +194,13 @@ def parse_arguments():
         help="Define if the generator will return masks for the text",
         default=0,
     )
+    parser.add_argument(
+        "-obb",
+        "--output_bboxes",
+        type=int,
+        help="Define if the generator will return bounding boxes for the text, 1: Bounding box file, 2: Tesseract format",
+        default=0
+    )
     parser.add_argument(
         "-d",
         "--distorsion",
@@ -455,6 +462,7 @@ def main():
                 [args.stroke_width] * string_count,
                 [args.stroke_fill] * string_count,
                 [args.image_mode] * string_count,
+                [args.output_bboxes] * string_count,
             ),
         ),
         total=args.count,
diff --git a/trdg/utils.py b/trdg/utils.py
@@ -3,7 +3,8 @@
 """
 
 import os
-
+import numpy as np
+from PIL import Image
 
 def load_dict(lang):
     """Read the dictionnary file and returns all words in it.
@@ -38,3 +39,45 @@ def load_fonts(lang):
                 os.path.join(os.path.dirname(__file__), "fonts/latin")
             )
         ]
+
+def mask_to_bboxes(mask, tess=False):
+    """Process the mask and turns it into a list of AABB bounding boxes
+    """
+
+    mask_arr = np.array(mask)
+
+    bboxes = []
+
+    i = 0
+    space_thresh = 1
+    while True:
+        try:
+            color_tuple = ((i + 1) // (255 * 255), (i + 1) // 255, (i + 1) % 255)
+            letter = np.where(np.all(mask_arr == color_tuple, axis=-1))
+            if space_thresh == 0 and letter:
+                x1 = min(bboxes[-1][2] + 1, np.min(letter[1]) - 1)
+                y1 = min(bboxes[-1][3] + 1, np.min(letter[0]) - 1) if not tess else min(mask_arr.shape[0] - np.min(letter[0]) + 2, bboxes[-1][1] - 1)
+                x2 = max(bboxes[-1][2] + 1, np.min(letter[1]) - 2)
+                y2 = max(bboxes[-1][3] + 1, np.min(letter[0]) - 2) if not tess else max(mask_arr.shape[0] - np.min(letter[0]) + 2, bboxes[-1][1] - 1)
+                bboxes.append((x1, y1, x2, y2))
+                space_thresh += 1
+            bboxes.append((
+                max(0, np.min(letter[1]) - 1),
+                max(0, np.min(letter[0]) - 1) if not tess else max(0, mask_arr.shape[0] - np.max(letter[0]) - 1),
+                min(mask_arr.shape[1] - 1, np.max(letter[1]) + 1),
+                min(mask_arr.shape[0] - 1, np.max(letter[0]) + 1) if not tess else min(mask_arr.shape[0] - 1, mask_arr.shape[0] - np.min(letter[0]) + 1)
+            ))
+            i += 1
+        except Exception as ex:
+            if space_thresh == 0:
+                break
+            space_thresh -= 1
+            i += 1
+
+    return bboxes        
+
+def draw_bounding_boxes(img, bboxes, color="green"):
+    d = ImageDraw.Draw(img)
+
+    for bbox in bboxes:
+        d.rectangle(bbox, outline=color)
