fix: docai_utilities.py to return Optional (#176)

* fix: `docai_utilities.py` to return `Optional`

- Should resolve customer reported issue in support case #47169701 relating to duplicate/inaccurate elements in hOCR output
- Followup to:
  - #161
  - #169

* Increase test coverage

* Addressed review comments

Author: holtskinner

google/cloud/documentai_toolbox/utilities/docai_utilities.py

@@ -15,15 +15,15 @@
 #
 """Utilities for Document AI"""
 
-from typing import Tuple
+from typing import Optional, Tuple
 
 from google.cloud import documentai
 
 
 def get_bounding_box(
     bounding_poly: documentai.BoundingPoly,
     page_dimension: documentai.Document.Page.Dimension,
-) -> Tuple[int, int, int, int]:
+) -> Optional[Tuple[int, int, int, int]]:
     r"""Returns the bounding box of an element from the element bounding_poly and page dimensions.
 
     Args:
@@ -35,10 +35,10 @@ def get_bounding_box(
     Returns:
         Tuple[int, int, int, int]:
             Bounding box coordinates in order (top, left, bottom, right).
-            Returns `0, 0, 0, 0` if `bounding_poly.normalized_vertices` is empty.
+            Returns `None` if `bounding_poly` or `bounding_poly.normalized_vertices` is empty.
     """
-    if not bounding_poly.normalized_vertices:
-        return 0, 0, 0, 0
+    if not bounding_poly or not bounding_poly.normalized_vertices:
+        return None
 
     vertices = [
         (

google/cloud/documentai_toolbox/wrappers/entity.py

@@ -95,9 +95,11 @@ def crop_image(
         if not documentai_page.image:
             raise ValueError("Document does not contain images.")
 
-        top, left, bottom, right = docai_utilities.get_bounding_box(
+        bbox = docai_utilities.get_bounding_box(
             bounding_poly=self.documentai_object.page_anchor.page_refs[0].bounding_poly,
             page_dimension=documentai_page.dimension,
         )
+        if bbox is None:
+            return None
         doc_image = Image.open(BytesIO(documentai_page.image.content))
-        return doc_image.crop((top, left, bottom, right))
+        return doc_image.crop(bbox)

google/cloud/documentai_toolbox/wrappers/page.py

@@ -16,7 +16,7 @@
 """Wrappers for Document AI Page type."""
 
 import dataclasses
-from typing import List, Optional, Union, cast
+from typing import cast, List, Optional, Union
 
 import pandas as pd
 
@@ -301,7 +301,7 @@ def _table_rows_from_documentai_table_rows(
 def _get_hocr_bounding_box(
     element_with_layout: ElementWithLayout,
     page_dimension: documentai.Document.Page.Dimension,
-) -> str:
+) -> Optional[str]:
     r"""Returns a hOCR bounding box string.
 
     Args:
@@ -311,13 +311,21 @@ def _get_hocr_bounding_box(
             Required. Page dimension.
 
     Returns:
-        str:
+        Optional[str]:
             hOCR bounding box sring.
     """
-    min_x, min_y, max_x, max_y = docai_utilities.get_bounding_box(
+    if not element_with_layout.layout.bounding_poly:
+        return None
+
+    bbox = docai_utilities.get_bounding_box(
         bounding_poly=element_with_layout.layout.bounding_poly,
         page_dimension=page_dimension,
     )
+
+    if not bbox:
+        return None
+
+    min_x, min_y, max_x, max_y = bbox
     return f"bbox {min_x} {min_y} {max_x} {max_y}"
 
 

tests/unit/test_document.py

@@ -105,6 +105,19 @@ def get_bytes_missing_shard_mock():
         yield byte_factory
 
 
+def create_document_with_images_without_bbox(get_bytes_images_mock):
+    doc = document.Document.from_gcs(
+        gcs_bucket_name="test-directory", gcs_prefix="documentai/output/123456789/0"
+    )
+
+    del (
+        doc.entities[0]
+        .documentai_object.page_anchor.page_refs[0]
+        .bounding_poly.normalized_vertices
+    )
+    return doc
+
+
 def test_get_shards_with_gcs_uri_contains_file_type():
     with pytest.raises(ValueError, match="gcs_prefix cannot contain file types"):
         document._get_shards(
@@ -299,6 +312,13 @@ def test_document_from_document_path_with_single_shard():
     assert len(actual.pages) == 1
 
 
+def test_document_from_document_path_with_directory():
+    actual = document.Document.from_document_path(
+        document_path="tests/unit/resources/0/"
+    )
+    assert len(actual.pages) == 1
+
+
 def test_document_from_documentai_document_with_single_shard():
     with open(
         "tests/unit/resources/0/toolbox_invoice_test-0.json", "r", encoding="utf-8"
@@ -626,6 +646,7 @@ def test_export_images(get_bytes_images_mock):
     output_path = "resources/output/"
     if os.path.exists(output_path):
         shutil.rmtree(output_path)
+        assert not os.path.exists(output_path)
 
     doc = document.Document.from_gcs(
         gcs_bucket_name="test-directory", gcs_prefix="documentai/output/123456789/0"
@@ -648,6 +669,20 @@ def test_export_images(get_bytes_images_mock):
     shutil.rmtree(output_path)
 
 
+def test_export_images_empty_bounding_box(get_bytes_images_mock):
+    output_path = "resources/output/"
+
+    doc = create_document_with_images_without_bbox(get_bytes_images_mock)
+    actual = doc.export_images(
+        output_path=output_path,
+        output_file_prefix="exported_photo",
+        output_file_extension="png",
+    )
+    get_bytes_images_mock.assert_called_once()
+
+    assert not actual
+
+
 def test_export_hocr_str():
     wrapped_document = document.Document.from_document_path(
         document_path="tests/unit/resources/0/toolbox_invoice_test-0.json"

tests/unit/test_entity.py

@@ -102,3 +102,15 @@ def test_crop_image_without_page_image(docproto):
         match="Document does not contain images.",
     ):
         doc.entities[0].crop_image(documentai_page=docproto.pages[0])
+
+
+def test_crop_image_empty_bounding_box(docproto):
+    doc = document.Document.from_documentai_document(docproto)
+    del (
+        doc.entities[0]
+        .documentai_object.page_anchor.page_refs[0]
+        .bounding_poly.normalized_vertices
+    )
+
+    actual = doc.entities[0].crop_image(documentai_page=docproto.pages[0])
+    assert actual is None

tests/unit/test_page.py

@@ -40,6 +40,12 @@ def docproto_form_parser():
         return documentai.Document.from_json(f.read())
 
 
+@pytest.fixture
+def docproto_blank_document():
+    with open("tests/unit/resources/blank_document.json", "r", encoding="utf-8") as f:
+        return documentai.Document.from_json(f.read())
+
+
 def test_table_to_csv(docproto):
     docproto_page = docproto.pages[0]
     table = page.Table(
@@ -160,6 +166,15 @@ def test_get_hocr_bounding_box(docproto):
     assert hocr_bounding_box_with_vertices == "bbox 1310 220 1534 282"
 
 
+def test_get_hocr_bounding_box_with_blank_document(docproto_blank_document):
+    hocr_bounding_box_normalized = page._get_hocr_bounding_box(
+        element_with_layout=docproto_blank_document.pages[0],
+        page_dimension=docproto_blank_document.pages[0].dimension,
+    )
+
+    assert hocr_bounding_box_normalized is None
+
+
 # Class init Tests