fix straighten pages (#1697)

Author: felixdittrich92

doctr/models/builder.py

@@ -266,7 +266,7 @@ def _build_blocks(
                 Line([
                     Word(
                         *word_preds[idx],
-                        tuple([tuple(pt) for pt in boxes[idx].tolist()]),  # type: ignore[arg-type]
+                        tuple(tuple(pt) for pt in boxes[idx].tolist()),  # type: ignore[arg-type]
                         float(objectness_scores[idx]),
                         crop_orientations[idx],
                     )
@@ -500,7 +500,7 @@ def _build_blocks(  # type: ignore[override]
             Prediction(
                 value=word_preds[idx][0],
                 confidence=word_preds[idx][1],
-                geometry=tuple([tuple(pt) for pt in boxes[idx].tolist()]),  # type: ignore[arg-type]
+                geometry=tuple(tuple(pt) for pt in boxes[idx].tolist()),  # type: ignore[arg-type]
                 objectness_score=float(objectness_scores[idx]),
                 crop_orientation=crop_orientations[idx],
             )

doctr/models/kie_predictor/pytorch.py

@@ -99,6 +99,9 @@ def forward(
             origin_pages_orientations = None
         if self.straighten_pages:
             pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore
+            # update page shapes after straightening
+            origin_page_shapes = [page.shape[:2] for page in pages]
+
             # Forward again to get predictions on straight pages
             loc_preds = self.det_predictor(pages, **kwargs)
 

doctr/models/kie_predictor/tensorflow.py

@@ -99,6 +99,9 @@ def __call__(
             origin_pages_orientations = None
         if self.straighten_pages:
             pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)
+            # update page shapes after straightening
+            origin_page_shapes = [page.shape[:2] for page in pages]
+
             # Forward again to get predictions on straight pages
             loc_preds = self.det_predictor(pages, **kwargs)  # type: ignore[assignment]
 

doctr/models/modules/vision_transformer/pytorch.py

@@ -20,7 +20,7 @@ def __init__(self, input_shape: Tuple[int, int, int], embed_dim: int, patch_size
         channels, height, width = input_shape
         self.patch_size = patch_size
         self.interpolate = True if patch_size[0] == patch_size[1] else False
-        self.grid_size = tuple([s // p for s, p in zip((height, width), self.patch_size)])
+        self.grid_size = tuple(s // p for s, p in zip((height, width), self.patch_size))
         self.num_patches = self.grid_size[0] * self.grid_size[1]
 
         self.cls_token = nn.Parameter(torch.randn(1, 1, embed_dim))

doctr/models/modules/vision_transformer/tensorflow.py

@@ -22,7 +22,7 @@ def __init__(self, input_shape: Tuple[int, int, int], embed_dim: int, patch_size
         height, width, _ = input_shape
         self.patch_size = patch_size
         self.interpolate = True if patch_size[0] == patch_size[1] else False
-        self.grid_size = tuple([s // p for s, p in zip((height, width), self.patch_size)])
+        self.grid_size = tuple(s // p for s, p in zip((height, width), self.patch_size))
         self.num_patches = self.grid_size[0] * self.grid_size[1]
 
         self.cls_token = self.add_weight(shape=(1, 1, embed_dim), initializer="zeros", trainable=True, name="cls_token")

doctr/models/predictor/base.py

@@ -101,8 +101,8 @@ def _straighten_pages(
             ]
         )
         return [
-            # We exapnd if the page is wider than tall and the angle is 90 or -90
-            rotate_image(page, angle, expand=page.shape[1] > page.shape[0] and abs(angle) == 90)
+            # expand if height and width are not equal
+            rotate_image(page, angle, expand=page.shape[0] != page.shape[1])
             for page, angle in zip(pages, origin_pages_orientations)
         ]
 

doctr/models/predictor/pytorch.py

@@ -97,6 +97,9 @@ def forward(
             origin_pages_orientations = None
         if self.straighten_pages:
             pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)  # type: ignore
+            # update page shapes after straightening
+            origin_page_shapes = [page.shape[:2] for page in pages]
+
             # Forward again to get predictions on straight pages
             loc_preds = self.det_predictor(pages, **kwargs)
 

doctr/models/predictor/tensorflow.py

@@ -97,6 +97,9 @@ def __call__(
             origin_pages_orientations = None
         if self.straighten_pages:
             pages = self._straighten_pages(pages, seg_maps, general_pages_orientations, origin_pages_orientations)
+            # update page shapes after straightening
+            origin_page_shapes = [page.shape[:2] for page in pages]
+
             # forward again to get predictions on straight pages
             loc_preds_dict = self.det_predictor(pages, **kwargs)  # type: ignore[assignment]