diff --git a/.github/workflows/test_package.yml b/.github/workflows/test_package.yml index 10d65dc..e2d97e2 100644 --- a/.github/workflows/test_package.yml +++ b/.github/workflows/test_package.yml @@ -28,7 +28,7 @@ jobs: python -m pip install --upgrade pip pip install pytest if [ -f requirements.txt ]; then pip install -r requirements.txt; fi - + - name: Install DLClibrary run: pip install . diff --git a/LICENSE b/LICENSE index 341c30b..65c5ca8 100644 --- a/LICENSE +++ b/LICENSE @@ -163,4 +163,3 @@ whether future versions of the GNU Lesser General Public License shall apply, that proxy's public statement of acceptance of any version is permanent authorization for you to choose that version for the Library. - diff --git a/README.md b/README.md index e0a3b42..55d4ce7 100644 --- a/README.md +++ b/README.md @@ -35,8 +35,8 @@ model_dir.mkdir() download_huggingface_model("superanimal_quadruped", model_dir) ``` -PyTorch models available for a given dataset (compatible with DeepLabCut>=3.0) can be -listed using the `dlclibrary.get_available_detectors` and +PyTorch models available for a given dataset (compatible with DeepLabCut>=3.0) can be +listed using the `dlclibrary.get_available_detectors` and `dlclibrary.get_available_models` methods. The datasets for which models are available can be listed using `dlclibrary.get_available_datasets`. Example use: @@ -57,7 +57,7 @@ can be listed using `dlclibrary.get_available_datasets`. Example use: ### TensorFlow models -Pick a good model_name. Follow the (novel) naming convention (modeltype_species), e.g. ```superanimal_topviewmouse```. +Pick a good model_name. Follow the (novel) naming convention (modeltype_species), e.g. ```superanimal_topviewmouse```. 1. Add the model_name with path and commit ID to: https://github.com/DeepLabCut/DLClibrary/blob/main/dlclibrary/dlcmodelzoo/modelzoo_urls.yaml @@ -88,8 +88,8 @@ i.e. for this example 3 models would be available: `my_cool_dataset_detector_nam To add a new model for `deeplabcut >= 3.0.0`, simply: - add a new line under detectors or pose models if the dataset is already defined -- add the structure if the model was trained on a new dataset +- add the structure if the model was trained on a new dataset The models will then be listed when calling `dlclibrary.get_available_detectors` or -`dlclibrary.get_available_models`! You can list the datasets for which models are +`dlclibrary.get_available_models`! You can list the datasets for which models are available using `dlclibrary.get_available_datasets`. diff --git a/dlclibrary/__init__.py b/dlclibrary/__init__.py index 27c39db..84b5d51 100644 --- a/dlclibrary/__init__.py +++ b/dlclibrary/__init__.py @@ -16,4 +16,4 @@ get_available_models, parse_available_supermodels, ) -from dlclibrary.version import __version__, VERSION +from dlclibrary.version import VERSION, __version__ diff --git a/dlclibrary/dlcmodelzoo/modelzoo_download.py b/dlclibrary/dlcmodelzoo/modelzoo_download.py index 49162a5..bdc6d51 100644 --- a/dlclibrary/dlcmodelzoo/modelzoo_download.py +++ b/dlclibrary/dlcmodelzoo/modelzoo_download.py @@ -12,8 +12,8 @@ import json import os -import tarfile import shutil +import tarfile import tempfile from pathlib import Path @@ -35,8 +35,8 @@ ] -def _get_dlclibrary_path(): - """Get path of where dlclibrary (this repo) is currently running""" +def _get_dlclibrary_path() -> str: + """Get path of where dlclibrary (this repo) is currently running.""" import importlib.util return os.path.split(importlib.util.find_spec("dlclibrary").origin)[0] @@ -46,24 +46,19 @@ 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 + return YAML(pure=True).load(file) def _load_pytorch_dataset_models(dataset: str) -> dict[str, dict[str, str]]: """Load URLs and commit hashes for available models.""" 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())}" - ) + if dataset not in models: + raise ValueError(f"Could not find any models for {dataset}. Models are available for {list(models.keys())}") return models[dataset] -def _load_model_names(): +def _load_model_names() -> dict[str, str | list[str]]: """Load URLs and commit hashes for available models.""" fn = os.path.join(_get_dlclibrary_path(), "dlcmodelzoo", "modelzoo_urls.yaml") with open(fn) as file: @@ -78,12 +73,12 @@ def _load_model_names(): return model_names -def parse_available_supermodels(): +def parse_available_supermodels() -> dict[str, str]: + """Return each SuperAnimal model name mapped to its config file name in ``superanimal_configs``.""" libpath = _get_dlclibrary_path() json_path = os.path.join(libpath, "dlcmodelzoo", "superanimal_models.json") with open(json_path) as file: - super_animal_models = json.load(file) - return super_animal_models + return json.load(file) def get_available_datasets() -> list[str]: @@ -113,10 +108,7 @@ def get_available_models(dataset: str) -> list[str]: return list(_load_pytorch_dataset_models(dataset)["pose_models"].keys()) - -def _handle_downloaded_file( - file_path: str, target_dir: str, rename_mapping: dict | None = None -): +def _handle_downloaded_file(file_path: str, target_dir: str, rename_mapping: dict | None = None) -> None: """Handle the downloaded file from HuggingFace cache and place the final artifact in target_dir.""" file_name = os.path.basename(file_path) @@ -158,9 +150,8 @@ def download_huggingface_model( model_name: str, target_dir: str = ".", rename_mapping: str | dict | None = None, -): - """ - Downloads a DeepLabCut Model Zoo Project from Hugging Face. +) -> None: + """Downloads a DeepLabCut Model Zoo Project from Hugging Face. Args: model_name (str): @@ -189,16 +180,11 @@ def download_huggingface_model( """ net_urls = _load_model_names() if model_name not in net_urls: - raise ValueError( - f"`modelname={model_name}` should be one of: {', '.join(net_urls)}." - ) + raise ValueError(f"`modelname={model_name}` should be one of: {', '.join(net_urls)}.") print("Loading....", model_name) urls = net_urls[model_name] - if isinstance(urls, CommentedBase): - urls = list(urls) - else: - urls = [urls] + urls = list(urls) if isinstance(urls, CommentedBase) else [urls] if not os.path.isabs(target_dir): target_dir = os.path.abspath(target_dir) @@ -215,9 +201,6 @@ def download_huggingface_model( cache_dir=hf_cache_dir, ) - if isinstance(rename_mapping, str): - mapping = {targzfn: rename_mapping} - else: - mapping = rename_mapping + mapping = {targzfn: rename_mapping} if isinstance(rename_mapping, str) else rename_mapping _handle_downloaded_file(downloaded, target_dir, mapping) diff --git a/setup.py b/setup.py index 700c675..957ae77 100644 --- a/setup.py +++ b/setup.py @@ -10,7 +10,7 @@ # import setuptools -with open("README.md", "r") as fh: +with open("README.md") as fh: long_description = fh.read() setuptools.setup( diff --git a/tests/test_modeldownload.py b/tests/test_modeldownload.py index 2e0d7bc..8d08065 100644 --- a/tests/test_modeldownload.py +++ b/tests/test_modeldownload.py @@ -11,7 +11,6 @@ from __future__ import annotations import io -import os import tarfile from pathlib import Path @@ -22,7 +21,7 @@ from dlclibrary.dlcmodelzoo.modelzoo_download import MODELOPTIONS -def _fake_model_names(): +def _fake_model_names() -> dict[str, str]: """ Return a deterministic fake URL for each model. Alternate between tar.gz and .pt to test both branches. @@ -34,7 +33,7 @@ def _fake_model_names(): return mapping -def _write_fake_tar_gz(path: Path): +def _write_fake_tar_gz(path: Path) -> None: """ Create a fake tar.gz archive with the files the downloader expects for archive-based DLC models. @@ -55,7 +54,7 @@ def _write_fake_tar_gz(path: Path): tar.addfile(info, io.BytesIO(content)) -def _write_fake_pt(path: Path): +def _write_fake_pt(path: Path) -> None: """ Create a fake .pt / .pth weight file. """ @@ -64,7 +63,7 @@ def _write_fake_pt(path: Path): @pytest.fixture -def mock_modelzoo(monkeypatch): +def mock_modelzoo(monkeypatch: pytest.MonkeyPatch) -> dict[str, str]: """ Patch both: - model name resolution @@ -76,7 +75,7 @@ def mock_modelzoo(monkeypatch): monkeypatch.setattr(modelzoo_download, "_load_model_names", lambda: fake_names) - def fake_hf_hub_download(repo_id, filename, cache_dir): + def fake_hf_hub_download(repo_id: str, filename: str, cache_dir: str) -> str: cache_dir = Path(cache_dir) hf_folder = cache_dir / f"models--{repo_id.replace('/', '--')}" snapshot_dir = hf_folder / "snapshots" / "fakecommit123" @@ -84,7 +83,7 @@ def fake_hf_hub_download(repo_id, filename, cache_dir): if filename.endswith(".tar.gz"): _write_fake_tar_gz(returned_file) - elif filename.endswith(".pt") or filename.endswith(".pth"): + elif filename.endswith((".pt", ".pth")): _write_fake_pt(returned_file) else: raise AssertionError(f"Unexpected mocked filename: {filename}") @@ -96,7 +95,7 @@ def fake_hf_hub_download(repo_id, filename, cache_dir): return fake_names -def _assert_download_success(folder: Path, model: str): +def _assert_download_success(folder: Path, model: str) -> None: """ Shared assertion helper for download_huggingface_model. """ @@ -117,7 +116,8 @@ def _assert_download_success(folder: Path, model: str): assert not any(name.startswith("models--") for name in files) -def test_download_huggingface_model_tar_or_pt(tmp_path, mock_modelzoo): +@pytest.mark.usefixtures("mock_modelzoo") +def test_download_huggingface_model_tar_or_pt(tmp_path: Path) -> None: folder = tmp_path / "download_one" folder.mkdir() @@ -126,25 +126,28 @@ def test_download_huggingface_model_tar_or_pt(tmp_path, mock_modelzoo): _assert_download_success(folder, "full_cat") -def test_download_huggingface_wrong_model(mock_modelzoo): +@pytest.mark.usefixtures("mock_modelzoo") +def test_download_huggingface_wrong_model() -> None: with pytest.raises(ValueError): dlclibrary.download_huggingface_model("wrong_model_name") -def test_parse_superanimal_models(): +def test_parse_superanimal_models() -> None: dict_ = dlclibrary.parse_available_supermodels() assert "superanimal_quadruped" in dict_ assert "superanimal_topviewmouse" in dict_ +@pytest.mark.usefixtures("mock_modelzoo") @pytest.mark.parametrize("model", MODELOPTIONS) -def test_download_all_models(tmp_path, mock_modelzoo, model): +def test_download_all_models(tmp_path: Path, model: str) -> None: folder = tmp_path / model folder.mkdir() _assert_download_success(folder, model) -def test_download_with_rename_mapping_for_pt(tmp_path, mock_modelzoo): +@pytest.mark.usefixtures("mock_modelzoo") +def test_download_with_rename_mapping_for_pt(tmp_path: Path) -> None: """ Explicitly test rename_mapping for a .pt model. """ diff --git a/tests/test_pytorch_models.py b/tests/test_pytorch_models.py index 0c04793..6d61492 100644 --- a/tests/test_pytorch_models.py +++ b/tests/test_pytorch_models.py @@ -8,10 +8,8 @@ # # Licensed under GNU Lesser General Public License v3.0 # -import os import pytest -import dlclibrary import dlclibrary.dlcmodelzoo.modelzoo_download as modelzoo @@ -21,9 +19,9 @@ ("superanimal_bird", ["ssdlite"]), ("superanimal_topviewmouse", ["fasterrcnn_resnet50_fpn_v2"]), ("superanimal_quadruped", ["fasterrcnn_resnet50_fpn_v2"]), - ] + ], ) -def test_get_super_animal_detectors(detector_data: tuple[str, list[str]]): +def test_get_super_animal_detectors(detector_data: tuple[str, list[str]]) -> None: dataset, expected_detectors = detector_data detectors = modelzoo.get_available_detectors(dataset) assert len(detectors) >= len(expected_detectors) @@ -38,11 +36,11 @@ def test_get_super_animal_detectors(detector_data: tuple[str, list[str]]): ("superanimal_topviewmouse", ["hrnet_w32"]), ("superanimal_quadruped", ["hrnet_w32"]), ("superanimal_humanbody", ["rtmpose_x"]), - ] + ], ) -def test_get_super_animal_pose_models(posemodel_data: tuple[str, list[str]]): +def test_get_super_animal_pose_models(posemodel_data: tuple[str, list[str]]) -> None: dataset, expected_pose_models = posemodel_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 + assert pose_model in pose_models diff --git a/tools/update_license_headers.py b/tools/update_license_headers.py index 2562b29..06e952c 100644 --- a/tools/update_license_headers.py +++ b/tools/update_license_headers.py @@ -14,34 +14,35 @@ configuration, see the instructions in NOTICE.yml. """ -import tempfile -import glob -import yaml import fnmatch +import glob import subprocess +import tempfile +from collections.abc import Iterable, Iterator + +import yaml -def load_config(filename): - with open(filename, "r") as fh: - config = yaml.safe_load(fh) - return config +def load_config(filename: str) -> list[dict]: + """Load the list of header entries from a NOTICE.yml file.""" + with open(filename) as fh: + return yaml.safe_load(fh) -def walk_directory(entry): - """Walk the directory""" +def walk_directory(entry: dict) -> list[str]: + """List the files an entry includes, minus those it excludes.""" if "header" not in entry: raise ValueError("Current entry does not have a header.") if "include" not in entry: raise ValueError("Current entry does not have an include list.") - def _list_include(): + def _list_include() -> Iterator[str]: """List all files specified in the include list.""" for include_pattern in entry["include"]: - for filename in glob.iglob(include_pattern, recursive=True): - yield filename + yield from glob.iglob(include_pattern, recursive=True) - def _filter_exclude(iterable): + def _filter_exclude(iterable: Iterable[str]) -> Iterator[str]: """Filter filenames from an iterator by the exclude patterns.""" for filename in iterable: for exclude_pattern in entry.get("exclude", []): @@ -54,7 +55,8 @@ def _filter_exclude(iterable): return list(files) -def main(input_file="NOTICE.yml"): +def main(input_file: str = "NOTICE.yml") -> None: + """Write each entry's header into the files it covers, with licenseheaders.""" config = load_config(input_file) for entry in config: filelist = list(walk_directory(entry)) @@ -62,7 +64,7 @@ def main(input_file="NOTICE.yml"): header_file.write(entry["header"]) header_file.flush() header_file.seek(0) - command = ["licenseheaders", "-t", str(header_file.name), "-f"] + filelist + command = ["licenseheaders", "-t", str(header_file.name), "-f", *filelist] result = subprocess.run(command, capture_output=True) if result.returncode != 0: print(result.stdout.decode())