diplomat.frontends.sleap.load_model.load_models

diplomat.frontends.sleap.load_model.load_models(config: List[PathLike | str] | PathLike | str, batch_size: int | None = None, num_outputs: int | None = None, gpu_index: int | None = None, refinement_kernel_size: int = 5, use_cpu: bool = False) → Tuple[ModelInfo, Callable][source]

Run DIPLOMAT tracking on videos using a SLEAP trained network.

Parameters:

  • 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.

  • batch_size – The batch size to use while processing. Defaults to None, which uses the default batch size for the project.

  • 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.

  • gpu_index – Integer index of the GPU to use for inference (in tensorflow) defaults to 0, or selecting the first detected GPU if available.

  • refinement_kernel_size – Size of refinement kernel used for computing offsets if an offset map is not generated by the model. Defaults to 5.

  • use_cpu – If True, run on cpu even if a gpu is available. Defaults to False.

Returns:

A model info dictionary, and a sleap model wrapper that can be used to estimate poses from video frames.