From 6dcb313b462ca519b2b5015f7e7508b894877df4 Mon Sep 17 00:00:00 2001 From: felix Date: Tue, 27 Oct 2026 07:31:39 +0100 Subject: [PATCH] [models] Update hook handling --- doctr/models/kie_predictor/pytorch.py | 20 +++++-- doctr/models/predictor/base.py | 85 +++++++++++++++++++++++++++ doctr/models/predictor/pytorch.py | 7 +-- tests/common/test_models.py | 47 +++++++++++++++ tests/pytorch/test_models_zoo_pt.py | 65 ++++++++++++++++++++ 5 files changed, 216 insertions(+), 8 deletions(-) diff --git a/doctr/models/kie_predictor/pytorch.py b/doctr/models/kie_predictor/pytorch.py index d00cdf929f..9dc929c5a9 100644 --- a/doctr/models/kie_predictor/pytorch.py +++ b/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: diff --git a/doctr/models/predictor/base.py b/doctr/models/predictor/base.py index acacfe92de..0f27152f5a 100644 --- a/doctr/models/predictor/base.py +++ b/doctr/models/predictor/base.py @@ -313,7 +313,92 @@ def _remap_to_original_coords( 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 diff --git a/doctr/models/predictor/pytorch.py b/doctr/models/predictor/pytorch.py index 1db1636226..db70ffbc07 100644 --- a/doctr/models/predictor/pytorch.py +++ b/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: diff --git a/tests/common/test_models.py b/tests/common/test_models.py index c35b7b3a73..f640606ae0 100644 --- a/tests/common/test_models.py +++ b/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) diff --git a/tests/pytorch/test_models_zoo_pt.py b/tests/pytorch/test_models_zoo_pt.py index 276ff367d3..9ba8652b41 100644 --- a/tests/pytorch/test_models_zoo_pt.py +++ b/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),