import argparse
import json
from pathlib import Path
from google.api_core.client_options import ClientOptions
from google.api_core.exceptions import GoogleAPICallError
from google.cloud import documentai
def anchor_text(text, anchor):
return "".join(text[int(segment.start_index):int(segment.end_index)]
for segment in anchor.text_segments)
def process_document(path, project, location, processor, mime_type):
content = path.read_bytes()
if not content:
raise ValueError("Document is empty")
client = documentai.DocumentProcessorServiceClient(
client_options=ClientOptions(
api_endpoint=f"{location}-documentai.googleapis.com"
)
)
try:
name = client.processor_path(project, location, processor)
request = documentai.ProcessRequest(
name=name,
raw_document=documentai.RawDocument(
content=content, mime_type=mime_type
),
)
document = client.process_document(request=request, timeout=120).document
return {
"text": document.text,
"pages": [{
"page_number": page.page_number,
"width": page.dimension.width,
"height": page.dimension.height,
"unit": page.dimension.unit,
"text": anchor_text(document.text, page.layout.text_anchor),
} for page in document.pages],
"entities": [{
"type": entity.type_,
"mention_text": entity.mention_text,
"anchored_text": anchor_text(document.text, entity.text_anchor),
"confidence": entity.confidence,
} for entity in document.entities],
}
finally:
client.transport.close()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Document AI online processing")
parser.add_argument("file", type=Path)
parser.add_argument("--project", required=True)
parser.add_argument("--location", required=True)
parser.add_argument("--processor", required=True)
parser.add_argument("--mime", required=True,
choices=["application/pdf", "image/jpeg", "image/png", "image/tiff"])
args = parser.parse_args()
try:
result = process_document(args.file, args.project, args.location,
args.processor, args.mime)
print(json.dumps(result, indent=2, ensure_ascii=False))
except (OSError, ValueError, GoogleAPICallError) as error:
parser.exit(1, f"Document processing failed: {error}\n")