Add centralized get_available_backends function.

The function is called in __init__ and emits a warning if neither tf nor torch are installed.

Added todo comments in tests that should test both tensorflow and pytorch installation.

Author: deruyter92

dlclive/__init__.py

@@ -5,6 +5,11 @@
 Licensed under GNU Lesser General Public License v3.0
 """
 
+# Check which backends are installed and get available backends
+# (Emits a warning if neither TensorFlow nor PyTorch is installed)
+from dlclive.utils import get_available_backends
+_AVAILABLE_BACKENDS = get_available_backends()
+
 from dlclive.display import Display
 from dlclive.dlclive import DLCLive
 from dlclive.processor.processor import Processor

dlclive/check_install/check_install.py

@@ -16,6 +16,7 @@
 from dlclive.utils import download_file
 from dlclive.benchmark import benchmark_videos
 from dlclive.engine import Engine
+from dlclive.utils import get_available_backends
 
 MODEL_NAME = "superanimal_quadruped"
 SNAPSHOT_NAME = "snapshot-700000.pb"
@@ -93,4 +94,15 @@ def main():
 
 
 if __name__ == "__main__":
+
+    # Get available backends (emits a warning if neither TensorFlow nor PyTorch is installed)
+    available_backends: list[Engine] = get_available_backends()
+    print(f"Available backends: {[b.value for b in available_backends]}")
+
+    # TODO: JR add support for PyTorch in check_install.py (requires some exported pytorch model to be downloaded)
+    if not Engine.TENSORFLOW in available_backends:
+        raise NotImplementedError(
+            "TensorFlow is not installed. Currently check_install.py only supports testing the TensorFlow installation."
+        )
+
     main()

dlclive/utils.py

@@ -12,6 +12,7 @@
 import urllib.error
 
 from dlclive.exceptions import DLCLiveWarning
+from dlclive.engine import Engine
 
 try:
     import skimage
@@ -214,6 +215,41 @@ def decode_fourcc(cc):
     return decoded
 
 
+def get_available_backends() -> list[Engine]:
+    """
+    Check which backends (TensorFlow or PyTorch) are installed.
+    
+    Returns:
+        list[str]: List of installed backends. Possible values: ["tensorflow"], ["pytorch"], 
+                   or ["tensorflow", "pytorch"]. Returns an empty list if neither is installed.
+    
+    Warns:
+        DLCLiveWarning: If neither TensorFlow nor PyTorch is installed.
+    """
+    backends = []
+    
+    try:
+        import tensorflow
+        backends.append(Engine.TENSORFLOW)
+    except (ImportError, ModuleNotFoundError):
+        pass
+    
+    try:
+        import torch
+        backends.append(Engine.PYTORCH)
+    except (ImportError, ModuleNotFoundError):
+        pass
+    
+    if not backends:
+        warnings.warn(
+            "Neither TensorFlow nor PyTorch is installed. One of these is required to use DLCLive!"
+            "Install with: pip install deeplabcut-live[tf] or pip install deeplabcut-live[pytorch]",
+            DLCLiveWarning,
+        )
+    
+    return backends
+
+
 def download_file(url: str, filepath: str, chunk_size: int = 8192) -> None:
     """
     Download a file from a URL with progress bar and error handling.

tests/test_benchmark_script.py

@@ -4,6 +4,7 @@
 from dlclive.engine import Engine
 
 
+# TODO: JR include separate functional tests for torch and tf backends
 @pytest.mark.functional
 def test_benchmark_script_runs(tmp_path):
     datafolder = tmp_path / "Data-DLC-live-benchmark"