Add PyTorch support and refactor check_install

Refactor check_install.py to support both PyTorch and TensorFlow backends and improve temporary file handling. Introduces TMP_DIR, MODELS_FOLDER, and separate run_pytorch_test/run_tensorflow_test helpers; PyTorch now exports and benchmarks an exported .pt checkpoint, TensorFlow model download logic is preserved. Replaces --nodisplay with --display, centralizes video download and assertions, tightens error handling for downloads, and ensures proper cleanup of temporary files. Also updates imports (urllib.error, export_modelzoo_model) and updates backend availability checks to require at least one backend.

Author: C-Achard

dlclive/check_install/check_install.py

@@ -8,101 +8,163 @@
 import argparse
 import shutil
 import warnings
+import urllib.error
 from pathlib import Path
 
 from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model
 
-import dlclive
 from dlclive.utils import download_file
 from dlclive.benchmark import benchmark_videos
 from dlclive.engine import Engine
+from dlclive.modelzoo.pytorch_model_zoo_export import export_modelzoo_model
 from dlclive.utils import get_available_backends
 
 MODEL_NAME = "superanimal_quadruped"
 SNAPSHOT_NAME = "snapshot-700000.pb"
+TMP_DIR = Path(__file__).parent / "dlc-live-tmp"
 
+MODELS_FOLDER = TMP_DIR / "test_models"
+TORCH_MODEL = "resnet_50"
+TORCH_CONFIG = {
+    "checkpoint": MODELS_FOLDER / f"exported_quadruped_{TORCH_MODEL}.pt",
+    "super_animal": "superanimal_quadruped",
+}
+TF_MODEL_DIR = TMP_DIR / "DLC_Dog_resnet_50_iteration-0_shuffle-0"
 
-def main():
-    parser = argparse.ArgumentParser(
-        description="Test DLC-Live installation by downloading and evaluating a demo DLC project!"
-    )
-    parser.add_argument(
-        "--nodisplay",
-        action="store_false",
-        help="Run the test without displaying tracking",
-    )
-    args = parser.parse_args()
-    display = args.nodisplay
-
-    if not display:
-        print("Running without displaying video")
+MODELS_FOLDER.mkdir(parents=True, exist_ok=True)
 
-    # make temporary directory
-    print("\nCreating temporary directory...\n")
-    tmp_dir = Path(dlclive.__file__).parent / "check_install" / "dlc-live-tmp"
-    tmp_dir.mkdir(mode=0o775, exist_ok=True)
 
-    video_file = str(tmp_dir / "dog_clip.avi")
-    model_dir = tmp_dir / "DLC_Dog_resnet_50_iteration-0_shuffle-0"
+def run_pytorch_test(video_file: str, display: bool = False):
+    if Engine.PYTORCH not in get_available_backends():
+        raise NotImplementedError(
+            "PyTorch backend is not available. Please ensure PyTorch is installed to run the PyTorch test."
+        )
+    # Download model from the DeepLabCut Model Zoo
+    export_modelzoo_model(
+        export_path=TORCH_CONFIG["checkpoint"],
+        super_animal=TORCH_CONFIG["super_animal"],
+        model_name=TORCH_MODEL,
+    )
+    assert TORCH_CONFIG["checkpoint"].exists(), (
+        f"Failed to export {TORCH_CONFIG['super_animal']} model"
+    )
+    assert TORCH_CONFIG["checkpoint"].stat().st_size > 0, (
+        f"Exported {TORCH_CONFIG['super_animal']} model is empty"
+    )
+    benchmark_videos(
+        model_path=str(TORCH_CONFIG["checkpoint"]),
+        model_type="pytorch",
+        video_path=video_file,
+        display=display,
+        resize=0.5,
+        pcutoff=0.25,
+        pixels=1000,
+    )
 
-    # download dog test video from github:
-    # Use raw.githubusercontent.com for direct file access
-    if not Path(video_file).exists():
-        print(f"Downloading Video to {video_file}")
-        url_link = "https://raw.githubusercontent.com/DeepLabCut/DeepLabCut-live/master/check_install/dog_clip.avi"
-        try:
-            download_file(url_link, video_file)
-        except (urllib.error.URLError, IOError) as e:
-            raise RuntimeError(f"Failed to download video file: {e}") from e
-    else:
-        print(f"Video file already exists at {video_file}, skipping download.")
 
-    # download model from the DeepLabCut Model Zoo
+def run_tensorflow_test(video_file: str, display: bool = False):
+    if Engine.TENSORFLOW not in get_available_backends():
+        raise NotImplementedError(
+            "TensorFlow backend is not available. Please ensure TensorFlow is installed to run the TensorFlow test."
+        )
+    model_dir = TF_MODEL_DIR
+    model_dir.mkdir(parents=True, exist_ok=True)
+    assert model_dir.exists(), f"Model directory {model_dir} does not exist"
     if Path(model_dir / SNAPSHOT_NAME).exists():
         print("Model already downloaded, using cached version")
     else:
-        print("Downloading superanimal_quadruped model from the DeepLabCut Model Zoo...")
-        download_huggingface_model(MODEL_NAME, model_dir)
+        print(
+            "Downloading superanimal_quadruped model from the DeepLabCut Model Zoo..."
+        )
+        download_huggingface_model(MODEL_NAME, str(model_dir))
 
-    # assert these things exist so we can give informative error messages
-    assert Path(video_file).exists(), f"Missing video file {video_file}"
-    assert Path(
-        model_dir / SNAPSHOT_NAME
-    ).exists(), f"Missing model file {model_dir / SNAPSHOT_NAME}"
+    assert Path(model_dir / SNAPSHOT_NAME).exists(), (
+        f"Missing model file {model_dir / SNAPSHOT_NAME}"
+    )
 
-    # run benchmark videos
-    print("\n Running inference...\n")
     benchmark_videos(
         model_path=str(model_dir),
-        model_type="base" if Engine.from_model_path(model_dir) == Engine.TENSORFLOW else "pytorch",
+        model_type="base",
         video_path=video_file,
         display=display,
         resize=0.5,
-        pcutoff=0.25
+        pcutoff=0.25,
+        pixels=1000,
     )
 
-    # deleting temporary files
-    print("\n Deleting temporary files...\n")
+
+def main():
+    tmp_dir = None
     try:
-        shutil.rmtree(tmp_dir)
-    except PermissionError:
-        warnings.warn(
-            f"Could not delete temporary directory {str(tmp_dir)} due to a permissions error, but otherwise dlc-live seems to be working fine!"
+        parser = argparse.ArgumentParser(
+            description="Test DLC-Live installation by downloading and evaluating a demo DLC project!"
+        )
+        parser.add_argument(
+            "--display",
+            action="store_true",
+            help="Run the test and display tracking",
         )
+        args = parser.parse_args()
+        display = args.display
+
+        if not display:
+            print("Running without displaying video")
+
+        # make temporary directory
+        print("\nCreating temporary directory...\n")
+        tmp_dir = TMP_DIR
+        tmp_dir.mkdir(mode=0o775, exist_ok=True)
+
+        video_file = str(tmp_dir / "dog_clip.avi")
+
+        # download dog test video from github:
+        # Use raw.githubusercontent.com for direct file access
+        if not Path(video_file).exists():
+            print(f"Downloading Video to {video_file}")
+            url_link = "https://raw.githubusercontent.com/DeepLabCut/DeepLabCut-live/master/check_install/dog_clip.avi"
+            try:
+                download_file(url_link, video_file)
+            except (urllib.error.URLError, IOError) as e:
+                raise RuntimeError(f"Failed to download video file: {e}") from e
+        else:
+            print(f"Video file already exists at {video_file}, skipping download.")
+
+        # assert these things exist so we can give informative error messages
+        assert Path(video_file).exists(), f"Missing video file {video_file}"
+
+        for backend in get_available_backends():
+            if backend == Engine.PYTORCH:
+                print("\nRunning PyTorch test...\n")
+                run_pytorch_test(video_file, display=display)
+            elif backend == Engine.TENSORFLOW:
+                print("\nRunning TensorFlow test...\n")
+                run_tensorflow_test(video_file, display=display)
+            else:
+                warnings.warn(
+                    f"Unrecognized backend {backend}, skipping...", UserWarning
+                )
+
+    finally:
+        # deleting temporary files
+        print("\n Deleting temporary files...\n")
+        try:
+            if tmp_dir is not None and tmp_dir.exists():
+                shutil.rmtree(tmp_dir)
+        except PermissionError:
+            warnings.warn(
+                f"Could not delete temporary directory {str(tmp_dir)} due to a permissions error, but otherwise dlc-live seems to be working fine!"
+            )
 
-    print("\nDone!\n")
+        print("\nDone!\n")
 
 
 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:
+    if len(available_backends) == 0:
         raise NotImplementedError(
-            "TensorFlow is not installed. Currently check_install.py only supports testing the TensorFlow installation."
+            "Neither TensorFlow nor PyTorch is installed. Please install at least one of these frameworks to run the installation test."
         )
 
     main()