diff --git a/google/cloud/documentai_toolbox/__init__.py b/google/cloud/documentai_toolbox/__init__.py index 4ff1b84d..9d1770b8 100644 --- a/google/cloud/documentai_toolbox/__init__.py +++ b/google/cloud/documentai_toolbox/__init__.py @@ -24,8 +24,13 @@ entity, ) +from .converters import ( + converters, +) + __all__ = ( document, page, entity, + converters, ) diff --git a/google/cloud/documentai_toolbox/converters/converters.py b/google/cloud/documentai_toolbox/converters/converters.py new file mode 100644 index 00000000..a00b21af --- /dev/null +++ b/google/cloud/documentai_toolbox/converters/converters.py @@ -0,0 +1,62 @@ +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Document.proto converters.""" + +from typing import List +from google.cloud.vision import AnnotateFileResponse, ImageAnnotationContext +from google.cloud.vision import AnnotateImageResponse + +from google.cloud.documentai_toolbox.wrappers import page + +from google.cloud.documentai_toolbox.converters.vision_helpers import ( + _convert_document_page, + _get_text_anchor_substring, + PageInfo, +) + + +def _convert_to_vision_annotate_file_response(text: str, pages: List[page.Page]): + """Convert OCR data from Document proto to AnnotateFileResponse proto (Vision API). + + Args: + text (str): + Contents of document. + List[Page]: + A list of Pages. + + Returns: + AnnotateFileResponse proto with a TextAnnotation per page. + """ + responses = [] + vision_file_response = AnnotateFileResponse() + page_idx = 0 + while page_idx < len(pages): + page_info = PageInfo(pages[page_idx].documentai_page, text) + page_vision_annotation = _convert_document_page(page_info) + page_vision_annotation.text = _get_text_anchor_substring( + text, pages[page_idx].documentai_page.layout.text_anchor + ) + responses.append( + AnnotateImageResponse( + full_text_annotation=page_vision_annotation, + context=ImageAnnotationContext(page_number=page_idx + 1), + ) + ) + page_idx += 1 + + vision_file_response.responses = responses + + return vision_file_response diff --git a/google/cloud/documentai_toolbox/converters/vision_helpers.py b/google/cloud/documentai_toolbox/converters/vision_helpers.py new file mode 100644 index 00000000..9562d51e --- /dev/null +++ b/google/cloud/documentai_toolbox/converters/vision_helpers.py @@ -0,0 +1,326 @@ +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +"""Helper functions for docproto to vision conversion.""" + +import dataclasses +from typing import List, Union + +import immutabledict + +from google.cloud.documentai import Document +from google.cloud.vision import TextAnnotation, Symbol, Word, Paragraph, Block, Page +from google.cloud import vision + + +_BREAK_TYPE_MAP = immutabledict.immutabledict( + { + Document.Page.Token.DetectedBreak.Type.SPACE: ( + TextAnnotation.DetectedBreak.BreakType.SPACE + ), + Document.Page.Token.DetectedBreak.Type.WIDE_SPACE: ( + TextAnnotation.DetectedBreak.BreakType.LINE_BREAK + ), + Document.Page.Token.DetectedBreak.Type.HYPHEN: ( + TextAnnotation.DetectedBreak.BreakType.HYPHEN + ), + } +) + + +ElementWithLayout = Union[ + Document.Page.Paragraph, + Document.Page, + Document.Page.Token, + Document.Page.Block, + Document.Page.Symbol, +] + + +@dataclasses.dataclass +class PageInfo: + page: Document.Page + text: str = "" + page_idx: int = 0 + block_idx: int = 0 + paragraph_idx: int = 0 + token_idx: int = 0 + symbol_idx: int = 0 + + +def _get_text_anchor_substring(text: str, text_anchor: Document.TextAnchor) -> str: + """Gets text corresponding to the TextAnchor.""" + return "".join( + [ + text[text_segment.start_index : text_segment.end_index] + for text_segment in text_anchor.text_segments + ] + ) + + +def _has_text_segment(layout: Document.Page.Layout) -> bool: + """Checks if the layout has not empty text_segments.""" + if layout.text_anchor.text_segments: + return True + return False + + +def _convert_common_info( + src: ElementWithLayout, dest: ElementWithLayout, page_info: PageInfo +): + """Copy common data such as languages and bounding poly from src to dest. + + Args: + src: + source message that contains common data such as languages and bounding + poly. + dest: + destination message that common data converts to. + page_info (PageInfo): + provides dimension information to help convert between vertices + and normalized vertices. + """ + dest.confidence = src.layout.confidence + for language in src.detected_languages: + dest.property.detected_languages.append( + TextAnnotation.DetectedLanguage( + language_code=language.language_code, confidence=language.confidence + ) + ) + + if src.layout.bounding_poly.vertices: + for scr_vertex in src.layout.bounding_poly.vertices: + dest.bounding_box.vertices.append( + {"x": int(scr_vertex.x), "y": int(scr_vertex.y)} + ) + else: + for scr_normalized_vertex in src.layout.bounding_poly.normalized_vertices: + dest.bounding_box.vertices.append( + { + "x": int(scr_normalized_vertex.x * page_info.page.dimension.width), + "y": int(scr_normalized_vertex.y * page_info.page.dimension.height), + } + ) + for scr_normalized_vertex in src.layout.bounding_poly.normalized_vertices: + dest.bounding_box.normalized_vertices.append( + {"x": scr_normalized_vertex.x, "y": scr_normalized_vertex.y} + ) + + +def _is_layout_included( + inner: Document.Page.Layout, + outer: Document.Page.Layout, +) -> bool: + """Checks if the inner layout is within the scope of the outer layout.""" + if not _has_text_segment(inner) or not _has_text_segment(outer): + return False + return ( + inner.text_anchor.text_segments[0].start_index + >= outer.text_anchor.text_segments[0].start_index + and inner.text_anchor.text_segments[0].end_index + <= outer.text_anchor.text_segments[0].end_index + ) + + +def _convert_document_symbol( + page_info: PageInfo, + break_type: Document.Page.Token.DetectedBreak.Type, +) -> List[Symbol]: + """Converts document symbols to vision symbols. + + Args: + page_info (PageInfo): + Current page information, including document page to be converted, + its text, and the position of reading cursor. + break_type (Document.Page.Token.DetectedBreak.Type): + The break type of the current word that needs to be added to the + last vision symbol. + + Returns: + List[Symbol]: + Symbols filled with OCR data that are within + current document token. + """ + vision_symbols = [] + doc_symbols = page_info.page.symbols + cur_doc_token = page_info.page.tokens[page_info.token_idx] + while page_info.symbol_idx < len(doc_symbols) and _is_layout_included( + doc_symbols[page_info.symbol_idx].layout, cur_doc_token.layout + ): + vision_symbols.append( + Symbol( + text=_get_text_anchor_substring( + page_info.text, + doc_symbols[page_info.symbol_idx].layout.text_anchor, + ) + ) + ) + _convert_common_info( + doc_symbols[page_info.symbol_idx], vision_symbols[-1], page_info + ) + + page_info.symbol_idx += 1 + # Add the break_type to the last symbol. + if ( + vision_symbols + and break_type != Document.Page.Token.DetectedBreak.Type.TYPE_UNSPECIFIED + ): + vision_symbols[-1].property.detected_break.type = ( + _BREAK_TYPE_MAP[break_type] + if break_type in _BREAK_TYPE_MAP + else TextAnnotation.DetectedBreak.BreakType.UNKNOWN + ) + return vision_symbols + + +def _convert_document_token( + page_info: PageInfo, +) -> List[Word]: + """Converts document tokens to vision words. + + Args: + page_info (PageInfo): + Current page information, including document page to be converted, + its text, and the position of reading cursor. + + Returns: + List[Word]: + Word filled with OCR data that are within + current document paragraph. + """ + vision_words = [] + doc_tokens = page_info.page.tokens + cur_doc_paragraph = page_info.page.paragraphs[page_info.paragraph_idx] + while page_info.token_idx < len(doc_tokens) and _is_layout_included( + doc_tokens[page_info.token_idx].layout, cur_doc_paragraph.layout + ): + doc_break_type = doc_tokens[page_info.token_idx].detected_break.type + vision_break_type = ( + _BREAK_TYPE_MAP[doc_break_type] + if doc_break_type in _BREAK_TYPE_MAP + else TextAnnotation.DetectedBreak.BreakType.UNKNOWN + ) + vision_words.append( + Word( + symbols=_convert_document_symbol(page_info, doc_break_type), + property=TextAnnotation.TextProperty( + detected_break=TextAnnotation.DetectedBreak(type=vision_break_type) + ), + ) + ) + _convert_common_info( + doc_tokens[page_info.token_idx], vision_words[-1], page_info + ) + page_info.token_idx += 1 + return vision_words + + +def _convert_document_paragraph( + page_info: PageInfo, +) -> List[Paragraph]: + """Converts document paragraphs to vision paragraphs. + + Args: + page_info (PageInfo): + Current page information, including document page to be converted, + its text, and the position of reading cursor. + + Returns: + List[Paragraph]: + Paragraph filled with OCR data that are within + current document block. + """ + vision_paragraphs = [] + doc_paragraphs = page_info.page.paragraphs + cur_doc_block = page_info.page.blocks[page_info.block_idx] + while page_info.paragraph_idx < len(doc_paragraphs) and _is_layout_included( + doc_paragraphs[page_info.paragraph_idx].layout, cur_doc_block.layout + ): + vision_paragraphs.append(Paragraph(words=_convert_document_token(page_info))) + _convert_common_info( + doc_paragraphs[page_info.paragraph_idx], + vision_paragraphs[-1], + page_info, + ) + + page_info.paragraph_idx += 1 + return vision_paragraphs + + +def _convert_document_block( + page_info: PageInfo, +) -> List[Block]: + """Converts document blocks to vision blocks. + + Args: + page_info (PageInfo): Current page information, including document page to be converted, + its text, and the position of reading cursor. + + Returns: + List[Block]: + Block filled with OCR data that are within the + current document page. + """ + vision_blocks = [] + while page_info.block_idx < len(page_info.page.blocks): + vision_blocks.append( + Block( + block_type=Block.BlockType.TEXT, + paragraphs=_convert_document_paragraph(page_info), + ) + ) + _convert_common_info( + page_info.page.blocks[page_info.block_idx], vision_blocks[-1], page_info + ) + + page_info.block_idx += 1 + return vision_blocks + + +def _convert_document_page( + page_info: PageInfo, +) -> TextAnnotation: + """Extracts OCR related data in `page` and converts it to TextAnnotation. + + Args: + page_info (PageInfo): Current page information, including document page to be converted, + its text, and the position of reading cursor. + + Returns: + TextAnnotation: + Proto that only contains one page OCR data. + """ + detected_languages = [] + for language in page_info.page.detected_languages: + detected_languages.append( + vision.TextAnnotation.DetectedLanguage( + language_code=language.language_code, confidence=language.confidence + ) + ) + + text_property = TextAnnotation.TextProperty(detected_languages=detected_languages) + + page = Page( + width=int(page_info.page.dimension.width), + height=int(page_info.page.dimension.height), + confidence=page_info.page.layout.confidence, + blocks=_convert_document_block(page_info), + property=text_property, + ) + + text_annotation = TextAnnotation() + text_annotation.pages = [page] + + return text_annotation diff --git a/google/cloud/documentai_toolbox/wrappers/document.py b/google/cloud/documentai_toolbox/wrappers/document.py index 8a834c75..5a9affbf 100644 --- a/google/cloud/documentai_toolbox/wrappers/document.py +++ b/google/cloud/documentai_toolbox/wrappers/document.py @@ -30,6 +30,12 @@ from google.cloud.documentai_toolbox.wrappers.page import Page from google.cloud.documentai_toolbox.wrappers.page import FormField from google.cloud.documentai_toolbox.wrappers.entity import Entity +from google.cloud.documentai_toolbox.converters.converters import ( + _convert_to_vision_annotate_file_response, +) + +from google.cloud.vision import AnnotateFileResponse + from pikepdf import Pdf @@ -157,6 +163,17 @@ def _get_shards(gcs_bucket_name: str, gcs_prefix: str) -> List[documentai.Docume return shards +def _text_from_shards(shards: List[documentai.Document]) -> str: + total_text = "" + for shard in shards: + if total_text == "": + total_text = shard.text + elif total_text != shard.text: + total_text += shard.text + + return total_text + + def print_gcs_document_tree(gcs_bucket_name: str, gcs_prefix: str) -> None: r"""Prints a tree of filenames in Cloud Storage folder. @@ -242,10 +259,12 @@ class Document: pages: List[Page] = dataclasses.field(init=False, repr=False) entities: List[Entity] = dataclasses.field(init=False, repr=False) + text: str = dataclasses.field(init=False, repr=False) def __post_init__(self): self.pages = _pages_from_shards(shards=self.shards) self.entities = _entities_from_shards(shards=self.shards) + self.text = _text_from_shards(shards=self.shards) @classmethod def from_document_path( @@ -474,3 +493,13 @@ def split_pdf(self, pdf_path: str, output_path: str) -> List[str]: ) output_files.append(output_filename) return output_files + + def convert_document_to_annotate_file_response(self) -> AnnotateFileResponse: + """Convert OCR data from Document proto to AnnotateFileResponse proto (Vision API). + + Args: + None. + Returns: + AnnotateFileResponse proto with a TextAnnotation per page. + """ + return _convert_to_vision_annotate_file_response(self.text, self.pages) diff --git a/samples/snippets/convert_document_to_vision_sample.py b/samples/snippets/convert_document_to_vision_sample.py new file mode 100644 index 00000000..0ecb5755 --- /dev/null +++ b/samples/snippets/convert_document_to_vision_sample.py @@ -0,0 +1,46 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + + +# [START documentai_toolbox_document_to_vision] + +from google.cloud.documentai_toolbox import document + +# TODO(developer): Uncomment these variables before running the sample. +# Given a document.proto or sharded document.proto in path gs://bucket/path/to/folder +# gcs_bucket_name = "bucket" +# gcs_prefix = "path/to/folder" + + +def convert_document_to_vision_sample( + gcs_bucket_name: str, + gcs_prefix: str, +) -> None: + wrapped_document = document.Document.from_gcs( + gcs_bucket_name=gcs_bucket_name, gcs_prefix=gcs_prefix + ) + + # Converting wrapped_document to vision AnnotateFileResponse + annotate_file_response = ( + wrapped_document.convert_document_to_annotate_file_response() + ) + + print("Document converted to AnnotateFileResponse!") + print( + f"Number of Pages : {len(annotate_file_response.responses[0].full_text_annotation.pages)}" + ) + + +# [END documentai_toolbox_document_to_vision] diff --git a/samples/snippets/test_convert_document_to_vision_sample.py b/samples/snippets/test_convert_document_to_vision_sample.py new file mode 100644 index 00000000..0f782fc4 --- /dev/null +++ b/samples/snippets/test_convert_document_to_vision_sample.py @@ -0,0 +1,34 @@ +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +import os + +import pytest +from samples.snippets import convert_document_to_vision_sample + +location = "us" +project_id = os.environ["GOOGLE_CLOUD_PROJECT"] +gcs_bucket_name = "documentai_toolbox_samples" +gcs_input_uri = "output/123456789/0" + + +def test_quickstart_sample(capsys: pytest.CaptureFixture) -> None: + convert_document_to_vision_sample.convert_document_to_vision_sample( + gcs_bucket_name=gcs_bucket_name, gcs_prefix=gcs_input_uri + ) + out, _ = capsys.readouterr() + + assert "Document converted to AnnotateFileResponse!" in out + assert "Number of Pages : 1" in out diff --git a/setup.py b/setup.py index 3c944367..9791bdf9 100644 --- a/setup.py +++ b/setup.py @@ -52,8 +52,10 @@ "google-cloud-bigquery >= 3.5.0, < 4.0.0dev", "google-cloud-documentai >= 1.2.1, < 3.0.0dev", "google-cloud-storage >= 1.31.0, < 3.0.0dev", + "google-cloud-vision >= 2.7.0, < 4.0.0dev ", "numpy >= 1.18.1", "pikepdf >= 6.2.9, < 8.0.0", + "immutabledict >= 2.0.0, < 3.0.0dev", ), python_requires=">=3.7", classifiers=[ diff --git a/tests/unit/test_document.py b/tests/unit/test_document.py index 8d7e2746..9e0bd86b 100644 --- a/tests/unit/test_document.py +++ b/tests/unit/test_document.py @@ -30,6 +30,8 @@ from google.cloud import documentai from google.cloud import storage +from google.cloud.vision import AnnotateFileResponse + def get_bytes(file_name): result = [] @@ -68,6 +70,23 @@ def get_bytes_splitter_mock(): yield byte_factory +@mock.patch("google.cloud.documentai_toolbox.wrappers.document.storage") +def test_get_bytes(mock_storage): + client = mock_storage.Client.return_value + mock_bucket = mock.Mock() + client.Bucket.return_value = mock_bucket + mock_blob1 = mock.Mock(name=[]) + mock_blob1.name.ends_with.return_value = True + mock_blob1.download_as_bytes.return_value = ( + "gs://test-directory/1/test-annotations.json" + ) + client.list_blobs.return_value = [mock_blob1] + + actual = document._get_bytes(gcs_bucket_name="test-directory", gcs_prefix="1") + + assert actual == ["gs://test-directory/1/test-annotations.json"] + + def test_get_shards_with_gcs_uri_contains_file_type(): with pytest.raises(ValueError, match="gcs_prefix cannot contain file types"): document._get_shards( @@ -427,3 +446,13 @@ def test_split_pdf(mock_Pdf, get_bytes_splitter_mock): "procurement_multi_document_pg5_restaurant_statement.pdf", "procurement_multi_document_pg6-7_other.pdf", ] + + +def test_convert_document_to_annotate_file_response(): + doc = document.Document.from_document_path( + document_path="tests/unit/resources/0/toolbox_invoice_test-0.json" + ) + + actual = doc.convert_document_to_annotate_file_response() + + assert actual != AnnotateFileResponse() diff --git a/tests/unit/test_vision_helper.py b/tests/unit/test_vision_helper.py new file mode 100644 index 00000000..82ce258e --- /dev/null +++ b/tests/unit/test_vision_helper.py @@ -0,0 +1,338 @@ +# -*- coding: utf-8 -*- +# Copyright 2023 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from google.cloud import documentai +from google.cloud.documentai_toolbox.converters import vision_helpers +from google.cloud.vision import Paragraph + + +def test_get_text_anchor_substring(): + text = "Testing this function" + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + actual = vision_helpers._get_text_anchor_substring( + text=text, text_anchor=text_anchor + ) + + assert actual == "Testing" + + +def test_has_text_segment_true(): + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + layout = documentai.Document.Page.Layout(text_anchor=text_anchor) + + assert vision_helpers._has_text_segment(layout) + + +def test_has_text_segment_false(): + layout = documentai.Document.Page.Layout(text_anchor=None) + + assert not vision_helpers._has_text_segment(layout) + + +def test_convert_common_info_src_with_vertices(): + test_src = documentai.Document.Page.Paragraph() + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + + vertices = [ + documentai.Vertex(x=1, y=1), + documentai.Vertex(x=2, y=1), + documentai.Vertex(x=2, y=2), + documentai.Vertex(x=1, y=2), + ] + + bounding_poly = documentai.BoundingPoly(vertices=vertices) + + detected_language = documentai.Document.Page.DetectedLanguage( + language_code="en-US", confidence=0.99 + ) + + layout = documentai.Document.Page.Layout( + text_anchor=text_anchor, + bounding_poly=bounding_poly, + confidence=0.98, + ) + test_src.layout = layout + test_src.detected_languages = [detected_language] + + page = documentai.Document.Page() + page.dimension.width = 1000 + page.dimension.height = 2500 + + test_dest = Paragraph() + + test_page_info = vision_helpers.PageInfo(page=page) + + vision_helpers._convert_common_info( + src=test_src, dest=test_dest, page_info=test_page_info + ) + + assert test_dest.bounding_box != Paragraph().bounding_box + + +def test_convert_common_info_src_with_normalized_vertices(): + test_src = documentai.Document.Page.Paragraph() + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + + normalized = [ + documentai.NormalizedVertex(x=0.1, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.2), + documentai.NormalizedVertex(x=0.1, y=0.2), + ] + + bounding_poly = documentai.BoundingPoly(normalized_vertices=normalized) + + layout = documentai.Document.Page.Layout( + text_anchor=text_anchor, bounding_poly=bounding_poly, confidence=0.98 + ) + test_src.layout = layout + + page = documentai.Document.Page() + page.dimension.width = 1000 + page.dimension.height = 2500 + + test_dest = Paragraph() + + test_page_info = vision_helpers.PageInfo(page=page) + + vision_helpers._convert_common_info( + src=test_src, dest=test_dest, page_info=test_page_info + ) + + assert test_dest.bounding_box != Paragraph().bounding_box + + +def test_is_layout_included_return_true(): + in_text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=1, end_index=4 + ) + in_text_anchor = documentai.Document.TextAnchor(text_segments=[in_text_segment]) + in_layout = documentai.Document.Page.Layout(text_anchor=in_text_anchor) + + out_text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + out_text_anchor = documentai.Document.TextAnchor(text_segments=[out_text_segment]) + out_layout = documentai.Document.Page.Layout(text_anchor=out_text_anchor) + + assert vision_helpers._is_layout_included(inner=in_layout, outer=out_layout) + + +def test_is_layout_included_return_false(): + in_layout = documentai.Document.Page.Layout(text_anchor=None) + out_layout = documentai.Document.Page.Layout(text_anchor=None) + + assert not vision_helpers._is_layout_included(inner=in_layout, outer=out_layout) + + +def test_convert_document_symbol(): + text = "Testing this function" + + vertices = [ + documentai.Vertex(x=1, y=1), + documentai.Vertex(x=2, y=1), + documentai.Vertex(x=2, y=2), + documentai.Vertex(x=1, y=2), + ] + normalized = [ + documentai.NormalizedVertex(x=0.1, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.2), + documentai.NormalizedVertex(x=0.1, y=0.2), + ] + + bounding_poly = documentai.BoundingPoly( + vertices=vertices, normalized_vertices=normalized + ) + + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + layout = documentai.Document.Page.Layout( + text_anchor=text_anchor, bounding_poly=bounding_poly + ) + + symbol = documentai.Document.Page.Symbol(layout=layout) + token = documentai.Document.Page.Token(layout=layout) + + page = documentai.Document.Page(symbols=[symbol], tokens=[token]) + page.dimension.width = 1000 + page.dimension.height = 2500 + + page_info = vision_helpers.PageInfo(page=page, text=text) + + actual = vision_helpers._convert_document_symbol( + page_info=page_info, + break_type=documentai.Document.Page.Token.DetectedBreak.Type.SPACE, + ) + + assert len(actual) == 1 + assert actual[0].text == "Testing" + assert len(actual[0].bounding_box.vertices) == 4 + assert len(actual[0].bounding_box.normalized_vertices) == 4 + + +def test_convert_document_token(): + vertices = [ + documentai.Vertex(x=1, y=1), + documentai.Vertex(x=2, y=1), + documentai.Vertex(x=2, y=2), + documentai.Vertex(x=1, y=2), + ] + normalized = [ + documentai.NormalizedVertex(x=0.1, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.2), + documentai.NormalizedVertex(x=0.1, y=0.2), + ] + + bounding_poly = documentai.BoundingPoly( + vertices=vertices, normalized_vertices=normalized + ) + + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + layout = documentai.Document.Page.Layout( + text_anchor=text_anchor, bounding_poly=bounding_poly + ) + + paragraph = documentai.Document.Page.Paragraph(layout=layout) + token = documentai.Document.Page.Token(layout=layout) + + page = documentai.Document.Page(paragraphs=[paragraph], tokens=[token]) + page.dimension.width = 1000 + page.dimension.height = 2500 + + page_info = vision_helpers.PageInfo(page=page) + + actual = vision_helpers._convert_document_token(page_info=page_info) + + assert len(actual) == 1 + assert len(actual[0].bounding_box.vertices) == 4 + assert len(actual[0].bounding_box.normalized_vertices) == 4 + + +def test_convert_document_paragraph(): + vertices = [ + documentai.Vertex(x=1, y=1), + documentai.Vertex(x=2, y=1), + documentai.Vertex(x=2, y=2), + documentai.Vertex(x=1, y=2), + ] + normalized = [ + documentai.NormalizedVertex(x=0.1, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.2), + documentai.NormalizedVertex(x=0.1, y=0.2), + ] + + bounding_poly = documentai.BoundingPoly( + vertices=vertices, normalized_vertices=normalized + ) + + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + layout = documentai.Document.Page.Layout( + text_anchor=text_anchor, bounding_poly=bounding_poly + ) + + paragraph = documentai.Document.Page.Paragraph(layout=layout) + block = documentai.Document.Page.Block(layout=layout) + + page = documentai.Document.Page(paragraphs=[paragraph], blocks=[block]) + page.dimension.width = 1000 + page.dimension.height = 2500 + + page_info = vision_helpers.PageInfo(page=page) + + actual = vision_helpers._convert_document_paragraph(page_info=page_info) + + assert len(actual) == 1 + assert len(actual[0].bounding_box.vertices) == 4 + assert len(actual[0].bounding_box.normalized_vertices) == 4 + + +def test_convert_document_block(): + vertices = [ + documentai.Vertex(x=1, y=1), + documentai.Vertex(x=2, y=1), + documentai.Vertex(x=2, y=2), + documentai.Vertex(x=1, y=2), + ] + normalized = [ + documentai.NormalizedVertex(x=0.1, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.1), + documentai.NormalizedVertex(x=0.2, y=0.2), + documentai.NormalizedVertex(x=0.1, y=0.2), + ] + + bounding_poly = documentai.BoundingPoly( + vertices=vertices, normalized_vertices=normalized + ) + + text_segment = documentai.Document.TextAnchor.TextSegment( + start_index=0, end_index=7 + ) + text_anchor = documentai.Document.TextAnchor(text_segments=[text_segment]) + layout = documentai.Document.Page.Layout( + text_anchor=text_anchor, bounding_poly=bounding_poly + ) + + block = documentai.Document.Page.Block(layout=layout) + + page = documentai.Document.Page(blocks=[block]) + page.dimension.width = 1000 + page.dimension.height = 2500 + + page_info = vision_helpers.PageInfo(page=page) + + actual = vision_helpers._convert_document_block(page_info=page_info) + + assert len(actual) == 1 + assert len(actual[0].bounding_box.vertices) == 4 + assert len(actual[0].bounding_box.normalized_vertices) == 4 + assert actual[0].block_type == 1 + + +def test_convert_document_page(): + detected_language = documentai.Document.Page.DetectedLanguage( + language_code="en-us", confidence=1.0 + ) + page = documentai.Document.Page(detected_languages=[detected_language]) + + page_info = vision_helpers.PageInfo(page=page) + + actual = vision_helpers._convert_document_page(page_info=page_info) + + assert actual.pages[0].property.detected_languages[0].language_code == "en-us" + assert actual.pages[0].property.detected_languages[0].confidence == 1