Prototype: [io] Lazy page sources & chunked/streaming predictors for long documents and large image collections

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
 

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
 

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:

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

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
 
 

doctr/io/__init__.py

@@ -5,3 +5,4 @@
 from .image import *
 from .pdf import *
 from .reader import *
+from .sources import *

doctr/io/elements.py

@@ -3,6 +3,7 @@
 # This program is licensed under the Apache License 2.0.
 # See LICENSE or go to <https://opensource.org/licenses/Apache-2.0> 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"""
 
@@ -434,7 +457,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)
@@ -450,9 +473,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],
@@ -526,7 +551,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
@@ -539,9 +564,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],

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()