diff --git a/google/cloud/documentai_toolbox/__init__.py b/google/cloud/documentai_toolbox/__init__.py index 9f475e57..6c93a1c8 100644 --- a/google/cloud/documentai_toolbox/__init__.py +++ b/google/cloud/documentai_toolbox/__init__.py @@ -25,11 +25,11 @@ ) from .converters import ( - converters, + converter, ) from .utilities import ( utilities, ) -__all__ = (document, page, entity, converters, utilities) +__all__ = (document, page, entity, converter, utilities) diff --git a/google/cloud/documentai_toolbox/converters/config/__init__.py b/google/cloud/documentai_toolbox/converters/config/__init__.py new file mode 100644 index 00000000..89a37dc9 --- /dev/null +++ b/google/cloud/documentai_toolbox/converters/config/__init__.py @@ -0,0 +1,15 @@ +# -*- 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. +# diff --git a/google/cloud/documentai_toolbox/converters/config/bbox_conversion.py b/google/cloud/documentai_toolbox/converters/config/bbox_conversion.py new file mode 100644 index 00000000..3b9bab01 --- /dev/null +++ b/google/cloud/documentai_toolbox/converters/config/bbox_conversion.py @@ -0,0 +1,296 @@ +# -*- 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 typing import Callable +from intervaltree import intervaltree + +from google.cloud import documentai +from google.cloud.documentai_v1.types import geometry + + +def _midpoint_in_bpoly( + box_a: geometry.BoundingPoly, box_b: geometry.BoundingPoly +) -> bool: + """Returns whether the midpoint in box_a is inside box_b.""" + + # Calculate the midpoint of box_a. + mid_x_a = (_get_norm_x_max(box_a) + _get_norm_x_min(box_a)) / 2.0 + mid_y_a = (_get_norm_y_max(box_a) + _get_norm_y_min(box_a)) / 2.0 + + max_x_b = _get_norm_x_max(box_b) + min_x_b = _get_norm_x_min(box_b) + max_y_b = _get_norm_y_max(box_b) + min_y_b = _get_norm_y_min(box_b) + + return min_x_b < mid_x_a < max_x_b and min_y_b < mid_y_a < max_y_b + + +def _merge_text_anchors( + text_anchor_1: documentai.Document.TextAnchor, + text_anchor_2: documentai.Document.TextAnchor, +) -> documentai.Document.TextAnchor: + """Merges two TextAnchor objects into one ascending sorted TextAnchor.""" + merged_text_anchor = documentai.Document.TextAnchor() + intervals = [] + for text_segment in text_anchor_1.text_segments: + intervals.append( + intervaltree.Interval(text_segment.start_index, text_segment.end_index) + ) + for text_segment in text_anchor_2.text_segments: + intervals.append( + intervaltree.Interval(text_segment.start_index, text_segment.end_index) + ) + + interval_tree = intervaltree.IntervalTree(intervals) + interval_tree.merge_overlaps(strict=False) + ts = [] + for iv in sorted(interval_tree): + ts.append( + documentai.Document.TextAnchor.TextSegment( + start_index=iv.begin, end_index=iv.end + ) + ) + + merged_text_anchor.text_segments = ts + return merged_text_anchor + + +def _get_text_anchor_in_bbox( + bbox: documentai.BoundingPoly, + page: documentai.Document.Page, + token_in_bounding_box_function: Callable[ + [documentai.BoundingPoly, documentai.BoundingPoly], bool + ] = _midpoint_in_bpoly, +) -> documentai.Document.TextAnchor: + """Gets mergedTextAnchor of Tokens in `page` that fall inside the `bbox`.""" + + text_anchor = documentai.Document.TextAnchor() + for token in page.tokens: + if token_in_bounding_box_function(token.layout.bounding_poly, bbox): + text_anchor = _merge_text_anchors(text_anchor, token.layout.text_anchor) + return text_anchor + + +def _get_norm_x_max(bbox: geometry.BoundingPoly) -> float: + return max([vertex.x for vertex in bbox.normalized_vertices]) + + +def _get_norm_x_min(bbox: geometry.BoundingPoly) -> float: + return min([vertex.x for vertex in bbox.normalized_vertices]) + + +def _get_norm_y_max(bbox: geometry.BoundingPoly) -> float: + return max([vertex.y for vertex in bbox.normalized_vertices]) + + +def _get_norm_y_min(bbox: geometry.BoundingPoly) -> float: + return min([vertex.y for vertex in bbox.normalized_vertices]) + + +def _normalize_coordinates(x, y) -> float: + return round(float(x / y), 9) + + +def _convert_to_pixels(x: float, conversion_rate: float) -> float: + return x * conversion_rate + + +def _convert_bbox_units( + coordinate, input_bbox_units, width=None, height=None, multiplier=1 +) -> float: + r"""Returns a converted coordinate. + + Args: + coordinate (float): + Required.The coordinate from document.proto + input_bbox_units (str): + Required. The bounding box units. + width (float): + Optional. + height (float): + Optional. + multiplier (float): + Optional. + + Returns: + float: + A converted coordinate. + + """ + final_coordinate = coordinate + if input_bbox_units != "normalized": + if input_bbox_units == "pxl": + if width is None: + final_coordinate = _normalize_coordinates(coordinate, height) + else: + final_coordinate = _normalize_coordinates(coordinate, width) + if input_bbox_units == "inch": + x = _convert_to_pixels(coordinate, 96) + if width is None: + final_coordinate = _normalize_coordinates(x, height) + else: + final_coordinate = _normalize_coordinates(x, width) + if input_bbox_units == "cm": + x = _convert_to_pixels(coordinate, 37.795) + if width is None: + final_coordinate = _normalize_coordinates(x, height) + else: + final_coordinate = _normalize_coordinates(x, width) + + return final_coordinate * multiplier + + +def _get_multiplier( + docproto_coordinate: float, external_coordinate: float, input_bbox_units: str +) -> float: + r"""Returns a multiplier to use when converting bounding boxes. + + Args: + docproto_coordinate (float): + Required.The coordinate from document.proto + external_coordinate (float): + Required.The coordinate from external annotations. + input_bbox_units (str): + Required. The bounding box units. + Returns: + float: + multiplier to use when converting bounding boxes. + + """ + if input_bbox_units == "inch": + converted = _convert_to_pixels(external_coordinate, 96) + return docproto_coordinate / converted + elif input_bbox_units == "cm": + converted = _convert_to_pixels(external_coordinate, 37.795) + return docproto_coordinate / converted + else: + return docproto_coordinate / external_coordinate + + +def _convert_bbox_to_docproto_bbox(block) -> geometry.BoundingPoly: + r"""Returns a converted bounding box from Block. + + Args: + block (Block): + Required. + Returns: + geometry.BoundingPoly: + A geometry.BoundingPoly from bounding box. + + """ + merged_bbox = geometry.BoundingPoly() + x_multiplier = 1 + y_multiplier = 1 + coordinates = [] + nv = [] + + # _convert_bbox_units should check if external_bbox is list or not + coordinates_object = block.bounding_box + if coordinates_object == []: + return coordinates_object + + if block.page_width and block.page_height: + x_multiplier = _get_multiplier( + docproto_coordinate=block.docproto_width, + external_coordinate=block.page_width, + input_bbox_units=block.bounding_unit, + ) + y_multiplier = _get_multiplier( + docproto_coordinate=block.docproto_height, + external_coordinate=block.page_height, + input_bbox_units=block.bounding_unit, + ) + + if block.bounding_type == "1": + # Type 1 : bounding box has 4 (x,y) coordinates + + if type(block.bounding_box) == list: + for coordinate in coordinates_object: + x = _convert_bbox_units( + coordinate[f"{block.bounding_x}"], + input_bbox_units=block.bounding_unit, + width=block.docproto_width, + multiplier=x_multiplier, + ) + y = _convert_bbox_units( + coordinate[f"{block.bounding_y}"], + input_bbox_units=block.bounding_unit, + height=block.docproto_height, + multiplier=y_multiplier, + ) + + coordinates.append({"x": x, "y": y}) + + coordinates_object = coordinates + + elif block.bounding_type == "2": + # Type 2 : bounding box has 1 (x,y) coordinates for the top left corner + # and (width, height) + original_x = coordinates_object[f"{block.bounding_x}"] + original_y = coordinates_object[f"{block.bounding_y}"] + + x = _convert_bbox_units( + original_x, + input_bbox_units=block.bounding_unit, + width=block.page_width, + multiplier=x_multiplier, + ) + y = _convert_bbox_units( + original_y, + input_bbox_units=block.bounding_unit, + width=block.page_height, + multiplier=y_multiplier, + ) + + # x_min_y_min + coordinates.append({"x": x, "y": y}) + # x_max_y_min + coordinates.append({"x": (x + block.bounding_width), "y": y}) + # x_max_y_max + coordinates.append( + {"x": (x + block.bounding_width), "y": (y + block.bounding_height)} + ) + # x_min_y_max + coordinates.append({"x": x, "y": (y + block.bounding_height)}) + + coordinates_object = coordinates + elif block.bounding_type == "3": + # Type 2 : bounding box has 1 (x,y) coordinates for the top left corner + # and (width, height) + for idx in range(0, len(block.bounding_box), 2): + x = _convert_bbox_units( + block.bounding_box[idx], + input_bbox_units=block.bounding_unit, + width=block.docproto_width, + multiplier=x_multiplier, + ) + y = _convert_bbox_units( + block.bounding_box[idx + 1], + input_bbox_units=block.bounding_unit, + width=block.docproto_height, + multiplier=y_multiplier, + ) + + coordinates.append({"x": x, "y": y}) + + coordinates_object = coordinates + + for coordinates in coordinates_object: + nv.append(documentai.NormalizedVertex(x=coordinates["x"], y=coordinates["y"])) + + merged_bbox.normalized_vertices = nv + + return merged_bbox diff --git a/google/cloud/documentai_toolbox/converters/config/blocks.py b/google/cloud/documentai_toolbox/converters/config/blocks.py new file mode 100644 index 00000000..1dd89bf3 --- /dev/null +++ b/google/cloud/documentai_toolbox/converters/config/blocks.py @@ -0,0 +1,279 @@ +# -*- 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. +# + +import dataclasses +from typing import List +import json +from types import SimpleNamespace + +from google.cloud import documentai + + +@dataclasses.dataclass +class Block: + r"""Represents a Block from OCR data. + + Attributes: + bounding_box (str): + Required. + block_references: + Optional. + block_id: + Optional. + confidence: + Optional. + type_: + Required. + text: + Required. + page_number: + Optional. + """ + bounding_box: dataclasses.field(init=True, repr=False, default=None) + block_references: dataclasses.field(init=False, repr=False, default=None) + block_id: dataclasses.field(init=False, repr=False, default=None) + confidence: dataclasses.field(init=False, repr=False, default=None) + type_: dataclasses.field(init=True, repr=False, default=None) + text: dataclasses.field(init=True, repr=False, default=None) + page_number: dataclasses.field(init=False, repr=False, default=None) + page_width: dataclasses.field(init=False, repr=False, default=None) + page_height: dataclasses.field(init=False, repr=False, default=None) + bounding_width: dataclasses.field(init=False, repr=False, default=None) + bounding_height: dataclasses.field(init=False, repr=False, default=None) + bounding_type: dataclasses.field(init=False, repr=False, default=None) + bounding_unit: dataclasses.field(init=False, repr=False, default=None) + bounding_x: dataclasses.field(init=False, repr=False, default=None) + bounding_y: dataclasses.field(init=False, repr=False, default=None) + docproto_width: dataclasses.field(init=False, repr=False, default=None) + docproto_height: dataclasses.field(init=False, repr=False, default=None) + + @classmethod + def create( + self, + type_, + text, + bounding_box=None, + block_references=None, + block_id=None, + confidence=None, + page_number=None, + page_width=None, + page_height=None, + bounding_width=None, + bounding_height=None, + bounding_type=None, + bounding_unit=None, + bounding_x=None, + bounding_y=None, + docproto_width=None, + docproto_height=None, + ): + return Block( + bounding_box=bounding_box, + block_references=block_references, + block_id=block_id, + confidence=confidence, + type_=type_, + text=text, + page_number=page_number, + page_width=page_width, + page_height=page_height, + bounding_width=bounding_width, + bounding_height=bounding_height, + bounding_type=bounding_type, + bounding_unit=bounding_unit, + bounding_x=bounding_x, + bounding_y=bounding_y, + docproto_width=docproto_width, + docproto_height=docproto_height, + ) + + +def _get_target_object(json_data: any, target_object: str) -> SimpleNamespace: + r"""Returns SimpleNamespace of target_object. + + Args: + json_data (str): + Required. data from JSON.loads . + target_object (str): + Required. The path to the target object. + + Returns: + SimpleNamespace. + + """ + json_data_s = SimpleNamespace(**json_data) + + target_object_parts = target_object.split(".") + + if not hasattr(json_data_s, target_object_parts[0]): + return None + + current_object = json_data_s + for part in target_object_parts: + if type(current_object) == dict: + current_object = SimpleNamespace(**current_object) + elif type(current_object) == list and part.isnumeric(): + current_object = current_object[int(part)] + continue + current_object = getattr(current_object, part) + return current_object + + +def _load_blocks_from_schema( + input_data: bytes, input_config: bytes, base_docproto: documentai.Document +) -> List[Block]: + r"""Loads Blocks from original annotation data and provided config. + + Args: + input_data (bytes): + Required.The bytes of the annotated data. + input_config (bytes): + Required.The bytes of config data. + base_docproto (bytes): + Required. The bytes of the original pdf. + + Returns: + List[Block]: + From original annotation data and provided config. + + """ + objects = json.loads(input_data) + schema_json = json.loads(input_config, object_hook=lambda d: SimpleNamespace(**d)) + + entities = schema_json.entity_object + type_ = schema_json.entity.type_ + + mention_text = schema_json.entity.mention_text + + document_height = None + document_width = None + + id_ = schema_json.entity.id if hasattr(schema_json.entity, "id") else None + if hasattr(schema_json, "page"): + document_height = ( + schema_json.page.height if hasattr(schema_json.page, "height") else None + ) + document_width = ( + schema_json.page.width if hasattr(schema_json.page, "width") else None + ) + + confidence = ( + schema_json.entity.confidence + if hasattr(schema_json.entity, "confidence") + else None + ) + page_number = ( + schema_json.entity.page_number + if hasattr(schema_json.entity, "page_number") + else None + ) + normalized_vertices = ( + schema_json.entity.normalized_vertices.base + if hasattr(schema_json.entity.normalized_vertices, "base") + else None + ) + bounding_width = ( + schema_json.entity.normalized_vertices.width + if hasattr(schema_json.entity.normalized_vertices, "width") + else None + ) + bounding_height = ( + schema_json.entity.normalized_vertices.height + if hasattr(schema_json.entity.normalized_vertices, "height") + else None + ) + bounding_type = ( + schema_json.entity.normalized_vertices.type + if hasattr(schema_json.entity.normalized_vertices, "type") + else None + ) + bounding_unit = ( + schema_json.entity.normalized_vertices.unit + if hasattr(schema_json.entity.normalized_vertices, "unit") + else None + ) + bounding_x = ( + schema_json.entity.normalized_vertices.x + if hasattr(schema_json.entity.normalized_vertices, "x") + else None + ) + bounding_y = ( + schema_json.entity.normalized_vertices.y + if hasattr(schema_json.entity.normalized_vertices, "y") + else None + ) + + blocks = [] + ens = _get_target_object(objects, entities) + for i in ens: + entity = i + + block_text = "" + + if type_ == f"{entities}:self": + block_type = i + entity = _get_target_object(objects, f"{entities}.{i}") + else: + block_type = _get_target_object(entity, type_) + + if "||" in mention_text: + text_commands = mention_text.split("||") + for command in text_commands: + if command in entity: + block_text = _get_target_object(entity, command) + continue + else: + block_text = _get_target_object(entity, mention_text) + + b = Block.create( + type_=block_type, + text=block_text, + ) + + b.bounding_box = _get_target_object(entity, normalized_vertices) + + if id_: + b.id_ = _get_target_object(entity, id_) + if confidence: + b.confidence = _get_target_object(entity, confidence) + if page_number and page_number in entity: + b.page_number = _get_target_object(entity, page_number) + if bounding_width: + b.bounding_width = _get_target_object(b.bounding_box, bounding_width) + if bounding_height: + b.bounding_height = _get_target_object(b.bounding_box, bounding_height) + if document_height: + b.page_height = _get_target_object(objects, document_height) + if document_width: + b.page_width = _get_target_object(objects, document_width) + if bounding_type: + b.bounding_type = bounding_type + if bounding_unit: + b.bounding_unit = bounding_unit + if bounding_x: + b.bounding_x = bounding_x + if bounding_y: + b.bounding_y = bounding_y + + if b.page_number is None: + b.page_number = 0 + + b.docproto_width = base_docproto.pages[int(b.page_number)].dimension.width + b.docproto_height = base_docproto.pages[int(b.page_number)].dimension.height + + blocks.append(b) + return blocks diff --git a/google/cloud/documentai_toolbox/converters/config/converter_helpers.py b/google/cloud/documentai_toolbox/converters/config/converter_helpers.py new file mode 100644 index 00000000..b6788d71 --- /dev/null +++ b/google/cloud/documentai_toolbox/converters/config/converter_helpers.py @@ -0,0 +1,534 @@ +# -*- 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. +# + +import re +import time +from concurrent import futures +from typing import List, Tuple + +from google.cloud.documentai_toolbox.converters.config.bbox_conversion import ( + _convert_bbox_to_docproto_bbox, + _get_text_anchor_in_bbox, +) +from google.cloud.documentai_toolbox.converters.config.blocks import ( + Block, + _load_blocks_from_schema, +) + +from google.cloud.documentai_toolbox import document, constants +from google.cloud import documentai, storage + + +def _get_base_ocr( + project_id: str, location: str, processor_id: str, file_bytes: bytes, mime_type: str +) -> documentai.Document: + r"""Returns documentai.Document from OCR processor. + + Args: + project_id (str): + Required. + location (str): + Required. + processor_id (str): + Required. + file_bytes (bytes): + Required. The bytes of the original pdf. + mime_type (str): + Required. usually "application/pdf". + Returns: + documentai.Document: + A documentai.Document from OCR processor. + + """ + + client = documentai.DocumentProcessorServiceClient() + + name = client.processor_path(project_id, location, processor_id) + + # Load Binary Data into Document AI RawDocument Object + raw_document = documentai.RawDocument(content=file_bytes, mime_type=mime_type) + + # Configure the process request + request = documentai.ProcessRequest(name=name, raw_document=raw_document) + + result = client.process_document(request=request) + return result.document + + +def _get_entity_content( + blocks: List[Block], docproto: documentai.Document +) -> List[documentai.Document.Entity]: + r"""Returns a list of documentai.Document entities. + + Args: + blocks (List[Block]): + Required.List of blocks from original annotation. + docproto (documentai.Document): + Required.The ocr docproto. + Returns: + List[documentai.Document.Entity]: + A list of documentai.Document entities. + """ + entities = [] + entity_id = 0 + + for block in blocks: + + docai_entity = documentai.Document.Entity() + if block.confidence: + docai_entity.confidence = block.confidence + + docai_entity.type = block.type_ + docai_entity.mention_text = block.text + docai_entity.id = str(entity_id) + + entity_id += 1 + # Generates the text anchors from bounding boxes + if block.bounding_box: + # Converts external bounding box format to docproto bounding box + + b1 = _convert_bbox_to_docproto_bbox(block) + + if block.page_number: + docai_entity.text_anchor = _get_text_anchor_in_bbox( + b1, docproto.pages[int(block.page_number) - 1] + ) + else: + docai_entity.text_anchor = _get_text_anchor_in_bbox( + b1, docproto.pages[0] + ) + + docai_entity.text_anchor.content = block.text + + page_anchor = documentai.Document.PageAnchor() + page_ref = documentai.Document.PageAnchor.PageRef() + + page_ref.bounding_poly = b1 + + page_anchor.page_refs = [page_ref] + docai_entity.page_anchor = page_anchor + entities.append(docai_entity) + + return entities + + +def _convert_to_docproto_with_config( + annotated_bytes: bytes, + config_bytes: bytes, + document_bytes: bytes, + project_id: str, + location: str, + processor_id: str, + retry_number: int, + name: str = "", +) -> documentai.Document: + r"""Converts a single document to docproto. + + Args: + annotated_bytes (bytes): + Required.The bytes of the annotated data. + config_bytes (bytes): + Required.The bytes of config data. + document_bytes (bytes): + Required. The bytes of the original pdf. + project_id (str): + Required. + location (str): + Required. + processor_id (str): + Required. + retry_number (str): + Required. The number of seconds needed to wait if an error occured. + name (str): + Optional. Name of the document to be converted. This is used for logging. + + Returns: + documentai.Document: + documentai.Document object. + + TODO: Depending on input type you will need to modify load_blocks. + Depending on input format, if your annotated data is not separate from the base OCR data you will need to modify _get_entity_content + Depending on input BoundingBox, if the input BoundingBox object is like https://cloud.google.com/document-ai/docs/reference/rest/v1/Document#BoundingPoly then you will need to + modify _convert_bbox_to_docproto_bbox since the objects are different. + """ + try: + base_docproto = _get_base_ocr( + project_id=project_id, + location=location, + processor_id=processor_id, + file_bytes=document_bytes, + mime_type="application/pdf", + ) + + # Loads OCR data into Blocks + # blocks = load_blocks(ocr_object=doc_object) + blocks = _load_blocks_from_schema( + input_data=annotated_bytes, + input_config=config_bytes, + base_docproto=base_docproto, + ) + + # Gets List[documentai.Document.Entity] + entities = _get_entity_content(blocks=blocks, docproto=base_docproto) + + base_docproto.entities = entities + print("Converted : %s\r" % name, end="") + return base_docproto + + except Exception as e: + print(e) + print(f"Could Not Convert {name}\nretrying") + if retry_number == 6: + return None + else: + time.sleep(retry_number) + return _convert_to_docproto_with_config( + name=name, + annotated_bytes=annotated_bytes, + config_bytes=config_bytes, + document_bytes=document_bytes, + project_id=project_id, + location=location, + processor_id=processor_id, + retry_number=retry_number + 1, + ) + + +def _get_bytes( + bucket_name: str, + prefix: str, + annotation_file_prefix: str, + config_file_prefix: str, + config_path: str = None, +) -> List[bytes]: + r"""Downloads documents and returns them as bytes. + + Args: + bucket_name (str): + Required. The bucket name. + prefix (str): + Required. The prefix for the location of the output folder. + annotation_file_prefix (str): + Required. The prefix to search for annotation file. + config_file_prefix (str): + Required. The prefix to search for config file. + config_path (str): + Optional. The gcs path to a config file. This should be used when there is a single config file. + + Returns: + List[bytes]. + + """ + + storage_client = document._get_storage_client() + bucket = storage_client.bucket(bucket_name=bucket_name) + blobs = storage_client.list_blobs(bucket_or_name=bucket_name, prefix=prefix) + + metadata_blob = None + + try: + for blob in blobs: + if "DS_Store" in blob.name: + continue + if not blob.name.endswith("/"): + blob_name = blob.name + file_name = blob_name.split("/")[-1] + if annotation_file_prefix in file_name: + annotation_blob = blob + elif config_file_prefix in file_name: + metadata_blob = blob + elif "pdf" in file_name: + doc_blob = blob + + if metadata_blob and config_path: + metadata_blob = bucket.get_blob(config_path) + + print("Downloaded : %s\r" % prefix.split("/")[-1], end="") + return [ + annotation_blob.download_as_bytes(), + doc_blob.download_as_bytes(), + metadata_blob.download_as_bytes(), + prefix.split("/")[-1], + file_name.split(".")[0], + ] + except Exception as e: + raise e + + +def _upload_file( + bucket_name: str, + output_prefix: str, + file: str, +) -> None: + r"""Uploads the converted docproto to gcs. + + Args: + bucket_name (str): + Required. The bucket name. + output_prefix (str): + Required. The prefix for the location of the output folder. + file (str): + Required. The docproto file in string format. + + Returns: + None. + + """ + storage_client = document._get_storage_client() + bucket = storage_client.bucket(bucket_name) + blob = bucket.blob(output_prefix) + + print("Uploaded : %s\r" % output_prefix.split("/")[-1], end="") + blob.upload_from_string(file, content_type="application/json") + + +def _get_files( + blob_list: List[storage.blob.Blob], + input_bucket: str, + input_prefix: str, + config_path: str = None, +): + r"""Returns a list of Futures of documents as bytes. + + Args: + blob_list (List[storage.blob.Blob]): + Required. The list of Futures from _get_files. + input_bucket (str): + Required. The name of the input bucket. + input_prefix (str): + Required. The prefix for the location of the input folder. + config_path (str): + Required. The optional + Returns: + Tuple[dict, list, list]: + Converted document.proto, unique entity types and documents that were not converted. + + """ + download_pool = futures.ThreadPoolExecutor(10) + downloads = [] + prev = None + print("-------- Downloading Started --------") + for i, blob in enumerate(blob_list): + if "DS_Store" in blob.name: + continue + + file_path = blob.name.split("/") + file_path.pop() + doc_directory = file_path[-1] + file_path2 = "/".join(file_path) + if prev == doc_directory or f"{file_path2}/" == input_prefix: + continue + + download = download_pool.submit( + _get_bytes, + input_bucket, + file_path2, + "annotation", + "config", + config_path, + ) + downloads.append(download) + + prev = doc_directory + + return downloads + + +def _get_docproto_files( + f: List[futures.Future], + project_id: str, + location: str, + processor_id: str, +) -> Tuple[dict, list, list]: + r"""Returns converted document.proto, unique entity types and documents that were not converted. + + Args: + f (List[futures.Future]): + Required. The list of Futures from _get_files. + project_id (str): + Required. + location (str): + Required. + processor_id (str): + Required. + Returns: + Tuple[dict, list, list]: + Converted document.proto, unique entity types and documents that were not converted. + + """ + did_not_convert = [] + files = {} + unique_types = [] + for future in f: + blobs = future.result() + docproto = _convert_to_docproto_with_config( + annotated_bytes=blobs[0], + document_bytes=blobs[1], + config_bytes=blobs[2], + project_id=project_id, + location=location, + processor_id=processor_id, + retry_number=1, + name=blobs[3], + ) + + if docproto is None: + did_not_convert.append(f"{blobs[3]}") + continue + + for entity in docproto.entities: + if entity.type_ not in unique_types: + unique_types.append(entity.type_) + + files[blobs[3]] = str(documentai.Document.to_json(docproto)) + + return files, unique_types, did_not_convert + + +def _upload(files: dict, gcs_output_path: str) -> None: + r"""Upload converted document.proto to gcs location. + + Args: + files (dict): + Required. The document.proto files to upload. + gcs_output_path (str): + Required. The gcs path to the folder to upload the converted docproto documents to. + + Format: `gs://{bucket}/{optional_folder}` + Returns: + None. + + """ + match = re.match(r"gs://(.*?)/(.*)", gcs_output_path) + + if match is None: + raise ValueError("gcs_prefix does not match accepted format") + + output_bucket, output_prefix = match.groups() + + if output_prefix is None: + output_prefix = "/" + + file_check = re.match(constants.FILE_CHECK_REGEX, output_prefix) + + if file_check: + raise ValueError("gcs_prefix cannot contain file types") + + download_pool = futures.ThreadPoolExecutor(10) + uploads = [] + print("-------- Uploading Started --------") + for i, key in enumerate(files): + op = output_prefix.split("/") + op.pop() + if "config" not in key and "annotations" not in key: + upload = download_pool.submit( + _upload_file, + output_bucket, + f"{output_prefix}/{key}.json", + files[key], + ) + uploads.append(upload) + + futures.wait(uploads) + + +def _convert_documents_with_config( + gcs_input_path: str, + gcs_output_path: str, + project_id: str, + location: str, + processor_id: str, + config_path: str = None, +) -> None: + r"""Converts all documents in gcs_path to docproto. + + Args: + gcs_input_path (str): + Required. The gcs path to the folder containing all non docproto documents. + + Format: `gs://{bucket}/{optional_folder}` + gcs_output_path (str): + Required. The gcs path to the folder to upload the converted docproto documents to. + + Format: `gs://{bucket}/{optional_folder}` + project_id (str): + Required. + location (str): + Required. + processor_id (str): + Required. + config_path: + Optional. The gcs path to a single config file. This will work if all the documents in gcs_input_path are of the same config type. + + Format: `gs://{bucket}/{optional_folder}/config.json` + + Returns: + None. + + """ + match = re.match(r"gs://(.*?)/(.*)", gcs_input_path) + + if match is None: + raise ValueError("gcs_prefix does not match accepted format") + + input_bucket, input_prefix = match.groups() + + if input_prefix is None: + input_prefix = "/" + + file_check = re.match(constants.FILE_CHECK_REGEX, input_prefix) + + if file_check: + raise ValueError("gcs_prefix cannot contain file types") + + storage_client = document._get_storage_client() + + blob_list = storage_client.list_blobs(input_bucket, prefix=input_prefix) + + downloads = _get_files( + blob_list=blob_list, + input_prefix=input_prefix, + input_bucket=input_bucket, + config_path=config_path, + ) + + f, _ = futures.wait(downloads) + + print("-------- Finished Downloading --------") + + print("-------- Converting Started --------") + + files = [] + did_not_convert = [] + labels = [] + + files, labels, did_not_convert = _get_docproto_files( + f, project_id, location, processor_id + ) + + print("-------- Finished Converting --------") + if did_not_convert != []: + print(f"Did not convert {len(did_not_convert)} documents") + print(did_not_convert) + + _upload(files, gcs_output_path) + + print("-------- Finished Uploading --------") + print("-------- Schema Information --------") + print(f"Unique Entity Types: {labels}") + + +# [min,min],[max,min],[max,max],[min,max] diff --git a/google/cloud/documentai_toolbox/converters/converter.py b/google/cloud/documentai_toolbox/converters/converter.py new file mode 100644 index 00000000..c80dd9cd --- /dev/null +++ b/google/cloud/documentai_toolbox/converters/converter.py @@ -0,0 +1,63 @@ +# -*- 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 google.cloud.documentai_toolbox.converters.config.converter_helpers import ( + _convert_documents_with_config, +) + + +def convert_from_config( + project_id: str, + location: str, + processor_id: str, + gcs_input_path: str, + gcs_output_path: str, + config_path: str = None, +) -> None: + r"""Converts all documents in gcs_input_path to docproto using configs. + + Args: + project_id (str): + Required. + location (str): + Required. + processor_id (str): + Required. + gcs_input_path (str): + Required. The gcs path to the folder containing all non docproto documents. + + Format: `gs://{bucket}/{optional_folder}` + gcs_output_path (str): + Required. The gcs path to the folder to upload the converted docproto documents to. + + Format: `gs://{bucket}/{optional_folder}` + config_path: + Optional. The gcs path to a single config file. This will work if all the documents in gcs_input_path are of the same config type. + + Format: `gs://{bucket}/{optional_folder}/config.json` + Returns: + None. + + """ + _convert_documents_with_config( + project_id=project_id, + location=location, + processor_id=processor_id, + gcs_input_path=gcs_input_path, + gcs_output_path=gcs_output_path, + config_path=config_path, + ) diff --git a/google/cloud/documentai_toolbox/converters/converters.py b/google/cloud/documentai_toolbox/converters/converters.py deleted file mode 100644 index a00b21af..00000000 --- a/google/cloud/documentai_toolbox/converters/converters.py +++ /dev/null @@ -1,62 +0,0 @@ -# -*- 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/wrappers/document.py b/google/cloud/documentai_toolbox/wrappers/document.py index 4fe105a7..82a0c6c3 100644 --- a/google/cloud/documentai_toolbox/wrappers/document.py +++ b/google/cloud/documentai_toolbox/wrappers/document.py @@ -30,12 +30,17 @@ 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 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, +) from pikepdf import Pdf @@ -76,8 +81,8 @@ def _pages_from_shards(shards: List[documentai.Document]) -> List[Page]: result = [] for shard in shards: text = shard.text - for page in shard.pages: - result.append(Page(documentai_page=page, text=text)) + for shard_page in shard.pages: + result.append(Page(documentai_page=shard_page, text=text)) return result @@ -166,6 +171,15 @@ def _get_shards(gcs_bucket_name: str, gcs_prefix: str) -> List[documentai.Docume def _text_from_shards(shards: List[documentai.Document]) -> str: + r"""Gets text from shards. + + Args: + shards (List[google.cloud.documentai.Document]): + Required. List of document shards. + Returns: + str: + Text in all shards. + """ total_text = "" for shard in shards: if total_text == "": @@ -176,6 +190,40 @@ def _text_from_shards(shards: List[documentai.Document]) -> str: return total_text +def _convert_to_vision_annotate_file_response(text: str, pages: List[page.Page]): + r"""Convert OCR data from Document.proto to AnnotateFileResponse.proto for Vision API. + + Args: + text (str): + Required. Contents of document. + pages (List[Page]): + Required. A list of pages. + Returns: + AnnotateFileResponse: + Proto with TextAnnotations. + """ + 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 + + @dataclasses.dataclass class Document: r"""Represents a wrapped Document. @@ -299,12 +347,12 @@ def search_pages( ) found_pages = [] - for page in self.pages: - for paragraph in page.paragraphs: + for p in self.pages: + for paragraph in p.paragraphs: if (target_string and target_string in paragraph.text) or ( pattern and re.search(pattern, paragraph.text) ): - found_pages.append(page) + found_pages.append(p) break return found_pages @@ -321,8 +369,8 @@ def get_form_field_by_name(self, target_field: str) -> List[FormField]: """ found_fields = [] - for page in self.pages: - for form_field in page.form_fields: + for p in self.pages: + for form_field in p.form_fields: if target_field.lower() in form_field.field_name.lower(): found_fields.append(form_field) @@ -447,11 +495,12 @@ def split_pdf(self, pdf_path: str, output_path: str) -> List[str]: return output_files def convert_document_to_annotate_file_response(self) -> AnnotateFileResponse: - """Convert OCR data from Document proto to AnnotateFileResponse proto (Vision API). + r"""Convert OCR data from Document.proto to AnnotateFileResponse.proto for Vision API. Args: None. Returns: - AnnotateFileResponse proto with a TextAnnotation per page. + AnnotateFileResponse: + Proto with TextAnnotations. """ return _convert_to_vision_annotate_file_response(self.text, self.pages) diff --git a/samples/sample-converter-configs/AWS/AWS-config.json b/samples/sample-converter-configs/AWS/AWS-config.json new file mode 100644 index 00000000..d33c1d50 --- /dev/null +++ b/samples/sample-converter-configs/AWS/AWS-config.json @@ -0,0 +1,14 @@ +{ + "entity_object":"Blocks", + "entity": { + "type_":"BlockType", + "mention_text":"Text", + "normalized_vertices":{ + "type":"1", + "unit":"normalized", + "base":"Geometry.Polygon", + "x":"X", + "y":"Y" + } + } +} \ No newline at end of file diff --git a/samples/sample-converter-configs/Azure/form-config.json b/samples/sample-converter-configs/Azure/form-config.json new file mode 100644 index 00000000..19749112 --- /dev/null +++ b/samples/sample-converter-configs/Azure/form-config.json @@ -0,0 +1,18 @@ +{ + "entity_object":"analyzeResult.pageResults.0.keyValuePairs", + "page": { + "height":"analyzeResult.readResults.0.height", + "width":"analyzeResult.readResults.0.width" + }, + "entity": { + "type_":"key.text", + "mention_text":"value.text", + "normalized_vertices":{ + "type":"3", + "unit":"inch", + "base":"key.boundingBox", + "x":"x", + "y":"y" + } + } +} \ No newline at end of file diff --git a/samples/sample-converter-configs/Azure/invoice-config.json b/samples/sample-converter-configs/Azure/invoice-config.json new file mode 100644 index 00000000..3ec3468e --- /dev/null +++ b/samples/sample-converter-configs/Azure/invoice-config.json @@ -0,0 +1,18 @@ +{ + "entity_object":"analyzeResult.documentResults.0.fields", + "page": { + "height":"analyzeResult.readResults.0.height", + "width":"analyzeResult.readResults.0.width" + }, + "entity": { + "type_":"analyzeResult.documentResults.0.fields:self", + "mention_text":"text", + "normalized_vertices":{ + "type":"3", + "unit":"pxl", + "base":"boundingBox", + "x":"x", + "y":"y" + } + } +} \ No newline at end of file diff --git a/samples/snippets/convert_external_annotations_sample.py b/samples/snippets/convert_external_annotations_sample.py new file mode 100644 index 00000000..02fa0d15 --- /dev/null +++ b/samples/snippets/convert_external_annotations_sample.py @@ -0,0 +1,71 @@ +# 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_convert_external_annotations] + +from google.cloud.documentai_toolbox import converter + +# TODO(developer): Uncomment these variables before running the sample. +# This sample will convert external annotations to the Document.json format used by Document AI Workbench for training. +# To process this the external annotation must have these type of objects: +# 1) Type +# 2) Text +# 3) Bounding Box (bounding boxes must be 1 of the 3 optional types) +# +# This is the bare minimum requirement to convert the annotations but for better accuracy you will need to also have: +# 1) Document width & height +# +# Bounding Box Types: +# Type 1: +# bounding_box:[{"x":1,"y":2},{"x":2,"y":2},{"x":2,"y":3},{"x":1,"y":3}] +# Type 2: +# bounding_box:{ "Width": 1, "Height": 1, "Left": 1, "Top": 1} +# Type 3: +# bounding_box: [1,2,2,2,2,3,1,3] +# +# Note: If these types are not sufficient you can propose a feature request or contribute the new type and conversion functionality. +# +# Given a folders in gcs_input_path with the following structure : +# +# gs://path/to/input/folder +# ├──test_annotations.json +# ├──test_config.json +# └──test.pdf +# +# An example of the config is in sample-converter-configs/Azure/form-config.json +# +# location = "us", +# processor_id = "my_processor_id" +# gcs_input_path = "gs://path/to/input/folder" +# gcs_output_path = "gs://path/to/input/folder" + + +def convert_external_annotations_sample( + location: str, + processor_id: str, + project_id: str, + gcs_input_path: str, + gcs_output_path: str, +) -> None: + converter.convert_from_config( + project_id=project_id, + location=location, + processor_id=processor_id, + gcs_input_path=gcs_input_path, + gcs_output_path=gcs_output_path, + ) + + +# [END documentai_toolbox_convert_external_annotations] diff --git a/samples/snippets/noxfile.py b/samples/snippets/noxfile.py index a768bbc6..1765c2a8 100644 --- a/samples/snippets/noxfile.py +++ b/samples/snippets/noxfile.py @@ -206,9 +206,7 @@ def _session_tests( if os.path.exists("requirements-test.txt"): if os.path.exists("constraints-test.txt"): - session.install( - "-r", "requirements-test.txt", "-c", "constraints-test.txt" - ) + session.install("-r", "requirements-test.txt", "-c", "constraints-test.txt") else: session.install("-r", "requirements-test.txt") with open("requirements-test.txt") as rtfile: @@ -221,9 +219,9 @@ def _session_tests( post_install(session) if "pytest-parallel" in packages: - concurrent_args.extend(['--workers', 'auto', '--tests-per-worker', 'auto']) + concurrent_args.extend(["--workers", "auto", "--tests-per-worker", "auto"]) elif "pytest-xdist" in packages: - concurrent_args.extend(['-n', 'auto']) + concurrent_args.extend(["-n", "auto"]) session.run( "pytest", diff --git a/samples/snippets/quickstart_sample.py b/samples/snippets/quickstart_sample.py index 15fe7484..c3e41670 100644 --- a/samples/snippets/quickstart_sample.py +++ b/samples/snippets/quickstart_sample.py @@ -27,7 +27,9 @@ def quickstart_sample(gcs_bucket_name: str, gcs_prefix: str) -> None: print("Document structure in Cloud Storage") - utilities.print_gcs_document_tree(gcs_bucket_name=gcs_bucket_name, gcs_prefix=gcs_prefix) + utilities.print_gcs_document_tree( + gcs_bucket_name=gcs_bucket_name, gcs_prefix=gcs_prefix + ) wrapped_document = document.Document.from_gcs( gcs_bucket_name=gcs_bucket_name, gcs_prefix=gcs_prefix diff --git a/samples/snippets/test_convert_document_to_vision_sample.py b/samples/snippets/test_convert_document_to_vision_sample.py index 0f782fc4..668e1acd 100644 --- a/samples/snippets/test_convert_document_to_vision_sample.py +++ b/samples/snippets/test_convert_document_to_vision_sample.py @@ -24,7 +24,7 @@ gcs_input_uri = "output/123456789/0" -def test_quickstart_sample(capsys: pytest.CaptureFixture) -> None: +def test_convert_document_to_vision_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 ) diff --git a/samples/snippets/test_convert_external_annotations_sample.py b/samples/snippets/test_convert_external_annotations_sample.py new file mode 100644 index 00000000..63acecdb --- /dev/null +++ b/samples/snippets/test_convert_external_annotations_sample.py @@ -0,0 +1,35 @@ +# 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_external_annotations_sample + +location = "us" +project_id = os.environ["GOOGLE_CLOUD_PROJECT"] + + +def test_convert_external_annotations_sample(capsys: pytest.CaptureFixture) -> None: + convert_external_annotations_sample.convert_external_annotations_sample( + location=location, + processor_id="52a38e080c1a7296", + project_id="project_id", + gcs_input_path="gs://documentai_toolbox_samples/documentai_toolbox_samples/converter/azure", + gcs_output_path="gs://documentai_toolbox_samples/documentai_toolbox_samples/converter/output", + ) + out, _ = capsys.readouterr() + + assert "-------- Finished Converting --------" in out diff --git a/setup.py b/setup.py index 9791bdf9..b4ca2f9b 100644 --- a/setup.py +++ b/setup.py @@ -54,6 +54,7 @@ "google-cloud-storage >= 1.31.0, < 3.0.0dev", "google-cloud-vision >= 2.7.0, < 4.0.0dev ", "numpy >= 1.18.1", + "intervaltree >= 3.0.0", "pikepdf >= 6.2.9, < 8.0.0", "immutabledict >= 2.0.0, < 3.0.0dev", ), diff --git a/tests/unit/resources/converters/test_config_type_1.json b/tests/unit/resources/converters/test_config_type_1.json new file mode 100644 index 00000000..554bdaf9 --- /dev/null +++ b/tests/unit/resources/converters/test_config_type_1.json @@ -0,0 +1,15 @@ +{ + "entity_object":"pages.1.Entities", + "entity": { + "mention_text":"Text", + "type_":"Type", + "page_number": "page", + "normalized_vertices":{ + "type":"1", + "unit":"inch", + "base":"bBox", + "x":"x", + "y":"y" + } + } +} \ No newline at end of file diff --git a/tests/unit/resources/converters/test_config_type_2.json b/tests/unit/resources/converters/test_config_type_2.json new file mode 100644 index 00000000..5403f0f5 --- /dev/null +++ b/tests/unit/resources/converters/test_config_type_2.json @@ -0,0 +1,16 @@ +{ + "entity_object":"document.entities", + "entity": { + "type_":"type", + "mention_text":"mentionText", + "normalized_vertices":{ + "type":"2", + "unit":"normalized", + "base":"pageAnchor.pageRefs.0.boundingPoly.normalizedVertices", + "x":"left", + "y":"top", + "width":"width", + "height":"height" + } + } +} \ No newline at end of file diff --git a/tests/unit/resources/converters/test_config_type_3.json b/tests/unit/resources/converters/test_config_type_3.json new file mode 100644 index 00000000..8b45f0bd --- /dev/null +++ b/tests/unit/resources/converters/test_config_type_3.json @@ -0,0 +1,21 @@ +{ + "entity_object":"Entities", + "page": { + "height":"page_height", + "width":"page_width" + }, + "entity": { + "type_":"Entities:self", + "mention_text":"Text||normalizedText", + "normalized_vertices":{ + "type":"3", + "unit":"pxl", + "base":"bBox", + "x":"x", + "y":"y" + }, + "id":"id", + "confidence":"confidence", + "page_number":"page" + } +} \ No newline at end of file diff --git a/tests/unit/resources/converters/test_type_1.json b/tests/unit/resources/converters/test_type_1.json new file mode 100644 index 00000000..5fd4ef28 --- /dev/null +++ b/tests/unit/resources/converters/test_type_1.json @@ -0,0 +1,36 @@ +{ + "DocumentType": "ScannedPDF", + "NoOfPages": 1, + "pages": [ + {}, + { + "Entities": [ + { + "Type": "BusinessName", + "Text": "411 I.T. Group", + "id":0, + "bBox": [ + { + "x": 4.083333, + "y": 1.208333 + }, + { + "x": 5.8125, + "y": 1.208333 + }, + { + "x": 5.8125, + "y": 1.510416 + }, + { + "x": 4.083333, + "y": 1.510416 + } + ], + "page": "0", + "confidence": 0.9997831 + } + ] + } + ] +} \ No newline at end of file diff --git a/tests/unit/resources/converters/test_type_2.json b/tests/unit/resources/converters/test_type_2.json new file mode 100644 index 00000000..5ed2411b --- /dev/null +++ b/tests/unit/resources/converters/test_type_2.json @@ -0,0 +1,30 @@ +{ + "document": { + "uri": "", + "mimeType": "application/pdf", + "page_height":1000, + "page_width":1000, + "entities": [ + { + "type": "invoice_id", + "mentionText": "4748", + "confidence": 0.980109, + "pageAnchor": { + "pageRefs": [ + { + "boundingPoly": { + "normalizedVertices": { + "width": 0.03, + "height":0.01, + "left": 0.07906712, + "top": 0.36043957 + } + } + } + ] + }, + "id": "0" + } + ] + } +} \ No newline at end of file diff --git a/tests/unit/resources/converters/test_type_3.json b/tests/unit/resources/converters/test_type_3.json new file mode 100644 index 00000000..ebba08cd --- /dev/null +++ b/tests/unit/resources/converters/test_type_3.json @@ -0,0 +1,25 @@ +{ + "DocumentType": "ScannedPDF", + "NoOfPages": 1, + "page_height":1000, + "page_width":1000, + "Entities": { + "BusinessName": { + "Text": "411 I.T. Group", + "normalizedText":"normalized 411 I.T. Group", + "id":0, + "bBox": [ + 392, + 116, + 558, + 116, + 558, + 145, + 392, + 145 + ], + "page": "0", + "confidence": 0.9997831 + } + } + } \ No newline at end of file diff --git a/tests/unit/test_bbox_conversion.py b/tests/unit/test_bbox_conversion.py new file mode 100644 index 00000000..c701debd --- /dev/null +++ b/tests/unit/test_bbox_conversion.py @@ -0,0 +1,264 @@ +from google.cloud import documentai +from google.cloud.documentai_v1.types import geometry +from google.cloud.documentai_toolbox.converters.config import bbox_conversion, blocks + + +def test_midpoint_in_bpoly(): + vertex_a = geometry.NormalizedVertex(x=2, y=2) + box_a = geometry.BoundingPoly(normalized_vertices=[vertex_a]) + + vertex_b = geometry.NormalizedVertex(x=1, y=1) + vertex_b_max = geometry.NormalizedVertex(x=4, y=4) + box_b = geometry.BoundingPoly(normalized_vertices=[vertex_b, vertex_b_max]) + + actual = bbox_conversion._midpoint_in_bpoly(box_a=box_a, box_b=box_b) + assert actual + + +def test_merge_text_anchors(): + text_segment_1 = documentai.Document.TextAnchor.TextSegment( + start_index="0", end_index="100" + ) + text_anchor_1 = documentai.Document.TextAnchor(text_segments=[text_segment_1]) + + text_segment_2 = documentai.Document.TextAnchor.TextSegment( + start_index="100", end_index="200" + ) + text_anchor_2 = documentai.Document.TextAnchor(text_segments=[text_segment_2]) + + text_segment_3 = documentai.Document.TextAnchor.TextSegment( + start_index="0", end_index="200" + ) + expected = documentai.Document.TextAnchor(text_segments=[text_segment_3]) + actual = bbox_conversion._merge_text_anchors( + text_anchor_1=text_anchor_1, text_anchor_2=text_anchor_2 + ) + assert actual == expected + + +def test_get_text_anchor_in_bbox(): + vertex_a = geometry.NormalizedVertex(x=2, y=2) + vertex_a_max = geometry.NormalizedVertex(x=5, y=5) + box_a = geometry.BoundingPoly(normalized_vertices=[vertex_a, vertex_a_max]) + + vertex_b = geometry.NormalizedVertex(x=1, y=1) + vertex_b_max = geometry.NormalizedVertex(x=8, y=8) + box_b = geometry.BoundingPoly(normalized_vertices=[vertex_b, vertex_b_max]) + + text_segment_1 = documentai.Document.TextAnchor.TextSegment( + start_index="0", end_index="100" + ) + text_anchor_1 = documentai.Document.TextAnchor(text_segments=[text_segment_1]) + + text_segment_2 = documentai.Document.TextAnchor.TextSegment( + start_index="100", end_index="200" + ) + text_anchor_2 = documentai.Document.TextAnchor(text_segments=[text_segment_2]) + + layout1 = documentai.Document.Page.Layout( + bounding_poly=box_b, text_anchor=text_anchor_1 + ) + layout2 = documentai.Document.Page.Layout( + bounding_poly=box_b, text_anchor=text_anchor_2 + ) + + token1 = documentai.Document.Page.Token(layout=layout1) + token2 = documentai.Document.Page.Token(layout=layout2) + + page = documentai.Document.Page(tokens=[token1, token2]) + actual = bbox_conversion._get_text_anchor_in_bbox(bbox=box_a, page=page) + + text_segment_3 = documentai.Document.TextAnchor.TextSegment( + start_index="0", end_index="200" + ) + expected = documentai.Document.TextAnchor(text_segments=[text_segment_3]) + assert actual == expected + + +def test_get_norm_x_max(): + vertex_a_min = geometry.NormalizedVertex(x=2, y=2) + vertex_a_max = geometry.NormalizedVertex(x=4, y=4) + + bbox = geometry.BoundingPoly(normalized_vertices=[vertex_a_min, vertex_a_max]) + actual = bbox_conversion._get_norm_x_max(bbox=bbox) + assert actual == 4 + + +def test_get_norm_x_min(): + vertex_a_min = geometry.NormalizedVertex(x=2, y=2) + vertex_a_max = geometry.NormalizedVertex(x=4, y=4) + + bbox = geometry.BoundingPoly(normalized_vertices=[vertex_a_min, vertex_a_max]) + actual = bbox_conversion._get_norm_x_min(bbox=bbox) + assert actual == 2 + + +def test_get_norm_y_max(): + vertex_a_min = geometry.NormalizedVertex(x=2, y=2) + vertex_a_max = geometry.NormalizedVertex(x=4, y=4) + + bbox = geometry.BoundingPoly(normalized_vertices=[vertex_a_min, vertex_a_max]) + actual = bbox_conversion._get_norm_y_min(bbox=bbox) + assert actual == 2 + + +def test_get_norm_y_min(): + vertex_a_min = geometry.NormalizedVertex(x=2, y=2) + vertex_a_max = geometry.NormalizedVertex(x=4, y=4) + + bbox = geometry.BoundingPoly(normalized_vertices=[vertex_a_min, vertex_a_max]) + actual = bbox_conversion._get_norm_y_max(bbox=bbox) + assert actual == 4 + + +def test_normalize_coordinates(): + actual = bbox_conversion._normalize_coordinates(x=4.0, y=2.0) + assert actual == 2.0 + + +def test_convert_to_pixels(): + actual = bbox_conversion._convert_to_pixels(x=1, conversion_rate=96) + assert actual == 96 + + +def test_convert_bbox_units_with_normalized(): + actual = bbox_conversion._convert_bbox_units( + coordinate=0.5, input_bbox_units="normalized", width=2550, height=3300 + ) + assert actual == 0.5 + + +def test_convert_bbox_units_with_pxl(): + actual = bbox_conversion._convert_bbox_units( + coordinate=1, input_bbox_units="pxl", width=2550, height=3300 + ) + assert actual == 0.000392157 + + +def test_convert_bbox_units_with_inch(): + actual = bbox_conversion._convert_bbox_units( + coordinate=1, input_bbox_units="inch", width=2550, height=3300 + ) + assert actual == 0.037647059 + + +def test_convert_bbox_units_with_cm(): + actual = bbox_conversion._convert_bbox_units( + coordinate=1, input_bbox_units="cm", width=2550, height=3300 + ) + assert actual == 0.014821569 + + +def test_get_multiplier_pxl(): + actual = bbox_conversion._get_multiplier( + docproto_coordinate=1000, external_coordinate=1000, input_bbox_units="pxl" + ) + assert actual == 1.0 + + +def test_get_multiplier_inch(): + actual = bbox_conversion._get_multiplier( + docproto_coordinate=1000, external_coordinate=10.416, input_bbox_units="inch" + ) + assert actual == 1.000064004096262 + + +def test_get_multiplier_cm(): + actual = bbox_conversion._get_multiplier( + docproto_coordinate=1000, external_coordinate=26.4585, input_bbox_units="cm" + ) + assert actual == 1.000000992500985 + + +def test_convert_bbox_to_docproto_bbox_empty_coordinate(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_1.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_1.json", "r") as (f): + config = f.read() + b = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + b[0].bounding_box = [] + + actual = bbox_conversion._convert_bbox_to_docproto_bbox(block=(b[0])) + + assert actual == [] + + +def test_convert_bbox_to_docproto_bbox_type_1(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_1.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_1.json", "r") as (f): + config = f.read() + b = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + actual = bbox_conversion._convert_bbox_to_docproto_bbox(block=(b[0])) + + assert actual.normalized_vertices != [] + assert actual.vertices == [] + assert "x" in str(actual.normalized_vertices) + assert "y" in str(actual.normalized_vertices) + + +def test_convert_bbox_to_docproto_bbox_type_2(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_2.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_2.json", "r") as (f): + config = f.read() + b = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + actual = bbox_conversion._convert_bbox_to_docproto_bbox(block=(b[0])) + + assert actual.normalized_vertices != [] + assert actual.vertices == [] + assert "x" in str(actual.normalized_vertices) + assert "y" in str(actual.normalized_vertices) + + +def test_convert_bbox_to_docproto_bbox_type_3(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_3.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_3.json", "r") as (f): + config = f.read() + b = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + + print(b[0].bounding_type) + + actual = bbox_conversion._convert_bbox_to_docproto_bbox(block=(b[0])) + + assert actual.normalized_vertices != [] + assert actual.vertices == [] + assert "x" in str(actual.normalized_vertices) + assert "y" in str(actual.normalized_vertices) diff --git a/tests/unit/test_blocks.py b/tests/unit/test_blocks.py new file mode 100644 index 00000000..9e1ac238 --- /dev/null +++ b/tests/unit/test_blocks.py @@ -0,0 +1,126 @@ +from google.cloud import documentai +from google.cloud.documentai_toolbox.converters.config import blocks + + +def test_create(): + actual = blocks.Block.create( + type_="test_type", + text="test_text", + bounding_box="", + block_references="", + block_id="", + confidence="", + page_number="", + page_width="", + page_height="", + bounding_width="", + bounding_height="", + bounding_type="", + bounding_unit="", + bounding_x="", + bounding_y="", + docproto_width="", + docproto_height="", + ) + + assert actual.type_ == "test_type" + assert actual.text == "test_text" + + +def test_get_target_object(): + test_json_data = { + "document": {"entities": [{}, {"text": "test_text", "type": "test_type"}]} + } + + text = blocks._get_target_object( + json_data=test_json_data, target_object="document.entities.1.text" + ) + type = blocks._get_target_object( + json_data=test_json_data, target_object="document.entities.1.type" + ) + + assert text == "test_text" + assert type == "test_type" + + +def test_get_target_object_with_one_object(): + test_json_data = {"document": "document_test"} + + text = blocks._get_target_object(json_data=test_json_data, target_object="document") + + assert text == "document_test" + + +def test_get_target_object_without_target(): + test_json_data = { + "document": {"entities": [{}, {"text": "test_text", "type": "test_type"}]} + } + + text = blocks._get_target_object( + json_data=test_json_data, target_object="entities.text" + ) + + assert text is None + + +def test_load_blocks_from_scheme_type_1(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_1.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_1.json", "r") as (f): + config = f.read() + + actual = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + + assert actual[0].text == "411 I.T. Group" + assert actual[0].type_ == "BusinessName" + + +def test_load_blocks_from_scheme_type_2(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_2.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_2.json", "r") as (f): + config = f.read() + + actual = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + + assert actual[0].text == "4748" + assert actual[0].type_ == "invoice_id" + + +def test__load_blocks_from_schema_type_3(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_3.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_3.json", "r") as (f): + config = f.read() + + actual = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + + assert actual[0].text == "normalized 411 I.T. Group" + assert actual[0].type_ == "BusinessName" diff --git a/tests/unit/test_converter.py b/tests/unit/test_converter.py new file mode 100644 index 00000000..0be99429 --- /dev/null +++ b/tests/unit/test_converter.py @@ -0,0 +1,63 @@ +# -*- 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. +# + +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock + +from google.cloud.documentai_toolbox.converters import converter + + +@mock.patch("google.cloud.documentai_toolbox.wrappers.document.storage") +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._get_docproto_files", + return_value=(["file1"], ["test_label"], []), +) +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._upload", + return_value="Done", +) +def test__convert_documents_with_config( + mock_storage, mock_get_docproto_files, mock_upload, capfd +): + client = mock_storage.Client.return_value + mock_bucket = mock.Mock() + client.Bucket.return_value = mock_bucket + + mock_blob1 = mock.Mock(name="gs://test-directory/1/test-annotations.json") + mock_blob1.download_as_bytes.return_value = ( + "gs://test-directory/1/test-annotations.json" + ) + + mock_blob2 = mock.Mock(name="gs://test-directory/1/test-config.json") + mock_blob2.download_as_bytes.return_value = "gs://test-directory/1/test-config.json" + + mock_blob3 = mock.Mock(name="gs://test-directory/1/test.pdf") + mock_blob3.download_as_bytes.return_value = "gs://test-directory/1/test.pdf" + + client.list_blobs.return_value = [mock_blob1, mock_blob2, mock_blob3] + + converter.convert_from_config( + project_id="project-id", + location="location", + processor_id="project-id", + gcs_input_path="gs://test-directory/1", + gcs_output_path="gs://test-directory/1/output", + ) + + out, err = capfd.readouterr() + assert "test_label" in out diff --git a/tests/unit/test_converter_helpers.py b/tests/unit/test_converter_helpers.py new file mode 100644 index 00000000..845505c1 --- /dev/null +++ b/tests/unit/test_converter_helpers.py @@ -0,0 +1,490 @@ +# -*- 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. +# + +try: + from unittest import mock +except ImportError: # pragma: NO COVER + import mock + +from google.cloud.documentai_toolbox.converters.config import blocks, converter_helpers +from google.cloud import documentai +import pytest + + +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers.documentai" +) +def test_get_base_ocr(mock_docai): + mock_client = mock_docai.DocumentProcessorServiceClient.return_value + + mock_client.process_document.return_value.document = "Done" + + actual = converter_helpers._get_base_ocr( + project_id="project_id", + location="location", + processor_id="processor_id", + file_bytes="file", + mime_type="application/pdf", + ) + + mock_client.process_document.assert_called() + assert actual == "Done" + + +def test_get_entity_content_type_3(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_3.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_3.json", "r") as (f): + config = f.read() + + b = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + + actual = converter_helpers._get_entity_content(blocks=b, docproto=docproto) + + assert actual[0].type == "BusinessName" + assert actual[0].mention_text == "normalized 411 I.T. Group" + + +def test_get_entity_content_type_2(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_2.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_2.json", "r") as (f): + config = f.read() + + b = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + + actual = converter_helpers._get_entity_content(blocks=b, docproto=docproto) + + assert actual[0].type == "invoice_id" + assert actual[0].mention_text == "4748" + + +def test_get_entity_content_type_1(): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + with open("tests/unit/resources/converters/test_type_1.json", "r") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_1.json", "r") as (f): + config = f.read() + + b = blocks._load_blocks_from_schema( + input_data=invoice, input_config=config, base_docproto=docproto + ) + + actual = converter_helpers._get_entity_content(blocks=b, docproto=docproto) + + assert actual[0].type == "BusinessName" + assert actual[0].mention_text == "411 I.T. Group" + + +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._get_base_ocr" +) +def test_convert_to_docproto_with_config(mock_ocr): + docproto = documentai.Document() + page = documentai.Document.Page() + dimensions = documentai.Document.Page.Dimension() + dimensions.width = 2550 + dimensions.height = 3300 + page.dimension = dimensions + docproto.pages = [page] + mock_ocr.return_value = docproto + + with open("tests/unit/resources/converters/test_type_3.json", "rb") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_3.json", "rb") as (f): + config = f.read() + with open("tests/unit/resources/toolbox_invoice_test.pdf", "rb") as (f): + pdf = f.read() + + actual = converter_helpers._convert_to_docproto_with_config( + name="test_document", + annotated_bytes=invoice, + config_bytes=config, + document_bytes=pdf, + project_id="project_id", + processor_id="processor_id", + location="location", + retry_number=0, + ) + + assert len(actual.pages) == 1 + assert len(actual.entities) == 1 + assert actual.entities[0].type == "BusinessName" + assert actual.entities[0].mention_text == "normalized 411 I.T. Group" + + +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._get_base_ocr" +) +def test_convert_to_docproto_with_config_with_error(mock_ocr, capfd): + mock_ocr.return_value = None + + with open("tests/unit/resources/converters/test_type_3.json", "rb") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_3.json", "rb") as (f): + config = f.read() + with open("tests/unit/resources/toolbox_invoice_test.pdf", "rb") as (f): + pdf = f.read() + + actual = converter_helpers._convert_to_docproto_with_config( + name="test_document", + annotated_bytes=invoice, + config_bytes=config, + document_bytes=pdf, + project_id="project_id", + processor_id="processor_id", + location="location", + retry_number=6, + ) + + out, err = capfd.readouterr() + + assert actual is None + assert "Could Not Convert test_document" in out + + +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._get_base_ocr" +) +def test_convert_to_docproto_with_config_with_error_and_retry(mock_ocr, capfd): + mock_ocr.return_value = None + + with open("tests/unit/resources/converters/test_type_3.json", "rb") as (f): + invoice = f.read() + with open("tests/unit/resources/converters/test_config_type_3.json", "rb") as (f): + config = f.read() + with open("tests/unit/resources/toolbox_invoice_test.pdf", "rb") as (f): + pdf = f.read() + + actual = converter_helpers._convert_to_docproto_with_config( + name="test_document", + annotated_bytes=invoice, + config_bytes=config, + document_bytes=pdf, + project_id="project_id", + processor_id="processor_id", + location="location", + retry_number=5, + ) + + out, err = capfd.readouterr() + + assert actual is None + assert "Could Not Convert test_document" in out + + +@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_ds_store = mock.Mock(name=[]) + mock_ds_store.name = "DS_Store" + + mock_blob1 = mock.Mock(name=[]) + mock_blob1.name = "gs://test-directory/1/test-annotations.json" + mock_blob1.download_as_bytes.return_value = ( + "gs://test-directory/1/test-annotations.json" + ) + + mock_blob2 = mock.Mock(name=[]) + mock_blob2.name = "gs://test-directory/1/test-config.json" + mock_blob2.download_as_bytes.return_value = "gs://test-directory/1/test-config.json" + + mock_blob3 = mock.Mock(name=[]) + mock_blob3.name = "gs://test-directory/1/test.pdf" + mock_blob3.download_as_bytes.return_value = "gs://test-directory/1/test.pdf" + + client.list_blobs.return_value = [mock_ds_store, mock_blob1, mock_blob2, mock_blob3] + + actual = converter_helpers._get_bytes( + bucket_name="bucket", + prefix="prefix", + annotation_file_prefix="annotations", + config_file_prefix="config", + ) + + assert actual == [ + "gs://test-directory/1/test-annotations.json", + "gs://test-directory/1/test.pdf", + "gs://test-directory/1/test-config.json", + "prefix", + "test", + ] + + +@mock.patch("google.cloud.documentai_toolbox.wrappers.document.storage") +def test_get_bytes_with_error(mock_storage): + with pytest.raises(Exception, match="Fail"): + client = mock_storage.Client.return_value + mock_bucket = mock.Mock() + client.Bucket.return_value = mock_bucket + + mock_blob1 = mock.Mock(name=[]) + mock_blob1.name = "gs://test-directory/1/test-annotations.json" + mock_blob1.download_as_bytes.side_effect = Exception("Fail") + + client.list_blobs.return_value = [mock_blob1] + + converter_helpers._get_bytes( + bucket_name="bucket", + prefix="prefix", + annotation_file_prefix="annotations", + config_file_prefix="config", + ) + + +@mock.patch("google.cloud.documentai_toolbox.wrappers.document.storage") +def test_upload_file(mock_storage): + client = mock_storage.Client.return_value + + converter_helpers._upload_file( + bucket_name="bucket", output_prefix="prefix", file="file" + ) + client.bucket.return_value.blob.return_value.upload_from_string.assert_called_with( + "file", content_type="application/json" + ) + + +@mock.patch("google.cloud.documentai_toolbox.wrappers.document.storage") +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._get_bytes", + return_value="file_bytes", +) +def test_get_files(mock_storage, mock_get_bytes): + client = mock_storage.Client.return_value + mock_bucket = mock.Mock() + client.Bucket.return_value = mock_bucket + + mock_ds_store = mock.Mock(name=[]) + mock_ds_store.name = "DS_Store" + + mock_blob1 = mock.Mock(name=[]) + mock_blob1.name = "gs://test-directory/1/test-annotations.json" + mock_blob1.download_as_bytes.return_value = ( + "gs://test-directory/1/test-annotations.json" + ) + + mock_blob2 = mock.Mock(name=[]) + mock_blob2.name = "gs://test-directory/1/test-config.json" + mock_blob2.download_as_bytes.return_value = "gs://test-directory/1/test-config.json" + + mock_blob3 = mock.Mock(name=[]) + mock_blob3.name = "gs://test-directory/1/test.pdf" + mock_blob3.download_as_bytes.return_value = "gs://test-directory/1/test.pdf" + + blob_list = [mock_ds_store, mock_blob1, mock_blob2, mock_blob3] + + actual = converter_helpers._get_files( + blob_list=blob_list, input_prefix="", input_bucket="test-directory" + ) + + assert actual[0].result() == "file_bytes" + + +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._convert_to_docproto_with_config", +) +def test_get_docproto_files(mocked_convert_docproto): + + mock_result = mock.Mock() + mock_result.result.return_value = [ + "annotated_bytes", + "document_bytes", + "config_bytes", + "document_1", + ] + + document = documentai.Document() + entities = [documentai.Document.Entity(type_="test_type", mention_text="test_text")] + document.entities = entities + + mocked_convert_docproto.return_value = document + ( + actual_files, + actual_unique_types, + actual_did_not_convert, + ) = converter_helpers._get_docproto_files( + f=[mock_result], + project_id="project-id", + processor_id="processor-id", + location="us", + ) + assert "test_type" in actual_files["document_1"] + assert "test_text" in actual_files["document_1"] + assert "test_type" in actual_unique_types + mocked_convert_docproto.assert_called_with( + annotated_bytes="annotated_bytes", + document_bytes="document_bytes", + config_bytes="config_bytes", + project_id="project-id", + location="us", + processor_id="processor-id", + retry_number=1, + name="document_1", + ) + + +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._convert_to_docproto_with_config", +) +def test_get_docproto_files_with_no_docproto(mocked_convert_docproto): + + mock_result = mock.Mock() + mock_result.result.return_value = [ + "annotated_bytes", + "document_bytes", + "config_bytes", + "document_1", + ] + + mocked_convert_docproto.return_value = None + ( + actual_files, + actual_unique_types, + actual_did_not_convert, + ) = converter_helpers._get_docproto_files( + f=[mock_result], + project_id="project-id", + processor_id="processor-id", + location="us", + ) + assert "document_1" in actual_did_not_convert + mocked_convert_docproto.assert_called_with( + annotated_bytes="annotated_bytes", + document_bytes="document_bytes", + config_bytes="config_bytes", + project_id="project-id", + location="us", + processor_id="processor-id", + retry_number=1, + name="document_1", + ) + + +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._upload_file", +) +def test_upload(mock_upload_file): + files = {} + files["document_1"] = "Document" + converter_helpers._upload(files, gcs_output_path="gs://output/") + + mock_upload_file.assert_called_with("output", "/document_1.json", "Document") + + +def test_upload_with_format_error(): + with pytest.raises(ValueError, match="gcs_prefix does not match accepted format"): + files = {} + files["document_1"] = "Document" + converter_helpers._upload(files, gcs_output_path="output/path") + + +def test_upload_with_file_error(): + with pytest.raises(ValueError, match="gcs_prefix cannot contain file types"): + files = {} + files["document_1"] = "Document" + converter_helpers._upload(files, gcs_output_path="gs://output/path.json") + + +@mock.patch("google.cloud.documentai_toolbox.wrappers.document.storage") +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._get_docproto_files", + return_value=(["file1"], ["test_label"], ["document_2"]), +) +@mock.patch( + "google.cloud.documentai_toolbox.converters.config.converter_helpers._upload", + return_value="Done", +) +def test_convert_documents_with_config( + mock_storage, mock_get_docproto_files, mock_upload, capfd +): + client = mock_storage.Client.return_value + mock_bucket = mock.Mock() + client.Bucket.return_value = mock_bucket + + mock_blob1 = mock.Mock(name="gs://test-directory/1/test-annotations.json") + mock_blob1.download_as_bytes.return_value = ( + "gs://test-directory/1/test-annotations.json" + ) + + mock_blob2 = mock.Mock(name="gs://test-directory/1/test-config.json") + mock_blob2.download_as_bytes.return_value = "gs://test-directory/1/test-config.json" + + mock_blob3 = mock.Mock(name="gs://test-directory/1/test.pdf") + mock_blob3.download_as_bytes.return_value = "gs://test-directory/1/test.pdf" + + client.list_blobs.return_value = [mock_blob1, mock_blob2, mock_blob3] + + converter_helpers._convert_documents_with_config( + project_id="project-id", + location="location", + processor_id="project-id", + gcs_input_path="gs://test-directory/", + gcs_output_path="gs://test-directory-output/", + ) + + out, err = capfd.readouterr() + assert "test_label" in out + assert "Did not convert 1 documents" in out + assert "document_2" in out + + +def test_convert_documents_with_config_with_gcs_path_error(): + with pytest.raises(ValueError, match="gcs_prefix does not match accepted format"): + converter_helpers._convert_documents_with_config( + project_id="project-id", + location="location", + processor_id="project-id", + gcs_input_path="test-directory/1", + gcs_output_path="gs://test-directory/1/output", + ) + + +def test_convert_documents_with_config_with_file_error(): + with pytest.raises(ValueError, match="gcs_prefix cannot contain file types"): + converter_helpers._convert_documents_with_config( + project_id="project-id", + location="location", + processor_id="project-id", + gcs_input_path="gs://test-directory/1.json", + gcs_output_path="gs://test-directory/1/output", + ) diff --git a/tests/unit/test_document.py b/tests/unit/test_document.py index 06f538d9..d2e89e77 100644 --- a/tests/unit/test_document.py +++ b/tests/unit/test_document.py @@ -75,23 +75,6 @@ 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( @@ -253,6 +236,42 @@ def test_get_entity_by_type(get_bytes_single_file_mock): assert actual[0].mention_text == "222 Main Street\nAnytown, USA" +@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_ds_store = mock.Mock(name=[]) + mock_ds_store.name = "DS_Store" + + mock_blob1 = mock.Mock(name=[]) + mock_blob1.name = "gs://test-directory/1/test-annotations.json" + mock_blob1.download_as_bytes.return_value = ( + "gs://test-directory/1/test-annotations.json" + ) + + mock_blob2 = mock.Mock(name=[]) + mock_blob2.name = "gs://test-directory/1/test-config.json" + mock_blob2.download_as_bytes.return_value = "gs://test-directory/1/test-config.json" + + mock_blob3 = mock.Mock(name=[]) + mock_blob3.name = "gs://test-directory/1/test.pdf" + mock_blob3.download_as_bytes.return_value = "gs://test-directory/1/test.pdf" + + client.list_blobs.return_value = [mock_ds_store, mock_blob1, mock_blob2, mock_blob3] + + actual = document._get_bytes( + gcs_bucket_name="bucket", + gcs_prefix="prefix", + ) + + assert actual == [ + "gs://test-directory/1/test-annotations.json", + "gs://test-directory/1/test-config.json", + ] + + def test_get_form_field_by_name(get_bytes_form_parser_mock): doc = document.Document.from_gcs( gcs_bucket_name="test-directory", gcs_prefix="documentai/output/123456789/0"