From 4fc8ac1ad4cc4cf1ad9d470e293fe4f0f1953400 Mon Sep 17 00:00:00 2001 From: Mohammad Hijjawi Date: Wed, 7 Oct 2026 09:34:35 +0100 Subject: [PATCH] [models] Keep the objectness scores aligned with the boxes when empty crops are dropped `_prepare_crops` drops the boxes whose crop is empty before recognition, but the objectness scores were detached earlier and kept their original length. Every word after a dropped box was then built with the score of an earlier box (the dropped one, for the first word after it), in both `OCRPredictor` and `KIEPredictor`. The scores are now filtered with the same mask as the boxes. --- doctr/models/kie_predictor/pytorch.py | 3 ++- doctr/models/predictor/base.py | 7 +++++-- doctr/models/predictor/pytorch.py | 3 ++- tests/pytorch/test_models_zoo_pt.py | 28 +++++++++++++++++++++++++++ 4 files changed, 37 insertions(+), 4 deletions(-) diff --git a/doctr/models/kie_predictor/pytorch.py b/doctr/models/kie_predictor/pytorch.py index 75a56e0d86..bd315f6c31 100644 --- a/doctr/models/kie_predictor/pytorch.py +++ b/doctr/models/kie_predictor/pytorch.py @@ -149,9 +149,10 @@ def forward( # Crop images crops = {} for class_name in dict_loc_preds.keys(): - crops[class_name], dict_loc_preds[class_name] = self._prepare_crops( + crops[class_name], dict_loc_preds[class_name], objectness_scores[class_name] = self._prepare_crops( pages, dict_loc_preds[class_name], + objectness_scores[class_name], assume_straight_pages=self.assume_straight_pages, assume_horizontal=self._page_orientation_disabled, ) diff --git a/doctr/models/predictor/base.py b/doctr/models/predictor/base.py index 644beeb7c2..acacfe92de 100644 --- a/doctr/models/predictor/base.py +++ b/doctr/models/predictor/base.py @@ -180,9 +180,10 @@ def _generate_crops( def _prepare_crops( pages: list[np.ndarray], loc_preds: list[np.ndarray], + objectness_scores: list[np.ndarray], assume_straight_pages: bool = False, assume_horizontal: bool = False, - ) -> tuple[list[list[np.ndarray]], list[np.ndarray]]: + ) -> tuple[list[list[np.ndarray]], list[np.ndarray], list[np.ndarray]]: crops = _OCRPredictor._generate_crops(pages, loc_preds, assume_straight_pages, assume_horizontal) # Avoid sending zero-sized crops @@ -192,8 +193,10 @@ def _prepare_crops( for page_crops, page_kept in zip(crops, is_kept) ] loc_preds = [_boxes[_kept] for _boxes, _kept in zip(loc_preds, is_kept)] + # Drop the scores of the discarded boxes too, so that each kept box keeps its own score + objectness_scores = [_scores[_kept] for _scores, _kept in zip(objectness_scores, is_kept)] - return crops, loc_preds + return crops, loc_preds, objectness_scores def _rectify_crops( self, diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py index 8a18d3ac39..66dde94f16 100644 --- a/doctr/models/predictor/pytorch.py +++ b/doctr/models/predictor/pytorch.py @@ -157,9 +157,10 @@ def forward( loc_preds = hook(loc_preds) # Crop images - crops, loc_preds = self._prepare_crops( + crops, loc_preds, objectness_scores = self._prepare_crops( pages, loc_preds, + objectness_scores, assume_straight_pages=self.assume_straight_pages, assume_horizontal=self._page_orientation_disabled, ) diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index 2654d69736..be6af75b75 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/tests/pytorch/test_models_zoo_pt.py @@ -156,6 +156,34 @@ def test_predictors_on_empty_batch(mock_vocab): assert out.export() == {"pages": []} +@pytest.mark.parametrize("predictor_cls", [OCRPredictor, KIEPredictor]) +def test_predictors_keep_objectness_scores_aligned(predictor_cls): + # The middle box lies on the right border of the page: its crop is empty, so it is dropped before recognition + boxes = np.array( + [[0.1, 0.1, 0.3, 0.2, 0.9], [1.0, 0.5, 1.0, 0.6, 0.1], [0.5, 0.1, 0.7, 0.2, 0.8]], + dtype=np.float32, + ) + + class _FixedDetPredictor(nn.Module): + def forward(self, pages, **kwargs): + return [{CLASS_NAME: boxes.copy()} for _ in pages] + + class _DummyRecoPredictor(nn.Module): + def forward(self, crops, **kwargs): + return [(f"word{idx}", 1.0) for idx in range(len(crops))] + + predictor = predictor_cls(_FixedDetPredictor(), _DummyRecoPredictor(), assume_straight_pages=True) + page = predictor([np.full((100, 200, 3), 255, dtype=np.uint8)]).pages[0] + if predictor_cls is KIEPredictor: + words = page.predictions[CLASS_NAME] + else: + words = [word for block in page.blocks for line in block.lines for word in line.words] + + # Each remaining word keeps the objectness score of its own box + assert [word.geometry[0][0] for word in words] == pytest.approx([0.1, 0.5]) + assert [word.objectness_score for word in words] == pytest.approx([0.9, 0.8]) + + def test_ocrpredictor_layout(mock_pdf, mock_vocab, mock_payslip): det_predictor = DetectionPredictor( PreProcessor(output_size=(512, 512), batch_size=2),