94 lines
3.1 KiB
Python
94 lines
3.1 KiB
Python
import os
|
|
import io
|
|
from google.api_core.client_options import ClientOptions
|
|
from google.cloud import documentai
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
|
|
PROJECT_ID = os.getenv("GOOGLE_PROJECT_ID")
|
|
LOCATION = os.getenv("GOOGLE_LOCATION")
|
|
PROCESSOR_ID = os.getenv("GOOGLE_PROCESSOR_ID")
|
|
|
|
PAGE_LIMIT = 14 # stay under the 30 page limit
|
|
|
|
|
|
def get_client():
|
|
opts = ClientOptions(api_endpoint=f"{LOCATION}-documentai.googleapis.com")
|
|
return documentai.DocumentProcessorServiceClient(client_options=opts)
|
|
|
|
|
|
def split_pdf(content: bytes, page_limit: int = PAGE_LIMIT) -> list[bytes]:
|
|
"""Split PDF bytes into chunks of page_limit pages."""
|
|
import pypdf
|
|
reader = pypdf.PdfReader(io.BytesIO(content))
|
|
total_pages = len(reader.pages)
|
|
chunks = []
|
|
for start in range(0, total_pages, page_limit):
|
|
writer = pypdf.PdfWriter()
|
|
for i in range(start, min(start + page_limit, total_pages)):
|
|
writer.add_page(reader.pages[i])
|
|
buf = io.BytesIO()
|
|
writer.write(buf)
|
|
chunks.append(buf.getvalue())
|
|
return chunks
|
|
|
|
|
|
def ocr_chunk(client, name: str, content: bytes, mime_type: str, page_offset: int) -> list[dict]:
|
|
"""OCR a single chunk and return pages with corrected page numbers."""
|
|
raw_doc = documentai.RawDocument(content=content, mime_type=mime_type)
|
|
request = documentai.ProcessRequest(name=name, raw_document=raw_doc)
|
|
result = client.process_document(request=request)
|
|
doc = result.document
|
|
|
|
pages = []
|
|
for i, page in enumerate(doc.pages):
|
|
page_text_parts = []
|
|
for block in page.blocks:
|
|
seg = block.layout.text_anchor.text_segments
|
|
for s in seg:
|
|
start = int(s.start_index) if s.start_index else 0
|
|
end = int(s.end_index)
|
|
page_text_parts.append(doc.text[start:end])
|
|
page_text = "\n".join(page_text_parts).strip()
|
|
if page_text:
|
|
pages.append({"page_num": page_offset + i + 1, "text": page_text})
|
|
return pages
|
|
|
|
|
|
def ocr_file(file_path: str) -> list[dict]:
|
|
client = get_client()
|
|
name = client.processor_path(PROJECT_ID, LOCATION, PROCESSOR_ID)
|
|
|
|
ext = os.path.splitext(file_path)[1].lower()
|
|
mime_map = {
|
|
".pdf": "application/pdf",
|
|
".png": "image/png",
|
|
".jpg": "image/jpeg",
|
|
".jpeg": "image/jpeg",
|
|
".tiff": "image/tiff",
|
|
".tif": "image/tiff",
|
|
".gif": "image/gif",
|
|
".bmp": "image/bmp",
|
|
".webp": "image/webp",
|
|
}
|
|
mime_type = mime_map.get(ext, "application/pdf")
|
|
|
|
with open(file_path, "rb") as f:
|
|
content = f.read()
|
|
|
|
# Only PDFs need splitting
|
|
if mime_type == "application/pdf":
|
|
chunks = split_pdf(content)
|
|
print(f" Split into {len(chunks)} chunk(s) of up to {PAGE_LIMIT} pages")
|
|
else:
|
|
chunks = [content]
|
|
|
|
all_pages = []
|
|
for chunk_idx, chunk in enumerate(chunks):
|
|
page_offset = chunk_idx * PAGE_LIMIT
|
|
print(f" OCR chunk {chunk_idx + 1}/{len(chunks)}...")
|
|
pages = ocr_chunk(client, name, chunk, mime_type, page_offset)
|
|
all_pages.extend(pages)
|
|
|
|
return all_pages |