Source code for diplomat.frontends.sleap.load_model

from typing import Optional

import numpy as np

import diplomat.processing.type_casters as tc
from diplomat.frontends import ModelInfo, ModelLike
from diplomat.frontends.sleap.run_utils import _load_configs
from diplomat.frontends.sleap.sleap_providers import PredictorExtractor
from diplomat.utils.cli_tools import Flag
from .sleap_imports import ort


def _build_provider_ordering(device_index: Optional[int], use_cpu: bool):
    supported_devices = ort.get_available_providers()
    device_config = []

    def _add(val, extra=None):
        if extra is None:
            extra = {}
        if device_index is not None:
            extra["device_id"] = device_index
        return (val, extra)

    if not use_cpu:
        if "CUDAExecutionProvider" in supported_devices:
            device_config.append(_add("CUDAExecutionProvider"))
        if "ROCMExecutionProvider" in supported_devices:
            device_config.append(_add("ROCMExecutionProvider"))
        if "CoreMLExecutionProvider" in supported_devices:
            device_config.append("CoreMLExecutionProvider")

    # Fallback...
    device_config.append("CPUExecutionProvider")
    return device_config




@tc.typecaster_function
def load_models(
    config: tc.Union[tc.List[tc.PathLike], tc.PathLike],
    batch_size: tc.Optional[int] = None,
    num_outputs: tc.Optional[int] = None,
    gpu_index: tc.Optional[int] = None,
    refinement_kernel_size: int = 5,
    use_cpu: Flag = False,
) -> tc.Tuple[ModelInfo, ModelLike]:
    """
    Run DIPLOMAT tracking on videos using a SLEAP trained network.

    :param config: The path or list of paths to the SLEAP model folders or config files ("training_config.json").
                   Can also be a zip file containing a single or multiple sleap models, in which case DIPLOMAT will
                   automatically extract models from the zip file.
    :param batch_size: The batch size to use while processing. Defaults to None, which uses the default batch size for the project.
    :param num_outputs: The number of outputs, or bodies to track in the video. Defaults to the value specified in the DLC config, or None if one
                        is not specified.
    :param gpu_index: Integer index of the GPU to use for inference (in tensorflow) defaults to 0, or selecting the first detected GPU if available.
    :param refinement_kernel_size: Size of refinement kernel used for computing offsets if an offset map is not generated by the model. Defaults to 5.
    :param use_cpu: If True, run on cpu even if a gpu is available. Defaults to False.

    :return: A model info dictionary, and a sleap model wrapper that can be used to estimate poses from video frames.
    """
    configs = _load_configs(config)

    provider = PredictorExtractor(
        configs,
        refinement_kernel_size,
        providers=_build_provider_ordering(gpu_index, bool(use_cpu)),
    )

    meta = provider.get_metadata()

    if batch_size is None:
        batch_size = meta["batch_size"]

    return (
        ModelInfo(
            num_outputs=num_outputs,
            batch_size=batch_size,
            dotsize=int(np.ceil(meta["sigma"] / meta["input_scaling"])),
            colormap=None,
            shape_list=None,
            alphavalue=0.7,
            pcutoff=0.1,
            line_thickness=1,
            bp_names=meta["bp_names"],
            skeleton=meta["skeleton"],
            frontend="sleap",
        ),
        provider.extract,
    )