[felixdittrich92]

Closes #1240 (ref.: #1111)

Summary

Long PDFs (1000+ pages) and large collections of images can now be processed in one call, with a memory footprint that no longer grows with the number of pages, and without any custom batching logic on the user side.

Before this PR, every page was decoded into a list up front. The predictor then ran every pipeline stage on all pages at once: for example, all pages were preprocessed for detection in a single step. A 300-page A4 PDF was enough to get the process killed by the OOM killer on a 7 GB machine.

This PR adds:

• Lazy page sources (doctr.io.PDFSource, ImageSource, ConcatSource): a class API over generators, as proposed in [IO] Refactor io module from functional API to class API with generators #1240. They behave like a read-only list of pages, but a page is only decoded when it is accessed.
• Chunked prediction: OCRPredictor / KIEPredictor process any input longer than page_chunk_size (default: 32) chunk by chunk. The next chunk is decoded in a background thread while the models run on the current one.
• predictor.stream(...): yields the predicted pages one by one, as soon as they are processed.
• Lightweight results: pages predicted from a lazy source keep a reference to their image instead of the pixels, and page.page decodes it again on access.

Everything is opt-in and backward compatible: DocumentFile.from_* still returns lists by default, and model(list_of_pages) returns the same Document as before.

Usage

A long PDF in a single call

from doctr.io import DocumentFile
from doctr.models import ocr_predictor

model = ocr_predictor(pretrained=True)

doc = DocumentFile.from_pdf("path/to/1200_pages.pdf", lazy=True)  # PDFSource: nothing is rendered yet
print(len(doc), doc[0].shape)  # only the first page is rendered

result = model(doc)  # processed in chunks of `page_chunk_size` pages, one Document returned
print(result.pages[-1].page_idx)  # 1199: indices refer to the whole input

Stream the results with a constant memory footprint

with open("result.md", "w") as f:
    for page in model.stream(doc):
        f.write(page.export_as("md"))

Folders, mixed inputs and huge image collections

# 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, in order
doc = DocumentFile.from_files(["contracts/", "annex.pdf", "signature.png"], lazy=True)

for page in model.stream(doc):
    source_idx, page_idx = doc.locate(page.page_idx)  # which input / which page of it
    print(f"input #{source_idx}, page {page_idx + 1}: {page.render()[:80]}")

The sources behave like lists

from doctr.io import ConcatSource, ImageSource, PDFSource

pdf = PDFSource("report.pdf", scale=2, password=None)
images = ImageSource(["cover.png", b"<raw image bytes>"])

doc = pdf + images  # ConcatSource
first_pages = doc[:10]  # lazy view
for batch in doc.iter_batches(16):  # lists of 16 decoded pages
    ...

with PDFSource("report.pdf") as pdf:  # releases the file handle on exit
    page = pdf[-1]

Tuning the chunk size, and any iterable as input

model = ocr_predictor(pretrained=True, page_chunk_size=16)  # lower peak memory
result = model(frame for frame in video_frames())  # generators work too
pages = model.stream(doc, page_chunk_size=8, prefetch_pages=False)

CLI

doctr-cli --input_path 1200_pages.pdf scans/*.png --output result.md --page_chunk_size 16

Inputs are now opened lazily, and a per-page progress bar is shown.

Implementation notes

• doctr/io/sources.py (new)
  • PageSource(Sequence[np.ndarray]) supports len, indexing, slicing (lazy views), iteration, +, iter_batches, close() and use as a context manager.
  • PDFSource
    • Opens the PDF eagerly to count the pages, so invalid or missing files fail immediately, then releases the handle.
    • Renders pages on demand and closes the document after a full iteration.
    • Every pdfium call goes through a global lock, because pdfium is not thread-safe, not even across documents.
    • It can be pickled: the handle is re-opened on demand.
  • ImageSource
    • Accepts paths, folders and bytes.
    • Missing files fail eagerly.
    • Decoding errors name the faulty file.
  • ConcatSource
    • Flattens nested concatenations and exposes locate().
    • Keeps at most one PDF open at a time.
• DocumentFile
  • from_pdf, from_images and from_url gain lazy=False, with typed overloads.
  • New from_files for mixed inputs.
  • read_pdf now uses PDFSource internally, so there is a single rendering code path.
• Predictors
  • The former forward body moved to _predict_pages.
  • forward dispatches through _run: short lists take the existing single-pass path, everything else goes through stream().
  • stream() makes page_idx global, which also keeps the figure file names of the exports unique.
  • page_chunk_size is a new _OCRPredictor argument, also exposed in the CLI.
• Page.page / KIEPage.page
  • Each is now a property backed by either an array or a zero-argument callable.
  • Pages that the pipeline did not modify point back to their source.
  • Transformed pages (e.g. straighten_pages=True) are kept losslessly PNG-compressed.
• doctr.utils.multithreading: new chunked() and prefetch() helpers. Prefetching is disabled when DOCTR_MULTIPROCESSING_DISABLE is set.

Benchmark

2 CPUs / 7 GB RAM, A4 pages, db_mobilenet_v3_large + crnn_mobilenet_v3_small:

	main	This PR
Load 1,200 pages	killed by OOM	0.9 GB peak, 47 s
OCR on 300 pages	killed by OOM	3.5 GB peak, 337 s
OCR on 1,200 pages, model(doc)	—	4.1 GB peak, 20 min
OCR on 1,200 pages, model.stream(doc)	—	3.9 GB peak, 20 min

Most of the remaining peak is the working set of one 32-page chunk, and it decreases with a lower page_chunk_size.

Behaviour changes to be aware of

• Hooks are called once per chunk on inputs longer than page_chunk_size. Each call receives the location predictions of that chunk's pages. This is documented in the hooks section.
• model(...) now also accepts generators and other iterables of pages.

Tests & docs

• New tests/common/test_io_sources.py:
  • sources, slicing and concatenation;
  • file-handle management, pickling and thread safety;
  • DocumentFile lazy readers and from_files;
  • chunked / prefetch.
• New predictor tests in tests/pytorch/test_models_zoo_pt.py:
  • chunking, streaming and generator inputs, for OCR and KIE;
  • lazy sources and transformed pages;
  • eager vs. lazy PDF results being identical, with and without straighten_pages.
• CLI tests updated (--page_chunk_size, lazy loading).
• Docs:
  • new "Long documents & large collections" section in Loading documents;
  • updated Process many documents section;
  • API reference entries for the new classes and stream.
• ruff check, ruff format --check and mypy doctr/ pass.

[codecov]

Codecov Report

❌ Patch coverage is 97.63314% with 8 lines in your changes missing coverage. Please review.
✅ Project coverage is 97.71%. Comparing base (1a4a35b) to head (3a7b546).

Files with missing lines	Patch %	Lines
doctr/io/sources.py	97.48%	5 Missing ⚠️
doctr/io/reader.py	95.34%	2 Missing ⚠️
doctr/models/predictor/base.py	96.15%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2186      +/-   ##
==========================================
- Coverage   97.74%   97.71%   -0.03%     
==========================================
  Files         170      171       +1     
  Lines       10932    11241     +309     
==========================================
+ Hits        10685    10984     +299     
- Misses        247      257      +10     

Flag	Coverage Δ
unittests	97.71% <97.63%> (-0.03%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.