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
[docs]
@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
)