Skip to content

Commit

Permalink
[Fix] crop orientation KIE (mindee#1548)
Browse files Browse the repository at this point in the history
  • Loading branch information
felixdittrich92 authored Apr 15, 2024
1 parent 55dce5e commit 99eb5a3
Show file tree
Hide file tree
Showing 2 changed files with 8 additions and 2 deletions.
5 changes: 4 additions & 1 deletion doctr/models/kie_predictor/pytorch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down
5 changes: 4 additions & 1 deletion doctr/models/kie_predictor/tensorflow.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = {
Expand Down

0 comments on commit 99eb5a3

Please sign in to comment.