rebase and add requires

Author: felixdittrich92

.github/workflows/public_docker_images.yml

@@ -23,7 +23,7 @@ jobs:
       matrix:
         # Must match version at https://www.python.org/ftp/python/
         python: ["3.9.18", "3.10.13", "3.11.8"]
-        framework: ["tf", "torch", "tf,viz,html", "torch,viz,html"]
+        framework: ["tf", "torch", "tf,viz,html,contrib", "torch,viz,html,contrib"]
         system: ["cpu", "gpu"]
 
     # Sets the permissions granted to the `GITHUB_TOKEN` for the actions in this job.

docs/source/getting_started/installing.rst

@@ -38,16 +38,16 @@ We strive towards reducing framework-specific dependencies to a minimum, but som
         .. code:: bash
 
             pip install "python-doctr[tf]"
-            # or with preinstalled packages for visualization & html support
-            pip install "python-doctr[tf,viz,html]"
+            # or with preinstalled packages for visualization & html & contrib module support
+            pip install "python-doctr[tf,viz,html,contib]"
 
     .. tab:: PyTorch
 
         .. code:: bash
 
             pip install "python-doctr[torch]"
-            # or with preinstalled packages for visualization & html support
-            pip install "python-doctr[torch,viz,html]"
+            # or with preinstalled packages for visualization & html & contrib module support
+            pip install "python-doctr[torch,viz,html,contrib]"
 
 
 

docs/source/using_doctr/using_contrib_modules.rst

@@ -14,7 +14,9 @@ Available contribution modules
 
 .. code:: bash
 
-    pip install onnxruntime
+    pip install python-doctr[contrib]
+    # Or
+    pip install onnxruntime  # pip install onnxruntime-gpu
 
 Here are all contribution modules that are available through docTR:
 

doctr/contrib/artefacts.py

@@ -8,6 +8,8 @@
 import cv2
 import numpy as np
 
+from doctr.file_utils import requires_package
+
 from .base import _BasePredictor
 
 __all__ = ["ArtefactDetector"]
@@ -110,7 +112,7 @@ def show(self, **kwargs: Any) -> None:
         ----
             **kwargs: additional keyword arguments to be passed to `plt.show`
         """
-        # TODO: requires check for matplotlib
+        requires_package("matplotlib", "`.show()` requires matplotlib installed")
         import matplotlib.pyplot as plt
         from matplotlib.patches import Rectangle
 

doctr/contrib/base.py

@@ -7,6 +7,7 @@
 
 import numpy as np
 
+from doctr.file_utils import requires_package
 from doctr.utils.data import download_from_url
 
 
@@ -41,10 +42,9 @@ def _init_model(self, url: Optional[str] = None, model_path: Optional[str] = Non
 
         Returns:
         -------
-            str: the path to the model
+            Any: the ONNX loaded model
         """
-        # TODO: requires check
-
+        requires_package("onnxruntime", "`.contrib` module requires `onnxruntime` to be installed.")
         import onnxruntime as ort
 
         if not url and not model_path:

pyproject.toml

@@ -69,6 +69,9 @@ viz = [
     "matplotlib>=3.1.0",
     "mplcursors>=0.3",
 ]
+contrib = [
+    "onnxruntime>=1.11.0",
+]
 testing = [
     "pytest>=5.3.2",
     "coverage[toml]>=4.5.4",