diff --git a/doctr/models/kie_predictor/pytorch.py b/doctr/models/kie_predictor/pytorch.py index bd315f6c31..d00cdf929f 100644 --- a/doctr/models/kie_predictor/pytorch.py +++ b/doctr/models/kie_predictor/pytorch.py @@ -142,10 +142,6 @@ def forward( dict_loc_preds[class_name] = _loc_preds objectness_scores[class_name] = _scores - # Apply hooks to loc_preds if any - for hook in self.hooks: - dict_loc_preds = hook(dict_loc_preds) - # Crop images crops = {} for class_name in dict_loc_preds.keys(): @@ -156,6 +152,11 @@ def forward( assume_straight_pages=self.assume_straight_pages, 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/pytorch.py b/doctr/models/predictor/pytorch.py index 66dde94f16..1db1636226 100644 --- a/doctr/models/predictor/pytorch.py +++ b/doctr/models/predictor/pytorch.py @@ -152,10 +152,6 @@ def forward( # Detach objectness scores from loc_preds loc_preds, objectness_scores = detach_scores(loc_preds) - # Apply hooks to loc_preds if any - for hook in self.hooks: - loc_preds = hook(loc_preds) - # Crop images crops, loc_preds, objectness_scores = self._prepare_crops( pages, @@ -164,6 +160,11 @@ def forward( assume_straight_pages=self.assume_straight_pages, 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: