fix: handle empty batches in detection and end-to-end predictors

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

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.")

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.")

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