import logging from collections.abc import Iterable, Iterator from io import BytesIO from pathlib import Path from typing import TYPE_CHECKING, Optional, Union import pypdfium2 as pdfium from docling_core.types.doc import BoundingBox, CoordOrigin, Size from docling_core.types.doc.page import SegmentedPdfPage, TextCell from docling_parse.pdf_parser import ( ContentConfig, ContentLevel, DecodeConfig, DoclingPdfParser, DoclingThreadedPdfParser, PageParseResult, PdfDocument, RenderConfig, ThreadedPdfParserConfig, ) from PIL import Image from pypdfium2 import PdfPage from docling.backend.managed_pdfium_backend import ( ManagedPdfiumDocumentBackend, ManagedPdfiumPageBackend, ) from docling.backend.pdf_backend import PdfDocumentBackend, PdfPageBackend from docling.datamodel.accelerator_options import AcceleratorOptions from docling.datamodel.backend_options import ( PdfBackendOptions, ThreadedDoclingParseBackendOptions, ) from docling.datamodel.settings import DEFAULT_PAGE_RANGE from docling.exceptions import DocumentLoadError from docling.utils.locks import pypdfium2_lock if TYPE_CHECKING: from docling.datamodel.document import InputDocument _log = logging.getLogger(__name__) def _make_docling_parse_decode_config( *, release_native_memory_every_n_pages: int | None = None, ) -> DecodeConfig: config = DecodeConfig(enforce_same_font=True) if release_native_memory_every_n_pages is not None: config.release_native_memory_every_n_pages = release_native_memory_every_n_pages return config def _make_docling_parse_page_content_config( *, create_words: bool, create_textlines: bool, ) -> ContentConfig: compute = ContentLevel.COMPUTE materialize = ContentLevel.COMPUTE_AND_MATERIALIZE skip = ContentLevel.SKIP return ContentConfig( char_cells_content_level=compute if (create_words or create_textlines) else skip, word_cells_content_level=materialize if create_words else skip, line_cells_content_level=materialize if create_textlines else skip, shapes_content_level=skip, bitmaps_content_level=materialize, include_bitmap_bytes=False, # only need bitmap rectangles for OCR ) class DoclingParsePageBackend(ManagedPdfiumPageBackend): def __init__( self, *, dp_doc: PdfDocument, page_obj: PdfPage, page_no: int, create_words: bool = True, create_textlines: bool = True, keep_chars: bool = False, keep_lines: bool = False, keep_images: bool = True, ): super().__init__() self._ppage = page_obj self._dp_doc: Optional[PdfDocument] = dp_doc self._page_no = page_no self._create_words = create_words self._create_textlines = create_textlines self._keep_chars = keep_chars self._keep_lines = keep_lines self._keep_images = keep_images self._dpage: Optional[SegmentedPdfPage] = None self._unloaded = False self.valid = (self._ppage is not None) and (self._dp_doc is not None) @property def page_no(self) -> int: return self._page_no + 1 def _require_page(self) -> PdfPage: assert self._ppage is not None, "Page backend was unloaded." return self._ppage def _ensure_parsed(self) -> None: if self._dpage is not None: return content_config = _make_docling_parse_page_content_config( create_words=self._create_words, create_textlines=self._create_textlines, ) assert self._dp_doc is not None seg_page = self._dp_doc.get_page( self._page_no + 1, content_config=content_config, ) # In Docling, all TextCell instances are expected with top-left origin. [ tc.to_top_left_origin(seg_page.dimension.height) for tc in seg_page.textline_cells ] [tc.to_top_left_origin(seg_page.dimension.height) for tc in seg_page.char_cells] [tc.to_top_left_origin(seg_page.dimension.height) for tc in seg_page.word_cells] self._dpage = seg_page def is_valid(self) -> bool: return self.valid def get_text_in_rect(self, bbox: BoundingBox) -> str: self._ensure_parsed() assert self._dpage is not None # Find intersecting cells on the page text_piece = "" page_size = self.get_size() scale = ( 1 # FIX - Replace with param in get_text_in_rect across backends (optional) ) for i, cell in enumerate(self._dpage.textline_cells): cell_bbox = ( cell.rect.to_bounding_box() .to_top_left_origin(page_height=page_size.height) .scaled(scale) ) overlap_frac = cell_bbox.intersection_over_self(bbox) if overlap_frac > 0.5: if len(text_piece) > 0: text_piece += " " text_piece += cell.text return text_piece def get_segmented_page(self) -> Optional[SegmentedPdfPage]: self._ensure_parsed() return self._dpage def get_text_cells(self) -> Iterable[TextCell]: self._ensure_parsed() assert self._dpage is not None return self._dpage.textline_cells def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: self._ensure_parsed() assert self._dpage is not None AREA_THRESHOLD = 0 # 32 * 32 images = self._dpage.bitmap_resources for img in images: cropbox = img.rect.to_bounding_box().to_top_left_origin( self.get_size().height ) if cropbox.area() > AREA_THRESHOLD: cropbox = cropbox.scaled(scale=scale) yield cropbox def get_page_image( self, scale: float = 1, cropbox: Optional[BoundingBox] = None ) -> Image.Image: page_size = self.get_size() if not cropbox: cropbox = BoundingBox( l=0, r=page_size.width, t=0, b=page_size.height, coord_origin=CoordOrigin.TOPLEFT, ) padbox = BoundingBox( l=0, r=0, t=0, b=0, coord_origin=CoordOrigin.BOTTOMLEFT ) else: padbox = cropbox.to_bottom_left_origin(page_size.height).model_copy() padbox.r = page_size.width - padbox.r padbox.t = page_size.height - padbox.t with pypdfium2_lock: bitmap = self._ppage.render( scale=scale * 1.5, rotation=0, # no additional rotation crop=padbox.as_tuple(), ) image = bitmap.to_pil().copy() bitmap.close() # We resize the image from 1.5x the given scale to make it sharper. image = image.resize( size=(round(cropbox.width * scale), round(cropbox.height * scale)) ) return image def get_size(self) -> Size: with pypdfium2_lock: page = self._require_page() return Size(width=page.get_width(), height=page.get_height()) # TODO: Take width and height from docling-parse. # return Size( # width=self._dpage.dimension.width, # height=self._dpage.dimension.height, # ) def _close_native_page(self) -> None: if not self._unloaded and self._dp_doc is not None: self._dp_doc.unload_pages((self._page_no + 1, self._page_no + 2)) self._unloaded = True with pypdfium2_lock: if self._ppage is not None: self._ppage.close() self._ppage = None self._dpage = None self._dp_doc = None class DoclingParseDocumentBackend(ManagedPdfiumDocumentBackend): def __init__( self, in_doc: "InputDocument", path_or_stream: Union[BytesIO, Path], options: Optional[PdfBackendOptions] = None, ): if options is None: options = PdfBackendOptions() super().__init__(in_doc, path_or_stream, options) password = ( self.options.password.get_secret_value() if self.options.password else None ) self.dp_doc: Optional[PdfDocument] try: with pypdfium2_lock: self._pdoc = pdfium.PdfDocument(self.path_or_stream, password=password) self.parser = DoclingPdfParser(loglevel="fatal") decode_config = _make_docling_parse_decode_config() self.dp_doc = self.parser.load( path_or_stream=self.path_or_stream, password=password, decode_config=decode_config, ) except RuntimeError as e: # pypdfium2 (PdfiumError) and docling-parse both signal unreadable # bytes by raising RuntimeError; tag it as a load failure. detail = str(e).strip() if detail: raise DocumentLoadError( f"docling-parse could not load document {self.document_hash}: {detail}" ) from e raise DocumentLoadError( f"docling-parse could not load document {self.document_hash}." ) from e if self.dp_doc is None: raise DocumentLoadError( f"docling-parse could not load document {self.document_hash}." ) def page_count(self) -> int: # return len(self._pdoc) # To be replaced with docling-parse API len_1 = len(self._pdoc) assert self.dp_doc is not None len_2 = self.dp_doc.number_of_pages() if len_1 != len_2: _log.error(f"Inconsistent number of pages: {len_1}!={len_2}") return len_2 def load_page( self, page_no: int, create_words: bool = True, create_textlines: bool = True ) -> DoclingParsePageBackend: assert self.dp_doc is not None with pypdfium2_lock: ppage = self._pdoc[page_no] return DoclingParsePageBackend( dp_doc=self.dp_doc, page_obj=ppage, page_no=page_no, create_words=create_words, create_textlines=create_textlines, ) def is_valid(self) -> bool: return self.page_count() > 0 def _close_native_document(self) -> None: if self.dp_doc is not None: self.dp_doc.unload() self.dp_doc = None if self._pdoc is not None: with pypdfium2_lock: try: self._pdoc.close() except Exception: pass self._pdoc = None def _resolve_threaded_page_numbers( path_or_stream: Union[BytesIO, Path], password: Optional[str], page_range: tuple[int, int], ) -> list[int] | None: start_page, end_page = page_range if page_range == DEFAULT_PAGE_RANGE: return None with pypdfium2_lock: pdoc = pdfium.PdfDocument(path_or_stream, password=password) try: page_count = len(pdoc) finally: pdoc.close() clipped_end_page = min(end_page, page_count) if start_page > clipped_end_page: return [] return list(range(start_page, clipped_end_page + 1)) class ThreadedDoclingParsePageBackend(PdfPageBackend): def __init__(self, result: PageParseResult): self._result = result self._seg_page: Optional[SegmentedPdfPage] = None @property def page_no(self) -> int: return self._result.page_number def is_valid(self) -> bool: return self._result.success def get_text_in_rect(self, bbox: BoundingBox) -> str: segmented_page = self.get_segmented_page() if segmented_page is None: return "" text_piece = "" for cell in segmented_page.textline_cells: cell_bbox = cell.rect.to_bounding_box() overlap_frac = cell_bbox.intersection_over_self(bbox) if overlap_frac > 0.5: if text_piece: text_piece += " " text_piece += cell.text return text_piece def get_segmented_page(self) -> Optional[SegmentedPdfPage]: if not self.is_valid(): return None if self._seg_page is None: seg_page = self._result.get_page() page_height = seg_page.dimension.height for tc in seg_page.textline_cells: tc.to_top_left_origin(page_height) for tc in seg_page.char_cells: tc.to_top_left_origin(page_height) for tc in seg_page.word_cells: tc.to_top_left_origin(page_height) self._seg_page = seg_page return self._seg_page def get_text_cells(self) -> Iterable[TextCell]: segmented_page = self.get_segmented_page() if segmented_page is None: return [] return segmented_page.textline_cells def get_bitmap_rects(self, scale: float = 1) -> Iterable[BoundingBox]: segmented_page = self.get_segmented_page() if segmented_page is None: return [] page_height = self.get_size().height cropboxes: list[BoundingBox] = [] for image_resource in segmented_page.bitmap_resources: cropbox = image_resource.rect.to_bounding_box().to_top_left_origin( page_height ) if cropbox.area() > 0: cropboxes.append(cropbox.scaled(scale=scale)) return cropboxes def get_page_image( self, scale: float = 1, cropbox: Optional[BoundingBox] = None ) -> Image.Image: return self._result.get_image(scale=scale, cropbox=cropbox).convert("RGB") def get_size(self) -> Size: return Size(width=self._result.page_width, height=self._result.page_height) def unload(self) -> None: return None class ThreadedDoclingParseDocumentBackend(PdfDocumentBackend): supports_random_page_access = False def __init__( self, in_doc: "InputDocument", path_or_stream: Union[BytesIO, Path], options: Optional[PdfBackendOptions] = None, ): if options is None: options = PdfBackendOptions() super().__init__(in_doc, path_or_stream, options) self.options: PdfBackendOptions self._closed = False password = ( self.options.password.get_secret_value() if self.options.password else None ) requested_page_numbers = _resolve_threaded_page_numbers( self.path_or_stream, password, in_doc.limits.page_range, ) parser_threads = ( self.options.parser_threads if isinstance(self.options, ThreadedDoclingParseBackendOptions) and self.options.parser_threads is not None else AcceleratorOptions().num_threads ) render_config = RenderConfig() render_config.scale = 1.0 native_memory_release_interval = ( self.options.release_native_memory_every_n_pages if isinstance(self.options, ThreadedDoclingParseBackendOptions) else 128 ) decode_config = _make_docling_parse_decode_config( release_native_memory_every_n_pages=native_memory_release_interval, ) content_config = _make_docling_parse_page_content_config( create_words=True, create_textlines=True, ) self.parser = DoclingThreadedPdfParser( parser_config=ThreadedPdfParserConfig( loglevel="fatal", threads=parser_threads, render_config=render_config, page_content_config=content_config, ), decode_config=decode_config, ) self.doc_key = self.parser.load( self.path_or_stream, password=password, page_numbers=requested_page_numbers, ) def is_valid(self) -> bool: return not self._closed and self.page_count() > 0 def page_count(self) -> int: return self.parser.page_count(self.doc_key) def load_page(self, page_no: int) -> PdfPageBackend: raise NotImplementedError( "ThreadedDoclingParseDocumentBackend only supports iter_pages()." ) def iter_pages(self) -> Iterator[ThreadedDoclingParsePageBackend]: for result in self.parser.iterate_results(): yield ThreadedDoclingParsePageBackend(result) def unload(self) -> None: if self._closed: return self._closed = True self.parser.unload(self.doc_key) super().unload()