| from __future__ import annotations |
|
|
| import math |
| from pathlib import Path |
| from types import ModuleType |
| from typing import Any |
|
|
| from .errors import InferenceError |
|
|
|
|
| class PDFPageRenderer: |
| """Bounded, page-at-a-time PDFium renderer.""" |
|
|
| def __init__( |
| self, |
| pdf_path: Path, |
| work_dir: Path, |
| *, |
| dpi: int, |
| start_page: int, |
| max_pages: int, |
| max_page_pixels: int, |
| max_total_pixels: int, |
| max_page_rendered_bytes: int, |
| max_rendered_bytes: int, |
| ) -> None: |
| self.pdf_path = pdf_path |
| self.work_dir = work_dir |
| self.dpi = dpi |
| self.start_page = start_page |
| self.max_pages = max_pages |
| self.max_page_pixels = max_page_pixels |
| self.max_total_pixels = max_total_pixels |
| self.max_page_rendered_bytes = max_page_rendered_bytes |
| self.max_rendered_bytes = max_rendered_bytes |
| self._pdfium: ModuleType | None = None |
| self._document: Any = None |
| self.total_pages = 0 |
| self._rendered_pixels = 0 |
| self._rendered_bytes = 0 |
|
|
| def __enter__(self) -> PDFPageRenderer: |
| try: |
| import pypdfium2 as pdfium |
| except ImportError as exc: |
| raise InferenceError("pypdfium2 is missing; reinstall unlimited-ocr-rdna4") from exc |
|
|
| self._pdfium = pdfium |
| try: |
| self._document = pdfium.PdfDocument(str(self.pdf_path)) |
| self.total_pages = len(self._document) |
| except Exception as exc: |
| raise InferenceError(f"could not open PDF {self.pdf_path}: {exc}") from exc |
|
|
| if self.total_pages < 1: |
| self.close() |
| raise InferenceError(f"PDF has no pages: {self.pdf_path}") |
| if self.start_page > self.total_pages: |
| self.close() |
| raise InferenceError(f"start page {self.start_page} exceeds PDF page count {self.total_pages}") |
| return self |
|
|
| def __exit__(self, _exc_type: object, _exc: object, _traceback: object) -> None: |
| self.close() |
|
|
| @property |
| def page_numbers(self) -> range: |
| last_page = min(self.total_pages, self.start_page - 1 + self.max_pages) |
| return range(self.start_page, last_page + 1) |
|
|
| def render(self, page_number: int) -> Path: |
| if self._document is None: |
| raise InferenceError("PDF renderer is not open") |
| page = None |
| bitmap = None |
| image = None |
| output = self.work_dir / f"page-{page_number:06d}.png" |
| try: |
| page = self._document[page_number - 1] |
| width_points, height_points = page.get_size() |
| scale = self.dpi / 72 |
| width_pixels = math.ceil(width_points * scale) |
| height_pixels = math.ceil(height_points * scale) |
| pixels = width_pixels * height_pixels |
| if pixels <= 0 or pixels > self.max_page_pixels: |
| raise InferenceError( |
| f"PDF page {page_number} would render {pixels:,} pixels; limit is {self.max_page_pixels:,}" |
| ) |
| if self._rendered_pixels + pixels > self.max_total_pixels: |
| raise InferenceError( |
| f"PDF rendered-pixel total would exceed {self.max_total_pixels:,} on page {page_number}" |
| ) |
|
|
| bitmap = page.render(scale=scale) |
| image = bitmap.to_pil() |
| image.save(output, format="PNG", optimize=True) |
| rendered_bytes = output.stat().st_size |
| if rendered_bytes > self.max_page_rendered_bytes: |
| output.unlink(missing_ok=True) |
| raise InferenceError( |
| f"PDF page {page_number} rendered to {rendered_bytes:,} bytes; " |
| f"limit is {self.max_page_rendered_bytes:,}" |
| ) |
| if self._rendered_bytes + rendered_bytes > self.max_rendered_bytes: |
| output.unlink(missing_ok=True) |
| raise InferenceError( |
| f"PDF rendered-byte total would exceed {self.max_rendered_bytes:,} on page {page_number}" |
| ) |
| self._rendered_pixels += pixels |
| self._rendered_bytes += rendered_bytes |
| return output |
| except InferenceError: |
| raise |
| except Exception as exc: |
| output.unlink(missing_ok=True) |
| raise InferenceError(f"could not render PDF page {page_number}: {exc}") from exc |
| finally: |
| if image is not None: |
| image.close() |
| if bitmap is not None: |
| bitmap.close() |
| if page is not None: |
| page.close() |
|
|
| def close(self) -> None: |
| if self._document is not None: |
| self._document.close() |
| self._document = None |
|
|
