feat: ✨ tf mobilenet_v3_small_page_orientation checkpoint (#1636)

mindee · Jun 6, 2024 · 866557a · 866557a
1 parent 5eda559
commit 866557a
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 4 deletions.
diff --git a/doctr/models/classification/mobilenet/tensorflow.py b/doctr/models/classification/mobilenet/tensorflow.py
@@ -67,7 +67,7 @@
         "std": (0.299, 0.296, 0.301),
         "input_shape": (512, 512, 3),
         "classes": [0, -90, 180, 90],
-        "url": None,
+        "url": "https://doctr-static.mindee.com/models?id=v0.8.1/mobilenet_v3_small_page_orientation-aec9553e.zip&src=0",
     },
 }
 

diff --git a/tests/tensorflow/test_models_classification_tf.py b/tests/tensorflow/test_models_classification_tf.py
@@ -114,8 +114,6 @@ def test_crop_orientation_model(mock_text_box):
     assert all(isinstance(pred, float) for pred in classifier([text_box_0, text_box_270, text_box_180, text_box_90])[2])
 
 
-# TODO: uncomment when model is available
-"""
 def test_page_orientation_model(mock_payslip):
     text_box_0 = cv2.imread(mock_payslip)
     # rotates counter-clockwise
@@ -127,7 +125,6 @@ def test_page_orientation_model(mock_payslip):
     # 270 degrees is equivalent to -90 degrees
     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)