diff --git a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py index a7cabc1a82..b94230e24b 100644 --- a/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py +++ b/deeplabcut/pose_estimation_pytorch/apis/analyze_images.py @@ -40,7 +40,10 @@ get_filtered_coco_detector_inference_runner, ) from deeplabcut.pose_estimation_pytorch.data.ctd import CondFromModel -from deeplabcut.pose_estimation_pytorch.modelzoo.utils import update_config +from deeplabcut.pose_estimation_pytorch.modelzoo.utils import ( + COCO_PERSON_CATEGORY_ID, + update_config, +) from deeplabcut.pose_estimation_pytorch.task import Task from deeplabcut.pose_estimation_pytorch.utils import resolve_device from deeplabcut.utils import auxfun_videos, auxiliaryfunctions @@ -176,10 +179,9 @@ def superanimal_analyze_images( torchvision_detector_name = detector_name else: torchvision_detector_name = "fasterrcnn_mobilenet_v3_large_fpn" - COCO_PERSON = 1 # COCO class ID for person filtered_detector_config = { "torchvision_detector_name": torchvision_detector_name, - "category_id": COCO_PERSON, + "category_id": COCO_PERSON_CATEGORY_ID, } if customized_model_config is None: diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/__init__.py b/deeplabcut/pose_estimation_pytorch/modelzoo/__init__.py index e8232cd895..51459a5087 100644 --- a/deeplabcut/pose_estimation_pytorch/modelzoo/__init__.py +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/__init__.py @@ -8,11 +8,55 @@ # # Licensed under GNU Lesser General Public License v3.0 # -from deeplabcut.pose_estimation_pytorch.modelzoo.utils import ( - download_super_animal_snapshot, - get_snapshot_folder_path, - get_super_animal_model_config_path, - get_super_animal_project_config_path, - get_super_animal_snapshot_path, - load_super_animal_config, -) +"""Public API for PyTorch modelzoo. + +Exports are resolved lazily to avoid import cycles between helpers and package +initialization. +""" + +from importlib import import_module + +_EXPORTS = { + "download_super_animal_snapshot": ( + "deeplabcut.pose_estimation_pytorch.modelzoo.utils", + "download_super_animal_snapshot", + ), + "get_snapshot_folder_path": ( + "deeplabcut.pose_estimation_pytorch.modelzoo.utils", + "get_snapshot_folder_path", + ), + "get_super_animal_model_config_path": ( + "deeplabcut.pose_estimation_pytorch.modelzoo.utils", + "get_super_animal_model_config_path", + ), + "get_super_animal_project_config_path": ( + "deeplabcut.pose_estimation_pytorch.modelzoo.utils", + "get_super_animal_project_config_path", + ), + "get_super_animal_snapshot_path": ( + "deeplabcut.pose_estimation_pytorch.modelzoo.utils", + "get_super_animal_snapshot_path", + ), + "load_super_animal_config": ( + "deeplabcut.pose_estimation_pytorch.modelzoo.utils", + "load_super_animal_config", + ), + "create_superanimal_inference_runners": ( + "deeplabcut.pose_estimation_pytorch.modelzoo.inference_helpers", + "create_superanimal_inference_runners", + ), +} + +__all__ = sorted(_EXPORTS) + + +def __getattr__(name: str): + try: + module_name, attr_name = _EXPORTS[name] + except KeyError as exc: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from exc + return getattr(import_module(module_name), attr_name) + + +def __dir__(): + return sorted(set(globals()) | set(__all__)) diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py b/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py index 15c5004264..27b1dd5648 100644 --- a/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/inference.py @@ -25,6 +25,7 @@ get_filtered_coco_detector_inference_runner, ) from deeplabcut.pose_estimation_pytorch.modelzoo.utils import ( + COCO_PERSON_CATEGORY_ID, raise_warning_if_called_directly, ) from deeplabcut.utils.make_labeled_video import create_video @@ -110,10 +111,9 @@ def _video_inference_superanimal( if superanimal_name == "superanimal_humanbody": if not torchvision_detector_name: torchvision_detector_name = "fasterrcnn_mobilenet_v3_large_fpn" - COCO_PERSON = 1 # COCO class ID for person detector_runner = get_filtered_coco_detector_inference_runner( model_name=torchvision_detector_name, - category_id=COCO_PERSON, + category_id=COCO_PERSON_CATEGORY_ID, batch_size=detector_batch_size, max_individuals=max_individuals, model_config=model_cfg, diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/inference_helpers.py b/deeplabcut/pose_estimation_pytorch/modelzoo/inference_helpers.py new file mode 100644 index 0000000000..a9d9fdce6c --- /dev/null +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/inference_helpers.py @@ -0,0 +1,224 @@ +# +# DeepLabCut Toolbox (deeplabcut.org) +# © A. & M.W. Mathis Labs +# https://github.com/DeepLabCut/DeepLabCut +# +# Please see AUTHORS for contributors. +# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS +# +# Licensed under GNU Lesser General Public License v3.0 +# +"""PyTorch-specific helper entrypoints for model zoo inference.""" + +from __future__ import annotations + +import copy +import logging +from pathlib import Path + +import torch + +import deeplabcut.modelzoo.weight_initialization as weight_initialization +from deeplabcut.core.config import read_config_as_dict +from deeplabcut.pose_estimation_pytorch.apis.utils import ( + get_filtered_coco_detector_inference_runner, + get_inference_runners, + get_pose_inference_runner, +) +from deeplabcut.pose_estimation_pytorch.modelzoo.utils import ( + COCO_PERSON_CATEGORY_ID, + get_super_animal_snapshot_path, + load_super_animal_config, + update_config, +) +from deeplabcut.pose_estimation_pytorch.runners import InferenceRunner +from deeplabcut.pose_estimation_pytorch.task import Task + + +def _build_humanbody_inference_runners( + model_cfg: dict, + model_name: str, + detector_name: str | None, + max_individuals: int, + batch_size: int, + detector_batch_size: int, + customized_pose_checkpoint: str | Path | None, + customized_detector_checkpoint: str | Path | None, +) -> tuple[InferenceRunner, InferenceRunner, dict]: + if customized_detector_checkpoint is not None: + logging.warning( + "customized_detector_checkpoint is ignored for superanimal_humanbody. " + "A filtered torchvision detector runner is used instead." + ) + + torchvision_detector_name = ( + detector_name + if detector_name is not None + else "fasterrcnn_mobilenet_v3_large_fpn" + ) + + pose_snapshot_path = customized_pose_checkpoint + if pose_snapshot_path is None: + pose_snapshot_path = get_super_animal_snapshot_path( + dataset="superanimal_humanbody", + model_name=model_name, + download=True, + ) + + detector_runner = get_filtered_coco_detector_inference_runner( + model_name=torchvision_detector_name, + category_id=COCO_PERSON_CATEGORY_ID, + batch_size=detector_batch_size, + max_individuals=max_individuals, + model_config=model_cfg, + ) + pose_runner = get_pose_inference_runner( + model_cfg, + snapshot_path=pose_snapshot_path, + batch_size=batch_size, + max_individuals=max_individuals, + ) + return pose_runner, detector_runner, model_cfg + + +def create_superanimal_inference_runners( + superanimal_name: str, + model_name: str, + detector_name: str | None = None, + max_individuals: int = 10, + batch_size: int = 1, + detector_batch_size: int = 1, + device: str | None = "auto", + customized_model_config: str | Path | dict | None = None, + customized_pose_checkpoint: str | Path | None = None, + customized_detector_checkpoint: str | Path | None = None, +) -> tuple[InferenceRunner, InferenceRunner | None, dict]: + """Create SuperAnimal inference runners for in-memory batched inference. + + This helper is intended for Model Zoo inference pipelines that run directly on + arrays. It prepares pose/detector runners and returns them with the resolved + model config. + + Args: + superanimal_name: Name of the SuperAnimal dataset, e.g. + ``"superanimal_quadruped"``. + model_name: Pose model architecture name, e.g. ``"hrnet_w32"``. + detector_name: Detector architecture name. For top-down SuperAnimal models, + use detector names such as ``"fasterrcnn_resnet50_fpn_v2"``. For + ``superanimal_humanbody``, this is interpreted as a torchvision detector + name (default: ``"fasterrcnn_mobilenet_v3_large_fpn"``). Can be ``None`` + for bottom-up models. + max_individuals: Maximum number of individuals to keep per frame. + batch_size: Batch size for pose inference. + detector_batch_size: Batch size for detector inference. + device: Device for inference. If ``"auto"`` or ``None``, resolves to CUDA + when available, else CPU. + customized_model_config: Optional path or dict for a custom model config. + If not provided, uses the default SuperAnimal config. Note that this config + determines whether the model is top-down or bottom-up; for bottom-up models, + ``detector_runner`` will be ``None`` even if ``detector_name`` is set. + customized_pose_checkpoint: Optional custom pose checkpoint path. + customized_detector_checkpoint: Optional custom detector checkpoint path. + + Returns: + tuple: ``(pose_runner, detector_runner, model_cfg)`` where: + - ``pose_runner`` is the pose inference runner + - ``detector_runner`` is the detector inference runner or ``None`` if no + detector is configured + - ``model_cfg`` is the resolved model configuration dict + + Example: + >>> from pathlib import Path + >>> import numpy as np + >>> from PIL import Image + >>> from deeplabcut.pose_estimation_pytorch.modelzoo.inference_helpers import ( + ... create_superanimal_inference_runners, + ... ) + >>> + >>> img_paths = [ + ... "/path/to/images/frame_0000.png", + ... "/path/to/images/frame_0001.png", + ... "/path/to/images/frame_0002.png", + ... ] + >>> images = [np.asarray(Image.open(Path(p)).convert("RGB")) for p in img_paths] + >>> + >>> pose_runner, det_runner, model_cfg = create_superanimal_inference_runners( + ... superanimal_name="superanimal_quadruped", + ... model_name="hrnet_w32", + ... detector_name="fasterrcnn_resnet50_fpn_v2", + ... max_individuals=10, + ... batch_size=1, + ... detector_batch_size=1, + ... ) + >>> + >>> det_preds = det_runner.inference(images) if det_runner is not None else None + >>> pose_inputs = list(zip(images, det_preds)) if det_preds is not None else images + >>> pose_preds = pose_runner.inference(pose_inputs) + >>> print(len(pose_preds)) + """ + if model_name.lower().startswith("fmpose3d"): + raise NotImplementedError( + "FMPose3D is not supported in this helper. Use the FMPose3D inference API." + ) + + if device is None: + device = "auto" + + if customized_model_config is not None: + if isinstance(customized_model_config, (str, Path)): + model_cfg = read_config_as_dict(customized_model_config) + else: + model_cfg = copy.deepcopy(customized_model_config) + else: + model_cfg = load_super_animal_config( + super_animal=superanimal_name, + model_name=model_name, + detector_name=detector_name, + ) + model_cfg = update_config(model_cfg, max_individuals=max_individuals, device=device) + + if superanimal_name == "superanimal_humanbody": + return _build_humanbody_inference_runners( + model_cfg=model_cfg, + model_name=model_name, + detector_name=detector_name, + max_individuals=max_individuals, + batch_size=batch_size, + detector_batch_size=detector_batch_size, + customized_pose_checkpoint=customized_pose_checkpoint, + customized_detector_checkpoint=customized_detector_checkpoint, + ) + + # Top-down models typically need a detector for bbox generation. If no detector + # is configured, the returned detector_runner will be None and callers should + # provide bboxes in the pose input context. + if ( + Task(model_cfg["method"]) == Task.TOP_DOWN + and detector_name is None + and customized_detector_checkpoint is None + ): + logging.warning( + "Top-down model configured without a detector. " + "Returning detector_runner=None; pass bboxes in pose input context." + ) + + weight_init = weight_initialization.build_weight_init( + cfg=model_cfg, + super_animal=superanimal_name, + model_name=model_name, + detector_name=detector_name, + with_decoder=False, + memory_replay=False, + customized_pose_checkpoint=customized_pose_checkpoint, + customized_detector_checkpoint=customized_detector_checkpoint, + ) + + pose_runner, detector_runner = get_inference_runners( + model_config=model_cfg, + snapshot_path=weight_init.snapshot_path, + max_individuals=max_individuals, + batch_size=batch_size, + detector_batch_size=detector_batch_size, + detector_path=weight_init.detector_snapshot_path, + ) + return pose_runner, detector_runner, model_cfg diff --git a/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py b/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py index 3c8b09d14c..ede5b803ed 100644 --- a/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py +++ b/deeplabcut/pose_estimation_pytorch/modelzoo/utils.py @@ -21,6 +21,9 @@ from deeplabcut.pose_estimation_pytorch.config.make_pose_config import add_metadata from deeplabcut.utils import auxiliaryfunctions +# COCO category ID for the "person" class. +COCO_PERSON_CATEGORY_ID = 1 + def get_model_configs_folder_path() -> Path: """Returns: the folder containing the SuperAnimal model configuration files""" diff --git a/tests/pose_estimation_pytorch/modelzoo/test_inference_helpers.py b/tests/pose_estimation_pytorch/modelzoo/test_inference_helpers.py new file mode 100644 index 0000000000..c3495a5545 --- /dev/null +++ b/tests/pose_estimation_pytorch/modelzoo/test_inference_helpers.py @@ -0,0 +1,178 @@ +# +# DeepLabCut Toolbox (deeplabcut.org) +# © A. & M.W. Mathis Labs +# https://github.com/DeepLabCut/DeepLabCut +# +# Please see AUTHORS for contributors. +# https://github.com/DeepLabCut/DeepLabCut/blob/main/AUTHORS +# +# Licensed under GNU Lesser General Public License v3.0 +# + +from types import SimpleNamespace + +import pytest + +import deeplabcut.pose_estimation_pytorch.modelzoo.inference_helpers as helpers + + +def _dummy_cfg(method: str = "TD") -> dict: + return { + "method": method, + "metadata": {"bodyparts": ["nose"], "unique_bodyparts": []}, + } + + +def test_create_superanimal_inference_runners_uses_custom_config_path(monkeypatch): + cfg = _dummy_cfg("TD") + read_calls = [] + + def fake_read_config_as_dict(path): + read_calls.append(path) + return cfg + + monkeypatch.setattr(helpers, "read_config_as_dict", fake_read_config_as_dict) + monkeypatch.setattr( + helpers, "update_config", lambda config, max_individuals, device: config + ) + monkeypatch.setattr( + helpers, + "get_inference_runners", + lambda **kwargs: ("pose_runner", "det_runner"), + ) + + import deeplabcut.modelzoo.weight_initialization as wi + + monkeypatch.setattr( + wi, + "build_weight_init", + lambda **kwargs: SimpleNamespace( + snapshot_path="pose.pt", + detector_snapshot_path="det.pt", + ), + ) + + pose_runner, detector_runner, model_cfg = helpers.create_superanimal_inference_runners( + superanimal_name="superanimal_quadruped", + model_name="hrnet_w32", + detector_name="fasterrcnn_resnet50_fpn_v2", + customized_model_config="/tmp/custom_model_cfg.yaml", + ) + + assert read_calls == ["/tmp/custom_model_cfg.yaml"] + assert pose_runner == "pose_runner" + assert detector_runner == "det_runner" + assert model_cfg is cfg + + +def test_create_superanimal_inference_runners_uses_deepcopy_for_custom_dict(monkeypatch): + custom_cfg = _dummy_cfg("TD") + monkeypatch.setattr( + helpers, + "read_config_as_dict", + lambda path: pytest.fail("read_config_as_dict should not be called for dict input"), + ) + + def fake_update_config(config, max_individuals, device): + # Mutate nested structure; caller-owned dict should stay unchanged. + config["metadata"]["bodyparts"].append("tail") + return config + + monkeypatch.setattr(helpers, "update_config", fake_update_config) + monkeypatch.setattr( + helpers, + "get_inference_runners", + lambda **kwargs: ("pose_runner", None), + ) + + import deeplabcut.modelzoo.weight_initialization as wi + + monkeypatch.setattr( + wi, + "build_weight_init", + lambda **kwargs: SimpleNamespace( + snapshot_path="pose.pt", + detector_snapshot_path=None, + ), + ) + + _, _, model_cfg = helpers.create_superanimal_inference_runners( + superanimal_name="superanimal_quadruped", + model_name="hrnet_w32", + detector_name=None, + customized_model_config=custom_cfg, + ) + + assert custom_cfg["metadata"]["bodyparts"] == ["nose"] + assert model_cfg["metadata"]["bodyparts"] == ["nose", "tail"] + + +@pytest.mark.parametrize("input_device", ["auto", None]) +def test_create_superanimal_inference_runners_auto_device_selection( + monkeypatch, input_device +): + cfg = _dummy_cfg("TD") + captured = {} + + monkeypatch.setattr(helpers, "read_config_as_dict", lambda path: cfg) + + def fake_update_config(config, max_individuals, device): + captured["device"] = device + return config + + monkeypatch.setattr(helpers, "update_config", fake_update_config) + monkeypatch.setattr( + helpers, + "get_inference_runners", + lambda **kwargs: ("pose_runner", "det_runner"), + ) + + import deeplabcut.modelzoo.weight_initialization as wi + + monkeypatch.setattr( + wi, + "build_weight_init", + lambda **kwargs: SimpleNamespace( + snapshot_path="pose.pt", + detector_snapshot_path="det.pt", + ), + ) + + helpers.create_superanimal_inference_runners( + superanimal_name="superanimal_quadruped", + model_name="hrnet_w32", + detector_name="fasterrcnn_resnet50_fpn_v2", + customized_model_config="/tmp/custom_model_cfg.yaml", + device=input_device, + ) + assert captured["device"] == "auto" + + +def test_create_superanimal_inference_runners_raises_for_fmpose3d(): + with pytest.raises(NotImplementedError, match="FMPose3D"): + helpers.create_superanimal_inference_runners( + superanimal_name="superanimal_quadruped", + model_name="FMPose3D_resnet", + detector_name="fasterrcnn_resnet50_fpn_v2", + customized_model_config=_dummy_cfg("TD"), + ) + + +def test_create_superanimal_inference_runners_propagates_unsupported_dataset_error( + monkeypatch, +): + monkeypatch.setattr( + helpers, + "load_super_animal_config", + lambda **kwargs: (_ for _ in ()).throw( + ValueError("Unsupported dataset for model zoo config") + ), + ) + + with pytest.raises(ValueError, match="Unsupported dataset"): + helpers.create_superanimal_inference_runners( + superanimal_name="superanimal_unknown", + model_name="hrnet_w32", + detector_name="fasterrcnn_resnet50_fpn_v2", + customized_model_config=None, + ) diff --git a/tests/pose_estimation_pytorch/runners/test_filtered_detector_inference_runner.py b/tests/pose_estimation_pytorch/runners/test_filtered_detector_inference_runner.py index 5d52050026..84489e6626 100644 --- a/tests/pose_estimation_pytorch/runners/test_filtered_detector_inference_runner.py +++ b/tests/pose_estimation_pytorch/runners/test_filtered_detector_inference_runner.py @@ -11,6 +11,7 @@ FilteredDetector, ) from deeplabcut.pose_estimation_pytorch.modelzoo import load_super_animal_config +from deeplabcut.pose_estimation_pytorch.modelzoo.utils import COCO_PERSON_CATEGORY_ID def test_torchvision_detector(): @@ -34,14 +35,15 @@ def test_torchvision_detector(): print("Torchvision detector loaded successfully!") # Test loading the FilteredDetector - COCO_PERSON = 1 # COCO class ID for person - person_detector = FilteredDetector(coco_detector, class_id=COCO_PERSON) + person_detector = FilteredDetector( + coco_detector, class_id=COCO_PERSON_CATEGORY_ID + ) person_detector.eval() print("Filtered detector loaded successfully!") _ = get_filtered_coco_detector_inference_runner( model_name=detector_name, - category_id=COCO_PERSON, + category_id=COCO_PERSON_CATEGORY_ID, batch_size=1, model_config=superanimal_config, )