fix: make sure that expand in rotate_image is keeping the same image …

…ratio
mindee · fg-mindee · Dec 6, 2021 · Jul 19, 2021 · Jul 19, 2021 · Sep 21, 2021
commit 8525b14471a8f9e426cdc04eb598e304cf50940f
diff --git a/doctr/utils/geometry.py b/doctr/utils/geometry.py
@@ -240,9 +240,8 @@ def rotate_image(
     Args:
         image: numpy tensor to rotate
         angle: rotation angle in degrees, between -90 and +90
-        expand: whether the image should be padded before the rotation
-        preserve_origin_shape: whether the image should be resized to the original image size after the rotation.
-            Only useful if expand is True.
+        expand: whether the image should be padded before the rotation,
+        preserve_origin_shape: if expand is set to True, resizes the final output to the original image size
 
     Returns:
         Rotated array, padded by 0 by default.
@@ -261,16 +260,16 @@ def rotate_image(
     rot_mat = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
     rot_img = cv2.warpAffine(exp_img, rot_mat, (width, height))
     if expand:
+        # Pad to get the same aspect ratio
+        if (image.shape[0] / image.shape[1]) != (rot_img.shape[0] / rot_img.shape[1]):
+            # Pad width
+            if (rot_img.shape[0] / rot_img.shape[1]) > (image.shape[0] / image.shape[1]):
+                h_pad, w_pad = 0, int(rot_img.shape[0] * image.shape[1] / image.shape[0] - rot_img.shape[1])
+            # Pad height
+            else:
+                h_pad, w_pad = int(rot_img.shape[1] * image.shape[0] / image.shape[1] - rot_img.shape[0]), 0
+            rot_img = np.pad(rot_img, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))
         if preserve_origin_shape:
-            # Pad to get the same aspect ratio
-            if (image.shape[0] / image.shape[1]) != (rot_img.shape[0] / rot_img.shape[1]):
-                # Pad width
-                if (rot_img.shape[0] / rot_img.shape[1]) > (image.shape[0] / image.shape[1]):
-                    h_pad, w_pad = 0, int(rot_img.shape[0] * image.shape[1] / image.shape[0] - rot_img.shape[1])
-                # Pad height
-                else:
-                    h_pad, w_pad = int(rot_img.shape[1] * image.shape[0] / image.shape[1] - rot_img.shape[0]), 0
-                rot_img = np.pad(rot_img, ((h_pad // 2, h_pad - h_pad // 2), (w_pad // 2, w_pad - w_pad // 2), (0, 0)))
             # rescale
             rot_img = cv2.resize(rot_img, image.shape[:-1][::-1], interpolation=cv2.INTER_LINEAR)
 

diff --git a/tests/common/test_utils_geometry.py b/tests/common/test_utils_geometry.py
@@ -124,7 +124,7 @@ def test_rotate_image():
 
     # Expand
     rotated = geometry.rotate_image(img, 30., expand=True)
-    assert rotated.shape[:-1] == (60, 72)
+    assert rotated.shape[:-1] == (60, 120)
     assert rotated[0, :, 0].sum() <= 1
 
     # Expand
@@ -134,5 +134,5 @@ def test_rotate_image():
 
     # Expand with 90° rotation
     rotated = geometry.rotate_image(img, 90., expand=True)
-    assert rotated.shape[:-1] == (64, 64)
+    assert rotated.shape[:-1] == (64, 128)
     assert rotated[0, :, 0].sum() <= 1