From 99eb5a33aaab2a7a0e037963810c7304a9cbd956 Mon Sep 17 00:00:00 2001 From: Felix Dittrich Date: Mon, 15 Apr 2024 10:37:53 +0200 Subject: [PATCH] [Fix] crop orientation KIE (#1548) --- doctr/models/kie_predictor/pytorch.py | 5 ++++- doctr/models/kie_predictor/tensorflow.py | 5 ++++- 2 files changed, 8 insertions(+), 2 deletions(-) diff --git a/doctr/models/kie_predictor/pytorch.py b/doctr/models/kie_predictor/pytorch.py index d8c30f826..0fa183f5b 100644 --- a/doctr/models/kie_predictor/pytorch.py +++ b/doctr/models/kie_predictor/pytorch.py @@ -123,9 +123,12 @@ def forward( crop_orientations: Any = {} if not self.assume_straight_pages: for class_name in dict_loc_preds.keys(): - crops[class_name], dict_loc_preds[class_name], crop_orientations[class_name] = self._rectify_crops( + crops[class_name], dict_loc_preds[class_name], word_orientations = self._rectify_crops( crops[class_name], dict_loc_preds[class_name] ) + crop_orientations[class_name] = [ + {"value": orientation[0], "confidence": orientation[1]} for orientation in word_orientations + ] # Identify character sequences word_preds = { diff --git a/doctr/models/kie_predictor/tensorflow.py b/doctr/models/kie_predictor/tensorflow.py index 8284f2ee5..52b1211dd 100644 --- a/doctr/models/kie_predictor/tensorflow.py +++ b/doctr/models/kie_predictor/tensorflow.py @@ -118,9 +118,12 @@ def __call__( crop_orientations: Any = {} if not self.assume_straight_pages: for class_name in dict_loc_preds.keys(): - crops[class_name], dict_loc_preds[class_name], crop_orientations[class_name] = self._rectify_crops( + crops[class_name], dict_loc_preds[class_name], word_orientations = self._rectify_crops( crops[class_name], dict_loc_preds[class_name] ) + crop_orientations[class_name] = [ + {"value": orientation[0], "confidence": orientation[1]} for orientation in word_orientations + ] # Identify character sequences word_preds = {