Merge pull request #21 from jeylau/superquad_hrnet

Handle .pt model download & add superanimal_quadruped_HRNetw32

Author: MMathisLab

dlclibrary/dlcmodelzoo/modelzoo_download.py

@@ -22,6 +22,7 @@
     "full_macaque",
     "superanimal_topviewmouse",
     "superanimal_quadruped",
+    "superanimal_quadruped_HRNetw32",
 ]
 
 
@@ -85,11 +86,14 @@ def download_huggingface_model(modelname, target_dir=".", remove_hf_folder=True)
     )
 
     filename = os.path.join(target_dir, hf_path)
-    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))
+    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))
 
     if remove_hf_folder:
         import shutil

dlclibrary/dlcmodelzoo/modelzoo_urls.yaml

@@ -36,3 +36,6 @@ superanimal_topviewmouse_commit: 1781f384b56ddb9b58c0d4acae98993b9a17e901
 
 superanimal_quadruped: mwmathis/DeepLabCutModelZoo-SuperAnimal-Quadruped/DLC_ma_superquadruped_resnet_50_iteration-0_shuffle-1.tar.gz
 superanimal_quadruped_commit: 673140e6dd9f7be492d77cab957f31c73a192f67
+
+superanimal_quadruped_HRNetw32: mwmathis/SuperAnimal_Quadruped_pt/blob/main/SuperAnimal_Quadruped_pt_HRNetw32.pth
+superanimal_quadruped_HRNetw32_commit: b3d7ddf5f82922a57fb7d7388c292e7b357add18