[models] Update hook handling

doctr/models/kie_predictor/pytorch.py

@@ -142,6 +142,22 @@ def forward(
             dict_loc_preds[class_name] = _loc_preds
             objectness_scores[class_name] = _scores
 
+        # Apply hooks to loc_preds if any, before the crops are extracted so that the hooked boxes are read
+        if self.hooks:
+            # Pass a copy, so that a hook modifying its input in place does not alter the reference
+            hooked_preds: Any = {class_name: list(preds) for class_name, preds in dict_loc_preds.items()}
+            for hook in self.hooks:
+                hooked_preds = hook(hooked_preds)
+            if not isinstance(hooked_preds, dict) or set(hooked_preds) != set(dict_loc_preds):
+                raise ValueError(
+                    "hooks of the KIE predictor must return a dict with the same classes as their input: "
+                    f"{sorted(dict_loc_preds)}"
+                )
+            for class_name in dict_loc_preds.keys():
+                dict_loc_preds[class_name], objectness_scores[class_name] = self._align_hook_output(
+                    dict_loc_preds[class_name], hooked_preds[class_name], objectness_scores[class_name]
+                )
+
         # Crop images
         crops = {}
         for class_name in dict_loc_preds.keys():
@@ -153,10 +169,6 @@ def forward(
                 assume_horizontal=self._page_orientation_disabled,
             )
 
-        # Apply hooks to loc_preds if any
-        for hook in self.hooks:
-            dict_loc_preds = hook(dict_loc_preds)
-
         # Rectify crop orientation
         crop_orientations: Any = {}
         if not self.assume_straight_pages:

doctr/models/predictor/base.py

@@ -313,7 +313,92 @@
     def add_hook(self, hook: Callable) -> None:
         """Add a hook to the predictor
 
+        Hooks run after text detection and before the word crops are extracted, in the order they were added.
+        Each hook receives the location predictions, one array of boxes per page in relative coordinates
+        (`(N, 4)` straight boxes or `(N, 4, 2)` polygons), and returns them in the same format. It can move,
+        resize, add or remove boxes: only the returned boxes are cropped and read by the recognition model.
+
         Args:
             hook: a callable that takes as input the `loc_preds` and returns the modified `loc_preds`
         """
         self.hooks.append(hook)
+
+    def _apply_hooks(
+        self, loc_preds: list[np.ndarray], objectness_scores: list[np.ndarray]
+    ) -> tuple[list[np.ndarray], list[np.ndarray]]:
+        """Run the hooks on the location predictions, before the word crops are extracted
+
+        Args:
+            loc_preds: one array of boxes per page, in relative coordinates
+            objectness_scores: one array of objectness scores per page
+
+        Returns:
+            the location predictions returned by the hooks, and the matching objectness scores
+        """
+        if not self.hooks:
+            return loc_preds, objectness_scores
+        # Pass a copy of the list, so that a hook replacing the arrays in place does not alter the reference
+        hooked_preds: Any = list(loc_preds)
+        for hook in self.hooks:
+            hooked_preds = hook(hooked_preds)
+        return self._align_hook_output(loc_preds, hooked_preds, objectness_scores)
+
+    @staticmethod
+    def _align_hook_output(
+        loc_preds: list[np.ndarray], hooked_preds: Any, objectness_scores: list[np.ndarray]
+    ) -> tuple[list[np.ndarray], list[np.ndarray]]:
+        """Check the location predictions returned by the hooks, and give each box an objectness score
+
+        A box keeps the score of the closest predicted box of its page, so that moved, resized, reordered or
+        filtered boxes keep their own score. A box added on a page without any prediction gets a score of 1.
+
+        Args:
+            loc_preds: the location predictions passed to the hooks, one array per page
+            hooked_preds: the location predictions returned by the hooks
+            objectness_scores: the objectness scores of `loc_preds`, one array per page
+
+        Returns:
+            the location predictions returned by the hooks as arrays, and their objectness scores
+        """
+        if not isinstance(hooked_preds, (list, tuple)) or len(hooked_preds) != len(loc_preds):
+            received = type(hooked_preds).__name__
+            if isinstance(hooked_preds, (list, tuple)):
+                received += f" of length {len(hooked_preds)}"
+            raise ValueError(
+                f"hooks must return one array of boxes per page: expected a list of {len(loc_preds)} arrays, "
+                f"got a {received}"
+            )
+
+        out_preds: list[np.ndarray] = []
+        out_scores: list[np.ndarray] = []
+        for page_idx, (boxes, new_boxes, scores) in enumerate(zip(loc_preds, hooked_preds, objectness_scores)):
+            new_boxes = np.asarray(new_boxes, dtype=boxes.dtype)
+            if new_boxes.size == 0:
+                new_boxes = new_boxes.reshape((0, *boxes.shape[1:]))
+            if new_boxes.shape[1:] != boxes.shape[1:]:
+                raise ValueError(
+                    "hooks must keep the format of the boxes: expected an array of shape "
+                    f"(N, {', '.join(map(str, boxes.shape[1:]))}) for page {page_idx}, got {new_boxes.shape}"
+                )
+
+            if new_boxes.shape == boxes.shape and np.array_equal(new_boxes, boxes):
+                new_scores = scores
+            elif len(new_boxes) == 0:
+                new_scores = scores[:0]
+            elif len(boxes) == 0:
+                new_scores = np.ones(len(new_boxes), dtype=scores.dtype)
+            else:
+                # Match each returned box with the closest predicted box (L1 distance between the corners),
+                # by chunks to bound the memory used on dense pages
+                ref = boxes.reshape(len(boxes), -1)
+                flat = new_boxes.reshape(len(new_boxes), -1)
+                idcs = np.concatenate([
+                    np.abs(flat[start : start + 256, None, :] - ref[None, :, :]).sum(-1).argmin(axis=1)
+                    for start in range(0, len(flat), 256)
+                ])
+                new_scores = scores[idcs]
+
+            out_preds.append(new_boxes)
+            out_scores.append(new_scores)
+
+        return out_preds, out_scores

doctr/models/predictor/pytorch.py

@@ -152,6 +152,9 @@ def forward(
         # Detach objectness scores from loc_preds
         loc_preds, objectness_scores = detach_scores(loc_preds)
 
+        # Apply hooks to loc_preds if any, before the crops are extracted so that the hooked boxes are read
+        loc_preds, objectness_scores = self._apply_hooks(loc_preds, objectness_scores)
+
         # Crop images
         crops, loc_preds, objectness_scores = self._prepare_crops(
             pages,
@@ -161,10 +164,6 @@ def forward(
             assume_horizontal=self._page_orientation_disabled,
         )
 
-        # Apply hooks to loc_preds if any
-        for hook in self.hooks:
-            loc_preds = hook(loc_preds)
-
         # Rectify crop orientation and get crop orientation predictions
         crop_orientations: Any = []
         if not self.assume_straight_pages:

tests/common/test_models.py

@@ -13,6 +13,7 @@
     invert_data_structure,
     mask_boxes,
 )
+from doctr.models.predictor.base import _OCRPredictor
 from doctr.utils import geometry
 
 
@@ -185,3 +186,49 @@ def test_aggregate_confidence_edge_cases():
     for invalid in ["average", ["mean"], {"mean"}, None, 1]:
         with pytest.raises(ValueError, match="Unknown confidence aggregation"):
             aggregate_confidence([0.5], invalid)
+
+
+def test_align_hook_output():
+    boxes = [
+        np.array([[0.1, 0.1, 0.2, 0.2], [0.5, 0.5, 0.6, 0.6], [0.8, 0.1, 0.9, 0.2]], dtype=np.float32),
+        np.zeros((0, 4), dtype=np.float32),
+    ]
+    scores = [np.array([0.9, 0.3, 0.6], dtype=np.float32), np.zeros((0,), dtype=np.float32)]
+
+    # Unchanged boxes keep their scores
+    out_boxes, out_scores = _OCRPredictor._align_hook_output(boxes, [b.copy() for b in boxes], scores)
+    assert all(np.array_equal(a, b) for a, b in zip(out_boxes, boxes))
+    assert out_scores[0] == pytest.approx([0.9, 0.3, 0.6])
+
+    # Moved, reordered and filtered boxes keep the score of their closest box; lists are converted to arrays,
+    # and a box added on a page without predictions gets a score of 1
+    hooked = [np.array([[0.81, 0.1, 0.91, 0.2], [0.11, 0.1, 0.21, 0.2]]), [[0.0, 0.0, 0.1, 0.1]]]
+    out_boxes, out_scores = _OCRPredictor._align_hook_output(boxes, hooked, scores)
+    assert out_boxes[0].dtype == np.float32 and out_boxes[0].shape == (2, 4)
+    assert out_boxes[1].shape == (1, 4)
+    assert out_scores[0] == pytest.approx([0.6, 0.9])
+    assert out_scores[1] == pytest.approx([1.0])
+
+    # All the boxes can be removed
+    out_boxes, out_scores = _OCRPredictor._align_hook_output(boxes, [[], []], scores)
+    assert [b.shape for b in out_boxes] == [(0, 4), (0, 4)]
+    assert [len(s) for s in out_scores] == [0, 0]
+
+    # Rotated boxes
+    polys = [
+        np.array([
+            [[0.1, 0.1], [0.2, 0.1], [0.2, 0.2], [0.1, 0.2]],
+            [[0.5, 0.5], [0.6, 0.5], [0.6, 0.6], [0.5, 0.6]],
+        ])
+    ]
+    out_boxes, out_scores = _OCRPredictor._align_hook_output(polys, [polys[0][1:]], [np.array([0.4, 0.7])])
+    assert out_boxes[0].shape == (1, 4, 2)
+    assert out_scores[0] == pytest.approx([0.7])
+
+    # Invalid outputs
+    with pytest.raises(ValueError, match="one array of boxes per page"):
+        _OCRPredictor._align_hook_output(boxes, boxes[:1], scores)
+    with pytest.raises(ValueError, match="one array of boxes per page"):
+        _OCRPredictor._align_hook_output(boxes, boxes[0], scores)
+    with pytest.raises(ValueError, match="format of the boxes"):
+        _OCRPredictor._align_hook_output(boxes, [boxes[0][:, :2], boxes[1]], scores)

tests/pytorch/test_models_zoo_pt.py

@@ -204,6 +204,71 @@ def forward(self, crops, **kwargs):
     assert [word.objectness_score for word in words] == pytest.approx([0.9, 0.8])
 
 
+@pytest.mark.parametrize("assume_straight_pages", [True, False])
+@pytest.mark.parametrize("predictor_cls", [OCRPredictor, KIEPredictor])
+def test_predictors_hooks_run_before_cropping(predictor_cls, assume_straight_pages):
+    # Each box covers a region of the page filled with its own gray level, read back by the recognition stub
+    levels = [10, 100, 200]
+    page = np.full((100, 200, 3), 255, dtype=np.uint8)
+    straight_boxes = [[0.1, 0.1, 0.3, 0.2], [0.4, 0.1, 0.6, 0.2], [0.7, 0.1, 0.9, 0.2]]
+    for (xmin, ymin, xmax, ymax), level in zip(straight_boxes, levels):
+        page[int(ymin * 100) : int(ymax * 100), int(xmin * 200) : int(xmax * 200)] = level
+    scores = [0.9, 0.8, 0.7]
+    if assume_straight_pages:
+        boxes = np.array([[*box, score] for box, score in zip(straight_boxes, scores)], dtype=np.float32)
+    else:
+        boxes = np.array(
+            [
+                [[x0, y0], [x1, y0], [x1, y1], [x0, y1], [score, score]]
+                for (x0, y0, x1, y1), score in zip(straight_boxes, scores)
+            ],
+            dtype=np.float32,
+        )
+
+    class _FixedDetPredictor(nn.Module):
+        def forward(self, pages, **kwargs):
+            return [{CLASS_NAME: boxes.copy()} for _ in pages]
+
+    class _GrayLevelRecoPredictor(nn.Module):
+        def forward(self, crops, **kwargs):
+            return [(str(int(np.median(crop))), 1.0) for crop in crops]
+
+    predictor = predictor_cls(
+        _FixedDetPredictor(),
+        _GrayLevelRecoPredictor(),
+        assume_straight_pages=assume_straight_pages,
+        disable_page_orientation=True,
+        disable_crop_orientation=True,
+    )
+
+    def _words(out):
+        page_out = out.pages[0]
+        if predictor_cls is KIEPredictor:
+            return page_out.predictions[CLASS_NAME]
+        return [word for block in page_out.blocks for line in block.lines for word in line.words]
+
+    # A hook removing the first box: only the other boxes are read, each with its own text and score
+    if predictor_cls is KIEPredictor:
+        predictor.add_hook(lambda preds: {name: [p[1:] for p in pages] for name, pages in preds.items()})
+    else:
+        predictor.add_hook(lambda preds: [p[1:] for p in preds])
+    words = sorted(_words(predictor([page])), key=lambda word: np.asarray(word.geometry)[..., 0].min())
+    assert [word.value for word in words] == ["100", "200"]
+    assert [word.objectness_score for word in words] == pytest.approx([0.8, 0.7])
+
+    # Hooks run in order: the second one only keeps the box of the third region
+    if predictor_cls is KIEPredictor:
+        predictor.add_hook(lambda preds: {name: [p[1:] for p in pages] for name, pages in preds.items()})
+    else:
+        predictor.add_hook(lambda preds: [p[1:] for p in preds])
+    assert [word.value for word in _words(predictor([page]))] == ["200"]
+
+    # Invalid outputs are reported
+    predictor.hooks = [lambda preds: preds[0] if isinstance(preds, list) else list(preds.values())]
+    with pytest.raises(ValueError):
+        predictor([page])
+
+
 def test_ocrpredictor_layout(mock_pdf, mock_vocab, mock_payslip):
     det_predictor = DetectionPredictor(
         PreProcessor(output_size=(512, 512), batch_size=2),