deeplabcut.pose_estimation_pytorch.modelzoo.inference_helpers

PyTorch-specific helper entrypoints for model zoo inference.

Functions:

Name	Description
create_superanimal_inference_runners	Create SuperAnimal inference runners for in-memory batched inference.

create_superanimal_inference_runners

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: PoseConfig | dict | str | Path | None = None,
    customized_pose_checkpoint: str | Path | None = None,
    customized_detector_checkpoint: str | Path | None = None,
) -> tuple[InferenceRunner, InferenceRunner | None, PoseConfig]

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.

Parameters:

Name	Type	Description	Default
superanimal_name	str	Name of the SuperAnimal dataset, e.g. "superanimal_quadruped".	required
model_name	str	Pose model architecture name, e.g. "hrnet_w32".	required
detector_name	str | None	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.	None
max_individuals	int	Maximum number of individuals to keep per frame.	10
batch_size	int	Batch size for pose inference.	1
detector_batch_size	int	Batch size for detector inference.	1
device	str | None	Device for inference. If "auto" or None, resolves to CUDA when available, else CPU.	'auto'
customized_model_config	PoseConfig | dict | str | Path | None	Optional pose configuration for a custom SuperAnimal model. If not provided, uses the default SuperAnimal configuration. This 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.	None
customized_pose_checkpoint	str | Path | None	Optional custom pose checkpoint path.	None
customized_detector_checkpoint	str | Path | None	Optional custom detector checkpoint path.	None

Returns:

Name	Type	Description
tuple	tuple[InferenceRunner, InferenceRunner | None, PoseConfig]	(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 PoseConfig

Examples:

Run inference on a SuperAnimal model:

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))

Source code in deeplabcut/pose_estimation_pytorch/modelzoo/inference_helpers.py

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: PoseConfig | dict | str | Path | None = None,
    customized_pose_checkpoint: str | Path | None = None,
    customized_detector_checkpoint: str | Path | None = None,
) -> tuple[InferenceRunner, InferenceRunner | None, PoseConfig]:
    """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 (PoseConfig | dict | str | Path | None, optional):
            Optional pose configuration for a custom SuperAnimal model. If not
            provided, uses the default SuperAnimal configuration. This 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 ``PoseConfig``

    Examples:
        Run inference on a SuperAnimal model:

            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:
        model_cfg = PoseConfig.from_any(customized_model_config)
    else:
        model_cfg = load_super_animal_config(
            super_animal=superanimal_name,
            model_name=model_name,
            detector_name=detector_name,
            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