""" InvoiceForge AI — utils/pdf_utils.py PDF to image conversion and multi-page data merging. Uses: - pdf2image (poppler wrapper) for high-DPI page rendering - PyMuPDF (fitz) as a fallback if poppler is not installed Provides: - PDFProcessor: convert PDF bytes/path to list of BGR numpy arrays - merge_page_results(): combine extraction results from all pages """ from __future__ import annotations import io import logging import os from pathlib import Path from typing import Generator, Optional import cv2 import numpy as np from PIL import Image logger = logging.getLogger(__name__) DEFAULT_DPI = 300 # High-quality render for OCR MAX_PAGES = 50 # Safety cap for very long documents # ───────────────────────────────────────────────────────────────────────────── # PDF PROCESSOR # ───────────────────────────────────────────────────────────────────────────── class PDFProcessor: """ Converts PDF files to sequences of BGR images for OCR processing. Tries pdf2image first (requires poppler-utils system package), falls back to PyMuPDF (fitz) if poppler is unavailable. """ @staticmethod def pdf_to_images( pdf_bytes: bytes, dpi: int = DEFAULT_DPI, max_pages: int = MAX_PAGES, ) -> list[np.ndarray]: """ Convert PDF bytes to a list of BGR numpy arrays (one per page). Args: pdf_bytes: Raw PDF file contents. dpi: Rendering resolution (300 recommended for OCR). max_pages: Maximum pages to process. Returns: List of BGR numpy arrays, one per page. Raises: ValueError: If the PDF cannot be decoded. """ images: list[np.ndarray] = [] # Try pdf2image (poppler) try: images = PDFProcessor._convert_with_pdf2image( pdf_bytes, dpi, max_pages ) logger.info("PDF converted via pdf2image: %d pages.", len(images)) return images except Exception as exc: logger.warning("pdf2image failed (%s), trying PyMuPDF …", exc) # Fallback: PyMuPDF try: images = PDFProcessor._convert_with_pymupdf( pdf_bytes, dpi, max_pages ) logger.info("PDF converted via PyMuPDF: %d pages.", len(images)) return images except Exception as exc: logger.error("PyMuPDF also failed: %s", exc) raise ValueError( "Could not render PDF. Ensure poppler-utils or PyMuPDF is installed." ) from exc @staticmethod def _convert_with_pdf2image( pdf_bytes: bytes, dpi: int, max_pages: int, ) -> list[np.ndarray]: """Use pdf2image (poppler) to render pages.""" from pdf2image import convert_from_bytes # type: ignore pil_pages: list[Image.Image] = convert_from_bytes( pdf_bytes, dpi=dpi, first_page=1, last_page=max_pages, fmt="JPEG", ) bgr_pages: list[np.ndarray] = [] for pil_img in pil_pages: rgb = np.array(pil_img.convert("RGB")) bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) bgr_pages.append(bgr) return bgr_pages @staticmethod def _convert_with_pymupdf( pdf_bytes: bytes, dpi: int, max_pages: int, ) -> list[np.ndarray]: """Use PyMuPDF (fitz) to render pages.""" import fitz # type: ignore # PyMuPDF pdf_doc = fitz.open(stream=pdf_bytes, filetype="pdf") bgr_pages: list[np.ndarray] = [] zoom = dpi / 72.0 # 72 DPI is PDF default mat = fitz.Matrix(zoom, zoom) for page_num in range(min(len(pdf_doc), max_pages)): page = pdf_doc.load_page(page_num) pixmap = page.get_pixmap(matrix=mat, colorspace=fitz.csRGB) img_bytes = pixmap.tobytes("jpeg") nparr = np.frombuffer(img_bytes, np.uint8) img_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR) if img_bgr is not None: bgr_pages.append(img_bgr) pdf_doc.close() return bgr_pages @staticmethod def pdf_from_request(request) -> bytes: """ Extract PDF bytes from a Flask request. Accepts multipart/form-data with key 'file'. Raises: ValueError: If no PDF found. """ if request.files: file_obj = ( request.files.get("file") or request.files.get("pdf") or request.files.get("image") ) if file_obj: return file_obj.read() raise ValueError( "No PDF found in request. " "Send a PDF as 'file' in multipart/form-data." ) # ───────────────────────────────────────────────────────────────────────────── # MULTI-PAGE RESULT MERGER # ───────────────────────────────────────────────────────────────────────────── def merge_page_results(page_results: list[dict]) -> dict: """ Merge extraction results from multiple pages into a single document result. Strategy: - header: taken from the first page that has non-empty fields - items: concatenated from all pages - totals: accumulated (subtotal, gstTotal, grandTotal are summed, but if any single page has a "grandTotal" that looks like the real document total it takes precedence) - validation: combined errors from all pages - meta: page count, per-page confidence Args: page_results: List of result dicts, one per page. Returns: Merged result dict. """ if not page_results: return _empty_result() merged_header: dict = {} merged_items: list[dict] = [] merged_errors: list[str] = [] per_page_confidence: list[float] = [] subtotal: float = 0.0 gst_total: float = 0.0 grand_total: float = 0.0 doc_grand_total: Optional[float] = None for page_idx, result in enumerate(page_results): header = result.get("header", {}) items = result.get("items", []) totals = result.get("totals", {}) validation = result.get("validation", {}) meta = result.get("meta", {}) # ── Header: prefer the first page with real data ────────────────── if not merged_header or not merged_header.get("vendorName"): merged_header = _merge_header(merged_header, header) # ── Items: accumulate across pages ──────────────────────────────── for item in items: item["_page"] = page_idx + 1 merged_items.append(item) # ── Totals ───────────────────────────────────────────────────────── page_sub = float(totals.get("subtotal", 0.0)) page_gst = float(totals.get("gstTotal", 0.0)) page_gt = float(totals.get("grandTotal", 0.0)) subtotal += page_sub gst_total += page_gst # Last page with a non-zero grandTotal is the document total if page_gt > 0: doc_grand_total = page_gt # ── Errors & confidence ──────────────────────────────────────────── errors = validation.get("errors", []) merged_errors.extend( [f"[Page {page_idx + 1}] {e}" for e in errors] ) per_page_confidence.append(float(validation.get("confidence", 0.85))) grand_total = doc_grand_total if doc_grand_total else round(subtotal + gst_total, 2) mean_conf = float(np.mean(per_page_confidence)) if per_page_confidence else 0.85 return { "header": merged_header, "items": merged_items, "totals": { "subtotal": round(subtotal, 2), "gstTotal": round(gst_total, 2), "grandTotal": round(grand_total, 2), }, "validation": { "isValid": len(merged_errors) == 0, "confidence": round(mean_conf, 4), "errors": merged_errors, "warnings": [], }, "meta": { "documentType": page_results[0].get("meta", {}).get("documentType", "unknown"), "pageCount": len(page_results), "perPageConfidence": per_page_confidence, }, } def _merge_header(existing: dict, new_header: dict) -> dict: """Fill missing header fields from new_header into existing.""" merged = dict(existing) for key, value in new_header.items(): if value and not merged.get(key): merged[key] = value return merged def _empty_result() -> dict: """Return an empty result structure.""" return { "header": { "vendorName": "", "vendorAddress": "", "vendorGSTIN": "", "buyerName": "", "buyerGSTIN": "", "invoiceNumber": "", "invoiceDate": "", "poNumber": "", }, "items": [], "totals": {"subtotal": 0.0, "gstTotal": 0.0, "grandTotal": 0.0}, "validation": { "isValid": False, "confidence": 0.0, "errors": ["Empty PDF"], "warnings": [], }, "meta": {"documentType": "unknown", "pageCount": 0}, }