rebase & update

Author: felixdittrich92

docs/source/modules/models.rst

@@ -27,6 +27,8 @@ doctr.models.classification
 
 .. autofunction:: doctr.models.classification.mobilenet_v3_small_crop_orientation
 
+.. autofunction:: doctr.models.classification.mobilenet_v3_small_page_orientation
+
 .. autofunction:: doctr.models.classification.magc_resnet31
 
 .. autofunction:: doctr.models.classification.vit_s
@@ -41,6 +43,8 @@ doctr.models.classification
 
 .. autofunction:: doctr.models.classification.crop_orientation_predictor
 
+.. autofunction:: doctr.models.classification.page_orientation_predictor
+
 
 doctr.models.detection
 ----------------------

tests/pytorch/test_models_classification_pt.py

@@ -136,18 +136,16 @@ def test_crop_orientation_model(mock_text_box):
 
 
 def test_page_orientation_model(mock_payslip):
-    image_rot_0 = cv2.imread(mock_payslip)
+    text_box_0 = cv2.imread(mock_payslip)
     # rotates counter-clockwise
-    image_rot_90 = np.rot90(image_rot_0, 3)
-    image_rot_180 = np.rot90(image_rot_0, 2)
-    image_rot_270 = np.rot90(image_rot_0, 1)
-    classifier = classification.page_orientation_predictor("mobilenet_v3_small_page_orientation", pretrained=True)
-    assert classifier([image_rot_0, image_rot_270, image_rot_180, image_rot_90])[0] == [0, 1, 2, 3]
+    text_box_270 = np.rot90(text_box_0, 1)
+    text_box_180 = np.rot90(text_box_0, 2)
+    text_box_90 = np.rot90(text_box_0, 3)
+    classifier = classification.crop_orientation_predictor("mobilenet_v3_small_page_orientation", pretrained=True)
+    assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[0] == [0, 1, 2, 3]
     # 270 degrees is equivalent to -90 degrees
-    assert classifier([image_rot_0, image_rot_270, image_rot_180, image_rot_90])[1] == [0, -90, 180, 90]
-    assert all(
-        isinstance(pred, float) for pred in classifier([image_rot_0, image_rot_270, image_rot_180, image_rot_90])[2]
-    )
+    assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[1] == [0, -90, 180, 90]
+    assert all(isinstance(pred, float) for pred in classifier([text_box_0, text_box_270, text_box_180, text_box_90])[2])
 
 
 @pytest.mark.parametrize(

tests/tensorflow/test_models_classification_tf.py

@@ -116,18 +116,16 @@ def test_crop_orientation_model(mock_text_box):
 
 
 def test_page_orientation_model(mock_payslip):
-    image_rot_0 = cv2.imread(mock_payslip)
+    text_box_0 = cv2.imread(mock_payslip)
     # rotates counter-clockwise
-    image_rot_90 = np.rot90(image_rot_0, 3)
-    image_rot_180 = np.rot90(image_rot_0, 2)
-    image_rot_270 = np.rot90(image_rot_0, 1)
-    classifier = classification.page_orientation_predictor("mobilenet_v3_small_page_orientation", pretrained=True)
-    assert classifier([image_rot_0, image_rot_270, image_rot_180, image_rot_90])[0] == [0, 1, 2, 3]
+    text_box_270 = np.rot90(text_box_0, 1)
+    text_box_180 = np.rot90(text_box_0, 2)
+    text_box_90 = np.rot90(text_box_0, 3)
+    classifier = classification.crop_orientation_predictor("mobilenet_v3_small_page_orientation", pretrained=True)
+    assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[0] == [0, 1, 2, 3]
     # 270 degrees is equivalent to -90 degrees
-    assert classifier([image_rot_0, image_rot_270, image_rot_180, image_rot_90])[1] == [0, -90, 180, 90]
-    assert all(
-        isinstance(pred, float) for pred in classifier([image_rot_0, image_rot_270, image_rot_180, image_rot_90])[2]
-    )
+    assert classifier([text_box_0, text_box_270, text_box_180, text_box_90])[1] == [0, -90, 180, 90]
+    assert all(isinstance(pred, float) for pred in classifier([text_box_0, text_box_270, text_box_180, text_box_90])[2])
 
 
 # temporarily fix to avoid killing the CI (tf2onnx v1.14 memory leak issue)