[models] Keep the objectness scores aligned with the boxes when empty crops are dropped

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

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,

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

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