Skip to content

Conversation

@deruyter92
Copy link
Collaborator

Motivation:
It would be helpful to have easy access to the torch superanimal models without having to load export them first in DeepLabCut. This PR adds the functionality for direct downloading and exporting the pytorch models from the modelzoo.

Implementation:
The current implementation duplicates code from the DeepLabCut main repository (as is the case for other models). Other solutions are also possible (e.g. hosting the readily exported models online).

  • Modelzoo configs are copied into this repository
  • New utility functions are added for parsing the config files and exporting models (based on DeepLabCut main repository)
  • Modelzoo tests are copied from DeepLabCut
  • Added new functional benchmarking test for pytorch models with model zoo as well

Example usage:

  • The main entry point is dlclive.modelzoo.pytorch_model_zoo_export.py:
from dlclive import modelzoo

# Output path for the exported model
checkpoint = "exported_models/exported_quadruped_resnet_50.pt"

modelzoo.export_modelzoo_model(
    export_path=checkpoint ,
    super_animal="superanimal_quadruped",
    model_name="resnet_50",
    detector_name="ssdlite",
)
  • Some utility functions are defined in dlclive.modelzoo.utils.py, such as:
modelzoo.list_available_combinations()
>>> ['superanimal_bird_dlcrnet', 'superanimal_bird_fasterrcnn_mobilenet_v3_large_fpn',  ..., superanimal_topviewmouse_rtmpose_x', 'superanimal_topviewmouse_ssdlite']

@deruyter92 deruyter92 changed the title Jaap/add modelzoo functionality add modelzoo functionality Jan 23, 2026
Comment on lines +33 to +40
def test_download_huggingface_model(tmp_path_factory, model="full_cat"):
folder = tmp_path_factory.mktemp("temp")
dlclibrary.download_huggingface_model(model, str(folder))

assert os.path.exists(folder / "pose_cfg.yaml")
assert any(f.startswith("snapshot-") for f in os.listdir(folder))
# Verify that the Hugging Face folder was removed
assert not any(f.startswith("models--") for f in os.listdir(folder))
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

looks like we're just testing dlclibrary code here?

@sneakers-the-rat
Copy link
Collaborator

Other solutions are also possible (e.g. hosting the readily exported models online).

strongly prefer not duplicating code or configs unless it's unavoidable, getting in the habit of duplicating code trades off with just making a proper dependency tree so that things have a single place they are known to live. DLClibrary looks like pretty much all of its contents are modelzoo things, so basically everything except the export function that requires pytorch looks like a good enhancement for that package!

not trying to be a busybody, i just have enjoyed how this package has been very solid and stable and i can just reach into my back pocket and use it whenever. the big code duplication of the pytorch models imo would also have been better as a separate package, but would really like to avoid any further code duplication, as that's a pretty quick path to bloat (and also confusing to downstream users, where rather than "use the model zoo package to do model zoo things" now there are model zoo things in 3 packages and they're all slightly different)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants