diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index 27ea85a..25c9a92 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -68,7 +68,10 @@ jobs: run: poetry lock --no-cache - name: Install project dependencies - run: poetry install --no-root + run: poetry install --with dev - name: Run DLC Live Tests run: poetry run dlc-live-test --nodisplay + + - name: Run Functional Benchmark Test + run: poetry run pytest tests/test_benchmark_script.py diff --git a/.gitignore b/.gitignore index 7fa3a49..2ce1454 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,8 @@ benchmarking/results* **DS_Store* *vscode* +**/__MACOSX/ + # Byte-compiled / optimized / DLL files __pycache__/ *.py[cod] diff --git a/dlclive/benchmark.py b/dlclive/benchmark.py index 1d53042..9643b09 100644 --- a/dlclive/benchmark.py +++ b/dlclive/benchmark.py @@ -37,42 +37,38 @@ def download_benchmarking_data( target_dir=".", - url="http://deeplabcut.rowland.harvard.edu/datasets/dlclivebenchmark.tar.gz", + url="https://huggingface.co/datasets/mwmathis/DLCspeed_benchmarking/resolve/main/Data-DLC-live-benchmark.zip", ): """ - Downloads a DeepLabCut-Live benchmarking Data (videos & DLC models). + Downloads and extracts DeepLabCut-Live benchmarking data (videos & DLC models). """ + import os import urllib.request - import tarfile from tqdm import tqdm + import zipfile - def show_progress(count, block_size, total_size): - pbar.update(block_size) - - def tarfilenamecutting(tarf): - """' auxfun to extract folder path - ie. /xyz-trainsetxyshufflez/ - """ - for memberid, member in enumerate(tarf.getmembers()): - if memberid == 0: - parent = str(member.path) - l = len(parent) + 1 - if member.path.startswith(parent): - member.path = member.path[l:] - yield member - - response = urllib.request.urlopen(url) - print( - "Downloading the benchmarking data from the DeepLabCut server @Harvard -> Go Crimson!!! {}....".format( - url - ) - ) - total_size = int(response.getheader("Content-Length")) - pbar = tqdm(unit="B", total=total_size, position=0) - filename, _ = urllib.request.urlretrieve(url, reporthook=show_progress) - with tarfile.open(filename, mode="r:gz") as tar: - tar.extractall(target_dir, members=tarfilenamecutting(tar)) + # Avoid nested folder issue + if os.path.basename(os.path.normpath(target_dir)) == "Data-DLC-live-benchmark": + target_dir = os.path.dirname(os.path.normpath(target_dir)) + os.makedirs(target_dir, exist_ok=True) # Ensure target directory exists + + zip_path = os.path.join(target_dir, "Data-DLC-live-benchmark.zip") + + if os.path.exists(zip_path): + print(f"{zip_path} already exists. Skipping download.") + else: + def show_progress(count, block_size, total_size): + pbar.update(block_size) + + print(f"Downloading the benchmarking data from {url} ...") + pbar = tqdm(unit="B", total=0, position=0, desc="Downloading") + + filename, _ = urllib.request.urlretrieve(url, filename=zip_path, reporthook=show_progress) + pbar.close() + print(f"Extracting {zip_path} to {target_dir} ...") + with zipfile.ZipFile(zip_path, 'r') as zip_ref: + zip_ref.extractall(target_dir) def get_system_info() -> dict: """ Return summary info for system running benchmark diff --git a/pyproject.toml b/pyproject.toml index 9fb4eed..7451d7e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -36,6 +36,7 @@ torch = ">=1.10,<3.0" dlclibrary = ">=0.0.6" pandas = "^1.3" tables = "^3.6" +pytest = "^8.0" # OS-specific TensorFlow packages tensorflow = [ @@ -44,7 +45,6 @@ tensorflow = [ ] tensorflow-macos = { version = ">=2.7.0,<2.12", markers = "sys_platform == 'darwin'" } - [tool.poetry.group.dev.dependencies] [build-system] diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..d08b10e --- /dev/null +++ b/pytest.ini @@ -0,0 +1,3 @@ +[pytest] +markers = + functional: functional tests diff --git a/tests/test_benchmark_script.py b/tests/test_benchmark_script.py new file mode 100644 index 0000000..ff3a3c2 --- /dev/null +++ b/tests/test_benchmark_script.py @@ -0,0 +1,33 @@ +import os +import glob +import pathlib +import pytest +from dlclive import benchmark_videos, download_benchmarking_data + +@pytest.mark.functional +def test_benchmark_script_runs(tmp_path): + datafolder = tmp_path / "Data-DLC-live-benchmark" + download_benchmarking_data(str(datafolder)) + + dog_models = glob.glob(str(datafolder / "dog" / "*[!avi]")) + dog_video = glob.glob(str(datafolder / "dog" / "*.avi"))[0] + mouse_models = glob.glob(str(datafolder / "mouse_lick" / "*[!avi]")) + mouse_video = glob.glob(str(datafolder / "mouse_lick" / "*.avi"))[0] + + out_dir = tmp_path / "results" + out_dir.mkdir(exist_ok=True) + + pixels = [100, 400] #[2500, 10000] + n_frames = 5 + + for m in dog_models: + print(f"Running dog model: {m}") + result = benchmark_videos(m, dog_video, output=str(out_dir), n_frames=n_frames, pixels=pixels) + print("Dog model result:", result) + + for m in mouse_models: + print(f"Running mouse model: {m}") + result = benchmark_videos(m, mouse_video, output=str(out_dir), n_frames=n_frames, pixels=pixels) + print("Mouse model result:", result) + + assert any(out_dir.iterdir())