deeplabcut.pose_estimation_pytorch.modelzoo.utils

Functions:

Name	Description
download_super_animal_snapshot	Downloads a SuperAnimal snapshot.
get_gpu_memory_map	Get the current gpu usage.
get_model_configs_folder_path	Returns: the folder containing the SuperAnimal model configuration files
get_project_configs_folder_path	Returns: the folder containing the SuperAnimal project configuration files
get_snapshot_folder_path	Returns: the path to the folder containing the SuperAnimal model snapshots
get_super_animal_model_config_path	Gets the path to the configuration file for a SuperAnimal model.
get_super_animal_project_config_path	Gets the path to a SuperAnimal project configuration file.
get_super_animal_snapshot_path	Gets the path to the snapshot containing SuperAnimal model weights.
load_super_animal_config	Loads the model configuration file for a model, detector and SuperAnimal.
update_config	Loads the model configuration file for a model, detector and SuperAnimal.

download_super_animal_snapshot

download_super_animal_snapshot(dataset: str, model_name: str) -> Path

Downloads a SuperAnimal snapshot.

Parameters:

Name	Type	Description	Default
dataset	str	The name of the SuperAnimal dataset for which to download a snapshot.	required
model_name	str	The name of the model for which to download a snapshot.	required

Returns:

Type	Description
Path	The path to the downloaded snapshot.

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

def download_super_animal_snapshot(dataset: str, model_name: str) -> Path:
    """Downloads a SuperAnimal snapshot.

    Args:
        dataset: The name of the SuperAnimal dataset for which to download a snapshot.
        model_name: The name of the model for which to download a snapshot.

    Returns:
        The path to the downloaded snapshot.

    Raises:
        RuntimeError if the model fails to download.
    """
    snapshot_dir = get_snapshot_folder_path()
    model_name = f"{dataset}_{model_name}"
    model_filename = f"{model_name}.pt"
    model_path = snapshot_dir / model_filename

    download_huggingface_model(
        model_name,
        target_dir=str(snapshot_dir),
        rename_mapping={model_filename: model_filename},
    )
    if not model_path.exists():
        raise RuntimeError(f"Failed to download {model_name} to {model_path}")

    return snapshot_dir / f"{model_name}.pt"

get_gpu_memory_map

get_gpu_memory_map()

Get the current gpu usage.

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

def get_gpu_memory_map():
    """Get the current gpu usage."""
    result = subprocess.check_output(
        ["nvidia-smi", "--query-gpu=memory.free", "--format=csv,nounits,noheader"],
        encoding="utf-8",
    )
    gpu_memory = [int(x) for x in result.strip().split("\n")]
    gpu_memory_map = dict(zip(range(len(gpu_memory)), gpu_memory, strict=False))

    return gpu_memory_map

get_model_configs_folder_path

get_model_configs_folder_path() -> Path

Returns: the folder containing the SuperAnimal model configuration files

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

def get_model_configs_folder_path() -> Path:
    """Returns: the folder containing the SuperAnimal model configuration files"""
    return Path(auxiliaryfunctions.get_deeplabcut_path()) / "modelzoo" / "model_configs"

get_project_configs_folder_path

get_project_configs_folder_path() -> Path

Returns: the folder containing the SuperAnimal project configuration files

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

def get_project_configs_folder_path() -> Path:
    """Returns: the folder containing the SuperAnimal project configuration files"""
    return Path(auxiliaryfunctions.get_deeplabcut_path()) / "modelzoo" / "project_configs"

get_snapshot_folder_path

get_snapshot_folder_path() -> Path

Returns: the path to the folder containing the SuperAnimal model snapshots

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

def get_snapshot_folder_path() -> Path:
    """Returns: the path to the folder containing the SuperAnimal model snapshots"""
    return Path(auxiliaryfunctions.get_deeplabcut_path()) / "modelzoo" / "checkpoints"

get_super_animal_model_config_path

get_super_animal_model_config_path(model_name: str) -> Path

Gets the path to the configuration file for a SuperAnimal model.

Parameters:

Name	Type	Description	Default
model_name	str	The name of the model for which to get the path.	required

Returns:

Type	Description
Path	The path to the config file for a SuperAnimal model.

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

def get_super_animal_model_config_path(model_name: str) -> Path:
    """Gets the path to the configuration file for a SuperAnimal model.

    Args:
        model_name: The name of the model for which to get the path.

    Returns:
        The path to the config file for a SuperAnimal model.
    """
    return get_model_configs_folder_path() / f"{model_name}.yaml"

get_super_animal_project_config_path

get_super_animal_project_config_path(super_animal: str) -> Path

Gets the path to a SuperAnimal project configuration file.

Parameters:

Name	Type	Description	Default
super_animal	str	The name of the SuperAnimal for which to get the config path.	required

Returns:

Type	Description
Path	The path to the config file for a SuperAnimal project.

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

def get_super_animal_project_config_path(super_animal: str) -> Path:
    """Gets the path to a SuperAnimal project configuration file.

    Args:
        super_animal: The name of the SuperAnimal for which to get the config path.

    Returns:
        The path to the config file for a SuperAnimal project.
    """
    return get_project_configs_folder_path() / f"{super_animal}.yaml"

get_super_animal_snapshot_path

get_super_animal_snapshot_path(dataset: str, model_name: str, download: bool = True) -> Path

Gets the path to the snapshot containing SuperAnimal model weights.

Parameters:

Name	Type	Description	Default
dataset	str	The name of the SuperAnimal dataset.	required
model_name	str	The name of the model.	required
download	bool	Whether to download the weights if they aren't already there.	True

Returns:

Type	Description
Path	The path to the weights for a SuperAnimal model.

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

def get_super_animal_snapshot_path(
    dataset: str,
    model_name: str,
    download: bool = True,
) -> Path:
    """Gets the path to the snapshot containing SuperAnimal model weights.

    Args:
        dataset: The name of the SuperAnimal dataset.
        model_name: The name of the model.
        download: Whether to download the weights if they aren't already there.

    Returns:
        The path to the weights for a SuperAnimal model.
    """
    model_path = get_snapshot_folder_path() / f"{dataset}_{model_name}.pt"
    if download and not model_path.exists():
        download_super_animal_snapshot(dataset, model_name)

    return model_path

load_super_animal_config

load_super_animal_config(
    super_animal: str,
    model_name: str,
    detector_name: str | None = None,
    max_individuals: int = 30,
    device: str | None = None,
) -> dict

Loads the model configuration file for a model, detector and SuperAnimal.

Parameters:

Name	Type	Description	Default
super_animal	str	The name of the SuperAnimal for which to create the model config.	required
model_name	str	The name of the model for which to create the model config.	required
detector_name	str | None	The name of the detector for which to create the model config.	None
max_individuals	int	The maximum number of detections to make in an image	30
device	str | None	The device to use to train/run inference on the model	None

Returns:

Type	Description
dict	The model configuration for a SuperAnimal-pretrained model.

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

def load_super_animal_config(
    super_animal: str,
    model_name: str,
    detector_name: str | None = None,
    max_individuals: int = 30,
    device: str | None = None,
) -> dict:
    """Loads the model configuration file for a model, detector and SuperAnimal.

    Args:
        super_animal: The name of the SuperAnimal for which to create the model config.
        model_name: The name of the model for which to create the model config.
        detector_name: The name of the detector for which to create the model config.
        max_individuals: The maximum number of detections to make in an image
        device: The device to use to train/run inference on the model

    Returns:
        The model configuration for a SuperAnimal-pretrained model.
    """
    project_cfg_path = get_super_animal_project_config_path(super_animal=super_animal)
    project_config = read_config_as_dict(project_cfg_path)

    model_cfg_path = get_super_animal_model_config_path(model_name=model_name)
    model_config = read_config_as_dict(model_cfg_path)
    model_config = add_metadata(project_config, model_config, model_cfg_path)
    model_config = update_config(model_config, max_individuals, device)

    if detector_name is None and super_animal != "superanimal_humanbody":
        model_config["method"] = "BU"
    else:
        model_config["method"] = "TD"
        if super_animal != "superanimal_humanbody":
            detector_cfg_path = get_super_animal_model_config_path(model_name=detector_name)
            detector_cfg = read_config_as_dict(detector_cfg_path)
            model_config["detector"] = detector_cfg
    return model_config

update_config

update_config(config: dict, max_individuals: int, device: str)

Loads the model configuration file for a model, detector and SuperAnimal.

Parameters:

Name	Type	Description	Default
config	dict	The default model configuration file.	required
max_individuals	int	The maximum number of detections to make in an image	required
device	str	The device to use to train/run inference on the model	required

Returns:

Type	Description
	The model configuration for a SuperAnimal-pretrained model.

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

def update_config(config: dict, max_individuals: int, device: str):
    """Loads the model configuration file for a model, detector and SuperAnimal.

    Args:
        config: The default model configuration file.
        max_individuals: The maximum number of detections to make in an image
        device: The device to use to train/run inference on the model

    Returns:
        The model configuration for a SuperAnimal-pretrained model.
    """
    config = config_utils.replace_default_values(
        config,
        num_bodyparts=len(config["metadata"]["bodyparts"]),
        num_individuals=max_individuals,
        backbone_output_channels=config["model"]["backbone_output_channels"],
    )
    config["metadata"]["individuals"] = [f"animal{i}" for i in range(max_individuals)]

    config["device"] = device
    if config.get("detector", None) is not None:
        config["detector"]["device"] = device

    return config