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..5e7ccdf935 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -124,6 +124,37 @@ def test_ocrpredictor( assert out.pages[0].orientation["value"] == orientation +def test_predictors_on_empty_batch(mock_vocab): + 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),