# Copyright 2022 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. """This module defines a CLI that uses Document AI to annotate a PDF document""" import argparse import os import sys from typing import Optional import google.auth from google.cloud.documentai_v1beta3 import DocumentProcessorServiceClient from google.cloud.documentai_v1beta3 import Processor import pikepdf DEFAULT_MULTI_REGION_LOCATION = "us" DEFAULT_PROCESSOR_TYPE = "FORM_PARSER_PROCESSOR" def main(args): """This functions annotates a PDF document using the Document AI API""" if not args.project_id: _, project_id = google.auth.default() args.project_id = project_id parent = f"projects/{args.project_id}/locations/{args.multi_region_location}" client = DocumentProcessorServiceClient() processor_id = find_processor_id_of_type(client, parent, args.form_processor_type) if processor_id is None: print( f"no form processor found. " f'creating new processor of type "{args.form_processor_type}"', ) processor_id = create_processor(client, parent, args.form_processor_type) if not os.path.isfile(os.path.abspath(args.input)): print(f"could not find file at {os.path.abspath(args.input)}") return 1 # If a output path is not specified, use input directory if not args.output: args.output = f'{os.path.abspath(args.input).rstrip(".pdf")}_annotated.pdf' print("Calling Document AI API...", end="") with open(args.input, "rb") as pdf_file: document = client.process_document( request={ "name": f"{parent}/processors/{processor_id}", "raw_document": { "content": pdf_file.read(), "mime_type": "application/pdf", }, } ).document original_pdf = pikepdf.Pdf.open(os.path.abspath(args.input)) annotated_pdf = pikepdf.Pdf.new() for page_num, page_info in enumerate(document.pages): annotated_pdf.pages.append(original_pdf.pages[page_num]) print( f"Found { len(page_info.form_fields)} form fields on page {page_num + 1}:" ) # Calculate the max "x" and "y" coordinate values for the PDF # this uses the PDF's own built in measuring units which need # to be used to place annotations page_max_x = float(annotated_pdf.pages[page_num].trimbox[2]) page_max_y = float(annotated_pdf.pages[page_num].trimbox[3]) page_annotations = [] for field in page_info.form_fields: # Use the normalized vertices of the form fields and the max # "x" and "y" coordinates to calculate the position of the # annotation using the PDF's built in measuring units coord1 = field.field_name.bounding_poly.normalized_vertices[0] coord2 = field.field_name.bounding_poly.normalized_vertices[1] rect = [ coord1.x * page_max_x, page_max_y - coord1.y * page_max_y, coord2.x * page_max_x, page_max_y - coord2.y * page_max_y, ] # Extract the parsed name and values of each field # as determined by Document AI's API name = layout_to_text(field.field_name, document.text) value = layout_to_text(field.field_value, document.text) annotation_text = f"{name}: {value}" # Create a PDF annotation for this field name value pair page_annotations.append( pikepdf.Dictionary( Type=pikepdf.Name.Annot, Subtype=pikepdf.Name.Text, Rect=rect, Name=pikepdf.Name.Note, Contents=pikepdf.String(annotation_text), Open=False, ) ) print(f"adding annotation: {annotation_text}") # Add all the annotations for this page annotated_pdf.pages[page_num].Annots = annotated_pdf.make_indirect( pikepdf.Array(page_annotations) ) print(f"Saving annotated PDF to {args.output}.") annotated_pdf.save( os.path.join(args.output), min_version=original_pdf.pdf_version, # Disable annotation modification encryption=pikepdf.Encryption( owner="", user="", allow=pikepdf.Permissions(modify_annotation=False) ), ) print("Done.") return 0 def layout_to_text(layout: dict, text: str) -> str: """ Document AI identifies form fields by their offsets in the entirety of the document's text. This function converts offsets to a string. """ response = "" # If a text segment spans several lines, it will # be stored in different text segments. for segment in layout.text_anchor.text_segments: # type: ignore start_index = ( int(segment.start_index) if segment in layout.text_anchor.text_segments # type: ignore else 0 ) end_index = int(segment.end_index) response += text[start_index:end_index] # remove whitespace response = "".join(response.split("\n")) response = "".join(response.split(":")) return response.strip() def create_processor( client: DocumentProcessorServiceClient, parent: str, processor_type: str ) -> str: """Create a processor for a given processor type.""" processor = client.create_processor( parent=parent, processor=Processor(display_name=processor_type, type_=processor_type), ) return processor.name.split("/")[-1] def find_processor_id_of_type( client: DocumentProcessorServiceClient, parent: str, tartget_processor_type: str ) -> Optional[str]: """Searches for a processor ID for a given processor type.""" for processor in client.list_processors(parent=parent).processors: if processor.type_ == tartget_processor_type: return processor.name.split("/")[-1] return None if __name__ == "__main__": parser = argparse.ArgumentParser(description="Annotate a PDF document.") parser.add_argument( "-i", "--input", help="filepath of input PDF to annotate", required=True ) parser.add_argument("--output", help="path to save annotated PDF") parser.add_argument( "--project-id", help="Project ID to use to call the Document AI API" ) parser.add_argument( "--multi-region-location", help="multi-regional location for document storage and processing", default=DEFAULT_MULTI_REGION_LOCATION, ) parser.add_argument( "--form-processor-type", help='type of form processor e.g. "FORM_W9_PROCESSOR"', default=DEFAULT_PROCESSOR_TYPE, ) sys.exit(main(parser.parse_args()))