Merge pull request #24 from DeepLabCut/update_urls

Add DLC3 detectors and pose estimation models

Author: MMathisLab

dlclibrary/dlcmodelzoo/modelzoo_download.py

@@ -8,6 +8,8 @@
 #
 # Licensed under GNU Lesser General Public License v3.0
 #
+from __future__ import annotations
+
 import json
 import os
 
@@ -20,10 +22,10 @@
     "mouse_pupil_vclose",
     "horse_sideview",
     "full_macaque",
-    "superanimal_topviewmouse",
-    "superanimal_quadruped",
-    "superanimal_quadruped_HRNetw32",
-    "superanimal_quadruped_ATP",
+    "superanimal_topviewmouse_dlcrnet",
+    "superanimal_quadruped_dlcrnet",
+    "superanimal_topviewmouse_hrnetw32",
+    "superanimal_quadruped_hrnetw32",
 ]
 
 
@@ -50,7 +52,9 @@ def parse_available_supermodels():
         return json.load(file)
 
 
-def download_huggingface_model(modelname, target_dir=".", remove_hf_folder=True):
+def download_huggingface_model(
+    modelname, target_dir=".", remove_hf_folder=True, rename_mapping: dict | None = None
+):
     """
     Download a DeepLabCut Model Zoo Project from Hugging Face
 
@@ -62,41 +66,58 @@ def download_huggingface_model(modelname, target_dir=".", remove_hf_folder=True)
         Directory where to store the model weights and pose_cfg.yaml file
     remove_hf_folder : bool, default True
         Whether to remove the directory structure provided by HuggingFace after downloading and decompressing data into DeepLabCut format.
+    rename_mapping : dict, default None
+        Dictionary to rename the downloaded file. If None, the original filename is used.
     """
     from huggingface_hub import hf_hub_download
     import tarfile
     from pathlib import Path
+    from ruamel.yaml.comments import CommentedBase
 
     neturls = _load_model_names()
     if modelname not in neturls:
         raise ValueError(f"`modelname` should be one of: {', '.join(modelname)}.")
 
     print("Loading....", modelname)
-    url = neturls[modelname].split("/")
-    repo_id, targzfn = url[0] + "/" + url[1], str(url[-1])
-
-    hf_hub_download(repo_id, targzfn, cache_dir=str(target_dir))
-
-    # Create a new subfolder as indicated below, unzipping from there and deleting this folder
-    hf_folder = f"models--{url[0]}--{url[1]}"
-    hf_path = os.path.join(
-        hf_folder,
-        "snapshots",
-        str(neturls[modelname + "_commit"]),
-        targzfn,
-    )
-
-    filename = os.path.join(target_dir, hf_path)
-    try:
-        with tarfile.open(filename, mode="r:gz") as tar:
-            for member in tar:
-                if not member.isdir():
-                    fname = Path(member.name).name
-                    tar.makefile(member, os.path.join(target_dir, fname))
-    except tarfile.ReadError:  # The model is a .pt file
-        os.rename(filename, os.path.join(target_dir, targzfn))
+    urls = neturls[modelname]
+    if isinstance(urls, CommentedBase):
+        urls = list(urls)
+    else:
+        urls = [urls]
+
+    if not os.path.isabs(target_dir):
+        target_dir = os.path.abspath(target_dir)
+
+    for url in urls:
+        url = url.split("/")
+        repo_id, targzfn = url[0] + "/" + url[1], str(url[-1])
+
+        hf_hub_download(repo_id, targzfn, cache_dir=str(target_dir))
+
+        # Create a new subfolder as indicated below, unzipping from there and deleting this folder
+        hf_folder = f"models--{url[0]}--{url[1]}"
+        path_ = os.path.join(target_dir, hf_folder, "snapshots")
+        commit = os.listdir(path_)[0]
+        filename = os.path.join(path_, commit, targzfn)
+        try:
+            with tarfile.open(filename, mode="r:gz") as tar:
+                for member in tar:
+                    if not member.isdir():
+                        fname = Path(member.name).name
+                        tar.makefile(member, os.path.join(target_dir, fname))
+        except tarfile.ReadError:  # The model is a .pt file
+            if rename_mapping is not None:
+                targzfn = rename_mapping.get(targzfn, targzfn)
+            if os.path.islink(filename):
+                filename_ = os.readlink(filename)
+                if not os.path.isabs(filename_):
+                    filename_ = os.path.abspath(os.path.join(os.path.dirname(filename), filename_))
+                filename = filename_
+            os.rename(filename, os.path.join(target_dir, targzfn))
 
     if remove_hf_folder:
         import shutil
 
         shutil.rmtree(os.path.join(target_dir, hf_folder))
+
+'../../blobs/6c9c66d48f25cac9f8adaea7a485b07f4bd781ba656785bc4e077d9064e8e5df'

dlclibrary/dlcmodelzoo/modelzoo_urls.yaml

@@ -9,36 +9,29 @@
 ## Licensed under GNU Lesser General Public License v3.0
 ##
 
-
-
 #Model Zoo from HuggingFace
 full_human: mwmathis/DeepLabCutModelZoo-DLC_human_fullbody_resnet_101/DLC_human_dancing_resnet_101_iteration-0_shuffle-1.tar.gz
-full_human_commit: 01acf989bfec7a25110b94677449d9e861b5ea81
+
 full_cat: AlexEMG/DeepLabCutModelZoo-cat/DLC_Cat_resnet_50_iteration-0_shuffle-0.tar.gz
-full_cat_commit: d39b3922a04d7077a84a05d141243c044cc18af6
+
 full_dog: AlexEMG/DeepLabCutModelZoo-dog/DLC_Dog_resnet_50_iteration-0_shuffle-0.tar.gz
-full_dog_commit: b420a859f881b3cabdfd0cc4a0eb2d1c7fb592c5
 
 primate_face: mwmathis/DeepLabCutModelZoo-primate_face/DLC_primate_face_resnet_50_iteration-1_shuffle-1.tar.gz
-primate_face_commit: df7e8f6ad38899489758e5a890b58905cfb4d893
 
 mouse_pupil_vclose: mwmathis/DeepLabCutModelZoo-mouse_pupil_vclose/DLC_mouse_pupil_vclose_resnet_50_iteration-0_shuffle-1.tar.gz
-mouse_pupil_vclose_commit: 93557f0e74390596a2b9fddaab7628a1cc79fb2d
 
 horse_sideview: mwmathis/DeepLabCutModelZoo-horse_sideview/DLC_Horses_resnet_50_iteration-1_shuffle-1.tar.gz
-horse_sideview_commit: fd0329b2ffc8fe7a5e6eb3d4850ebca75987e92c
 
 full_macaque: mwmathis/DeepLabCutModelZoo-macaque_full/DLC_macaque_full_resnet50.tar.gz
-full_macaque_commit: 4c7ebf2628d5b7eb0483356595256fb01b7e1a9e
 
-superanimal_topviewmouse: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/DLC_ma_supertopview5k_resnet_50_iteration-0_shuffle-1.tar.gz
-superanimal_topviewmouse_commit: 1781f384b56ddb9b58c0d4acae98993b9a17e901
+superanimal_topviewmouse_dlcrnet: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/DLC_ma_supertopview5k_resnet_50_iteration-0_shuffle-1.tar.gz
 
-superanimal_quadruped: mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/DLC_ma_superquadruped_resnet_50_iteration-0_shuffle-1.tar.gz
-superanimal_quadruped_commit: 673140e6dd9f7be492d77cab957f31c73a192f67
+superanimal_quadruped_dlcrnet: mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped_DLC2/DLC_ma_superquadruped_resnet_50_iteration-0_shuffle-1.tar.gz
 
-superanimal_quadruped_HRNetw32: mwmathis/SuperAnimal_Quadruped_pt/blob/main/SuperAnimal_Quadruped_pt_HRNetw32.pth
-superanimal_quadruped_HRNetw32_commit: b3d7ddf5f82922a57fb7d7388c292e7b357add18
+superanimal_topviewmouse_hrnetw32:
+  - mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/pose_model.pth
+  - mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/detector.pt
 
-superanimal_quadruped_ATP: mwmathis/SuperAnimal_quadruped_ATP_pt/blob/main/AnimalTokenPose_SplitD.pth
-superanimal_quadruped_ATP_commit: 41186b920187e92b897a6483b2a4f053a8ef3700
+superanimal_quadruped_hrnetw32:
+  - mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/pose_model.pth
+  - mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/detector.pt