pytorch_model_zoo_export.py

import warnings
from pathlib import Path
from collections import OrderedDict

import torch

from dlclive.modelzoo.utils import load_super_animal_config, download_super_animal_snapshot


def export_modelzoo_model(
    export_path: str | Path,
    super_animal: str,
    model_name: str,
    detector_name: str | None = None,
) -> None:
    """

    """
    Path(export_path).parent.mkdir(parents=True, exist_ok=True)
    if Path(export_path).exists():
        warnings.warn(f"Export path {export_path} already exists, skipping export", UserWarning)
        return

    model_cfg = load_super_animal_config(
        super_animal=super_animal,
        model_name=model_name,
        detector_name=detector_name,
    )

    def _load_model_weights(model_name: str, super_animal: str = super_animal) -> OrderedDict:
        """Download the model weights from huggingface and load them in torch state dict"""
        checkpoint: Path = download_super_animal_snapshot(dataset=super_animal, model_name=model_name)
        return torch.load(checkpoint, map_location="cpu", weights_only=True)["model"]
    
    # Skip downloading the detector weights for humanbody models, as they are not on huggingface
    skip_detector_download = (detector_name is None) or (super_animal == "superanimal_humanbody")
    export_dict = {
        "config": model_cfg,
        "pose": _load_model_weights(model_name),
        "detector": None if skip_detector_download else _load_model_weights(detector_name),
    }
    torch.save(export_dict, export_path)


if __name__ == "__main__":
    """Example usage"""	
    from utils import _MODELZOO_PATH
    
    model_name = "resnet_50"
    super_animal = "superanimal_quadruped"

    export_modelzoo_model(
        export_path=_MODELZOO_PATH / 'exported_models' / f'exported_{super_animal}_{model_name}.pt',
        super_animal=super_animal,
        model_name=model_name,
    )