OCR-Viewer (#1)

Co-authored-by: nathan <nathan.kehler@gmail.com>
Reviewed-on: #1

GoogleDocumentOCR/ocr.py

@@ -0,0 +1,94 @@
+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