better separation of models

Author: n-poulsen

dlclibrary/dlcmodelzoo/modelzoo_download.py

@@ -16,6 +16,7 @@
 from pathlib import Path
 
 from huggingface_hub import hf_hub_download
+from ruamel.yaml import YAML
 from ruamel.yaml.comments import CommentedBase
 
 # just expand this list when adding new models:
@@ -40,20 +41,64 @@ def _get_dlclibrary_path():
     return os.path.split(importlib.util.find_spec("dlclibrary").origin)[0]
 
 
-def _load_model_names():
+def _load_pytorch_models() -> dict[str, dict[str, dict[str, str]]]:
+    """Load URLs and commit hashes for available models."""
+    urls = Path(_get_dlclibrary_path()) / "dlcmodelzoo" / "modelzoo_urls_pytorch.yaml"
+    with open(urls) as file:
+        data = YAML(pure=True).load(file)
+
+    return data
+
+
+def _load_pytorch_dataset_models(dataset: str) -> dict[str, dict[str, str]]:
     """Load URLs and commit hashes for available models."""
-    from ruamel.yaml import YAML
+    models = _load_pytorch_models()
+    if not dataset in models:
+        raise ValueError(
+            f"Could not find any models for {dataset}. Models are available for "
+            f"{list(models.keys())}"
+        )
 
+    return models[dataset]
+
+
+def _load_model_names():
+    """Load URLs and commit hashes for available models."""
     fn = os.path.join(_get_dlclibrary_path(), "dlcmodelzoo", "modelzoo_urls.yaml")
     with open(fn) as file:
-        return YAML().load(file)
+        model_names = YAML().load(file)
+
+    # add PyTorch models
+    for dataset, model_types in _load_pytorch_models().items():
+        for model_type, models in model_types.items():
+            for model, url in models.items():
+                model_names[f"{dataset}_{model}"] = url
+
+    return model_names
 
 
 def parse_available_supermodels():
     libpath = _get_dlclibrary_path()
     json_path = os.path.join(libpath, "dlcmodelzoo", "superanimal_models.json")
     with open(json_path) as file:
-        return json.load(file)
+        super_animal_models = json.load(file)
+    return super_animal_models
+
+
+def get_available_detectors(dataset: str) -> list[str]:
+    """
+    Returns:
+        The detectors available for the dataset.
+    """
+    return list(_load_pytorch_dataset_models(dataset)["detectors"].keys())
+
+
+def get_available_models(dataset: str) -> list[str]:
+    """
+    Returns:
+        The pose models available for the dataset.
+    """
+    return list(_load_pytorch_dataset_models(dataset)["pose_models"].keys())
 
 
 def _handle_downloaded_file(

dlclibrary/dlcmodelzoo/modelzoo_urls.yaml

@@ -36,14 +36,6 @@ superanimal_quadruped_hrnetw32:
   - mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/pose_model.pth
   - mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/detector.pt
 
-# Updated format - split detectors and pose models
-superanimal_bird_resnet_50: DeepLabCut/DeepLabCutModelZoo-SuperAnimal-Bird/superanimal_bird_resnet_50.pt
-superanimal_bird_ssdlite: DeepLabCut/DeepLabCutModelZoo-SuperAnimal-Bird/superanimal_bird_ssdlite.pt
-superanimal_topviewmouse_hrnet_w32: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/pose_model.pth
-superanimal_topviewmouse_fasterrcnn_resnet50_fpn_v2: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/detector.pt
-superanimal_quadruped_hrnet_w32: mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/pose_model.pth
-superanimal_quadruped_fasterrcnn_resnet50_fpn_v2: mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/detector.pt
-
 # DeepLabCut 2.X backwards compatibility
 superanimal_topviewmouse: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/DLC_ma_supertopview5k_resnet_50_iteration-0_shuffle-1.tar.gz
 

dlclibrary/dlcmodelzoo/modelzoo_urls_pytorch.yaml

@@ -0,0 +1,19 @@
+# DeepLabCut 3.0: SuperAnimal detectors and pose model URLS
+
+superanimal_bird:
+  detectors:
+    ssdlite: DeepLabCut/DeepLabCutModelZoo-SuperAnimal-Bird/superanimal_bird_ssdlite.pt
+  pose_models:
+    resnet_50: DeepLabCut/DeepLabCutModelZoo-SuperAnimal-Bird/superanimal_bird_resnet_50.pt
+
+superanimal_topviewmouse:
+  detectors:
+    fasterrcnn_resnet50_fpn_v2: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/detector.pt
+  pose_models:
+    hrnet_w32: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/pose_model.pth
+
+superanimal_quadruped:
+  detectors:
+    fasterrcnn_resnet50_fpn_v2: mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/detector.pt
+  pose_models:
+    hrnet_w32: mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/pose_model.pth

setup.py

@@ -32,6 +32,7 @@
             "dlclibrary",
             [
                 "dlclibrary/dlcmodelzoo/modelzoo_urls.yaml",
+                "dlclibrary/dlcmodelzoo/modelzoo_urls_pytorch.yaml",
                 "dlclibrary/dlcmodelzoo/superanimal_models.json",
                 "dlclibrary/dlcmodelzoo/superanimal_configs/superquadruped.yaml",
                 "dlclibrary/dlcmodelzoo/superanimal_configs/supertopview.yaml",

tests/test_pytorch_models.py

@@ -0,0 +1,47 @@
+#
+# DeepLabCut Toolbox (deeplabcut.org)
+# © A. & M.W. Mathis Labs
+# https://github.com/DeepLabCut/DeepLabCut
+#
+# Please see AUTHORS for contributors.
+# https://github.com/DeepLabCut/DeepLabCut/blob/master/AUTHORS
+#
+# Licensed under GNU Lesser General Public License v3.0
+#
+import os
+import pytest
+
+import dlclibrary
+import dlclibrary.dlcmodelzoo.modelzoo_download as modelzoo
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        ("superanimal_bird", ["ssdlite"]),
+        ("superanimal_topviewmouse", ["fasterrcnn_resnet50_fpn_v2"]),
+        ("superanimal_quadruped", ["fasterrcnn_resnet50_fpn_v2"]),
+    ]
+)
+def test_get_super_animal_detectors(data: tuple[str, list[str]]):
+    dataset, expected_detectors = data
+    detectors = modelzoo.get_available_detectors(dataset)
+    assert len(detectors) >= len(expected_detectors)
+    for det in expected_detectors:
+        assert det in detectors
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        ("superanimal_bird", ["resnet_50"]),
+        ("superanimal_topviewmouse", ["hrnet_w32"]),
+        ("superanimal_quadruped", ["hrnet_w32"]),
+    ]
+)
+def test_get_super_animal_pose_models(data: tuple[str, list[str]]):
+    dataset, expected_pose_models = data
+    pose_models = modelzoo.get_available_models(dataset)
+    assert len(pose_models) >= len(expected_pose_models)
+    for pose_model in expected_pose_models:
+        assert pose_model in pose_models