update

Author: felixdittrich92

doctr/models/_utils.py

@@ -98,7 +98,7 @@ def estimate_orientation(
     if general_page_orientation and general_page_orientation[1] >= min_confidence:
         estimated_angle = estimated_angle + general_page_orientation[0]
 
-    return -estimated_angle  # return the clockwise angle
+    return estimated_angle  # return the clockwise angle
 
 
 def rectify_crops(
@@ -107,10 +107,12 @@ def rectify_crops(
 ) -> List[np.ndarray]:
     """Rotate each crop of the list according to the predicted orientation:
     0: already straight, no rotation
-    1: 90 cw, rotate 3 times ccw
-    2: 180 cw, rotate 2 times ccw
-    3: 270 cw, rotate 1 times ccw
+    1: 90 ccw, rotate 3 times ccw
+    2: 180, rotate 2 times ccw
+    3: 270 ccw, rotate 1 time ccw
     """
+    # Inverse predictions (if angle of +90 is detected, rotate by -90)
+    orientations = [4 - pred if pred != 0 else 0 for pred in orientations]
     return (
         [crop if orientation == 0 else np.rot90(crop, orientation) for orientation, crop in zip(orientations, crops)]
         if len(orientations) > 0
@@ -128,8 +130,7 @@ def rectify_loc_preds(
     return (
         np.stack(
             [
-                # shift clockwise
-                np.roll(page_loc_pred, -orientation, axis=0)
+                np.roll(page_loc_pred, orientation, axis=0)
                 for orientation, page_loc_pred in zip(orientations, page_loc_preds)
             ],
             axis=0,

doctr/models/kie_predictor/pytorch.py

@@ -97,7 +97,7 @@ def forward(
             general_pages_orientations = None
             origin_pages_orientations = None
         if self.straighten_pages:
-            pages = self._get_straightened_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore[arg-type]
+            pages = self._get_straightened_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore
             # Forward again to get predictions on straight pages
             loc_preds = self.det_predictor(pages, **kwargs)
 

doctr/models/predictor/base.py

@@ -77,7 +77,7 @@ def _get_pages_orientations(
     ) -> Tuple[List[Tuple[int, float]], List[int]]:
         general_pages_orientations = self._get_general_pages_orientations(pages)
         origin_page_orientations = [
-            estimate_orientation(seq_map, general_orientation)
+            -estimate_orientation(seq_map, general_orientation)
             for seq_map, general_orientation in zip(seg_maps, general_pages_orientations)
         ]
         return general_pages_orientations, origin_page_orientations
@@ -115,7 +115,6 @@ def _get_straightened_pages(
             ]
         )
         # TODO: test the -> expand if page page width is larger than height
-        # rotate image counter clockwise
         return [
             rotate_image(page, angle, expand=page.shape[1] > page.shape[0])
             for page, angle in zip(pages, origin_pages_orientations)

doctr/models/predictor/pytorch.py

@@ -95,7 +95,7 @@ def forward(
             general_pages_orientations = None
             origin_pages_orientations = None
         if self.straighten_pages:
-            pages = self._get_straightened_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore[assignment]
+            pages = self._get_straightened_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore
             # Forward again to get predictions on straight pages
             loc_preds = self.det_predictor(pages, **kwargs)