From 78b6c190602fbdb7c1478bbe3fbf03f8ce02ac00 Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 27 Oct 2026 11:56:19 +0100 Subject: [PATCH 1/6] [io] Lazy page sources & chunked/streaming predictors for long documents and large image collections --- docs/source/api/io.md | 14 + docs/source/api/models.md | 4 +- docs/source/user_guide/loading_documents.md | 61 +++ docs/source/user_guide/ocr_predictor.md | 24 +- doctr/cli/main.py | 58 ++- doctr/io/__init__.py | 1 + doctr/io/elements.py | 35 +- doctr/io/pdf.py | 11 +- doctr/io/reader.py | 164 +++++++- doctr/io/sources.py | 392 ++++++++++++++++++++ doctr/models/kie_predictor/base.py | 7 + doctr/models/kie_predictor/pytorch.py | 23 ++ doctr/models/predictor/base.py | 72 +++- doctr/models/predictor/pytorch.py | 23 ++ doctr/models/zoo.py | 6 +- doctr/utils/multithreading.py | 68 +++- tests/common/test_cli.py | 14 +- tests/common/test_io_sources.py | 334 +++++++++++++++++ tests/pytorch/test_models_zoo_pt.py | 155 +++++++- 19 files changed, 1409 insertions(+), 57 deletions(-) create mode 100644 doctr/io/sources.py create mode 100644 tests/common/test_io_sources.py diff --git a/docs/source/api/io.md b/docs/source/api/io.md index 5957cf70c9..82b84bc8e0 100644 --- a/docs/source/api/io.md +++ b/docs/source/api/io.md @@ -14,6 +14,20 @@ Read documents into pages, represent the predictions as a tree of elements, and .. automethod:: from_pdf .. automethod:: from_images .. automethod:: from_url + .. automethod:: from_files + +.. autoclass:: PageSource + + .. automethod:: iter_batches + .. automethod:: close + +.. autoclass:: PDFSource + +.. autoclass:: ImageSource + +.. autoclass:: ConcatSource + + .. automethod:: locate .. autofunction:: read_pdf diff --git a/docs/source/api/models.md b/docs/source/api/models.md index 00bbdf32b4..091f0efff6 100644 --- a/docs/source/api/models.md +++ b/docs/source/api/models.md @@ -15,10 +15,10 @@ directly. See the [model zoo](../models/index.md) for benchmarks. .. autofunction:: doctr.models.kie_predictor .. autoclass:: doctr.models.predictor.OCRPredictor - :members: add_hook + :members: add_hook, stream .. autoclass:: doctr.models.kie_predictor.KIEPredictor - :members: add_hook + :members: add_hook, stream .. autoclass:: doctr.models.builder.DocumentBuilder diff --git a/docs/source/user_guide/loading_documents.md b/docs/source/user_guide/loading_documents.md index 5798f20eb6..4b10e68726 100644 --- a/docs/source/user_guide/loading_documents.md +++ b/docs/source/user_guide/loading_documents.md @@ -114,6 +114,67 @@ pages = DocumentFile.from_pdf("cover_letter.pdf") + DocumentFile.from_images(["i result = predictor(pages) # one Document with all the pages, in order ``` +{meth}`~doctr.io.DocumentFile.from_files` does the same in a single call, for any mix of PDFs, images, folders and +URLs (folders are expanded to the PDFs and images they contain, sorted by name): + +```python +pages = DocumentFile.from_files(["cover_letter.pdf", "id_front.jpg", "id_back.jpg", "path/to/scans/"]) +``` + +## Long documents & large collections + +Every page of a list is decoded in memory: at the default resolution, an A4 page weighs about 11 MB, so a 1,000-page +PDF would need more than 10 GB. Pass `lazy=True` to get a **page source** instead: it behaves like the list (`len`, +indexing, slicing, iteration), but a page is only decoded when it is accessed. + +```python +from doctr.io import DocumentFile +from doctr.models import ocr_predictor + +predictor = ocr_predictor(pretrained=True) + +doc = DocumentFile.from_pdf("path/to/1000_pages.pdf", lazy=True) +print(len(doc), doc[0].shape) # only the first page is rendered + +result = predictor(doc) # one call, a bounded memory footprint +``` + +The predictor processes a page source chunk by chunk (`page_chunk_size` pages at a time, 32 by default), and decodes +the next chunk in a background thread while the models run on the current one. The pages of the returned +{class}`~doctr.io.Document` only keep a reference to their image: `result.pages[i].page` renders it again on access, +so `show()` and the figure exports keep working. + +The same goes for images, folders and mixed inputs: + +```python +# Every image of a folder (sorted by name), decoded on demand +doc = DocumentFile.from_images("path/to/100k_scans/", lazy=True) + +# Any mix of PDFs, images, folders and URLs +doc = DocumentFile.from_files(["contracts/", "annex.pdf", "signature.png"], lazy=True) +source_idx, page_idx = doc.locate(42) # the input that the 43rd page comes from +``` + +To write the results as you go instead of keeping them all, iterate over +{meth}`~doctr.models.predictor.OCRPredictor.stream`: it yields the pages one by one, as soon as they are processed. + +```python +with open("result.txt", "w") as f: + for page in predictor.stream(doc): + f.write(f"--- page {page.page_idx + 1} ---\n{page.render()}\n") +``` + +| Class | Pages | +| ---------------------------- | ------------------------------------------------------------------ | +| {class}`~doctr.io.PDFSource` | The pages of a PDF, rendered on demand (`scale`, `password`, `rgb_mode`) | +| {class}`~doctr.io.ImageSource` | Image files, folders or bytes, decoded on demand | +| {class}`~doctr.io.ConcatSource` | The concatenation of several sources (also obtained with `+`) | + +:::{tip} +Missing files and invalid PDFs are reported when the source is created. Images, on the other hand, are only decoded +when they are processed: an unreadable image raises a `ValueError` naming the file. +::: + ## Low-level readers `DocumentFile` is a thin wrapper around these functions, which you can use directly when you need more control: diff --git a/docs/source/user_guide/ocr_predictor.md b/docs/source/user_guide/ocr_predictor.md index f43b950039..20c54cc098 100644 --- a/docs/source/user_guide/ocr_predictor.md +++ b/docs/source/user_guide/ocr_predictor.md @@ -286,22 +286,30 @@ Good to know: - **Reuse the predictor**: building it is the expensive part. - **Batch your pages**: one call on a list of pages is faster than one call per page, as the models process several pages and word crops at once (see `det_bs` and `reco_bs` above). -- **Free memory between documents**: the predictor returns plain Python objects, so drop the results you no - longer need. +- **Open long documents lazily**: with `lazy=True`, pages are only decoded when they are processed, and the predictor + works chunk by chunk (`page_chunk_size` pages at a time). A 1,000-page PDF or a folder of 100k images can then be + passed in a single call (see [](loading_documents.md#long-documents-large-collections)). +- **Stream the results**: {meth}`~doctr.models.predictor.OCRPredictor.stream` yields the pages as soon as they are + processed, so that the results can be written as you go. ```python -from pathlib import Path - from doctr.io import DocumentFile from doctr.models import ocr_predictor -predictor = ocr_predictor(pretrained=True) +predictor = ocr_predictor(pretrained=True, page_chunk_size=32) -for path in sorted(Path("invoices").glob("*.pdf")): - result = predictor(DocumentFile.from_pdf(path)) - path.with_suffix(".txt").write_text(result.render()) +# Every PDF and image of a folder, as one lazy document +doc = DocumentFile.from_files("invoices/", lazy=True) + +for page in predictor.stream(doc): + source_idx, page_idx = doc.locate(page.page_idx) + print(f"input #{source_idx}, page {page_idx + 1}: {page.render()[:80]}") ``` +Inputs longer than `page_chunk_size` pages, page sources and iterators (e.g. a generator of frames) are all processed +chunk by chunk, and `page_idx` is always the index of the page in the whole input. Hooks are therefore called once +per chunk, with the location predictions of its pages. + ## Next steps ::::{grid} 1 2 2 2 diff --git a/doctr/cli/main.py b/doctr/cli/main.py index e40c1bd2a1..6dc17768e0 100644 --- a/doctr/cli/main.py +++ b/doctr/cli/main.py @@ -11,10 +11,10 @@ from pathlib import Path from typing import Any -import numpy as np import torch +from tqdm import tqdm -from doctr.io import DocumentFile +from doctr.io import ConcatSource, Document, DocumentFile, PageSource from doctr.io.figures import IMAGE_FORMATS, IMAGE_MODES, FigureEncoder from doctr.models import ocr_predictor from doctr.models._utils import _CONFIDENCE_AGGREGATIONS @@ -95,26 +95,26 @@ def _resolve_format(args: argparse.Namespace) -> str: return SUFFIX_ALIASES.get(Path(args.output).suffix.lower(), "json") -def _load_document(args: argparse.Namespace) -> list[np.ndarray]: - """Load every input into a single list of pages +def _load_document(args: argparse.Namespace) -> ConcatSource: + """Open every input as a single lazy document: pages are only decoded when they are processed Args: args: the parsed command-line arguments Returns: - the list of pages decoded as numpy ndarray of shape H x W x 3 + the concatenation of the inputs, decoding its pages on demand """ pdf_kwargs: dict[str, Any] = {"scale": args.pdf_scale, "password": args.pdf_password} - pages: list[np.ndarray] = [] + sources: list[PageSource] = [] for input_path in args.input_path: try: if input_path.lower().startswith(("http://", "https://")): - pages.extend(DocumentFile.from_url(input_path, **pdf_kwargs)) + sources.append(DocumentFile.from_url(input_path, lazy=True, **pdf_kwargs)) elif input_path.lower().endswith(".pdf"): - pages.extend(DocumentFile.from_pdf(input_path, **pdf_kwargs)) + sources.append(DocumentFile.from_pdf(input_path, lazy=True, **pdf_kwargs)) else: - pages.extend(DocumentFile.from_images(input_path)) + sources.append(DocumentFile.from_images(input_path, lazy=True)) logger.info(f"Document loaded successfully from {input_path}") except FileNotFoundError: logger.error(f"File not found: {input_path}") @@ -126,7 +126,31 @@ def _load_document(args: argparse.Namespace) -> list[np.ndarray]: logger.error(f"Error occurred while loading the document: {e}") sys.exit(1) - return pages + return ConcatSource(sources) + + +def _predict(model: Any, doc: ConcatSource, args: argparse.Namespace) -> Document: + """Run the predictor on the document chunk by chunk, with a progress bar + + Args: + model: the predictor + doc: the lazy document + args: the parsed command-line arguments + + Returns: + the predicted document + """ + try: + pages = list( + tqdm(model.stream(doc), total=len(doc), unit="page", desc="OCR", disable=args.quiet, file=sys.stderr) + ) + except (FileNotFoundError, ValueError) as e: + # Images are only decoded when they are processed + logger.error(f"File could not be read as a valid image or PDF: {e}") + sys.exit(1) + finally: + doc.close() + return Document(pages=pages) def _resolve_device(device: str | None) -> torch.device: @@ -227,6 +251,7 @@ def _build_predictor(args: argparse.Namespace) -> Any: "disable_page_orientation": args.disable_page_orientation, "disable_crop_orientation": args.disable_crop_orientation, "preserve_original_coords": args.preserve_original_coords, + "page_chunk_size": args.page_chunk_size, # `DocumentBuilder` keyword args "resolve_lines": args.resolve_lines, "resolve_blocks": args.resolve_blocks, @@ -368,8 +393,8 @@ def main(argv=None): doc = _load_document(args) # perform OCR - logger.info("Performing OCR...") - result = model(doc) + logger.info(f"Performing OCR on {len(doc)} page(s)...") + result = _predict(model, doc, args) # save results to the requested format _save_results(result, fmt, args) @@ -459,6 +484,12 @@ def _parse_args(argv=None): ) parser.add_argument("--det_bs", type=int, default=2, help="batch size for detection") parser.add_argument("--reco_bs", type=int, default=128, help="batch size for recognition") + parser.add_argument( + "--page_chunk_size", + type=int, + default=32, + help="number of pages decoded & processed at once: bounds the memory footprint on long documents", + ) parser.add_argument("--detect_orientation", action="store_true", help="automatically detect page orientation") parser.add_argument("--detect_language", action="store_true", help="detect language of the text") parser.add_argument("--detect_layout", action="store_true", help="attach the detected layout regions to each page") @@ -573,6 +604,9 @@ def _parse_args(argv=None): if value is not None and not 0.0 <= value <= 1.0: parser.error(f"--{name} must be between 0 and 1") + if args.page_chunk_size < 1: + parser.error("--page_chunk_size must be a positive integer") + return args diff --git a/doctr/io/__init__.py b/doctr/io/__init__.py index 1744c3b3d8..00241df425 100644 --- a/doctr/io/__init__.py +++ b/doctr/io/__init__.py @@ -5,3 +5,4 @@ from .image import * from .pdf import * from .reader import * +from .sources import * diff --git a/doctr/io/elements.py b/doctr/io/elements.py index dd560e5612..e13dc9fd67 100644 --- a/doctr/io/elements.py +++ b/doctr/io/elements.py @@ -3,6 +3,7 @@ # This program is licensed under the Apache License 2.0. # See LICENSE or go to for full license details. +from collections.abc import Callable from typing import Any import numpy as np @@ -41,6 +42,28 @@ def _empty_page_image(page: np.ndarray | None) -> np.ndarray: return page if page is not None else np.zeros((0, 0, 3), dtype=np.uint8) +# A page image, or a zero-argument callable decoding it on demand. Predictions made on a lazy page source +# (cf. `doctr.io.PageSource`) only keep a reference to their page, so that their memory footprint does not grow +# with the image of every page of a long document. +PageImage = np.ndarray | Callable[[], np.ndarray] + + +def _get_page_image(self: Any) -> np.ndarray: + image = self._page + return image() if callable(image) else image + + +def _set_page_image(self: Any, image: PageImage) -> None: + self._page = image + + +_page_image_property = property( + _get_page_image, + _set_page_image, + doc="the page image as a numpy array (decoded on access when the page was predicted from a lazy page source)", +) + + class Element(NestedObject): """Implements an abstract document element with exporting and text rendering capabilities""" @@ -435,7 +458,7 @@ class Page(PageExportsMixin, Element): """Implements a page element as a collection of blocks Args: - page: image encoded as a numpy array in uint8 + page: image encoded as a numpy array in uint8, or a callable returning it (decoded on access) blocks: list of block elements page_idx: the index of the page in the input raw document dimensions: the page size in pixels in format (height, width) @@ -451,9 +474,11 @@ class Page(PageExportsMixin, Element): layout: list[LayoutElement] = [] tables: list[Table] = [] + page = _page_image_property + def __init__( self, - page: np.ndarray, + page: PageImage, blocks: list[Block], page_idx: int, dimensions: tuple[int, int], @@ -527,7 +552,7 @@ class KIEPage(KIEPageExportsMixin, Element): Args: predictions: Dictionary with list of block elements for each detection class - page: image encoded as a numpy array in uint8 + page: image encoded as a numpy array in uint8, or a callable returning it (decoded on access) page_idx: the index of the page in the input raw document dimensions: the page size in pixels in format (height, width) orientation: a dictionary with the value of the rotation angle in degress and confidence of the prediction @@ -540,9 +565,11 @@ class KIEPage(KIEPageExportsMixin, Element): predictions: dict[str, list[Prediction]] = {} layout: list[LayoutElement] = [] + page = _page_image_property + def __init__( self, - page: np.ndarray, + page: PageImage, predictions: dict[str, list[Prediction]], page_idx: int, dimensions: tuple[int, int], diff --git a/doctr/io/pdf.py b/doctr/io/pdf.py index ee5542da68..432711f039 100644 --- a/doctr/io/pdf.py +++ b/doctr/io/pdf.py @@ -6,10 +6,11 @@ from typing import Any import numpy as np -import pypdfium2 as pdfium from doctr.utils.common_types import AbstractFile +from .sources import PDFSource + __all__ = ["read_pdf"] @@ -35,9 +36,9 @@ def read_pdf( Returns: the list of pages decoded as numpy ndarray of shape H x W x C """ - # Rasterise pages to numpy ndarrays with pypdfium2 - pdf = pdfium.PdfDocument(file, password=password) + # Rasterise pages to numpy ndarrays with pypdfium2 (use `PDFSource` directly to render them on demand) + source = PDFSource(file, scale=scale, rgb_mode=rgb_mode, password=password, **kwargs) try: - return [page.render(scale=scale, rev_byteorder=rgb_mode, **kwargs).to_numpy() for page in pdf] + return list(source) finally: - pdf.close() + source.close() diff --git a/doctr/io/reader.py b/doctr/io/reader.py index 090b2c73c2..764b4137c9 100644 --- a/doctr/io/reader.py +++ b/doctr/io/reader.py @@ -5,6 +5,7 @@ from collections.abc import Sequence from pathlib import Path +from typing import Any, Literal, overload import numpy as np @@ -12,33 +13,67 @@ from doctr.utils.common_types import AbstractFile from .html import read_html -from .image import read_img_as_numpy -from .pdf import read_pdf +from .sources import IMAGE_EXTENSIONS, ConcatSource, ImageSource, PageSource, PDFSource, _is_pdf, _list_files __all__ = ["DocumentFile"] +def _materialize(source: PageSource) -> list[np.ndarray]: + """Decode every page of a source into memory""" + try: + return list(source) + finally: + source.close() + + class DocumentFile: - """Read a document from multiple extensions""" + """Read a document from multiple extensions + + Every reader decodes all the pages into a list of numpy arrays by default. Pass `lazy=True` to get a + :class:`~doctr.io.PageSource` instead: it behaves like that list, but a page is only decoded when it is + accessed. Predictors process page sources chunk by chunk, so that very long documents (1000+ pages) or + huge collections of images can be passed in a single call with a bounded memory footprint. + """ + + @overload + @classmethod + def from_pdf(cls, file: AbstractFile, lazy: Literal[False] = ..., **kwargs: Any) -> list[np.ndarray]: ... + + @overload + @classmethod + def from_pdf(cls, file: AbstractFile, lazy: Literal[True], **kwargs: Any) -> PDFSource: ... @classmethod - def from_pdf(cls, file: AbstractFile, **kwargs) -> list[np.ndarray]: + def from_pdf(cls, file: AbstractFile, lazy: bool = False, **kwargs: Any) -> list[np.ndarray] | PDFSource: """Read a PDF file >>> from doctr.io import DocumentFile >>> doc = DocumentFile.from_pdf("path/to/your/doc.pdf") + >>> # Pages are rendered on demand: suited for very long documents + >>> doc = DocumentFile.from_pdf("path/to/your/long_doc.pdf", lazy=True) Args: file: the path to the PDF file or a binary stream - **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render` + lazy: if True, return a :class:`~doctr.io.PDFSource` rendering the pages on demand + **kwargs: additional parameters to :class:`~doctr.io.PDFSource` (`scale`, `rgb_mode`, `password`) + and :meth:`pypdfium2.PdfPage.render` Returns: - the list of pages decoded as numpy ndarray of shape H x W x 3 + the list of pages decoded as numpy ndarray of shape H x W x 3, or a lazy source if `lazy` is True """ - return read_pdf(file, **kwargs) + source = PDFSource(file, **kwargs) + return source if lazy else _materialize(source) + @overload @classmethod - def from_url(cls, url: str, **kwargs) -> list[np.ndarray]: + def from_url(cls, url: str, lazy: Literal[False] = ..., **kwargs: Any) -> list[np.ndarray]: ... + + @overload + @classmethod + def from_url(cls, url: str, lazy: Literal[True], **kwargs: Any) -> PDFSource: ... + + @classmethod + def from_url(cls, url: str, lazy: bool = False, **kwargs: Any) -> list[np.ndarray] | PDFSource: """Interpret a web page as a PDF document >>> from doctr.io import DocumentFile @@ -46,10 +81,11 @@ def from_url(cls, url: str, **kwargs) -> list[np.ndarray]: Args: url: the URL of the target web page - **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render` + lazy: if True, return a :class:`~doctr.io.PDFSource` rendering the pages on demand + **kwargs: additional parameters to :class:`~doctr.io.PDFSource` and :meth:`pypdfium2.PdfPage.render` Returns: - the list of pages decoded as numpy ndarray of shape H x W x 3 + the list of pages decoded as numpy ndarray of shape H x W x 3, or a lazy source if `lazy` is True """ requires_package( "weasyprint", @@ -57,23 +93,117 @@ def from_url(cls, url: str, **kwargs) -> list[np.ndarray]: + "Installation instructions: https://doc.courtbouillon.org/weasyprint/stable/first_steps.html#installation", ) pdf_stream = read_html(url) - return cls.from_pdf(pdf_stream, **kwargs) + source = PDFSource(pdf_stream, **kwargs) + return source if lazy else _materialize(source) + + @overload + @classmethod + def from_images( + cls, files: Sequence[AbstractFile] | AbstractFile, lazy: Literal[False] = ..., **kwargs: Any + ) -> list[np.ndarray]: ... + @overload @classmethod - def from_images(cls, files: Sequence[AbstractFile] | AbstractFile, **kwargs) -> list[np.ndarray]: + def from_images( + cls, files: Sequence[AbstractFile] | AbstractFile, lazy: Literal[True], **kwargs: Any + ) -> ImageSource: ... + + @classmethod + def from_images( + cls, files: Sequence[AbstractFile] | AbstractFile, lazy: bool = False, **kwargs: Any + ) -> list[np.ndarray] | ImageSource: """Read an image file (or a collection of image files) and convert it into an image in numpy format >>> from doctr.io import DocumentFile >>> pages = DocumentFile.from_images(["path/to/your/page1.png", "path/to/your/page2.png"]) + >>> # Every image of a folder (sorted by name), decoded on demand + >>> pages = DocumentFile.from_images("path/to/your/folder", lazy=True) Args: - files: the path to the image file or a binary stream, or a collection of those + files: the path to an image file, to a directory of images, or an image as bytes, or a collection of those + lazy: if True, return a :class:`~doctr.io.ImageSource` decoding the images on demand **kwargs: additional parameters to :meth:`doctr.io.image.read_img_as_numpy` Returns: - the list of pages decoded as numpy ndarray of shape H x W x 3 + the list of pages decoded as numpy ndarray of shape H x W x 3, or a lazy source if `lazy` is True """ - if isinstance(files, (str, Path, bytes)): - files = [files] + source = ImageSource(files, **kwargs) + return source if lazy else _materialize(source) + + @overload + @classmethod + def from_files( + cls, + files: Sequence[AbstractFile] | AbstractFile, + lazy: Literal[False] = ..., + pdf_kwargs: dict[str, Any] | None = ..., + image_kwargs: dict[str, Any] | None = ..., + ) -> list[np.ndarray]: ... + + @overload + @classmethod + def from_files( + cls, + files: Sequence[AbstractFile] | AbstractFile, + lazy: Literal[True], + pdf_kwargs: dict[str, Any] | None = ..., + image_kwargs: dict[str, Any] | None = ..., + ) -> ConcatSource: ... + + @classmethod + def from_files( + cls, + files: Sequence[AbstractFile] | AbstractFile, + lazy: bool = False, + pdf_kwargs: dict[str, Any] | None = None, + image_kwargs: dict[str, Any] | None = None, + ) -> list[np.ndarray] | ConcatSource: + """Read any mix of PDFs, images, directories and URLs as a single document, in the given order - return [read_img_as_numpy(file, **kwargs) for file in files] + >>> from doctr.io import DocumentFile + >>> doc = DocumentFile.from_files(["cover.png", "report.pdf", "path/to/scans/"], lazy=True) + >>> source_idx, page_idx = doc.locate(42) # the input that the 43rd page comes from + + Args: + files: a path (to a PDF, an image or a directory), a URL, a PDF or an image as bytes, + or a collection of those. Directories are expanded (non-recursively) to the PDFs and images they + contain, sorted by name. + lazy: if True, return a :class:`~doctr.io.ConcatSource` decoding the pages on demand + pdf_kwargs: keyword arguments of :class:`~doctr.io.PDFSource`, applied to every PDF and URL + image_kwargs: keyword arguments of :meth:`doctr.io.image.read_img_as_numpy`, applied to every image + + Returns: + the list of pages decoded as numpy ndarray of shape H x W x 3, or a lazy source if `lazy` is True + """ + pdf_kwargs = pdf_kwargs or {} + image_kwargs = image_kwargs or {} + + inputs: list[AbstractFile] = [] + for file in [files] if isinstance(files, (str, Path, bytes)) else files: + if isinstance(file, (str, Path)) and Path(file).is_dir(): + inputs.extend(_list_files(Path(file), (".pdf", *IMAGE_EXTENSIONS))) + else: + inputs.append(file) + + sources: list[PageSource] = [] + images: list[AbstractFile] = [] + + def _flush_images() -> None: + # Consecutive images are grouped into a single source + if images: + sources.append(ImageSource(list(images), **image_kwargs)) + images.clear() + + for file in inputs: + if isinstance(file, str) and file.lower().startswith(("http://", "https://")): + _flush_images() + sources.append(cls.from_url(file, lazy=True, **pdf_kwargs)) + elif isinstance(file, (str, Path, bytes)) and _is_pdf(file): + _flush_images() + sources.append(PDFSource(file, **pdf_kwargs)) + else: + images.append(file) + _flush_images() + + source = ConcatSource(sources) + return source if lazy else _materialize(source) diff --git a/doctr/io/sources.py b/doctr/io/sources.py new file mode 100644 index 0000000000..12ba4e439f --- /dev/null +++ b/doctr/io/sources.py @@ -0,0 +1,392 @@ +# Copyright (C) 2021-2026, Mindee. + +# This program is licensed under the Apache License 2.0. +# See LICENSE or go to for full license details. + +import atexit +import threading +import weakref +from bisect import bisect_right +from collections.abc import Iterable, Iterator, Sequence +from itertools import accumulate +from pathlib import Path +from typing import Any, overload + +import cv2 +import numpy as np +import pypdfium2 as pdfium + +from doctr.utils.common_types import AbstractFile + +from .image import read_img_as_numpy + +__all__ = ["PageSource", "PDFSource", "ImageSource", "ConcatSource", "IMAGE_EXTENSIONS"] + +# File extensions recognized as images when a directory is expanded +IMAGE_EXTENSIONS: tuple[str, ...] = ( + ".jpg", + ".jpeg", + ".png", + ".bmp", + ".tif", + ".tiff", + ".webp", + ".jp2", + ".pbm", + ".pgm", + ".ppm", +) + +# pdfium is not thread-safe, not even across different documents: every call into it is serialized, +# which makes it safe to render pages in a background thread (cf. `OCRPredictor.stream`) +_PDFIUM_LOCK = threading.RLock() +# Sources holding an open PDF: they are closed at exit, before pypdfium2 reports the documents left open +_OPEN_PDF_SOURCES: "weakref.WeakSet[PDFSource]" = weakref.WeakSet() + + +@atexit.register +def _close_pdf_sources() -> None: # pragma: no cover + for source in list(_OPEN_PDF_SOURCES): + source.close() + + +class PageSource(Sequence[np.ndarray]): + """Abstract lazy collection of pages: a page is only decoded when it is accessed. + + A page source behaves like a read-only list of H x W x 3 numpy arrays (`len`, indexing, slicing, iteration), + but only keeps the encoded inputs in memory. It can be passed to the predictors as-is: they process it in + chunks, so that documents of any length are handled with a bounded memory footprint. + + >>> from doctr.io import DocumentFile + >>> doc = DocumentFile.from_pdf("path/to/your/doc.pdf", lazy=True) + >>> len(doc), doc[0].shape + >>> for batch in doc.iter_batches(16): + ... ... + """ + + def __len__(self) -> int: + raise NotImplementedError + + def _load(self, idx: int) -> np.ndarray: + """Decode the page at a given (non-negative, in-range) index""" + raise NotImplementedError + + @overload + def __getitem__(self, idx: int) -> np.ndarray: ... + + @overload + def __getitem__(self, idx: slice) -> "PageSource": ... + + def __getitem__(self, idx: int | slice) -> "np.ndarray | PageSource": + if isinstance(idx, slice): + return _PageView(self, range(len(self))[idx]) + if not isinstance(idx, (int, np.integer)): + raise TypeError(f"page indices must be integers or slices, not {type(idx).__name__}") + num_pages = len(self) + _idx = int(idx) + num_pages if idx < 0 else int(idx) + if not 0 <= _idx < num_pages: + raise IndexError(f"page index {idx} is out of range for a document of {num_pages} page(s)") + return self._load(_idx) + + def __iter__(self) -> Iterator[np.ndarray]: + for idx in range(len(self)): + yield self._load(idx) + + def __add__(self, other: Sequence[np.ndarray]) -> "ConcatSource": + if not isinstance(other, Sequence): + return NotImplemented + return ConcatSource([self, other]) + + def __radd__(self, other: Sequence[np.ndarray]) -> "ConcatSource": + if not isinstance(other, Sequence): + return NotImplemented + return ConcatSource([other, self]) + + def iter_batches(self, batch_size: int) -> Iterator[list[np.ndarray]]: + """Iterate over the pages by batches, only `batch_size` pages being decoded at a time + + Args: + batch_size: the number of pages per batch + + Returns: + an iterator of lists of pages + """ + if batch_size < 1: + raise ValueError("`batch_size` must be a positive integer") + batch: list[np.ndarray] = [] + for page in self: + batch.append(page) + if len(batch) == batch_size: + yield batch + batch = [] + if batch: + yield batch + + def close(self) -> None: + """Release the resources held by the source (e.g. open PDF handles). The source can still be used + afterwards: the resources are acquired again on demand.""" + + def __enter__(self) -> "PageSource": + return self + + def __exit__(self, *args: Any) -> None: + self.close() + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(num_pages={len(self)})" + + +class PDFSource(PageSource): + """Lazy PDF document: pages are rasterized on access, so a 1000+ page PDF only costs the memory of + the pages being processed. + + >>> from doctr.io import PDFSource + >>> doc = PDFSource("path/to/your/doc.pdf", scale=2) + >>> first_page = doc[0] + + Args: + file: the path to the PDF file or its content as bytes + scale: rendering scale (1 corresponds to 72dpi) + rgb_mode: if True, the output will be RGB, otherwise BGR + password: a password to unlock the document, if encrypted + **kwargs: additional parameters to :meth:`pypdfium2.PdfPage.render` + """ + + def __init__( + self, + file: AbstractFile, + scale: float = 2, + rgb_mode: bool = True, + password: str | None = None, + **kwargs: Any, + ) -> None: + self.file = file + self.scale = scale + self.rgb_mode = rgb_mode + self.password = password + self.render_kwargs = kwargs + self._pdf: pdfium.PdfDocument | None = None + # Open the document eagerly: invalid inputs fail here rather than in the middle of a prediction + with _PDFIUM_LOCK: + pdf = self._open() + self._num_pages = len(pdf) + # Do not hold a file handle per document: this matters when thousands of PDFs are concatenated + self.close() + + def _open(self) -> pdfium.PdfDocument: + if self._pdf is None: + file = str(self.file) if isinstance(self.file, Path) else self.file + self._pdf = pdfium.PdfDocument(file, password=self.password) + _OPEN_PDF_SOURCES.add(self) + return self._pdf + + def __len__(self) -> int: + return self._num_pages + + def _load(self, idx: int) -> np.ndarray: + with _PDFIUM_LOCK: + page = self._open()[idx] + try: + bitmap = page.render(scale=self.scale, rev_byteorder=self.rgb_mode, **self.render_kwargs) + # The bitmap buffer is allocated by Python: the array safely outlives the pdfium page & document + return bitmap.to_numpy() + finally: + page.close() + + def __iter__(self) -> Iterator[np.ndarray]: + was_open = self._pdf is not None + try: + for idx in range(len(self)): + yield self._load(idx) + finally: + # Release the file handle once the whole document has been read + if not was_open: + self.close() + + def close(self) -> None: + with _PDFIUM_LOCK: + if self._pdf is not None: + self._pdf.close() + self._pdf = None + _OPEN_PDF_SOURCES.discard(self) + + def __getstate__(self) -> dict[str, Any]: + # pdfium handles cannot be pickled: the document is re-opened on demand + state = self.__dict__.copy() + state["_pdf"] = None + return state + + def __repr__(self) -> str: + name = self.file if isinstance(self.file, (str, Path)) else "" + return f"{self.__class__.__name__}({name!s}, num_pages={len(self)})" + + +class ImageSource(PageSource): + """Lazy collection of images: each image is only decoded when it is accessed, so that tens of thousands + of files can be processed in a single call. + + >>> from doctr.io import ImageSource + >>> doc = ImageSource(["path/to/page1.png", "path/to/page2.png"]) + >>> doc = ImageSource("path/to/a/folder") # every image of the folder, sorted by name + + Args: + files: the path to an image file, to a directory of images, or an image as bytes, or a collection of those + **kwargs: additional parameters to :meth:`doctr.io.image.read_img_as_numpy` + """ + + def __init__(self, files: Sequence[AbstractFile] | AbstractFile, **kwargs: Any) -> None: + self.files: list[AbstractFile] = [] + for file in [files] if isinstance(files, (str, Path, bytes)) else files: + if isinstance(file, (str, Path)): + path = Path(file) + if path.is_dir(): + self.files.extend(_list_files(path, IMAGE_EXTENSIONS)) + continue + if not path.is_file(): + raise FileNotFoundError(f"unable to access {file}") + elif not isinstance(file, bytes): + raise TypeError(f"unsupported object type for an image file: {type(file).__name__}") + self.files.append(file) + self.read_kwargs = kwargs + + def __len__(self) -> int: + return len(self.files) + + def _load(self, idx: int) -> np.ndarray: + file = self.files[idx] + try: + return read_img_as_numpy(file, **self.read_kwargs) + except ValueError as e: + # Images are decoded lazily: point to the faulty file among possibly thousands of them + name = file if isinstance(file, (str, Path)) else f"" + raise ValueError(f"unable to read the image {name!s}: {e}") from e + + +class ConcatSource(PageSource): + """Concatenation of several page sources (or in-memory lists of pages) into a single lazy document. + + >>> from doctr.io import ConcatSource, PDFSource, ImageSource + >>> doc = ConcatSource([PDFSource("report.pdf"), ImageSource("scans/")]) + >>> source_idx, page_idx = doc.locate(42) # where does the 43rd page come from? + + Args: + sources: the sequences of pages to concatenate + """ + + def __init__(self, sources: Iterable[Sequence[np.ndarray]]) -> None: + self.sources: list[Sequence[np.ndarray]] = [] + for source in sources: + # Flatten nested concatenations + self.sources.extend(source.sources if isinstance(source, ConcatSource) else [source]) + self._offsets = [0, *accumulate(len(source) for source in self.sources)] + self._active: Sequence[np.ndarray] | None = None + + def __len__(self) -> int: + return self._offsets[-1] + + def locate(self, idx: int) -> tuple[int, int]: + """Map a page index of the concatenation to the index of its source and its page index in that source + + Args: + idx: the page index in the concatenation + + Returns: + a tuple (index of the source in `sources`, index of the page in that source) + """ + _idx = idx + len(self) if idx < 0 else idx + if not 0 <= _idx < len(self): + raise IndexError(f"page index {idx} is out of range for a document of {len(self)} page(s)") + source_idx = bisect_right(self._offsets, _idx) - 1 + return source_idx, _idx - self._offsets[source_idx] + + def _switch_to(self, source: Sequence[np.ndarray]) -> None: + # Keep at most one underlying document open at a time + if self._active is not None and self._active is not source and isinstance(self._active, PageSource): + self._active.close() + self._active = source + + def _load(self, idx: int) -> np.ndarray: + source_idx, page_idx = self.locate(idx) + source = self.sources[source_idx] + self._switch_to(source) + return source[page_idx] + + def __iter__(self) -> Iterator[np.ndarray]: + for source in self.sources: + self._switch_to(source) + yield from source + + def close(self) -> None: + for source in self.sources: + if isinstance(source, PageSource): + source.close() + self._active = None + + def __repr__(self) -> str: + return f"{self.__class__.__name__}(num_sources={len(self.sources)}, num_pages={len(self)})" + + +class _PageView(PageSource): + """Lazy view on a subset of the pages of another source (returned when slicing a source)""" + + def __init__(self, source: PageSource, indices: Sequence[int]) -> None: + self.source = source + self.indices = indices + + def __len__(self) -> int: + return len(self.indices) + + def _load(self, idx: int) -> np.ndarray: + return self.source[self.indices[idx]] + + def close(self) -> None: + self.source.close() + + +def _list_files(folder: Path, extensions: Sequence[str]) -> list[Path]: + """List the files of a directory (non-recursively) with one of the given extensions, sorted by name""" + return sorted( + (path for path in folder.iterdir() if path.is_file() and path.suffix.lower() in extensions), + key=lambda path: path.name, + ) + + +def _is_pdf(file: AbstractFile) -> bool: + """Whether a file (path or bytes) is a PDF""" + if isinstance(file, bytes): + return file.lstrip()[:5] == b"%PDF-" + return Path(file).suffix.lower() == ".pdf" + + +class _SourcePage: + """Picklable reference to a page of a source, decoded on demand (used to keep predictions lightweight)""" + + __slots__ = ("source", "idx") + + def __init__(self, source: PageSource, idx: int) -> None: + self.source = source + self.idx = idx + + def __call__(self) -> np.ndarray: + return self.source[self.idx] + + +class _EncodedImage: + """Losslessly compressed in-memory image, decoded on demand (used for pages that cannot be re-read from their + source as-is, e.g. straightened pages)""" + + __slots__ = ("data", "shape") + + def __init__(self, image: np.ndarray) -> None: + self.shape = image.shape + # Channel order is irrelevant here: the image is decoded exactly as it was encoded + ok, buffer = cv2.imencode(".png", image, [cv2.IMWRITE_PNG_COMPRESSION, 1]) + if not ok: # pragma: no cover + raise ValueError("unable to encode the page image") + self.data = buffer.tobytes() + + def __call__(self) -> np.ndarray: + image = cv2.imdecode(np.frombuffer(self.data, np.uint8), cv2.IMREAD_UNCHANGED) + if image is None: # pragma: no cover + raise ValueError("unable to decode the page image") + return image.reshape(self.shape) diff --git a/doctr/models/kie_predictor/base.py b/doctr/models/kie_predictor/base.py index 70b9f8fcdc..c65ba7abf9 100644 --- a/doctr/models/kie_predictor/base.py +++ b/doctr/models/kie_predictor/base.py @@ -6,6 +6,7 @@ from collections.abc import Collection from typing import Any +from doctr.io.elements import Document, KIEDocument from doctr.models.builder import KIEDocumentBuilder from ..classification.predictor import OrientationPredictor @@ -32,11 +33,15 @@ class _KIEPredictor(_OCRPredictor): be masked out (filled with black) before passing the pages to the detection/recognition modules. preserve_original_coords: if True and straighten_pages is True, bounding boxes are mapped back to the original page coordinates. Useful for redaction and annotation. + page_chunk_size: the maximum number of pages processed at once. Longer inputs, as well as lazy page sources + and iterators, are processed chunk by chunk so that the memory footprint does not grow with the number + of pages. kwargs: keyword args of `DocumentBuilder` """ crop_orientation_predictor: OrientationPredictor | None page_orientation_predictor: OrientationPredictor | None + _document_cls: type[Document] = KIEDocument def __init__( self, @@ -47,6 +52,7 @@ def __init__( detect_orientation: bool = False, ignore_regions: Collection[str] | None = None, preserve_original_coords: bool = False, + page_chunk_size: int = 32, **kwargs: Any, ) -> None: super().__init__( @@ -57,6 +63,7 @@ def __init__( detect_orientation, ignore_regions, preserve_original_coords=preserve_original_coords, + page_chunk_size=page_chunk_size, **kwargs, ) diff --git a/doctr/models/kie_predictor/pytorch.py b/doctr/models/kie_predictor/pytorch.py index 9dc929c5a9..a1741aab0c 100644 --- a/doctr/models/kie_predictor/pytorch.py +++ b/doctr/models/kie_predictor/pytorch.py @@ -3,6 +3,7 @@ # This program is licensed under the Apache License 2.0. # See LICENSE or go to for full license details. +from collections.abc import Iterable, Sequence from typing import Any import numpy as np @@ -39,6 +40,9 @@ class KIEPredictor(nn.Module, _KIEPredictor): preserve_original_coords: if True and straighten_pages is True, bounding boxes are mapped back to the original page coordinates. Useful for redaction and annotation. layout_predictor: optional layout detection module + page_chunk_size: the maximum number of pages processed at once (32 by default). Longer inputs, lazy page + sources (cf. `DocumentFile.from_pdf(..., lazy=True)`) and iterators are processed chunk by chunk, with a + memory footprint that does not grow with the number of pages. See also :meth:`stream`. **kwargs: keyword args of `DocumentBuilder` """ @@ -75,6 +79,25 @@ def __init__( @torch.inference_mode() def forward( + self, + pages: Sequence[np.ndarray] | Iterable[np.ndarray], + **kwargs: Any, + ) -> Document: + """Predict on a document + + Args: + pages: a list of pages (numpy arrays of shape H x W x C), a lazy page source + (cf. `DocumentFile.from_pdf(..., lazy=True)`) or any iterable of pages. Inputs longer than + `page_chunk_size` are processed chunk by chunk, with a bounded memory footprint. + **kwargs: keyword arguments forwarded to the sub-predictors + + Returns: + the predicted document + """ + return self._run(pages, **kwargs) + + @torch.inference_mode() + def _predict_pages( self, pages: list[np.ndarray], **kwargs: Any, diff --git a/doctr/models/predictor/base.py b/doctr/models/predictor/base.py index 0f27152f5a..12ae706d9c 100644 --- a/doctr/models/predictor/base.py +++ b/doctr/models/predictor/base.py @@ -3,12 +3,13 @@ # This program is licensed under the Apache License 2.0. # See LICENSE or go to for full license details. -from collections.abc import Callable, Collection +from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from typing import Any import numpy as np -from doctr.io.elements import Document, KIEPage +from doctr.io.elements import Document, KIEPage, Page +from doctr.io.sources import PageSource, _EncodedImage, _SourcePage from doctr.models.builder import DocumentBuilder from doctr.utils.geometry import ( extract_crops, @@ -17,6 +18,7 @@ rotate_image, straighten_page, ) +from doctr.utils.multithreading import chunked, prefetch from .._utils import estimate_orientation, mask_boxes, rectify_crops, rectify_loc_preds from ..classification import crop_orientation_predictor, page_orientation_predictor @@ -43,6 +45,9 @@ class _OCRPredictor: be masked out (filled with black) before passing the pages to the detection/recognition modules. preserve_original_coords: if True and straighten_pages is True, bounding boxes are mapped back to the original page coordinates. Useful for redaction and annotation. + page_chunk_size: the maximum number of pages processed at once. Longer inputs, as well as lazy page sources + (cf. `DocumentFile.from_pdf(..., lazy=True)`) and iterators, are processed chunk by chunk so that the + memory footprint does not grow with the number of pages. **kwargs: keyword args of `DocumentBuilder` """ @@ -58,8 +63,12 @@ def __init__( detect_orientation: bool = False, ignore_regions: Collection[str] | None = None, preserve_original_coords: bool = False, + page_chunk_size: int = 32, **kwargs: Any, ) -> None: + if page_chunk_size < 1: + raise ValueError("`page_chunk_size` must be a positive integer") + self.page_chunk_size = page_chunk_size self.assume_straight_pages = assume_straight_pages self.straighten_pages = straighten_pages self.preserve_original_coords = preserve_original_coords @@ -84,6 +93,65 @@ def __init__( self.hooks: list[Callable] = [] + # Type of the documents returned by the predictor + _document_cls: type[Document] = Document + + def _predict_pages(self, pages: list[np.ndarray], **kwargs: Any) -> Document: + """Run the full pipeline on a list of pages held in memory""" + raise NotImplementedError + + def _run(self, pages: Sequence[np.ndarray] | Iterable[np.ndarray], **kwargs: Any) -> Document: + """Predict on any input: short lists are processed at once, everything else chunk by chunk""" + if isinstance(pages, (list, tuple)) and len(pages) <= self.page_chunk_size: + return self._predict_pages(list(pages), **kwargs) + return self._document_cls(pages=list(self.stream(pages, **kwargs))) # type: ignore[arg-type] + + def stream( + self, + pages: Sequence[np.ndarray] | Iterable[np.ndarray], + page_chunk_size: int | None = None, + prefetch_pages: bool = True, + **kwargs: Any, + ) -> Iterator[Page | KIEPage]: + """Predict on a document of any length, yielding the pages one by one as soon as they are processed + + The pages are processed by chunks of `page_chunk_size`, and the next chunk is decoded in a background thread + while the current one goes through the models. Only the pages of the chunks in flight are held in memory, + so that a 1000+ page PDF or a folder of 100k images can be processed with a constant memory footprint. + + >>> from doctr.io import DocumentFile + >>> from doctr.models import ocr_predictor + >>> model = ocr_predictor(pretrained=True) + >>> doc = DocumentFile.from_pdf("path/to/your/long_doc.pdf", lazy=True) + >>> for page in model.stream(doc): + ... print(page.page_idx, page.render()) + + Args: + pages: a list of pages, a lazy page source (cf. `DocumentFile.from_pdf(..., lazy=True)`) or any + iterable of pages (e.g. a generator) + page_chunk_size: the number of pages processed at once (defaults to the `page_chunk_size` of the predictor) + prefetch_pages: whether the next chunk of pages is decoded in a background thread + **kwargs: keyword arguments forwarded to the sub-predictors (as for a regular call) + + Returns: + an iterator over the predicted pages, in the input order. `page_idx` is the index of the page in the + whole input. For lazy page sources, the predicted pages hold a reference to their image instead of the + image itself: `page.page` decodes it again on access. + """ + chunk_size = self.page_chunk_size if page_chunk_size is None else page_chunk_size + source = pages if isinstance(pages, PageSource) else None + offset = 0 + for chunk in prefetch(chunked(pages, chunk_size), enabled=prefetch_pages): + result = self._predict_pages(chunk, **kwargs) + for local_idx, (page, image) in enumerate(zip(result.pages, chunk)): + page.page_idx = offset + local_idx + if source is not None: + # Unchanged pages are decoded again from the source on access, transformed ones (e.g. + # straightened) are kept losslessly compressed + page.page = _SourcePage(source, page.page_idx) if page.page is image else _EncodedImage(page.page) + yield page + offset += len(chunk) + def _general_page_orientations( self, pages: list[np.ndarray], diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py index db70ffbc07..be78fb99fe 100644 --- a/doctr/models/predictor/pytorch.py +++ b/doctr/models/predictor/pytorch.py @@ -3,6 +3,7 @@ # This program is licensed under the Apache License 2.0. # See LICENSE or go to for full license details. +from collections.abc import Iterable, Sequence from typing import Any import cv2 @@ -44,6 +45,9 @@ class OCRPredictor(nn.Module, _OCRPredictor): table_predictor: optional table structure recognition module. Requires `layout_predictor`: table regions are located by the layout model, cropped, and passed to this module. Words falling inside a detected table are regrouped into a structured table and removed from the regular text output. + page_chunk_size: the maximum number of pages processed at once (32 by default). Longer inputs, lazy page + sources (cf. `DocumentFile.from_pdf(..., lazy=True)`) and iterators are processed chunk by chunk, with a + memory footprint that does not grow with the number of pages. See also :meth:`stream`. **kwargs: keyword args of `DocumentBuilder` """ @@ -89,6 +93,25 @@ def __init__( @torch.inference_mode() def forward( + self, + pages: Sequence[np.ndarray] | Iterable[np.ndarray], + **kwargs: Any, + ) -> Document: + """Predict on a document + + Args: + pages: a list of pages (numpy arrays of shape H x W x C), a lazy page source + (cf. `DocumentFile.from_pdf(..., lazy=True)`) or any iterable of pages. Inputs longer than + `page_chunk_size` are processed chunk by chunk, with a bounded memory footprint. + **kwargs: keyword arguments forwarded to the sub-predictors + + Returns: + the predicted document + """ + return self._run(pages, **kwargs) + + @torch.inference_mode() + def _predict_pages( self, pages: list[np.ndarray], **kwargs: Any, diff --git a/doctr/models/zoo.py b/doctr/models/zoo.py index 52c7690b17..9cc81c7273 100644 --- a/doctr/models/zoo.py +++ b/doctr/models/zoo.py @@ -169,7 +169,8 @@ def ocr_predictor( taking the 1D array of probabilities and returning a float. If None (default), the method of the recognition model is kept ("min" for CRNN, VIPTR, SAR and MASTER, "mean" for ViTSTR and PARSeq). A recognition model instance passed as `reco_arch` is modified. - kwargs: keyword args of `OCRPredictor` + kwargs: keyword args of `OCRPredictor` (e.g. `page_chunk_size`, the maximum number of pages processed + at once on long documents) Returns: OCR predictor @@ -327,7 +328,8 @@ def kie_predictor( taking the 1D array of probabilities and returning a float. If None (default), the method of the recognition model is kept ("min" for CRNN, VIPTR, SAR and MASTER, "mean" for ViTSTR and PARSeq). A recognition model instance passed as `reco_arch` is modified. - kwargs: keyword args of `OCRPredictor` + kwargs: keyword args of `OCRPredictor` (e.g. `page_chunk_size`, the maximum number of pages processed + at once on long documents) Returns: KIE predictor diff --git a/doctr/utils/multithreading.py b/doctr/utils/multithreading.py index 9020fe8712..3c388e137e 100644 --- a/doctr/utils/multithreading.py +++ b/doctr/utils/multithreading.py @@ -7,12 +7,14 @@ import multiprocessing as mp import os from collections.abc import Callable, Iterable, Iterator +from concurrent.futures import ThreadPoolExecutor +from itertools import islice from multiprocessing.pool import ThreadPool from typing import Any from doctr.file_utils import ENV_VARS_TRUE_VALUES -__all__ = ["multithread_exec"] +__all__ = ["multithread_exec", "chunked", "prefetch"] def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: int | None = None) -> Iterator[Any]: @@ -49,3 +51,67 @@ def multithread_exec(func: Callable[[Any], Any], seq: Iterable[Any], threads: in # That's why wrapping result in map to return iterator results = map(lambda x: x, tp.map(func, items)) # noqa: C417 return results + + +def chunked(seq: Iterable[Any], size: int) -> Iterator[list[Any]]: + """Split an iterable into consecutive lists of at most `size` elements, consuming it lazily + + >>> from doctr.utils.multithreading import chunked + >>> list(chunked(range(5), 2)) + [[0, 1], [2, 3], [4]] + + Args: + seq: iterable + size: the maximum number of elements per chunk + + Returns: + iterator of chunks + """ + if size < 1: + raise ValueError("the chunk size must be a positive integer") + iterator = iter(seq) + while chunk := list(islice(iterator, size)): + yield chunk + + +_EXHAUSTED = object() + + +def prefetch(seq: Iterable[Any], enabled: bool = True) -> Iterator[Any]: + """Iterate over an iterable while producing its next element in a background thread + + This overlaps the production of the next element (e.g. decoding the next pages of a document) with the + processing of the current one (e.g. running the models). At most one element is produced ahead, and the + iterable is only ever advanced by one thread at a time. + + >>> from doctr.utils.multithreading import prefetch + >>> for item in prefetch(expensive_generator()): + ... process(item) + + Args: + seq: iterable + enabled: whether the elements are produced in a background thread + + Returns: + iterator over the elements of the iterable + + Notes: + Prefetching is disabled when 'DOCTR_MULTIPROCESSING_DISABLE' is set to 'TRUE'. + """ + if not enabled or os.environ.get("DOCTR_MULTIPROCESSING_DISABLE", "").upper() in ENV_VARS_TRUE_VALUES: + yield from seq + return + + iterator = iter(seq) + try: + # On early exit, leaving the context waits for the element being produced, so that the iterable is + # never closed while it is advanced in the background thread + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="doctr-prefetch") as executor: + future = executor.submit(next, iterator, _EXHAUSTED) + while (item := future.result()) is not _EXHAUSTED: + future = executor.submit(next, iterator, _EXHAUSTED) + yield item + finally: + close = getattr(iterator, "close", None) + if callable(close): + close() diff --git a/tests/common/test_cli.py b/tests/common/test_cli.py index 9ace213579..a257724f44 100644 --- a/tests/common/test_cli.py +++ b/tests/common/test_cli.py @@ -12,7 +12,7 @@ import doctr.cli.main as cli from doctr import __version__ from doctr.datasets.generator.base import synthesize_text_img -from doctr.io import DocumentFile, FigureEncoder +from doctr.io import ConcatSource, DocumentFile, FigureEncoder from doctr.io.exporters import DocumentExportsMixin from doctr.models import ocr_predictor from doctr.models.builder import DocumentBuilder @@ -95,6 +95,7 @@ def test_parse_args_defaults(): assert args.symmetric_pad is True assert args.det_bs == 2 assert args.reco_bs == 128 + assert args.page_chunk_size == 32 assert args.detect_orientation is False assert args.detect_language is False assert args.detect_layout is False @@ -238,6 +239,7 @@ def test_parse_args_version(capsys): ["--box_thresh", "-0.1"], ["--bin_thresh", "high"], ["--confidence_aggregation", "average"], + ["--page_chunk_size", "0"], ], ) def test_parse_args_invalid_values(argv): @@ -421,6 +423,8 @@ def test_build_predictor_forwards_every_option(recorded_predictor): "4", "--reco_bs", "64", + "--page_chunk_size", + "8", "--confidence_aggregation", "geometric_mean", "--detect_orientation", @@ -461,6 +465,7 @@ def test_build_predictor_forwards_every_option(recorded_predictor): "disable_page_orientation": True, "disable_crop_orientation": True, "preserve_original_coords": True, + "page_chunk_size": 8, "resolve_lines": False, "resolve_blocks": True, "paragraph_break": 0.1, @@ -472,6 +477,7 @@ def test_build_predictor_forwards_every_option(recorded_predictor): assert model.layout_predictor is not None assert model.table_predictor is not None assert model.det_predictor.pre_processor.batch_size == 4 + assert model.page_chunk_size == 8 assert model.reco_predictor.pre_processor.batch_size == 64 assert model.reco_predictor.model.postprocessor.confidence_aggregation == "geometric_mean" @@ -559,7 +565,8 @@ def test_set_thresholds_without_postprocessor(monkeypatch, predictor, caplog): def test_load_document_concatenates_inputs(text_images, mock_pdf): pages = cli._load_document(cli._parse_args(["--input_path", *text_images, mock_pdf])) - # 2 images + the 2 pages of the mock PDF + # 2 images + the 2 pages of the mock PDF, decoded on demand + assert isinstance(pages, ConcatSource) assert len(pages) == 4 assert all(isinstance(page, np.ndarray) and page.ndim == 3 for page in pages) @@ -575,7 +582,7 @@ def _record_from_pdf(file, **kwargs): args = cli._parse_args(["--input_path", mock_pdf, "--pdf_scale", "4", "--pdf_password", "secret"]) assert len(cli._load_document(args)) == 1 - assert recorder == {"file": mock_pdf, "scale": 4, "password": "secret"} + assert recorder == {"file": mock_pdf, "lazy": True, "scale": 4, "password": "secret"} def test_load_document_from_url(monkeypatch): @@ -590,6 +597,7 @@ def _record_from_url(url, **kwargs): args = cli._parse_args(["--input_path", "https://www.mindee.com"]) assert len(cli._load_document(args)) == 1 assert recorder["url"] == "https://www.mindee.com" + assert recorder["lazy"] is True # Exporting the predictions diff --git a/tests/common/test_io_sources.py b/tests/common/test_io_sources.py new file mode 100644 index 0000000000..930b96552f --- /dev/null +++ b/tests/common/test_io_sources.py @@ -0,0 +1,334 @@ +import pickle +import threading +from pathlib import Path + +import cv2 +import numpy as np +import pytest + +from doctr import io +from doctr.io import sources +from doctr.io.elements import Page +from doctr.utils.multithreading import chunked, prefetch + + +@pytest.fixture(scope="module") +def image_folder(tmp_path_factory): + folder = tmp_path_factory.mktemp("scans") + # Written in a non-sorted order, with a non-image file that must be ignored + for name, value in [("page_2.png", 2), ("page_0.png", 0), ("page_1.jpg", 1)]: + cv2.imwrite(str(folder / name), np.full((20 + value, 30, 3), 40 * value, dtype=np.uint8)) + (folder / "notes.txt").write_text("not an image") + return folder + + +# Lazy sources + + +def test_pdf_source(mock_pdf): + eager = io.read_pdf(mock_pdf) + source = io.PDFSource(mock_pdf) + + assert isinstance(source, io.PageSource) + assert len(source) == 2 + assert repr(source).startswith("PDFSource(") + # pages are rendered on access, and identical to the eager reader + assert all(np.array_equal(page, ref) for page, ref in zip(source, eager)) + assert np.array_equal(source[-1], eager[-1]) + assert source[0].dtype == np.uint8 and source[0].ndim == 3 + # Path & bytes inputs + assert len(io.PDFSource(Path(mock_pdf))) == 2 + with open(mock_pdf, "rb") as f: + assert np.array_equal(io.PDFSource(f.read())[1], eager[1]) + # Rendering options + assert io.PDFSource(mock_pdf, scale=1)[0].shape[0] < eager[0].shape[0] + + with pytest.raises(IndexError): + source[2] + with pytest.raises(IndexError): + source[-3] + with pytest.raises(TypeError): + source["0"] + + +def test_pdf_source_fails_eagerly(tmp_path): + with pytest.raises(FileNotFoundError): + io.PDFSource("my_imaginary_file.pdf") + with pytest.raises(TypeError): + io.PDFSource(123) + corrupted = tmp_path / "corrupted.pdf" + corrupted.write_text("not a real pdf") + with pytest.raises(Exception): + io.PDFSource(corrupted) + + +def test_pdf_source_handles(mock_pdf): + source = io.PDFSource(mock_pdf) + # no file handle is held until a page is rendered + assert source._pdf is None + _ = source[0] + assert source._pdf is not None + source.close() + assert source._pdf is None + # the source can still be used after being closed + assert source[1].ndim == 3 + with source as doc: + assert doc[0].ndim == 3 + assert source._pdf is None + # a full iteration releases the handle it acquired + assert len(list(source)) == 2 + assert source._pdf is None + + +def test_pdf_source_pickle(mock_pdf): + source = io.PDFSource(mock_pdf) + _ = source[0] + restored = pickle.loads(pickle.dumps(source)) + assert restored._pdf is None + assert np.array_equal(restored[0], source[0]) + + +def test_pdf_source_thread_safety(mock_pdf): + source = io.PDFSource(mock_pdf, scale=1) + expected = [page.copy() for page in source] + errors = [] + + def _render(): + try: + for idx in range(20): + assert np.array_equal(source[idx % 2], expected[idx % 2]) + except Exception as e: # pragma: no cover + errors.append(e) + + threads = [threading.Thread(target=_render) for _ in range(4)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert not errors + + +def test_image_source(image_folder): + files = sorted(str(path) for path in image_folder.glob("page_*")) + source = io.ImageSource(files) + assert len(source) == 3 + assert [page.shape[0] for page in source] == [20, 21, 22] + + # a directory expands to its images, sorted by name + folder_source = io.ImageSource(image_folder) + assert [Path(file).name for file in folder_source.files] == ["page_0.png", "page_1.jpg", "page_2.png"] + assert [page.shape[0] for page in folder_source] == [20, 21, 22] + + # single file, bytes & reader options + assert len(io.ImageSource(files[0])) == 1 + raw = (image_folder / "page_0.png").read_bytes() + assert io.ImageSource([raw, raw])[1].shape == (20, 30, 3) + assert io.ImageSource(files, output_size=(8, 8))[2].shape == (8, 8, 3) + + # missing files and unsupported types fail when the source is built, not when the page is processed + with pytest.raises(FileNotFoundError): + io.ImageSource([files[0], "my_imaginary_file.png"]) + with pytest.raises(TypeError): + io.ImageSource([123]) + + +def test_image_source_invalid_image_names_the_file(image_folder): + source = io.ImageSource([image_folder / "page_0.png", image_folder / "notes.txt"]) + assert source[0].ndim == 3 + with pytest.raises(ValueError, match="notes.txt"): + source[1] + with pytest.raises(ValueError, match=""): + io.ImageSource([b"not an image"])[0] + + +def test_slicing(mock_pdf): + source = io.PDFSource(mock_pdf) + view = source[1:] + assert isinstance(view, io.PageSource) + assert len(view) == 1 + assert np.array_equal(view[0], source[1]) + assert len(source[::-1]) == 2 + assert np.array_equal(source[::-1][0], source[1]) + assert len(source[5:]) == 0 + view.close() + assert source._pdf is None + + +def test_iter_batches(image_folder): + source = io.ImageSource(image_folder) + assert [len(batch) for batch in source.iter_batches(2)] == [2, 1] + assert [len(batch) for batch in source.iter_batches(10)] == [3] + with pytest.raises(ValueError): + next(source.iter_batches(0)) + + +def test_concat_source(mock_pdf, image_folder): + pdf, images = io.PDFSource(mock_pdf), io.ImageSource(image_folder) + in_memory = [np.zeros((5, 5, 3), dtype=np.uint8)] + source = io.ConcatSource([pdf, images, in_memory]) + + assert len(source) == 6 + assert repr(source) == "ConcatSource(num_sources=3, num_pages=6)" + assert [page.shape for page in source][-1] == (5, 5, 3) + assert source.locate(0) == (0, 0) + assert source.locate(2) == (1, 0) + assert source.locate(-1) == (2, 0) + assert np.array_equal(source[3], images[1]) + with pytest.raises(IndexError): + source.locate(6) + + # nested concatenations are flattened, `+` concatenates + assert len(io.ConcatSource([source, pdf]).sources) == 4 + assert len(pdf + images) == 5 + assert len(in_memory + pdf) == 3 + assert isinstance(pdf + images, io.ConcatSource) + assert len(io.ConcatSource([])) == 0 + assert list(io.ConcatSource([])) == [] + + +def test_concat_source_keeps_one_pdf_open(mock_pdf): + pdfs = [io.PDFSource(mock_pdf) for _ in range(3)] + source = io.ConcatSource(pdfs) + for idx in range(len(source)): + _ = source[idx] + assert sum(pdf._pdf is not None for pdf in pdfs) == 1 + for _ in source: + assert sum(pdf._pdf is not None for pdf in pdfs) <= 1 + source.close() + assert all(pdf._pdf is None for pdf in pdfs) + + +# DocumentFile + + +def test_document_file_lazy(mock_pdf, image_folder): + assert isinstance(io.DocumentFile.from_pdf(mock_pdf), list) + assert isinstance(io.DocumentFile.from_pdf(mock_pdf, lazy=True), io.PDFSource) + assert isinstance(io.DocumentFile.from_images(image_folder), list) + lazy_images = io.DocumentFile.from_images(image_folder, lazy=True) + assert isinstance(lazy_images, io.ImageSource) + assert len(lazy_images) == 3 + # options are forwarded to the lazy source + assert io.DocumentFile.from_pdf(mock_pdf, lazy=True, scale=1).scale == 1 + + +def test_document_file_from_files(mock_pdf, image_folder, tmp_path): + pdf_bytes = Path(mock_pdf).read_bytes() + image = str(image_folder / "page_0.png") + + doc = io.DocumentFile.from_files([image, mock_pdf, pdf_bytes, image, image], lazy=True) + assert isinstance(doc, io.ConcatSource) + assert len(doc) == 7 + # consecutive images are grouped into a single source + assert [type(source).__name__ for source in doc.sources] == ["ImageSource", "PDFSource", "PDFSource", "ImageSource"] + assert doc.locate(6) == (3, 1) + + eager = io.DocumentFile.from_files([image, mock_pdf]) + assert isinstance(eager, list) and len(eager) == 3 + + # directories expand to their PDFs and images, sorted by name + folder = tmp_path / "mixed" + folder.mkdir() + (folder / "b.pdf").write_bytes(pdf_bytes) + (folder / "a.png").write_bytes((image_folder / "page_0.png").read_bytes()) + (folder / "c.txt").write_text("ignored") + doc = io.DocumentFile.from_files(folder, lazy=True, pdf_kwargs={"scale": 1}, image_kwargs={"output_size": (8, 8)}) + assert len(doc) == 3 + assert doc[0].shape == (8, 8, 3) + assert doc.sources[1].scale == 1 + + +# Lazy page images + + +def test_page_image_on_demand(mock_pdf): + source = io.PDFSource(mock_pdf) + page = Page(sources._SourcePage(source, 1), [], 1, source[1].shape[:2]) + assert np.array_equal(page.page, source[1]) + # assigning an array still works + page.page = np.zeros((2, 2, 3), dtype=np.uint8) + assert page.page.shape == (2, 2, 3) + # picklable, so that results can be sent across processes + restored = pickle.loads(pickle.dumps(sources._SourcePage(source, 0))) + assert np.array_equal(restored(), source[0]) + + +@pytest.mark.parametrize("shape", [(30, 40, 3), (30, 40), (30, 40, 4)]) +def test_encoded_image(shape): + image = np.random.randint(0, 255, shape, dtype=np.uint8) + encoded = sources._EncodedImage(image) + assert len(encoded.data) > 0 + assert np.array_equal(encoded(), image) + + +# Iteration helpers + + +def test_chunked(): + assert list(chunked(range(5), 2)) == [[0, 1], [2, 3], [4]] + assert list(chunked([], 2)) == [] + # lazily consumed + assert next(chunked(iter(range(10**9)), 3)) == [0, 1, 2] + with pytest.raises(ValueError): + next(chunked(range(3), 0)) + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_prefetch(enabled): + assert list(prefetch(range(5), enabled=enabled)) == list(range(5)) + assert list(prefetch([], enabled=enabled)) == [] + + +def test_prefetch_runs_ahead_in_background(): + main_thread = threading.get_ident() + produced = [] + + def _gen(): + for idx in range(3): + produced.append((idx, threading.get_ident() != main_thread)) + yield idx + + iterator = prefetch(_gen()) + assert next(iterator) == 0 + # the next element is already being produced in another thread + assert next(iterator) == 1 + assert all(in_background for _, in_background in produced) + assert list(iterator) == [2] + + +def test_prefetch_disabled_by_env(monkeypatch): + monkeypatch.setenv("DOCTR_MULTIPROCESSING_DISABLE", "TRUE") + main_thread = threading.get_ident() + threads = [] + + def _gen(): + for idx in range(2): + threads.append(threading.get_ident()) + yield idx + + assert list(prefetch(_gen())) == [0, 1] + assert set(threads) == {main_thread} + + +def test_prefetch_propagates_errors_and_closes(): + def _failing(): + yield 0 + raise RuntimeError("decoding failed") + + iterator = prefetch(_failing()) + assert next(iterator) == 0 + with pytest.raises(RuntimeError, match="decoding failed"): + next(iterator) + + closed = [] + + def _gen(): + try: + yield from range(100) + finally: + closed.append(True) + + for item in prefetch(_gen()): + if item == 2: + break + assert closed == [True] diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index 9ba8652b41..d8a0d88e64 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -7,7 +7,7 @@ from doctr import models from doctr.file_utils import CLASS_NAME from doctr.io import Document, DocumentFile, FigureEncoder, crop_layout_region -from doctr.io.elements import KIEDocument, LayoutElement, Table +from doctr.io.elements import KIEDocument, KIEPage, LayoutElement, Table from doctr.models import detection, layout, recognition from doctr.models.classification import mobilenet_v3_small_crop_orientation, mobilenet_v3_small_page_orientation from doctr.models.classification.zoo import crop_orientation_predictor, page_orientation_predictor @@ -1025,3 +1025,156 @@ def test_kie_predictor_straighten_with_preserve_original_coords(mock_tilted_pays assert geoms_on != geoms_off assert out_on.pages[0].page.shape[:2] == out_on.pages[0].dimensions assert out_on.pages[0].page.shape[:2] == doc[0].shape[:2] + + +# Long documents: chunking, streaming & lazy page sources + + +class _ValueDetPredictor(nn.Module): + """Fake detection: one box per page, and records the number of pages of each call""" + + def __init__(self): + super().__init__() + self.calls = [] + + def forward(self, pages, **kwargs): + self.calls.append(len(pages)) + return [{CLASS_NAME: np.array([[0.1, 0.1, 0.9, 0.9, 0.9]], dtype=np.float32)} for _ in pages] + + +class _ValueRecoPredictor(nn.Module): + """Fake recognition: reads the gray level of the crop, which identifies the page""" + + def forward(self, crops, **kwargs): + return [(str(round(float(crop.mean()))), 1.0) for crop in crops] + + +def _page_words(page): + if isinstance(page, KIEPage): + return [word.value for word in page.predictions[CLASS_NAME]] + return [word.value for block in page.blocks for line in block.lines for word in line.words] + + +def _value_pages(num_pages): + return [np.full((60 + idx, 80, 3), 10 * idx, dtype=np.uint8) for idx in range(num_pages)] + + +def _value_predictor(predictor_cls, **kwargs): + return predictor_cls(_ValueDetPredictor(), _ValueRecoPredictor(), assume_straight_pages=True, **kwargs) + + +@pytest.mark.parametrize( + "predictor_cls, doc_cls", + [(OCRPredictor, Document), (KIEPredictor, KIEDocument)], +) +def test_predictors_chunk_long_inputs(predictor_cls, doc_cls): + pages = _value_pages(7) + predictor = _value_predictor(predictor_cls, page_chunk_size=3) + + out = predictor(pages) + assert type(out) is doc_cls + # processed by chunks of at most 3 pages, the results being identical to a single pass + assert predictor.det_predictor.calls == [3, 3, 1] + assert [page.page_idx for page in out.pages] == list(range(7)) + assert [_page_words(page) for page in out.pages] == [[str(10 * idx)] for idx in range(7)] + assert [page.dimensions for page in out.pages] == [page.shape[:2] for page in pages] + # in-memory pages are referenced as-is + assert all(page.page is image for page, image in zip(out.pages, pages)) + + single_pass = _value_predictor(predictor_cls) + assert single_pass(pages).export() == out.export() + assert single_pass.det_predictor.calls == [7] + + # any iterable of pages, e.g. a generator, is supported + assert predictor(page for page in pages).export() == out.export() + assert predictor(tuple(pages)).export() == out.export() + + +@pytest.mark.parametrize("predictor_cls", [OCRPredictor, KIEPredictor]) +def test_predictors_stream(predictor_cls): + pages = _value_pages(5) + predictor = _value_predictor(predictor_cls, page_chunk_size=4) + + stream = predictor.stream(pages, page_chunk_size=2, prefetch_pages=False) + first = next(stream) + # results are yielded as soon as their chunk is processed + assert predictor.det_predictor.calls == [2] + assert (first.page_idx, _page_words(first)) == (0, ["0"]) + rest = list(stream) + assert [page.page_idx for page in rest] == [1, 2, 3, 4] + assert predictor.det_predictor.calls == [2, 2, 1] + + # same results with prefetching, and the predictor default chunk size + predictor.det_predictor.calls.clear() + streamed = list(predictor.stream(iter(pages))) + assert predictor.det_predictor.calls == [4, 1] + assert [page.export() for page in streamed] == [page.export() for page in [first, *rest]] + assert list(predictor.stream([])) == [] + + +@pytest.mark.parametrize( + "predictor_cls, doc_cls", + [(OCRPredictor, Document), (KIEPredictor, KIEDocument)], +) +def test_predictors_lazy_source(tmp_path, predictor_cls, doc_cls): + pages = _value_pages(5) + for idx, page in enumerate(pages): + cv2.imwrite(str(tmp_path / f"page_{idx}.png"), page) + source = DocumentFile.from_images(tmp_path, lazy=True) + + predictor = _value_predictor(predictor_cls, page_chunk_size=2) + out = predictor(source) + assert type(out) is doc_cls + assert predictor.det_predictor.calls == [2, 2, 1] + assert out.export() == _value_predictor(predictor_cls)(pages).export() + + # the predictions only keep a reference to their page, decoded again on access + assert all(type(page._page).__name__ == "_SourcePage" for page in out.pages) + assert all(np.array_equal(page.page, image) for page, image in zip(out.pages, pages)) + assert type(predictor(source[:0])) is doc_cls + + +def test_predictors_lazy_source_transformed_pages(tmp_path): + """Pages modified by the pipeline (e.g. straightened) cannot be read again from the source: they are kept + losslessly compressed""" + + class _TransformingPredictor(OCRPredictor): + def _predict_pages(self, pages, **kwargs): + out = super()._predict_pages(pages, **kwargs) + for page in out.pages: + page.page = np.ascontiguousarray(page.page[::-1]) + return out + + pages = _value_pages(3) + pages[1][:10] = 255 + for idx, page in enumerate(pages): + cv2.imwrite(str(tmp_path / f"page_{idx}.png"), page) + + predictor = _TransformingPredictor(_ValueDetPredictor(), _ValueRecoPredictor(), page_chunk_size=2) + out = predictor(DocumentFile.from_images(tmp_path, lazy=True)) + assert all(type(page._page).__name__ == "_EncodedImage" for page in out.pages) + assert all(np.array_equal(page.page, image[::-1]) for page, image in zip(out.pages, pages)) + + +def test_predictors_invalid_page_chunk_size(): + for predictor_cls in (OCRPredictor, KIEPredictor): + with pytest.raises(ValueError): + _value_predictor(predictor_cls, page_chunk_size=0) + + +@pytest.mark.parametrize("straighten_pages", [False, True]) +def test_ocr_predictor_lazy_pdf(mock_pdf, straighten_pages): + predictor = models.ocr_predictor( + "db_mobilenet_v3_large", + "crnn_mobilenet_v3_small", + pretrained=True, + straighten_pages=straighten_pages, + page_chunk_size=1, + ) + eager = predictor(DocumentFile.from_pdf(mock_pdf)) + lazy = predictor(DocumentFile.from_pdf(mock_pdf, lazy=True)) + + assert lazy.export() == eager.export() + assert lazy.render() == eager.render() + assert all(np.array_equal(lazy_page.page, page.page) for lazy_page, page in zip(lazy.pages, eager.pages)) + assert len(lazy.export_as_markdown(images="embedded")) > 0 From 3aad544cf519afc8be634e7b48321312a0916ac3 Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 27 Oct 2026 13:06:44 +0100 Subject: [PATCH 2/6] Fix imports --- tests/common/test_models.py | 6 +++--- tests/conftest.py | 4 ++-- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/common/test_models.py b/tests/common/test_models.py index f640606ae0..d861e7f51a 100644 --- a/tests/common/test_models.py +++ b/tests/common/test_models.py @@ -5,7 +5,7 @@ import pytest import requests -from doctr.io import reader +from doctr.io import read_img_as_numpy from doctr.models._utils import ( aggregate_confidence, estimate_orientation, @@ -24,7 +24,7 @@ def mock_image(tmpdir_factory): tmp_path = str(tmpdir_factory.mktemp("data").join("mock_bitmap.jpg")) with open(tmp_path, "wb") as f: f.write(file.getbuffer()) - image = reader.read_img_as_numpy(tmp_path) + image = read_img_as_numpy(tmp_path) return image @@ -52,7 +52,7 @@ def test_estimate_orientation(mock_image, mock_bitmap, mock_tilted_payslip): angle_rotated = estimate_orientation(rotated) assert abs(angle_rotated) == 0 - mock_tilted_payslip = reader.read_img_as_numpy(mock_tilted_payslip) + mock_tilted_payslip = read_img_as_numpy(mock_tilted_payslip) assert estimate_orientation(mock_tilted_payslip) == -30 rotated = geometry.rotate_image(mock_tilted_payslip, -30, expand=True) diff --git a/tests/conftest.py b/tests/conftest.py index 43636b8b44..b7fbb95df0 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -11,7 +11,7 @@ from PIL import Image from doctr.datasets.generator.base import synthesize_text_img -from doctr.io import reader +from doctr.io import read_img_as_numpy from doctr.utils import geometry @@ -85,7 +85,7 @@ def mock_figure_page(mock_image_stream, tmpdir_factory): @pytest.fixture(scope="session") def mock_tilted_payslip(mock_payslip, tmpdir_factory): - image = reader.read_img_as_numpy(mock_payslip) + image = read_img_as_numpy(mock_payslip) image = geometry.rotate_image(image, 30, expand=True) tmp_path = str(tmpdir_factory.mktemp("data").join("mock_tilted_payslip.jpg")) cv2.imwrite(tmp_path, image) From b27288aceaf86bb271978ce466f824c97ccded63 Mon Sep 17 00:00:00 2001 From: felix Date: Thu, 29 Oct 2026 09:37:17 +0100 Subject: [PATCH 3/6] Fixes --- docs/source/user_guide/loading_documents.md | 4 +- doctr/io/elements.py | 11 +++++- doctr/io/sources.py | 22 ----------- doctr/models/predictor/base.py | 29 ++++++++------ tests/common/test_io_elements.py | 38 ++++++++++++++++++ tests/common/test_io_sources.py | 31 +++++++++++---- tests/pytorch/test_models_zoo_pt.py | 44 +++++++++++++++++++-- 7 files changed, 133 insertions(+), 46 deletions(-) diff --git a/docs/source/user_guide/loading_documents.md b/docs/source/user_guide/loading_documents.md index 4b10e68726..569b83276e 100644 --- a/docs/source/user_guide/loading_documents.md +++ b/docs/source/user_guide/loading_documents.md @@ -142,7 +142,9 @@ result = predictor(doc) # one call, a bounded memory footprint The predictor processes a page source chunk by chunk (`page_chunk_size` pages at a time, 32 by default), and decodes the next chunk in a background thread while the models run on the current one. The pages of the returned {class}`~doctr.io.Document` only keep a reference to their image: `result.pages[i].page` renders it again on access, -so `show()` and the figure exports keep working. +so `show()` and the figure exports keep working. Pages transformed by the predictor (with `straighten_pages=True`) +cannot be rendered again from the source: as for a list of pages, each of them keeps its straightened image in memory. +Lists of pages already held in memory are simply processed chunk by chunk, without a background thread. The same goes for images, folders and mixed inputs: diff --git a/doctr/io/elements.py b/doctr/io/elements.py index e13dc9fd67..3b0b304c94 100644 --- a/doctr/io/elements.py +++ b/doctr/io/elements.py @@ -49,11 +49,20 @@ def _empty_page_image(page: np.ndarray | None) -> np.ndarray: def _get_page_image(self: Any) -> np.ndarray: - image = self._page + state = self.__dict__ + if "_page" in state: + image = state["_page"] + elif "page" in state: + # Pages pickled by a previous release store their image under `page` (a plain instance attribute back then) + image = state["page"] + else: + raise AttributeError(f"'{type(self).__name__}' object has no attribute 'page'") return image() if callable(image) else image def _set_page_image(self: Any, image: PageImage) -> None: + # Drop the legacy entry of a page unpickled from a previous release, so that it cannot shadow the new image + self.__dict__.pop("page", None) self._page = image diff --git a/doctr/io/sources.py b/doctr/io/sources.py index 12ba4e439f..97fd616284 100644 --- a/doctr/io/sources.py +++ b/doctr/io/sources.py @@ -12,7 +12,6 @@ from pathlib import Path from typing import Any, overload -import cv2 import numpy as np import pypdfium2 as pdfium @@ -369,24 +368,3 @@ def __init__(self, source: PageSource, idx: int) -> None: def __call__(self) -> np.ndarray: return self.source[self.idx] - - -class _EncodedImage: - """Losslessly compressed in-memory image, decoded on demand (used for pages that cannot be re-read from their - source as-is, e.g. straightened pages)""" - - __slots__ = ("data", "shape") - - def __init__(self, image: np.ndarray) -> None: - self.shape = image.shape - # Channel order is irrelevant here: the image is decoded exactly as it was encoded - ok, buffer = cv2.imencode(".png", image, [cv2.IMWRITE_PNG_COMPRESSION, 1]) - if not ok: # pragma: no cover - raise ValueError("unable to encode the page image") - self.data = buffer.tobytes() - - def __call__(self) -> np.ndarray: - image = cv2.imdecode(np.frombuffer(self.data, np.uint8), cv2.IMREAD_UNCHANGED) - if image is None: # pragma: no cover - raise ValueError("unable to decode the page image") - return image.reshape(self.shape) diff --git a/doctr/models/predictor/base.py b/doctr/models/predictor/base.py index 12ae706d9c..263e7a5a3a 100644 --- a/doctr/models/predictor/base.py +++ b/doctr/models/predictor/base.py @@ -9,7 +9,7 @@ import numpy as np from doctr.io.elements import Document, KIEPage, Page -from doctr.io.sources import PageSource, _EncodedImage, _SourcePage +from doctr.io.sources import PageSource, _SourcePage from doctr.models.builder import DocumentBuilder from doctr.utils.geometry import ( extract_crops, @@ -115,9 +115,11 @@ def stream( ) -> Iterator[Page | KIEPage]: """Predict on a document of any length, yielding the pages one by one as soon as they are processed - The pages are processed by chunks of `page_chunk_size`, and the next chunk is decoded in a background thread - while the current one goes through the models. Only the pages of the chunks in flight are held in memory, - so that a 1000+ page PDF or a folder of 100k images can be processed with a constant memory footprint. + The pages are processed by chunks of `page_chunk_size`. For lazy page sources and iterators, the next chunk + is decoded in a background thread while the current one goes through the models, and only the pages of the + chunks in flight are held in memory, so that a 1000+ page PDF or a folder of 100k images can be processed + with a bounded memory footprint. In-memory inputs (lists, tuples, arrays) have nothing to decode: they are + chunked directly, without a background thread. >>> from doctr.io import DocumentFile >>> from doctr.models import ocr_predictor @@ -130,25 +132,30 @@ def stream( pages: a list of pages, a lazy page source (cf. `DocumentFile.from_pdf(..., lazy=True)`) or any iterable of pages (e.g. a generator) page_chunk_size: the number of pages processed at once (defaults to the `page_chunk_size` of the predictor) - prefetch_pages: whether the next chunk of pages is decoded in a background thread + prefetch_pages: whether the next chunk of pages is decoded in a background thread (only applies to lazy + page sources and iterators: in-memory inputs are never prefetched) **kwargs: keyword arguments forwarded to the sub-predictors (as for a regular call) Returns: an iterator over the predicted pages, in the input order. `page_idx` is the index of the page in the whole input. For lazy page sources, the predicted pages hold a reference to their image instead of the - image itself: `page.page` decodes it again on access. + image itself: `page.page` decodes it again on access. Pages transformed by the pipeline (e.g. with + `straighten_pages=True`) cannot be read again from the source: they keep their transformed image, as + for in-memory inputs. """ chunk_size = self.page_chunk_size if page_chunk_size is None else page_chunk_size source = pages if isinstance(pages, PageSource) else None + # Pages already held in memory gain nothing from a background thread: it would only slice the list + in_memory = isinstance(pages, (list, tuple, np.ndarray)) offset = 0 - for chunk in prefetch(chunked(pages, chunk_size), enabled=prefetch_pages): + for chunk in prefetch(chunked(pages, chunk_size), enabled=prefetch_pages and not in_memory): result = self._predict_pages(chunk, **kwargs) for local_idx, (page, image) in enumerate(zip(result.pages, chunk)): page.page_idx = offset + local_idx - if source is not None: - # Unchanged pages are decoded again from the source on access, transformed ones (e.g. - # straightened) are kept losslessly compressed - page.page = _SourcePage(source, page.page_idx) if page.page is image else _EncodedImage(page.page) + # Unchanged pages are decoded again from the source on access. Transformed ones (e.g. straightened) + # keep their image, exactly as for in-memory inputs + if source is not None and page.page is image: + page.page = _SourcePage(source, page.page_idx) yield page offset += len(chunk) diff --git a/tests/common/test_io_elements.py b/tests/common/test_io_elements.py index 153dc38561..f9f517f6c7 100644 --- a/tests/common/test_io_elements.py +++ b/tests/common/test_io_elements.py @@ -1,4 +1,5 @@ import json +import pickle from xml.etree.ElementTree import ElementTree import numpy as np @@ -600,6 +601,43 @@ def test_kiepage(): assert img.shape == (*page_size, 3) +def _legacy_page(page_cls, image, children): + """Rebuild a page the way a previous release pickled it: its image stored as a plain `page` attribute""" + page = page_cls(image, children, 0, image.shape[:2], {"value": 0.0, "confidence": None}, None) + state = dict(page.__dict__) + state["page"] = state.pop("_page") + legacy = page_cls.__new__(page_cls) + legacy.__dict__.update(state) + return legacy + + +@pytest.mark.parametrize( + "page_cls, children", + [(elements.Page, []), (elements.KIEPage, {CLASS_NAME: []})], +) +def test_page_legacy_pickle(page_cls, children): + image = np.random.randint(0, 255, (30, 20, 3), dtype=np.uint8) + legacy = _legacy_page(page_cls, image, children) + # what `pickle` stores for such a page: the `page` entry of its instance dict, no `_page` + assert "page" in legacy.__reduce_ex__(4)[2] + assert "_page" not in legacy.__reduce_ex__(4)[2] + + restored = pickle.loads(pickle.dumps(legacy)) + assert np.array_equal(restored.page, image) + assert restored.synthesize().shape == (*image.shape[:2], 3) + + # setting the image afterwards replaces the legacy entry + new_image = np.zeros_like(image) + restored.page = new_image + assert "page" not in restored.__dict__ + assert restored.page is new_image + assert np.array_equal(pickle.loads(pickle.dumps(restored)).page, new_image) + + # a page without any image still behaves like a missing attribute + empty = page_cls.__new__(page_cls) + assert not hasattr(empty, "page") + + def test_document(): pages = _mock_pages() doc = elements.Document(pages) diff --git a/tests/common/test_io_sources.py b/tests/common/test_io_sources.py index 930b96552f..e19319e9ac 100644 --- a/tests/common/test_io_sources.py +++ b/tests/common/test_io_sources.py @@ -5,6 +5,7 @@ import cv2 import numpy as np import pytest +from PIL import Image from doctr import io from doctr.io import sources @@ -141,6 +142,28 @@ def test_image_source_invalid_image_names_the_file(image_folder): io.ImageSource([b"not an image"])[0] +def test_image_source_transparent_images(tmp_path): + """Lazy sources decode with `read_img_as_numpy`: transparent pixels land on white, exactly as when eager""" + rgba = np.zeros((20, 30, 4), dtype=np.uint8) # fully transparent black background + rgba[5:15, 5:25] = (200, 10, 10, 255) # opaque red box + Image.fromarray(rgba, "RGBA").save(tmp_path / "a_rgba.png") + Image.fromarray(rgba, "RGBA").save(tmp_path / "b_rgba.webp", lossless=True) + Image.fromarray(rgba, "RGBA").convert("P").save(tmp_path / "c_palette.png", transparency=0) + + lazy = io.DocumentFile.from_images(tmp_path, lazy=True) + eager = io.DocumentFile.from_images(sorted(tmp_path.iterdir())) + assert len(lazy) == 3 + for page, ref in zip(lazy, eager): + assert np.array_equal(page, ref) + assert page[0, 0].tolist() == [255, 255, 255] + assert page[10, 10].tolist() == [200, 10, 10] + + # bytes inputs, mixed documents and decoding options go through the same path + content = (tmp_path / "a_rgba.png").read_bytes() + assert io.DocumentFile.from_files([content], lazy=True)[0][0, 0].tolist() == [255, 255, 255] + assert io.ImageSource([content], rgb_output=False)[0][10, 10].tolist() == [10, 10, 200] + + def test_slicing(mock_pdf): source = io.PDFSource(mock_pdf) view = source[1:] @@ -253,14 +276,6 @@ def test_page_image_on_demand(mock_pdf): assert np.array_equal(restored(), source[0]) -@pytest.mark.parametrize("shape", [(30, 40, 3), (30, 40), (30, 40, 4)]) -def test_encoded_image(shape): - image = np.random.randint(0, 255, shape, dtype=np.uint8) - encoded = sources._EncodedImage(image) - assert len(encoded.data) > 0 - assert np.array_equal(encoded(), image) - - # Iteration helpers diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index d8a0d88e64..c4b9cea531 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -1134,9 +1134,44 @@ def test_predictors_lazy_source(tmp_path, predictor_cls, doc_cls): assert type(predictor(source[:0])) is doc_cls +def test_predictors_prefetch_only_lazy_inputs(tmp_path, monkeypatch): + """In-memory inputs are chunked directly: only lazy sources and iterators are prefetched in a background thread""" + from doctr.models.predictor import base as predictor_base + + calls = [] + _prefetch = predictor_base.prefetch + + def _recording_prefetch(seq, enabled=True): + calls.append(enabled) + return _prefetch(seq, enabled=enabled) + + monkeypatch.setattr(predictor_base, "prefetch", _recording_prefetch) + + pages = _value_pages(5) + for idx, page in enumerate(pages): + cv2.imwrite(str(tmp_path / f"page_{idx}.png"), page) + predictor = _value_predictor(OCRPredictor, page_chunk_size=2) + expected = _value_predictor(OCRPredictor)(pages).export() + + for in_memory in (pages, tuple(pages)): + calls.clear() + assert predictor(in_memory).export() == expected + assert calls == [False] + + for lazy in (iter(pages), DocumentFile.from_images(tmp_path, lazy=True)): + calls.clear() + assert predictor(lazy).export() == expected + assert calls == [True] + + # explicitly disabled + calls.clear() + list(predictor.stream(iter(pages), prefetch_pages=False)) + assert calls == [False] + + def test_predictors_lazy_source_transformed_pages(tmp_path): - """Pages modified by the pipeline (e.g. straightened) cannot be read again from the source: they are kept - losslessly compressed""" + """Pages modified by the pipeline (e.g. straightened) cannot be read again from the source: they keep their + transformed image, exactly as for in-memory inputs""" class _TransformingPredictor(OCRPredictor): def _predict_pages(self, pages, **kwargs): @@ -1152,8 +1187,11 @@ def _predict_pages(self, pages, **kwargs): predictor = _TransformingPredictor(_ValueDetPredictor(), _ValueRecoPredictor(), page_chunk_size=2) out = predictor(DocumentFile.from_images(tmp_path, lazy=True)) - assert all(type(page._page).__name__ == "_EncodedImage" for page in out.pages) + assert all(type(page._page) is np.ndarray for page in out.pages) assert all(np.array_equal(page.page, image[::-1]) for page, image in zip(out.pages, pages)) + # same pages as an in-memory run + eager = predictor(pages) + assert all(np.array_equal(lazy.page, ref.page) for lazy, ref in zip(out.pages, eager.pages)) def test_predictors_invalid_page_chunk_size(): From 9017c74bcac2a5a6d2a232a4937a5862b146ae14 Mon Sep 17 00:00:00 2001 From: felix Date: Thu, 29 Oct 2026 12:00:59 +0100 Subject: [PATCH 4/6] Add handling for corupted files --- docs/source/api/io.md | 2 + docs/source/user_guide/loading_documents.md | 9 ++-- docs/source/user_guide/ocr_predictor.md | 4 ++ doctr/cli/main.py | 12 +++--- doctr/io/sources.py | 38 +++++++++++++++-- doctr/models/predictor/base.py | 4 ++ tests/common/test_cli.py | 47 ++++++++++++++++++++- tests/common/test_io_sources.py | 24 +++++++++-- 8 files changed, 123 insertions(+), 17 deletions(-) diff --git a/docs/source/api/io.md b/docs/source/api/io.md index 82b84bc8e0..3b17f59507 100644 --- a/docs/source/api/io.md +++ b/docs/source/api/io.md @@ -29,6 +29,8 @@ Read documents into pages, represent the predictions as a tree of elements, and .. automethod:: locate +.. autoclass:: PageDecodeError + .. autofunction:: read_pdf .. autofunction:: read_img_as_numpy diff --git a/docs/source/user_guide/loading_documents.md b/docs/source/user_guide/loading_documents.md index 569b83276e..e5a33d6e64 100644 --- a/docs/source/user_guide/loading_documents.md +++ b/docs/source/user_guide/loading_documents.md @@ -123,8 +123,8 @@ pages = DocumentFile.from_files(["cover_letter.pdf", "id_front.jpg", "id_back.jp ## Long documents & large collections -Every page of a list is decoded in memory: at the default resolution, an A4 page weighs about 11 MB, so a 1,000-page -PDF would need more than 10 GB. Pass `lazy=True` to get a **page source** instead: it behaves like the list (`len`, +Every page of a list is decoded in memory: at the default resolution (`scale=2`, 1684 x 1191 pixels), an A4 page +weighs about 6 MB, so a 1,000-page PDF would need about 6 GB. Pass `lazy=True` to get a **page source** instead: it behaves like the list (`len`, indexing, slicing, iteration), but a page is only decoded when it is accessed. ```python @@ -173,8 +173,9 @@ with open("result.txt", "w") as f: | {class}`~doctr.io.ConcatSource` | The concatenation of several sources (also obtained with `+`) | :::{tip} -Missing files and invalid PDFs are reported when the source is created. Images, on the other hand, are only decoded -when they are processed: an unreadable image raises a `ValueError` naming the file. +Missing files, invalid PDFs and files that are not images are reported when the source is created (images are only +checked from their header, a few microseconds per file). Images are decoded when they are processed: one with a +valid header but corrupted content raises a {class}`~doctr.io.PageDecodeError` (a `ValueError`) naming the file. ::: ## Low-level readers diff --git a/docs/source/user_guide/ocr_predictor.md b/docs/source/user_guide/ocr_predictor.md index 20c54cc098..78d4179c8c 100644 --- a/docs/source/user_guide/ocr_predictor.md +++ b/docs/source/user_guide/ocr_predictor.md @@ -276,6 +276,10 @@ predictor.add_hook(pad_boxes) Good to know: - Hooks run in the order they were added, each one receiving the output of the previous one. +- Long inputs are processed chunk by chunk (`page_chunk_size` pages at a time, 32 by default, see + [](loading_documents.md#long-documents-large-collections)): hooks are then called once per chunk, with the pages of + that chunk only. Don't rely on the number of pages a hook receives, or on a list position as the page index in the + whole document. - Each returned box keeps the `objectness_score` of the closest detected box of its page; a box added on a page without any detection gets a score of `1`. - With the [KIE predictor](kie.md), hooks receive and return a dict with one such list per class. diff --git a/doctr/cli/main.py b/doctr/cli/main.py index 6dc17768e0..d1458a719a 100644 --- a/doctr/cli/main.py +++ b/doctr/cli/main.py @@ -14,7 +14,7 @@ import torch from tqdm import tqdm -from doctr.io import ConcatSource, Document, DocumentFile, PageSource +from doctr.io import ConcatSource, Document, DocumentFile, PageDecodeError, PageSource from doctr.io.figures import IMAGE_FORMATS, IMAGE_MODES, FigureEncoder from doctr.models import ocr_predictor from doctr.models._utils import _CONFIDENCE_AGGREGATIONS @@ -144,8 +144,9 @@ def _predict(model: Any, doc: ConcatSource, args: argparse.Namespace) -> Documen pages = list( tqdm(model.stream(doc), total=len(doc), unit="page", desc="OCR", disable=args.quiet, file=sys.stderr) ) - except (FileNotFoundError, ValueError) as e: - # Images are only decoded when they are processed + except (FileNotFoundError, PageDecodeError) as e: + # Invalid files are reported by `_load_document`: only a corrupted image (valid header, unreadable content) + # or an input deleted in the meantime gets here. Other errors (e.g. raised by the models) propagate as-is. logger.error(f"File could not be read as a valid image or PDF: {e}") sys.exit(1) finally: @@ -387,11 +388,12 @@ def main(argv=None): _setup_logging(args.quiet) fmt = _resolve_format(args) - model = _build_predictor(args) - # load the document + # open the inputs first (pages are decoded on demand): missing or invalid files fail before the models load doc = _load_document(args) + model = _build_predictor(args) + # perform OCR logger.info(f"Performing OCR on {len(doc)} page(s)...") result = _predict(model, doc, args) diff --git a/doctr/io/sources.py b/doctr/io/sources.py index 97fd616284..7b7a23e685 100644 --- a/doctr/io/sources.py +++ b/doctr/io/sources.py @@ -12,14 +12,16 @@ from pathlib import Path from typing import Any, overload +import cv2 import numpy as np import pypdfium2 as pdfium +from PIL import Image from doctr.utils.common_types import AbstractFile from .image import read_img_as_numpy -__all__ = ["PageSource", "PDFSource", "ImageSource", "ConcatSource", "IMAGE_EXTENSIONS"] +__all__ = ["PageSource", "PDFSource", "ImageSource", "ConcatSource", "PageDecodeError", "IMAGE_EXTENSIONS"] # File extensions recognized as images when a directory is expanded IMAGE_EXTENSIONS: tuple[str, ...] = ( @@ -49,6 +51,29 @@ def _close_pdf_sources() -> None: # pragma: no cover source.close() +class PageDecodeError(ValueError): + """Raised when a page of a source cannot be decoded (e.g. an unreadable or corrupted image file). + + It subclasses `ValueError`, so that code catching the errors of the eager readers keeps working. + """ + + +def _check_image_header(path: Path) -> None: + """Fail fast on files that are not images: only their header is read (a few microseconds per file), so that + an invalid file is reported when the source is created rather than after hours of processing. A file with a + valid header but a corrupted body is still only detected when it is decoded. + """ + if cv2.haveImageReader(str(path)): + return + # Formats read with Pillow by `read_img_as_numpy` (e.g. images with transparency) + try: + with Image.open(path): + return + except (OSError, ValueError, Image.DecompressionBombError): + pass + raise PageDecodeError(f"unable to read the image {path}: unknown or invalid image format") + + class PageSource(Sequence[np.ndarray]): """Abstract lazy collection of pages: a page is only decoded when it is accessed. @@ -228,6 +253,10 @@ class ImageSource(PageSource): >>> doc = ImageSource(["path/to/page1.png", "path/to/page2.png"]) >>> doc = ImageSource("path/to/a/folder") # every image of the folder, sorted by name + Missing files and files that are not images (checked from their header only) are reported when the source is + created. An image that still cannot be decoded raises a :class:`~doctr.io.PageDecodeError` naming the file + when it is accessed. + Args: files: the path to an image file, to a directory of images, or an image as bytes, or a collection of those **kwargs: additional parameters to :meth:`doctr.io.image.read_img_as_numpy` @@ -239,10 +268,13 @@ def __init__(self, files: Sequence[AbstractFile] | AbstractFile, **kwargs: Any) if isinstance(file, (str, Path)): path = Path(file) if path.is_dir(): - self.files.extend(_list_files(path, IMAGE_EXTENSIONS)) + for image_path in _list_files(path, IMAGE_EXTENSIONS): + _check_image_header(image_path) + self.files.append(image_path) continue if not path.is_file(): raise FileNotFoundError(f"unable to access {file}") + _check_image_header(path) elif not isinstance(file, bytes): raise TypeError(f"unsupported object type for an image file: {type(file).__name__}") self.files.append(file) @@ -258,7 +290,7 @@ def _load(self, idx: int) -> np.ndarray: except ValueError as e: # Images are decoded lazily: point to the faulty file among possibly thousands of them name = file if isinstance(file, (str, Path)) else f"" - raise ValueError(f"unable to read the image {name!s}: {e}") from e + raise PageDecodeError(f"unable to read the image {name!s}: {e}") from e class ConcatSource(PageSource): diff --git a/doctr/models/predictor/base.py b/doctr/models/predictor/base.py index 263e7a5a3a..dd4f024c13 100644 --- a/doctr/models/predictor/base.py +++ b/doctr/models/predictor/base.py @@ -393,6 +393,10 @@ def add_hook(self, hook: Callable) -> None: (`(N, 4)` straight boxes or `(N, 4, 2)` polygons), and returns them in the same format. It can move, resize, add or remove boxes: only the returned boxes are cropped and read by the recognition model. + Inputs longer than `page_chunk_size`, lazy page sources and iterators are processed chunk by chunk: hooks are + then called once per chunk, with the pages of that chunk only. A hook must therefore not rely on the number + of pages it receives, nor use a list position as the page index in the whole document. + Args: hook: a callable that takes as input the `loc_preds` and returns the modified `loc_preds` """ diff --git a/tests/common/test_cli.py b/tests/common/test_cli.py index a257724f44..56415cf9ff 100644 --- a/tests/common/test_cli.py +++ b/tests/common/test_cli.py @@ -12,7 +12,7 @@ import doctr.cli.main as cli from doctr import __version__ from doctr.datasets.generator.base import synthesize_text_img -from doctr.io import ConcatSource, DocumentFile, FigureEncoder +from doctr.io import ConcatSource, DocumentFile, FigureEncoder, Page from doctr.io.exporters import DocumentExportsMixin from doctr.models import ocr_predictor from doctr.models.builder import DocumentBuilder @@ -777,3 +777,48 @@ def test_main_one_invalid_input_among_several(text_images, tmp_path): "--output", str(tmp_path / "results.json"), ]) + + +def test_main_invalid_input_fails_before_loading_models(monkeypatch, text_images, tmp_path): + # inputs are opened (and images checked from their header) before the models are built + def _no_predictor(args): + raise AssertionError("the predictor must not be built") + + monkeypatch.setattr(cli, "_build_predictor", _no_predictor) + not_an_image = tmp_path / "scan.png" + not_an_image.write_text("not an image") + for invalid in ("non_existent_file.pdf", str(not_an_image)): + with pytest.raises(SystemExit) as exc_info: + cli.main(["--input_path", text_images[0], invalid, "--output", str(tmp_path / "results.json")]) + assert exc_info.value.code == 1 + + +class _StreamingModel: + """Minimal predictor: decodes the pages of the document as it streams them, then optionally fails""" + + def __init__(self, error=None): + self.error = error + + def stream(self, doc): + for page in doc: + if self.error is not None: + raise self.error + yield Page(page, [], 0, page.shape[:2]) + + +def test_predict_reports_corrupted_images(text_images, tmp_path, caplog): + # a valid header with corrupted content passes `_load_document` and is only detected while processing + truncated = tmp_path / "truncated.png" + truncated.write_bytes(Path(text_images[0]).read_bytes()[:60]) + args = cli._parse_args(["--input_path", text_images[0], str(truncated), "--quiet"]) + doc = cli._load_document(args) + with pytest.raises(SystemExit) as exc_info: + cli._predict(_StreamingModel(), doc, args) + assert exc_info.value.code == 1 + assert "truncated.png" in caplog.text + + # other errors (e.g. from the models) are not reported as unreadable files + doc = cli._load_document(cli._parse_args(["--input_path", text_images[0], "--quiet"])) + with pytest.raises(ValueError, match="model failure"): + cli._predict(_StreamingModel(ValueError("model failure")), doc, args) + assert len(cli._predict(_StreamingModel(), doc, args).pages) == 1 diff --git a/tests/common/test_io_sources.py b/tests/common/test_io_sources.py index e19319e9ac..09aec12212 100644 --- a/tests/common/test_io_sources.py +++ b/tests/common/test_io_sources.py @@ -133,12 +133,28 @@ def test_image_source(image_folder): io.ImageSource([123]) -def test_image_source_invalid_image_names_the_file(image_folder): - source = io.ImageSource([image_folder / "page_0.png", image_folder / "notes.txt"]) +def test_image_source_invalid_image_names_the_file(image_folder, tmp_path): + # files that are not images fail when the source is built, from their header only + with pytest.raises(io.PageDecodeError, match="notes.txt"): + io.ImageSource([image_folder / "page_0.png", image_folder / "notes.txt"]) + # including the files found in a directory + (tmp_path / "page_0.png").write_bytes((image_folder / "page_0.png").read_bytes()) + (tmp_path / "page_1.png").write_text("not an image") + with pytest.raises(io.PageDecodeError, match="page_1.png"): + io.ImageSource(tmp_path) + # still a `ValueError`, as raised by the eager readers + with pytest.raises(ValueError): + io.DocumentFile.from_images([image_folder / "notes.txt"]) + + # a valid header with corrupted content is only detected when the image is decoded + truncated = tmp_path / "page_1.png" + truncated.write_bytes((image_folder / "page_0.png").read_bytes()[:60]) + source = io.ImageSource([image_folder / "page_0.png", truncated]) assert source[0].ndim == 3 - with pytest.raises(ValueError, match="notes.txt"): + with pytest.raises(io.PageDecodeError, match="page_1.png"): source[1] - with pytest.raises(ValueError, match=""): + # bytes are not checked upfront + with pytest.raises(io.PageDecodeError, match=""): io.ImageSource([b"not an image"])[0] From 387c97e5de61c1a9e3e1c02f49c26307bde34059 Mon Sep 17 00:00:00 2001 From: felix Date: Thu, 29 Oct 2026 12:42:36 +0100 Subject: [PATCH 5/6] Fix several small issues --- docs/source/user_guide/loading_documents.md | 10 ++ doctr/io/reader.py | 24 +++ doctr/io/sources.py | 146 ++++++++++++----- doctr/models/predictor/base.py | 46 ++++-- doctr/utils/multithreading.py | 6 +- tests/common/test_io_sources.py | 169 ++++++++++++++++++-- tests/pytorch/test_models_zoo_pt.py | 33 ++++ 7 files changed, 366 insertions(+), 68 deletions(-) diff --git a/docs/source/user_guide/loading_documents.md b/docs/source/user_guide/loading_documents.md index e5a33d6e64..1df9b14d72 100644 --- a/docs/source/user_guide/loading_documents.md +++ b/docs/source/user_guide/loading_documents.md @@ -146,6 +146,16 @@ so `show()` and the figure exports keep working. Pages transformed by the predic cannot be rendered again from the source: as for a list of pages, each of them keeps its straightened image in memory. Lists of pages already held in memory are simply processed chunk by chunk, without a background thread. +A page source holds no open file between accesses: iterating over it opens each PDF once for the pass, while a +single page (`doc[i]`, or `result.pages[i].page`) opens the file for that page only. The results therefore stay tied +to their input files: if one is moved or deleted, accessing the image of its pages raises a `FileNotFoundError`. For +many random accesses to the same PDF, a `with` block keeps it open until the block ends: + +```python +with DocumentFile.from_pdf("path/to/1000_pages.pdf", lazy=True) as doc: + pages = [doc[i] for i in (0, 10, 500)] +``` + The same goes for images, folders and mixed inputs: ```python diff --git a/doctr/io/reader.py b/doctr/io/reader.py index 764b4137c9..0ada223f13 100644 --- a/doctr/io/reader.py +++ b/doctr/io/reader.py @@ -43,6 +43,10 @@ def from_pdf(cls, file: AbstractFile, lazy: Literal[False] = ..., **kwargs: Any) @classmethod def from_pdf(cls, file: AbstractFile, lazy: Literal[True], **kwargs: Any) -> PDFSource: ... + @overload + @classmethod + def from_pdf(cls, file: AbstractFile, lazy: bool, **kwargs: Any) -> list[np.ndarray] | PDFSource: ... + @classmethod def from_pdf(cls, file: AbstractFile, lazy: bool = False, **kwargs: Any) -> list[np.ndarray] | PDFSource: """Read a PDF file @@ -72,6 +76,10 @@ def from_url(cls, url: str, lazy: Literal[False] = ..., **kwargs: Any) -> list[n @classmethod def from_url(cls, url: str, lazy: Literal[True], **kwargs: Any) -> PDFSource: ... + @overload + @classmethod + def from_url(cls, url: str, lazy: bool, **kwargs: Any) -> list[np.ndarray] | PDFSource: ... + @classmethod def from_url(cls, url: str, lazy: bool = False, **kwargs: Any) -> list[np.ndarray] | PDFSource: """Interpret a web page as a PDF document @@ -108,6 +116,12 @@ def from_images( cls, files: Sequence[AbstractFile] | AbstractFile, lazy: Literal[True], **kwargs: Any ) -> ImageSource: ... + @overload + @classmethod + def from_images( + cls, files: Sequence[AbstractFile] | AbstractFile, lazy: bool, **kwargs: Any + ) -> list[np.ndarray] | ImageSource: ... + @classmethod def from_images( cls, files: Sequence[AbstractFile] | AbstractFile, lazy: bool = False, **kwargs: Any @@ -150,6 +164,16 @@ def from_files( image_kwargs: dict[str, Any] | None = ..., ) -> ConcatSource: ... + @overload + @classmethod + def from_files( + cls, + files: Sequence[AbstractFile] | AbstractFile, + lazy: bool, + pdf_kwargs: dict[str, Any] | None = ..., + image_kwargs: dict[str, Any] | None = ..., + ) -> list[np.ndarray] | ConcatSource: ... + @classmethod def from_files( cls, diff --git a/doctr/io/sources.py b/doctr/io/sources.py index 7b7a23e685..681fb08517 100644 --- a/doctr/io/sources.py +++ b/doctr/io/sources.py @@ -8,7 +8,8 @@ import weakref from bisect import bisect_right from collections.abc import Iterable, Iterator, Sequence -from itertools import accumulate +from contextlib import AbstractContextManager, contextmanager, nullcontext +from itertools import accumulate, groupby from pathlib import Path from typing import Any, overload @@ -112,9 +113,22 @@ def __getitem__(self, idx: int | slice) -> "np.ndarray | PageSource": raise IndexError(f"page index {idx} is out of range for a document of {num_pages} page(s)") return self._load(_idx) + def _session(self) -> AbstractContextManager[Any]: + """Context during which the resources needed to decode the pages (e.g. an open PDF) are kept. + + Iterations run within one, so that the resources are acquired once for the whole pass. Outside of it, an + access to a single page (`source[idx]`) acquires them for that page only, and releases them right after. + """ + return nullcontext() + + def _iter_indices(self, indices: Iterable[int]) -> Iterator[np.ndarray]: + """Decode the pages at the given (non-negative, in-range) indices, in this order""" + with self._session(): + for idx in indices: + yield self._load(idx) + def __iter__(self) -> Iterator[np.ndarray]: - for idx in range(len(self)): - yield self._load(idx) + return self._iter_indices(range(len(self))) def __add__(self, other: Sequence[np.ndarray]) -> "ConcatSource": if not isinstance(other, Sequence): @@ -148,14 +162,37 @@ def iter_batches(self, batch_size: int) -> Iterator[list[np.ndarray]]: def close(self) -> None: """Release the resources held by the source (e.g. open PDF handles). The source can still be used - afterwards: the resources are acquired again on demand.""" + afterwards: the resources are acquired again on demand. + + Resources are only held during an iteration or a `with` block: accessing a single page outside of those + (e.g. `source[idx]`, or `page.page` on a prediction made from this source) releases them right after. + """ def __enter__(self) -> "PageSource": + # Keep the resources for the whole block, e.g. for many random accesses to the pages of a PDF + session = self._session() + session.__enter__() + self.__dict__.setdefault("_entered", []).append(session) return self def __exit__(self, *args: Any) -> None: + entered = self.__dict__.get("_entered") + if entered: + entered.pop().__exit__(*args) + if not entered: + # Leaving the outermost block + self._release() + + def _release(self) -> None: + """Release the resources once the last `with` block on the source is left""" self.close() + def __getstate__(self) -> dict[str, Any]: + # The resources held by a `with` block are not part of the source + state = self.__dict__.copy() + state.pop("_entered", None) + return state + def __repr__(self) -> str: return f"{self.__class__.__name__}(num_pages={len(self)})" @@ -167,6 +204,11 @@ class PDFSource(PageSource): >>> from doctr.io import PDFSource >>> doc = PDFSource("path/to/your/doc.pdf", scale=2) >>> first_page = doc[0] + >>> with doc: # keep the document open for many random accesses + ... pages = [doc[idx] for idx in (0, 10, 500)] + + No file handle is held between accesses: an iteration opens the document once for the pass, and accessing a + single page outside of an iteration or a `with` block opens it for that page only. Args: file: the path to the PDF file or its content as bytes @@ -190,6 +232,8 @@ def __init__( self.password = password self.render_kwargs = kwargs self._pdf: pdfium.PdfDocument | None = None + # Number of sessions (iterations, `with` blocks) keeping the document open + self._holders = 0 # Open the document eagerly: invalid inputs fail here rather than in the middle of a prediction with _PDFIUM_LOCK: pdf = self._open() @@ -207,25 +251,37 @@ def _open(self) -> pdfium.PdfDocument: def __len__(self) -> int: return self._num_pages + @contextmanager + def _session(self) -> Iterator[None]: + with _PDFIUM_LOCK: + self._holders += 1 + try: + yield + finally: + with _PDFIUM_LOCK: + self._holders -= 1 + # Release the file handle once nobody reads the document anymore + if self._holders == 0: + self.close() + def _load(self, idx: int) -> np.ndarray: with _PDFIUM_LOCK: - page = self._open()[idx] + # Outside of any session (e.g. `source[idx]`, or the page image of a prediction accessed later on), the + # document is only opened for this page: no file handle is left behind. Within a session, it is kept open + # (and opened again if it was closed in the meantime). + transient = self._pdf is None and self._holders == 0 + pdf = self._open() try: - bitmap = page.render(scale=self.scale, rev_byteorder=self.rgb_mode, **self.render_kwargs) - # The bitmap buffer is allocated by Python: the array safely outlives the pdfium page & document - return bitmap.to_numpy() + page = pdf[idx] + try: + bitmap = page.render(scale=self.scale, rev_byteorder=self.rgb_mode, **self.render_kwargs) + # The bitmap buffer is allocated by Python: the array safely outlives the pdfium page & document + return bitmap.to_numpy() + finally: + page.close() finally: - page.close() - - def __iter__(self) -> Iterator[np.ndarray]: - was_open = self._pdf is not None - try: - for idx in range(len(self)): - yield self._load(idx) - finally: - # Release the file handle once the whole document has been read - if not was_open: - self.close() + if transient: + self.close() def close(self) -> None: with _PDFIUM_LOCK: @@ -234,10 +290,16 @@ def close(self) -> None: self._pdf = None _OPEN_PDF_SOURCES.discard(self) + def _release(self) -> None: + # The document is closed by the last session releasing it: an iteration still running in another thread + # keeps it open + pass + def __getstate__(self) -> dict[str, Any]: # pdfium handles cannot be pickled: the document is re-opened on demand - state = self.__dict__.copy() + state = super().__getstate__() state["_pdf"] = None + state["_holders"] = 0 return state def __repr__(self) -> str: @@ -310,7 +372,6 @@ def __init__(self, sources: Iterable[Sequence[np.ndarray]]) -> None: # Flatten nested concatenations self.sources.extend(source.sources if isinstance(source, ConcatSource) else [source]) self._offsets = [0, *accumulate(len(source) for source in self.sources)] - self._active: Sequence[np.ndarray] | None = None def __len__(self) -> int: return self._offsets[-1] @@ -330,35 +391,37 @@ def locate(self, idx: int) -> tuple[int, int]: source_idx = bisect_right(self._offsets, _idx) - 1 return source_idx, _idx - self._offsets[source_idx] - def _switch_to(self, source: Sequence[np.ndarray]) -> None: - # Keep at most one underlying document open at a time - if self._active is not None and self._active is not source and isinstance(self._active, PageSource): - self._active.close() - self._active = source - def _load(self, idx: int) -> np.ndarray: + # No state is shared between accesses, so that pages can be read from several threads at once (e.g. the + # page images of predictions while the predictor still iterates over the document in the background) source_idx, page_idx = self.locate(idx) - source = self.sources[source_idx] - self._switch_to(source) - return source[page_idx] + return self.sources[source_idx][page_idx] - def __iter__(self) -> Iterator[np.ndarray]: - for source in self.sources: - self._switch_to(source) - yield from source + def _iter_indices(self, indices: Iterable[int]) -> Iterator[np.ndarray]: + # Consecutive pages of the same input are read within one session of that input: a PDF is opened once for + # them and closed before moving on, so that at most one underlying document is open at a time + for source_idx, run in groupby(indices, key=lambda idx: bisect_right(self._offsets, idx) - 1): + source, offset = self.sources[source_idx], self._offsets[source_idx] + if isinstance(source, PageSource): + yield from source._iter_indices(idx - offset for idx in run) + else: + yield from (source[idx - offset] for idx in run) def close(self) -> None: for source in self.sources: if isinstance(source, PageSource): source.close() - self._active = None def __repr__(self) -> str: return f"{self.__class__.__name__}(num_sources={len(self.sources)}, num_pages={len(self)})" class _PageView(PageSource): - """Lazy view on a subset of the pages of another source (returned when slicing a source)""" + """Lazy view on a subset of the pages of another source (returned when slicing a source) + + A view holds no resources of its own: iterating over it holds those of its source for the pass only, and + closing it leaves its source (and the other views on it) untouched. + """ def __init__(self, source: PageSource, indices: Sequence[int]) -> None: self.source = source @@ -370,8 +433,12 @@ def __len__(self) -> int: def _load(self, idx: int) -> np.ndarray: return self.source[self.indices[idx]] - def close(self) -> None: - self.source.close() + def _session(self) -> AbstractContextManager[Any]: + # `with view:` keeps the resources of the source for the block + return self.source._session() + + def _iter_indices(self, indices: Iterable[int]) -> Iterator[np.ndarray]: + return self.source._iter_indices(self.indices[idx] for idx in indices) def _list_files(folder: Path, extensions: Sequence[str]) -> list[Path]: @@ -385,7 +452,8 @@ def _list_files(folder: Path, extensions: Sequence[str]) -> list[Path]: def _is_pdf(file: AbstractFile) -> bool: """Whether a file (path or bytes) is a PDF""" if isinstance(file, bytes): - return file.lstrip()[:5] == b"%PDF-" + # Only look at the beginning of the content: stripping a whole document would copy it + return file[:1024].lstrip()[:5] == b"%PDF-" return Path(file).suffix.lower() == ".pdf" diff --git a/doctr/models/predictor/base.py b/doctr/models/predictor/base.py index dd4f024c13..7b7f100b0c 100644 --- a/doctr/models/predictor/base.py +++ b/doctr/models/predictor/base.py @@ -3,6 +3,7 @@ # This program is licensed under the Apache License 2.0. # See LICENSE or go to for full license details. +import threading from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from typing import Any @@ -141,23 +142,46 @@ def stream( whole input. For lazy page sources, the predicted pages hold a reference to their image instead of the image itself: `page.page` decodes it again on access. Pages transformed by the pipeline (e.g. with `straighten_pages=True`) cannot be read again from the source: they keep their transformed image, as - for in-memory inputs. + for in-memory inputs. Stopping the iteration early (e.g. with `break`) stops decoding ahead after the + page in progress. """ chunk_size = self.page_chunk_size if page_chunk_size is None else page_chunk_size source = pages if isinstance(pages, PageSource) else None # Pages already held in memory gain nothing from a background thread: it would only slice the list in_memory = isinstance(pages, (list, tuple, np.ndarray)) + # Set when the consumer stops early (or on error), so that the background thread stops decoding ahead after + # the page in progress instead of completing a whole chunk that nobody will process + stop = threading.Event() + + def _pages_until_stopped() -> Iterator[np.ndarray]: + iterator = iter(pages) + try: + for page in iterator: + yield page + if stop.is_set(): + return + finally: + # Release the resources held by the iteration (e.g. an open PDF) as soon as it is over + close = getattr(iterator, "close", None) + if callable(close): + close() + + chunks = prefetch(chunked(_pages_until_stopped(), chunk_size), enabled=prefetch_pages and not in_memory) offset = 0 - for chunk in prefetch(chunked(pages, chunk_size), enabled=prefetch_pages and not in_memory): - result = self._predict_pages(chunk, **kwargs) - for local_idx, (page, image) in enumerate(zip(result.pages, chunk)): - page.page_idx = offset + local_idx - # Unchanged pages are decoded again from the source on access. Transformed ones (e.g. straightened) - # keep their image, exactly as for in-memory inputs - if source is not None and page.page is image: - page.page = _SourcePage(source, page.page_idx) - yield page - offset += len(chunk) + try: + for chunk in chunks: + result = self._predict_pages(chunk, **kwargs) + for local_idx, (page, image) in enumerate(zip(result.pages, chunk)): + page.page_idx = offset + local_idx + # Unchanged pages are decoded again from the source on access. Transformed ones (e.g. + # straightened) keep their image, exactly as for in-memory inputs + if source is not None and page.page is image: + page.page = _SourcePage(source, page.page_idx) + yield page + offset += len(chunk) + finally: + stop.set() + chunks.close() def _general_page_orientations( self, diff --git a/doctr/utils/multithreading.py b/doctr/utils/multithreading.py index 3c388e137e..9869c20989 100644 --- a/doctr/utils/multithreading.py +++ b/doctr/utils/multithreading.py @@ -6,7 +6,7 @@ import multiprocessing as mp import os -from collections.abc import Callable, Iterable, Iterator +from collections.abc import Callable, Generator, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor from itertools import islice from multiprocessing.pool import ThreadPool @@ -77,7 +77,7 @@ def chunked(seq: Iterable[Any], size: int) -> Iterator[list[Any]]: _EXHAUSTED = object() -def prefetch(seq: Iterable[Any], enabled: bool = True) -> Iterator[Any]: +def prefetch(seq: Iterable[Any], enabled: bool = True) -> Generator[Any, None, None]: """Iterate over an iterable while producing its next element in a background thread This overlaps the production of the next element (e.g. decoding the next pages of a document) with the @@ -93,7 +93,7 @@ def prefetch(seq: Iterable[Any], enabled: bool = True) -> Iterator[Any]: enabled: whether the elements are produced in a background thread Returns: - iterator over the elements of the iterable + generator over the elements of the iterable: closing it waits for the element being produced, if any Notes: Prefetching is disabled when 'DOCTR_MULTIPROCESSING_DISABLE' is set to 'TRUE'. diff --git a/tests/common/test_io_sources.py b/tests/common/test_io_sources.py index 09aec12212..d326ff3016 100644 --- a/tests/common/test_io_sources.py +++ b/tests/common/test_io_sources.py @@ -4,6 +4,7 @@ import cv2 import numpy as np +import pypdfium2 as pdfium import pytest from PIL import Image @@ -63,24 +64,82 @@ def test_pdf_source_fails_eagerly(tmp_path): io.PDFSource(corrupted) -def test_pdf_source_handles(mock_pdf): +def test_pdf_source_handles(mock_pdf, monkeypatch): + opened = [] + _open = io.PDFSource._open + + def _counting_open(self): + if self._pdf is None: + opened.append(self) + return _open(self) + + monkeypatch.setattr(io.PDFSource, "_open", _counting_open) source = io.PDFSource(mock_pdf) - # no file handle is held until a page is rendered + opened.clear() + + # no file handle is held outside of an iteration or a `with` block: random access opens the document for the + # page only assert source._pdf is None - _ = source[0] - assert source._pdf is not None - source.close() + assert source[0].ndim == 3 and source[1].ndim == 3 assert source._pdf is None - # the source can still be used after being closed - assert source[1].ndim == 3 + assert len(opened) == 2 + assert not sources._OPEN_PDF_SOURCES + + # a `with` block keeps it open for all the accesses of the block + opened.clear() with source as doc: - assert doc[0].ndim == 3 + assert doc[0].ndim == 3 and doc[1].ndim == 3 + assert source._pdf is not None + # nested blocks + with source: + assert doc[1].ndim == 3 + assert source._pdf is not None assert source._pdf is None - # a full iteration releases the handle it acquired - assert len(list(source)) == 2 + assert len(opened) == 1 + + # an iteration opens it once, and releases it at the end + opened.clear() + iterator = iter(source) + next(iterator) + assert source._pdf is not None + assert len(list(iterator)) == 1 + assert source._pdf is None + assert len(opened) == 1 + # also when stopped early + iterator = iter(source) + next(iterator) + iterator.close() + assert source._pdf is None + + # closing it within a `with` block or an iteration is not an issue: it is opened again on demand + with source: + source.close() + assert source[0].ndim == 3 + assert source._pdf is not None assert source._pdf is None +def test_pdf_source_handle_released_on_errors(mock_pdf, monkeypatch): + source = io.PDFSource(mock_pdf) + + def _failing_render(*args, **kwargs): + raise RuntimeError("render failure") + + # a page that cannot be loaded + with pytest.raises(Exception): + source._load(5) + assert source._pdf is None + # a page that cannot be rendered + monkeypatch.setattr(pdfium.PdfPage, "render", _failing_render) + with pytest.raises(RuntimeError, match="render failure"): + source[0] + assert source._pdf is None + with pytest.raises(RuntimeError, match="render failure"): + list(source) + assert source._pdf is None + assert not sources._OPEN_PDF_SOURCES + + def test_pdf_source_pickle(mock_pdf): source = io.PDFSource(mock_pdf) _ = source[0] @@ -189,8 +248,27 @@ def test_slicing(mock_pdf): assert len(source[::-1]) == 2 assert np.array_equal(source[::-1][0], source[1]) assert len(source[5:]) == 0 - view.close() + + # a view holds nothing of its own: closing it leaves its source and the other views untouched + with source: + _ = source[0] + view.close() + assert source._pdf is not None + assert np.array_equal(source[:1][0], source[0]) + assert source._pdf is None + # iterating over a view opens its source for the pass only, `with view:` for the block + assert len(list(source[::-1])) == 2 + assert source._pdf is None + with view: + _ = view[0] + assert source._pdf is not None assert source._pdf is None + # views of views + assert np.array_equal(source[::-1][1:][0], source[0]) + # a source used in a `with` block can still be pickled + with source: + assert pickle.loads(pickle.dumps(source))._pdf is None + assert len(pickle.loads(pickle.dumps(view))) == 1 def test_iter_batches(image_folder): @@ -227,14 +305,60 @@ def test_concat_source(mock_pdf, image_folder): def test_concat_source_keeps_one_pdf_open(mock_pdf): pdfs = [io.PDFSource(mock_pdf) for _ in range(3)] - source = io.ConcatSource(pdfs) + source = io.ConcatSource([pdfs[0], [np.zeros((4, 4, 3), dtype=np.uint8)], pdfs[1], pdfs[2]]) + # random access holds no handle for idx in range(len(source)): - _ = source[idx] - assert sum(pdf._pdf is not None for pdf in pdfs) == 1 + assert source[idx].ndim == 3 + assert all(pdf._pdf is None for pdf in pdfs) + # an iteration holds at most one, and none once it is over for _ in source: assert sum(pdf._pdf is not None for pdf in pdfs) <= 1 - source.close() assert all(pdf._pdf is None for pdf in pdfs) + # same for the views on a concatenation, in any order + pages = list(source) + view = source[::-1] + for page, ref in zip(view, pages[::-1]): + assert np.array_equal(page, ref) + assert sum(pdf._pdf is not None for pdf in pdfs) <= 1 + assert all(pdf._pdf is None for pdf in pdfs) + + +def test_concat_source_concurrent_access(mock_pdf, monkeypatch): + """Pages can be read (e.g. through the page images of predictions) while another thread iterates over the + concatenation: no state is shared, and the document being iterated is not closed under the iteration""" + pdfs = [io.PDFSource(mock_pdf) for _ in range(4)] + source = io.ConcatSource(pdfs) + expected = list(source) + # documents closed while an iteration was still reading them + closed_under_iteration = [] + _close = io.PDFSource.close + + def _recording_close(self): + if self._holders > 0 and self._pdf is not None: + closed_under_iteration.append(self) + return _close(self) + + monkeypatch.setattr(io.PDFSource, "close", _recording_close) + + errors = [] + + def _read_randomly(): + try: + for _ in range(5): + for idx in range(len(source)): + assert np.array_equal(source[idx], expected[idx]) + except Exception as e: # pragma: no cover + errors.append(e) + + reader = threading.Thread(target=_read_randomly) + reader.start() + iterated = list(source) + reader.join() + assert not errors + assert all(np.array_equal(page, ref) for page, ref in zip(iterated, expected)) + assert not closed_under_iteration + assert all(pdf._pdf is None for pdf in pdfs) + assert not sources._OPEN_PDF_SOURCES # DocumentFile @@ -277,6 +401,18 @@ def test_document_file_from_files(mock_pdf, image_folder, tmp_path): assert doc.sources[1].scale == 1 +def test_is_pdf(mock_pdf): + pdf_bytes = Path(mock_pdf).read_bytes() + assert sources._is_pdf(pdf_bytes) + assert sources._is_pdf(b"\n \r\n" + pdf_bytes) + assert not sources._is_pdf(b"\x89PNG\r\n") + assert not sources._is_pdf(b"") + # only the beginning of the content is looked at + assert not sources._is_pdf(b" " * 2048 + pdf_bytes) + assert sources._is_pdf(Path("doc.PDF")) and sources._is_pdf("doc.pdf") + assert not sources._is_pdf("doc.png") + + # Lazy page images @@ -284,6 +420,9 @@ def test_page_image_on_demand(mock_pdf): source = io.PDFSource(mock_pdf) page = Page(sources._SourcePage(source, 1), [], 1, source[1].shape[:2]) assert np.array_equal(page.page, source[1]) + # accessing the image of a prediction leaves no file handle behind + assert source._pdf is None + assert not sources._OPEN_PDF_SOURCES # assigning an array still works page.page = np.zeros((2, 2, 3), dtype=np.uint8) assert page.page.shape == (2, 2, 3) diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index c4b9cea531..e08032666c 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -1,3 +1,6 @@ +import threading +import time + import cv2 import numpy as np import pytest @@ -1134,6 +1137,36 @@ def test_predictors_lazy_source(tmp_path, predictor_cls, doc_cls): assert type(predictor(source[:0])) is doc_cls +def test_predictors_stream_early_exit_stops_decoding(): + """Stopping a stream early does not wait for the next chunk to be fully decoded in the background""" + pages = _value_pages(4) + decoded = [] + closed = threading.Event() + + def _slow_pages(): + try: + for idx in range(1000): + time.sleep(0.02) + decoded.append(idx) + yield pages[idx % len(pages)] + finally: + closed.set() + + predictor = _value_predictor(OCRPredictor, page_chunk_size=8) + stream = predictor.stream(_slow_pages()) + first = next(stream) + assert first.page_idx == 0 + stream.close() + # the first chunk (8 pages), plus at most the pages decoded ahead before the stop + assert len(decoded) < 16 + # no page is decoded after the stream is closed, and the input is closed + count = len(decoded) + time.sleep(0.1) + assert len(decoded) == count + assert closed.is_set() + assert not [thread for thread in threading.enumerate() if thread.name.startswith("doctr-prefetch")] + + def test_predictors_prefetch_only_lazy_inputs(tmp_path, monkeypatch): """In-memory inputs are chunked directly: only lazy sources and iterators are prefetched in a background thread""" from doctr.models.predictor import base as predictor_base From a331eb645f9a5f1bc3843625afba5e8a51be3c22 Mon Sep 17 00:00:00 2001 From: felix Date: Thu, 29 Oct 2026 14:08:11 +0100 Subject: [PATCH 6/6] rebase