Add figure layout prefix and suffix

aws-samples · May 10, 2024 · 03024da · 03024da
2 parents 44a002d + e4421e0
commit 03024da
Show file tree

Hide file tree

Showing 4 changed files with 51 additions and 0 deletions.
diff --git a/.../fixtures/saved_api_responses/test_figure_layout_prefixes_and_suffixes_in_text_words.json b/.../fixtures/saved_api_responses/test_figure_layout_prefixes_and_suffixes_in_text_words.json
diff --git a/tests/test_get_text_and_words.py b/tests/test_get_text_and_words.py
@@ -248,3 +248,39 @@ def test_table_prefixes_and_suffixes_in_words(self):
             "</value>",
         ]:
             self.assertTrue(token in words, f"{token} is not in text")
+
+    def test_figure_layout_prefixes_and_suffixes_in_text_words(self):
+        if os.environ.get("CALL_TEXTRACT"):
+            document = self.extractor.analyze_document(
+                os.path.join(self.current_directory, "fixtures/matrix.png"),
+                features=[
+                    TextractFeatures.LAYOUT,
+                ]
+            )
+            with open(get_fixture_path(), "w") as f:
+                json.dump(document.response, f)
+        else:
+            document = Document.open(get_fixture_path())
+
+        config = TextLinearizationConfig(
+            figure_layout_prefix = "<figure>",  #: Prefix for figure elements
+            figure_layout_suffix = "</figure>",  #: Suffix for figure elements
+            add_prefixes_and_suffixes_in_text=True,
+            add_prefixes_and_suffixes_as_words=True,
+        )
+
+        text, words = document.get_text_and_words(config)
+
+        for token in [
+            "<figure>",
+            "</figure>",
+        ]:
+            self.assertTrue(token in text, f"{token} is not in text")
+
+        words = [w.text for w in words]
+
+        for token in [
+            "<figure>",
+            "</figure>",
+        ]:
+            self.assertTrue(token in words, f"{token} is not in words")
diff --git a/textractor/data/text_linearization_config.py b/textractor/data/text_linearization_config.py
@@ -134,6 +134,10 @@ class TextLinearizationConfig:
 
     value_suffix: str = ""  #: Suffix for value elements
 
+    figure_layout_prefix: str = "" #: Prefix for figure layout elements 
+
+    figure_layout_suffix: str = "" #: Suffix for figure layout elements
+
     selection_element_selected: str = (
         "[X]"  #: Representation for selection element when selected
     )

diff --git a/textractor/entities/layout.py b/textractor/entities/layout.py
@@ -232,6 +232,10 @@ def get_text_and_words(
                     final_text = (
                         config.key_value_layout_prefix + final_text + config.key_value_layout_suffix
                     )
+                elif self.layout_type == LAYOUT_FIGURE:
+                    final_text = (
+                        config.figure_layout_prefix + final_text + config.figure_layout_suffix
+                    )
             if config.add_prefixes_and_suffixes_as_words:
                 if self.layout_type == LAYOUT_TABLE:
                     final_words = (
@@ -245,6 +249,12 @@ def get_text_and_words(
                         final_words + 
                         ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.key_value_layout_suffix, is_structure=True)] if config.key_value_layout_suffix else []) 
                     )
+                elif self.layout_type == LAYOUT_FIGURE:
+                    final_words = (
+                        ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.figure_layout_prefix, is_structure=True)] if config.figure_layout_prefix else []) + 
+                        final_words + 
+                        ([Word(str(uuid.uuid4()), BoundingBox.enclosing_bbox(final_words), config.figure_layout_suffix, is_structure=True)] if config.figure_layout_suffix else []) 
+                    )
 
         while (
             config.layout_element_separator * (config.max_number_of_consecutive_new_lines + 1) in final_text