From a11910cff1caac05ce163c337c977c6ad8c007ae Mon Sep 17 00:00:00 2001 From: linhongyu510 Date: Sat, 5 Sep 2026 16:40:09 +0800 Subject: [PATCH] fix: handle empty batches in detection and end-to-end predictors Passing an empty page list crashed three levels deep instead of returning an empty result. `RecognitionPredictor` has always short-circuited on an empty input, and `OrientationPredictor` gained the same behaviour in #2069, but the detection and end-to-end predictors did not, so filtering a batch down to nothing raised from internals that never mention the empty input: - `PreProcessor.batch_inputs` computes `num_batches == 0` correctly, then reads `samples[0]` to pick the tuple/tensor branch -> `IndexError` - `detach_scores` calls `zip(*(...))` over no boxes -> `ValueError: not enough values to unpack (expected 2, got 0)` - on the KIE path `invert_data_structure` reads `x[0]` -> `IndexError` Guard at the three public entry points rather than patching each internal, so the existing helpers keep their non-empty precondition. `DetectionPredictor` returns the shape its `return_maps` contract promises, and the end-to-end predictors return an empty document of their own type -- `Document` for `OCRPredictor`, `KIEDocument` for `KIEPredictor`, whose per-class page shape would otherwise be lost to the base class. Verified `ruff check`, `ruff format --check` and `mypy doctr/` clean; reverting the guards turns the new test red at `preprocessor/pytorch.py:73`. Co-authored-by: Claude --- doctr/models/detection/predictor/pytorch.py | 3 ++ doctr/models/kie_predictor/pytorch.py | 5 ++- doctr/models/predictor/pytorch.py | 3 ++ tests/pytorch/test_models_zoo_pt.py | 41 +++++++++++++++++++++ 4 files changed, 51 insertions(+), 1 deletion(-) diff --git a/doctr/models/detection/predictor/pytorch.py b/doctr/models/detection/predictor/pytorch.py index 69d7e7bd03..2de50c3ed6 100644 --- a/doctr/models/detection/predictor/pytorch.py +++ b/doctr/models/detection/predictor/pytorch.py @@ -40,6 +40,9 @@ def forward( return_maps: bool = False, **kwargs: Any, ) -> list[dict[str, np.ndarray]] | tuple[list[dict[str, np.ndarray]], list[np.ndarray]]: + if len(pages) == 0: + return ([], []) if return_maps else [] + # Extract parameters from the preprocessor preserve_aspect_ratio = self.pre_processor.resize.preserve_aspect_ratio symmetric_pad = self.pre_processor.resize.symmetric_pad diff --git a/doctr/models/kie_predictor/pytorch.py b/doctr/models/kie_predictor/pytorch.py index 6e71d4e5c0..75a56e0d86 100644 --- a/doctr/models/kie_predictor/pytorch.py +++ b/doctr/models/kie_predictor/pytorch.py @@ -9,7 +9,7 @@ import torch from torch import nn -from doctr.io.elements import Document +from doctr.io.elements import Document, KIEDocument from doctr.models._utils import get_language, invert_data_structure from doctr.models.detection.predictor import DetectionPredictor from doctr.models.layout.predictor import LayoutPredictor @@ -79,6 +79,9 @@ def forward( pages: list[np.ndarray], **kwargs: Any, ) -> Document: + if len(pages) == 0: + return KIEDocument(pages=[]) + # Dimension check if any(page.ndim != 3 for page in pages): raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.") diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py index c0ec204aca..8a18d3ac39 100644 --- a/doctr/models/predictor/pytorch.py +++ b/doctr/models/predictor/pytorch.py @@ -93,6 +93,9 @@ def forward( pages: list[np.ndarray], **kwargs: Any, ) -> Document: + if len(pages) == 0: + return Document(pages=[]) + # Dimension check if any(page.ndim != 3 for page in pages): raise ValueError("incorrect input shape: all pages are expected to be multi-channel 2D images.") diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index 3ba4a42fe2..927c5af2a0 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -124,6 +124,47 @@ def test_ocrpredictor( assert out.pages[0].orientation["value"] == orientation +def test_predictors_on_empty_batch(mock_vocab): + """An empty page list must yield an empty Document instead of raising. + + Filtering a batch down to nothing is ordinary caller code, and + `RecognitionPredictor` (and `OrientationPredictor` since #2069) already + return empty results for it. The detection and end-to-end predictors did + not, so they crashed deep in the stack -- `IndexError` from `samples[0]` in + `PreProcessor.batch_inputs`, then `ValueError` from `zip(*...)` in + `detach_scores`, then `IndexError` from `x[0]` in `invert_data_structure` + on the KIE path -- none of which names the empty input. + """ + det_predictor = DetectionPredictor( + PreProcessor(output_size=(512, 512), batch_size=2), + detection.db_mobilenet_v3_large(pretrained=False, pretrained_backbone=False, assume_straight_pages=True), + ) + reco_predictor = RecognitionPredictor( + PreProcessor(output_size=(32, 128), batch_size=32, preserve_aspect_ratio=True), + recognition.crnn_vgg16_bn(pretrained=False, pretrained_backbone=False, vocab=mock_vocab), + ) + + # Detection keeps the shape its `return_maps` contract promises. + assert det_predictor([]) == [] + assert det_predictor([], return_maps=True) == ([], []) + + # The recognition predictor already behaved; asserted here so the three + # stay consistent if one of them is touched again. + assert reco_predictor([]) == [] + + for predictor, expected_type in ( + (OCRPredictor(det_predictor, reco_predictor, assume_straight_pages=True), Document), + (KIEPredictor(det_predictor, reco_predictor, assume_straight_pages=True), KIEDocument), + ): + out = predictor([]) + # Exact type, not isinstance: KIEDocument subclasses Document, so an + # isinstance check would not notice the KIE path degrading to the base + # class and dropping the per-class prediction shape. + assert type(out) is expected_type + assert out.pages == [] + assert out.export() == {"pages": []} + + def test_ocrpredictor_layout(mock_pdf, mock_vocab, mock_payslip): det_predictor = DetectionPredictor( PreProcessor(output_size=(512, 512), batch_size=2),