test: Added unittest for PyTorch high-level zoo (#499)

* test: Renamed TF zoo test * test: Added PyTorch zoo unittest
mindee · Sep 29, 2021 · 3b2d375 · 3b2d375
1 parent e90edff
commit 3b2d375
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 0 deletions.
diff --git a/test/pytorch/test_models_zoo_pt.py b/test/pytorch/test_models_zoo_pt.py
@@ -0,0 +1,17 @@
+import pytest
+
+from doctr import models
+from doctr.models.predictor import OCRPredictor
+
+
+@pytest.mark.parametrize(
+    "det_arch, reco_arch",
+    [
+        ["db_mobilenet_v3_large", "crnn_mobilenet_v3_large"],
+    ],
+)
+def test_zoo_models(det_arch, reco_arch):
+    # Model
+    predictor = models.ocr_predictor(det_arch, reco_arch, pretrained=True)
+    # Output checks
+    assert isinstance(predictor, OCRPredictor)
diff --git a/test/tensorflow/test_models_tf.py → test/tensorflow/test_models_zoo_tf.py b/test/tensorflow/test_models_tf.py → test/tensorflow/test_models_zoo_tf.py