From 8e992a3b0c4b9d3dd73686ff484ad4b2702f395b Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Wed, 14 Dec 2022 08:51:04 +0100 Subject: [PATCH 1/5] Fix typo --- tools/update_license_headers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tools/update_license_headers.py b/tools/update_license_headers.py index 0e2d4d1..2562b29 100644 --- a/tools/update_license_headers.py +++ b/tools/update_license_headers.py @@ -28,7 +28,7 @@ def load_config(filename): def walk_directory(entry): - """Talk the directory""" + """Walk the directory""" if "header" not in entry: raise ValueError("Current entry does not have a header.") From fb1113363446b240f25393f85d736f160d69f644 Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Wed, 14 Dec 2022 09:19:38 +0100 Subject: [PATCH 2/5] Remove redundant test --- dlclibrary/dlcmodelzoo/modelzoo_download.py | 12 ------------ tests/test_modeldownload.py | 1 - 2 files changed, 13 deletions(-) diff --git a/dlclibrary/dlcmodelzoo/modelzoo_download.py b/dlclibrary/dlcmodelzoo/modelzoo_download.py index 51aa37e..724f5d1 100644 --- a/dlclibrary/dlcmodelzoo/modelzoo_download.py +++ b/dlclibrary/dlcmodelzoo/modelzoo_download.py @@ -97,15 +97,3 @@ def download_huggingface_model(modelname, target_dir=".", removeHFfolder=True): models = [fn for fn in neturls.keys()] print("Model does not exist: ", modelname) print("Pick one of the following: ", MODELOPTIONS) - - -if __name__ == "__main__": - print("Randomly downloading a model for testing...") - - import random - - # modelname = 'full_cat' - modelname = random.choice(MODELOPTIONS) - - target_dir = "/Users/alex/Downloads" # folder has to exist! - download_hugginface_model(modelname, target_dir) diff --git a/tests/test_modeldownload.py b/tests/test_modeldownload.py index da49649..40206cb 100644 --- a/tests/test_modeldownload.py +++ b/tests/test_modeldownload.py @@ -8,7 +8,6 @@ # # Licensed under GNU Lesser General Public License v3.0 # -import pytest def test_catdownload(tmp_path_factory): From e31fbb64b012633ec117ad2ad8ee6b4ab06e1601 Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Wed, 14 Dec 2022 09:55:21 +0100 Subject: [PATCH 3/5] Fix and test code --- dlclibrary/__init__.py | 2 +- dlclibrary/dlcmodelzoo/modelzoo_download.py | 80 +++++++++------------ tests/test_modeldownload.py | 15 ++-- 3 files changed, 46 insertions(+), 51 deletions(-) diff --git a/dlclibrary/__init__.py b/dlclibrary/__init__.py index b0db77e..a2473e8 100644 --- a/dlclibrary/__init__.py +++ b/dlclibrary/__init__.py @@ -9,5 +9,5 @@ # Licensed under GNU Lesser General Public License v3.0 # -from dlclibrary.dlcmodelzoo.modelzoo_download import download_hugginface_model +from dlclibrary.dlcmodelzoo.modelzoo_download import download_huggingface_model from dlclibrary.version import __version__, VERSION diff --git a/dlclibrary/dlcmodelzoo/modelzoo_download.py b/dlclibrary/dlcmodelzoo/modelzoo_download.py index 724f5d1..f7ae5c1 100644 --- a/dlclibrary/dlcmodelzoo/modelzoo_download.py +++ b/dlclibrary/dlcmodelzoo/modelzoo_download.py @@ -29,7 +29,7 @@ def _get_dlclibrary_path(): return os.path.split(importlib.util.find_spec("dlclibrary").origin)[0] -def _loadmodelnames(): +def _load_model_names(): """Loads URLs and commit hashes for available models.""" from ruamel.yaml import YAML @@ -38,7 +38,7 @@ def _loadmodelnames(): return YAML().load(file) -def download_huggingface_model(modelname, target_dir=".", removeHFfolder=True): +def download_huggingface_model(modelname, target_dir=".", remove_hf_folder=True): """ Downloads a DeepLabCut Model Zoo Project from Hugging Face @@ -48,52 +48,40 @@ def download_huggingface_model(modelname, target_dir=".", removeHFfolder=True): Name of the ModelZoo model. For visualizations see: http://www.mackenziemathislab.org/dlc-modelzoo target_dir : directory (as string) Directory where to store the model weigths and pose_cfg.yaml file - removeHFfolder : bool, default True + remove_hf_folder : bool, default True Whether to remove the directory structure provided by HuggingFace after downloading and decompressing data into DeepLabCut format. """ from huggingface_hub import hf_hub_download - import tarfile, os + import tarfile from pathlib import Path - neturls = _loadmodelnames() - - if modelname in neturls.keys(): - 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)) - # creates a new subfolder as indicated below, unzipping from there and deleting this folder - - # Building the HuggingFaceHub download path: - hf_path = ( - "models--" - + url[0] - + "--" - + url[1] - + "/snapshots/" - + str(neturls[modelname + "_commit"]) - + "/" - + targzfn - ) - - 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 # getting the filename - tar.makefile(member, target_dir + "/" + fname) - # tar.extractall(target_dir, members=tarfilenamecutting(tar)) - - if removeHFfolder: - # Removing folder - import shutil - - shutil.rmtree( - Path(os.path.join(target_dir, "models--" + url[0] + "--" + url[1])) - ) - - else: - models = [fn for fn in neturls.keys()] - print("Model does not exist: ", modelname) - print("Pick one of the following: ", MODELOPTIONS) + 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) + 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)) + + if remove_hf_folder: + import shutil + + shutil.rmtree(os.path.join(target_dir, hf_folder)) diff --git a/tests/test_modeldownload.py b/tests/test_modeldownload.py index 40206cb..e3a7c5c 100644 --- a/tests/test_modeldownload.py +++ b/tests/test_modeldownload.py @@ -8,14 +8,21 @@ # # Licensed under GNU Lesser General Public License v3.0 # +import dlclibrary +import os +import pytest -def test_catdownload(tmp_path_factory): - # TODO: just download the lightweight stuff.. - import dlclibrary, os - +def test_download_huggingface_model(tmp_path_factory): folder = tmp_path_factory.mktemp("cat") dlclibrary.download_huggingface_model("full_cat", str(folder)) assert os.path.exists(folder / "pose_cfg.yaml") assert os.path.exists(folder / "snapshot-75000.meta") + # Verify that the Hugging Face folder was removed + assert not any(f.startswith("models--") for f in os.listdir(folder)) + + +def test_download_huggingface_wrong_model(): + with pytest.raises(ValueError): + dlclibrary.download_huggingface_model("wrong_model_name") From d6f133f8c547b68b16655dcee4c133aeb1b023df Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Wed, 14 Dec 2022 09:59:51 +0100 Subject: [PATCH 4/5] Update mouse top view model name --- dlclibrary/dlcmodelzoo/modelzoo_download.py | 2 +- dlclibrary/modelzoo_urls.yaml | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/dlclibrary/dlcmodelzoo/modelzoo_download.py b/dlclibrary/dlcmodelzoo/modelzoo_download.py index f7ae5c1..1485341 100644 --- a/dlclibrary/dlcmodelzoo/modelzoo_download.py +++ b/dlclibrary/dlcmodelzoo/modelzoo_download.py @@ -18,7 +18,7 @@ "mouse_pupil_vclose", "horse_sideview", "full_macaque", - "superanimal_mouse", + "superanimal_mouse_topview", ] diff --git a/dlclibrary/modelzoo_urls.yaml b/dlclibrary/modelzoo_urls.yaml index 1f5bd66..1155d7c 100644 --- a/dlclibrary/modelzoo_urls.yaml +++ b/dlclibrary/modelzoo_urls.yaml @@ -31,5 +31,5 @@ horse_sideview_commit: fd0329b2ffc8fe7a5e6eb3d4850ebca75987e92c full_macaque: mwmathis/DeepLabCutModelZoo-macaque_full/DLC_macaque_full_resnet50.tar.gz full_macaque_commit: 4c7ebf2628d5b7eb0483356595256fb01b7e1a9e -superanimal_mouse: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/DLC_ma_supertopview5k_resnet_50_iteration-0_shuffle-1.tar.gz -superanimal_mouse_commit: a7d7df40c3307a3c7a0ceeb2593d46a783235b28 +superanimal_mouse_topview: mwmathis/DeepLabCutModelZoo-SuperAnimal-TopViewMouse/DLC_ma_supertopview5k_resnet_50_iteration-0_shuffle-1.tar.gz +superanimal_mouse_topview_commit: a7d7df40c3307a3c7a0ceeb2593d46a783235b28 From 18a0f99a94f7ae6d33ce09e3be0d5ab0c7b7d3ed Mon Sep 17 00:00:00 2001 From: Jessy Lauer <30733203+jeylau@users.noreply.github.com> Date: Wed, 14 Dec 2022 10:01:24 +0100 Subject: [PATCH 5/5] Fix typos --- dlclibrary/dlcmodelzoo/modelzoo_download.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/dlclibrary/dlcmodelzoo/modelzoo_download.py b/dlclibrary/dlcmodelzoo/modelzoo_download.py index 1485341..e231337 100644 --- a/dlclibrary/dlcmodelzoo/modelzoo_download.py +++ b/dlclibrary/dlcmodelzoo/modelzoo_download.py @@ -30,7 +30,7 @@ def _get_dlclibrary_path(): def _load_model_names(): - """Loads URLs and commit hashes for available models.""" + """Load URLs and commit hashes for available models.""" from ruamel.yaml import YAML fn = os.path.join(_get_dlclibrary_path(), "modelzoo_urls.yaml") @@ -40,14 +40,14 @@ def _load_model_names(): def download_huggingface_model(modelname, target_dir=".", remove_hf_folder=True): """ - Downloads a DeepLabCut Model Zoo Project from Hugging Face + Download a DeepLabCut Model Zoo Project from Hugging Face Parameters ---------- modelname : string Name of the ModelZoo model. For visualizations see: http://www.mackenziemathislab.org/dlc-modelzoo target_dir : directory (as string) - Directory where to store the model weigths and pose_cfg.yaml file + 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. """