|
7 | 7 |
|
8 | 8 | import argparse |
9 | 9 | import shutil |
| 10 | +import tempfile |
10 | 11 | import urllib |
| 12 | +import urllib.error |
11 | 13 | import warnings |
12 | 14 | from pathlib import Path |
13 | 15 |
|
14 | 16 | from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model |
15 | 17 |
|
16 | | -import dlclive |
17 | 18 | from dlclive.benchmark import benchmark_videos |
18 | 19 | from dlclive.engine import Engine |
19 | 20 | from dlclive.utils import download_file, get_available_backends |
20 | 21 |
|
21 | 22 | MODEL_NAME = "superanimal_quadruped" |
22 | 23 | SNAPSHOT_NAME = "snapshot-700000.pb" |
| 24 | +TMP_DIR = Path(tempfile.gettempdir()) / "dlc-live-tmp" |
| 25 | + |
| 26 | +MODELS_DIR = TMP_DIR / "test_models" |
| 27 | +TORCH_MODEL = "resnet_50" |
| 28 | +TORCH_CONFIG = { |
| 29 | + "checkpoint": MODELS_DIR / f"exported_quadruped_{TORCH_MODEL}.pt", |
| 30 | + "super_animal": "superanimal_quadruped", |
| 31 | +} |
| 32 | +TF_MODEL_DIR = MODELS_DIR / "DLC_Dog_resnet_50_iteration-0_shuffle-0" |
| 33 | + |
| 34 | + |
| 35 | +def run_pytorch_test(video_file: str, display: bool = False): |
| 36 | + from dlclive.modelzoo.pytorch_model_zoo_export import export_modelzoo_model |
| 37 | + |
| 38 | + if Engine.PYTORCH not in get_available_backends(): |
| 39 | + raise ImportError( |
| 40 | + "PyTorch backend is not available. Please ensure PyTorch is installed to run the PyTorch test." |
| 41 | + ) |
| 42 | + # Download model from the DeepLabCut Model Zoo |
| 43 | + export_modelzoo_model( |
| 44 | + export_path=TORCH_CONFIG["checkpoint"], |
| 45 | + super_animal=TORCH_CONFIG["super_animal"], |
| 46 | + model_name=TORCH_MODEL, |
| 47 | + ) |
| 48 | + if not TORCH_CONFIG["checkpoint"].exists(): |
| 49 | + raise FileNotFoundError(f"Failed to export {TORCH_CONFIG['super_animal']} model") |
| 50 | + if TORCH_CONFIG["checkpoint"].stat().st_size == 0: |
| 51 | + raise ValueError(f"Exported {TORCH_CONFIG['super_animal']} model is empty") |
| 52 | + benchmark_videos( |
| 53 | + model_path=str(TORCH_CONFIG["checkpoint"]), |
| 54 | + model_type="pytorch", |
| 55 | + video_path=video_file, |
| 56 | + display=display, |
| 57 | + pcutoff=0.25, |
| 58 | + pixels=1000, |
| 59 | + ) |
| 60 | + |
| 61 | + |
| 62 | +def run_tensorflow_test(video_file: str, display: bool = False): |
| 63 | + if Engine.TENSORFLOW not in get_available_backends(): |
| 64 | + raise ImportError( |
| 65 | + "TensorFlow backend is not available. Please ensure TensorFlow is installed to run the TensorFlow test." |
| 66 | + ) |
| 67 | + model_dir = TF_MODEL_DIR |
| 68 | + model_dir.mkdir(parents=True, exist_ok=True) |
| 69 | + if Path(model_dir / SNAPSHOT_NAME).exists(): |
| 70 | + print("Model already downloaded, using cached version") |
| 71 | + else: |
| 72 | + print("Downloading superanimal_quadruped model from the DeepLabCut Model Zoo...") |
| 73 | + download_huggingface_model(MODEL_NAME, str(model_dir)) |
| 74 | + |
| 75 | + if not Path(model_dir / SNAPSHOT_NAME).exists(): |
| 76 | + raise FileNotFoundError(f"Missing model file {model_dir / SNAPSHOT_NAME}") |
| 77 | + |
| 78 | + benchmark_videos( |
| 79 | + model_path=str(model_dir), |
| 80 | + model_type="base", |
| 81 | + video_path=video_file, |
| 82 | + display=display, |
| 83 | + pcutoff=0.25, |
| 84 | + pixels=1000, |
| 85 | + ) |
| 86 | + |
| 87 | + |
| 88 | +BACKEND_TESTS = { |
| 89 | + Engine.PYTORCH: run_pytorch_test, |
| 90 | + Engine.TENSORFLOW: run_tensorflow_test, |
| 91 | +} |
| 92 | +BACKEND_DISPLAY_NAMES = { |
| 93 | + Engine.PYTORCH: "PyTorch", |
| 94 | + Engine.TENSORFLOW: "TensorFlow", |
| 95 | +} |
23 | 96 |
|
24 | 97 |
|
25 | 98 | def main(): |
| 99 | + backend_results: dict[Engine, tuple[str, str | None]] = {} |
| 100 | + backend_failures: dict[Engine, Exception] = {} |
| 101 | + |
26 | 102 | parser = argparse.ArgumentParser( |
27 | 103 | description="Test DLC-Live installation by downloading and evaluating a demo DLC project!" |
28 | 104 | ) |
| 105 | + parser.add_argument( |
| 106 | + "--display", |
| 107 | + action="store_true", |
| 108 | + default=False, |
| 109 | + help="Run the test and display tracking", |
| 110 | + ) |
29 | 111 | parser.add_argument( |
30 | 112 | "--nodisplay", |
31 | 113 | action="store_false", |
32 | | - help="Run the test without displaying tracking", |
| 114 | + dest="display", |
| 115 | + help=argparse.SUPPRESS, |
| 116 | + default=False, |
33 | 117 | ) |
| 118 | + |
34 | 119 | args = parser.parse_args() |
35 | | - display = args.nodisplay |
| 120 | + display = args.display |
36 | 121 |
|
37 | 122 | if not display: |
38 | 123 | print("Running without displaying video") |
39 | 124 |
|
40 | | - # make temporary directory |
41 | 125 | print("\nCreating temporary directory...\n") |
42 | | - tmp_dir = Path(dlclive.__file__).parent / "check_install" / "dlc-live-tmp" |
| 126 | + tmp_dir = TMP_DIR |
43 | 127 | tmp_dir.mkdir(mode=0o775, exist_ok=True) |
| 128 | + MODELS_DIR.mkdir(parents=True, exist_ok=True) |
44 | 129 |
|
45 | 130 | video_file = str(tmp_dir / "dog_clip.avi") |
46 | | - model_dir = tmp_dir / "DLC_Dog_resnet_50_iteration-0_shuffle-0" |
47 | | - |
48 | | - # download dog test video from github: |
49 | | - # Use raw.githubusercontent.com for direct file access |
50 | | - if not Path(video_file).exists(): |
51 | | - print(f"Downloading Video to {video_file}") |
52 | | - url_link = "https://raw.githubusercontent.com/DeepLabCut/DeepLabCut-live/master/check_install/dog_clip.avi" |
53 | | - try: |
54 | | - download_file(url_link, video_file) |
55 | | - except (OSError, urllib.error.URLError) as e: |
56 | | - raise RuntimeError(f"Failed to download video file: {e}") from e |
57 | | - else: |
58 | | - print(f"Video file already exists at {video_file}, skipping download.") |
59 | | - |
60 | | - # download model from the DeepLabCut Model Zoo |
61 | | - if Path(model_dir / SNAPSHOT_NAME).exists(): |
62 | | - print("Model already downloaded, using cached version") |
63 | | - else: |
64 | | - print("Downloading superanimal_quadruped model from the DeepLabCut Model Zoo...") |
65 | | - download_huggingface_model(MODEL_NAME, model_dir) |
66 | | - |
67 | | - # assert these things exist so we can give informative error messages |
68 | | - assert Path(video_file).exists(), f"Missing video file {video_file}" |
69 | | - assert Path(model_dir / SNAPSHOT_NAME).exists(), f"Missing model file {model_dir / SNAPSHOT_NAME}" |
70 | 131 |
|
71 | | - # run benchmark videos |
72 | | - print("\n Running inference...\n") |
73 | | - benchmark_videos( |
74 | | - model_path=str(model_dir), |
75 | | - model_type="base" if Engine.from_model_path(model_dir) == Engine.TENSORFLOW else "pytorch", |
76 | | - video_path=video_file, |
77 | | - display=display, |
78 | | - resize=0.5, |
79 | | - pcutoff=0.25, |
80 | | - ) |
81 | | - |
82 | | - # deleting temporary files |
83 | | - print("\n Deleting temporary files...\n") |
84 | 132 | try: |
85 | | - shutil.rmtree(tmp_dir) |
86 | | - except PermissionError: |
87 | | - warnings.warn( |
88 | | - f"Could not delete temporary directory {str(tmp_dir)} due to a permissions error.", |
89 | | - stacklevel=2, |
90 | | - ) |
91 | | - |
92 | | - print("\nDone!\n") |
| 133 | + # download dog test video from github: |
| 134 | + # Use raw.githubusercontent.com for direct file access |
| 135 | + if not Path(video_file).exists(): |
| 136 | + print(f"Downloading Video to {video_file}") |
| 137 | + url_link = "https://raw.githubusercontent.com/DeepLabCut/DeepLabCut-live/master/check_install/dog_clip.avi" |
| 138 | + try: |
| 139 | + download_file(url_link, video_file) |
| 140 | + except (OSError, urllib.error.URLError) as e: |
| 141 | + raise RuntimeError(f"Failed to download video file: {e}") from e |
| 142 | + else: |
| 143 | + print(f"Video file already exists at {video_file}, skipping download.") |
| 144 | + |
| 145 | + if not Path(video_file).exists(): |
| 146 | + raise FileNotFoundError(f"Missing video file {video_file}") |
| 147 | + |
| 148 | + any_backend_succeeded = False |
| 149 | + |
| 150 | + available_backends = get_available_backends() |
| 151 | + if not available_backends: |
| 152 | + raise RuntimeError( |
| 153 | + "No available backends to test. " |
| 154 | + "Please ensure that at least one of the supported backends " |
| 155 | + "(TensorFlow or PyTorch) is installed." |
| 156 | + ) |
| 157 | + |
| 158 | + for backend in available_backends: |
| 159 | + test_func = BACKEND_TESTS.get(backend) |
| 160 | + if test_func is None: |
| 161 | + warnings.warn( |
| 162 | + f"No test function defined for backend {backend}, skipping...", |
| 163 | + UserWarning, |
| 164 | + stacklevel=2, |
| 165 | + ) |
| 166 | + continue |
| 167 | + |
| 168 | + try: |
| 169 | + print(f"\nRunning {BACKEND_DISPLAY_NAMES.get(backend, backend.value)} test...\n") |
| 170 | + test_func(video_file, display=display) |
| 171 | + any_backend_succeeded = True |
| 172 | + backend_results[backend] = ("SUCCESS", None) |
| 173 | + |
| 174 | + except Exception as e: |
| 175 | + backend_results[backend] = ("ERROR", str(e)) |
| 176 | + backend_failures[backend] = e |
| 177 | + warnings.warn( |
| 178 | + f"Error while running test for backend {backend}: {e}. " |
| 179 | + "Continuing to test other available backends.", |
| 180 | + UserWarning, |
| 181 | + stacklevel=2, |
| 182 | + ) |
| 183 | + |
| 184 | + print("\n---\nBackend test summary:") |
| 185 | + for backend in BACKEND_TESTS.keys(): |
| 186 | + status, _ = backend_results.get(backend, ("SKIPPED", None)) |
| 187 | + print(f"{backend.value:<11} [{status}]") |
| 188 | + print("---") |
| 189 | + |
| 190 | + for backend, (status, error) in backend_results.items(): |
| 191 | + if status == "ERROR": |
| 192 | + backend_name = BACKEND_DISPLAY_NAMES.get(backend, backend.value) |
| 193 | + print(f"{backend_name} error:\n{error}\n") |
| 194 | + |
| 195 | + if not any_backend_succeeded and backend_failures: |
| 196 | + failure_messages = "; ".join(f"{b}: {e}" for b, e in backend_failures.items()) |
| 197 | + raise RuntimeError(f"All backend tests failed. Details: {failure_messages}") |
| 198 | + |
| 199 | + finally: |
| 200 | + # deleting temporary files |
| 201 | + print("\n Deleting temporary files...\n") |
| 202 | + try: |
| 203 | + if tmp_dir.exists(): |
| 204 | + shutil.rmtree(tmp_dir) |
| 205 | + except PermissionError: |
| 206 | + warnings.warn( |
| 207 | + f"Could not delete temporary directory {str(tmp_dir)} due to a permissions error.", |
| 208 | + stacklevel=2, |
| 209 | + ) |
93 | 210 |
|
94 | 211 |
|
95 | 212 | if __name__ == "__main__": |
96 | 213 | # Get available backends (emits a warning if neither TensorFlow nor PyTorch is installed) |
97 | 214 | available_backends: list[Engine] = get_available_backends() |
98 | 215 | print(f"Available backends: {[b.value for b in available_backends]}") |
99 | | - |
100 | | - # TODO: JR add support for PyTorch in check_install.py (requires some exported pytorch model to be downloaded) |
101 | | - if Engine.TENSORFLOW not in available_backends: |
| 216 | + if len(available_backends) == 0: |
102 | 217 | raise NotImplementedError( |
103 | | - "TensorFlow is not installed. Currently check_install.py only supports testing the TensorFlow installation." |
| 218 | + "Neither TensorFlow nor PyTorch is installed. " |
| 219 | + "Please install at least one of these frameworks to run the installation test." |
104 | 220 | ) |
105 | 221 |
|
106 | 222 | main() |
0 commit comments