Update benchmark.py (#123)

Author: MMathisLab

.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

.gitignore

@@ -8,6 +8,8 @@ benchmarking/results*
 **DS_Store*
 *vscode*
 
+**/__MACOSX/
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

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

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]

pytest.ini

@@ -0,0 +1,3 @@
+[pytest]
+markers =
+    functional: functional tests

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())