sleap.nn.inference

Inference pipelines and utilities.

This module contains the classes and high level APIs for predicting instances on new data using trained models.

The inference logic is implemented at two levels:

  • Low-level InferenceModel`s which subclass `tf.keras.Model and implement the core TensorFlow operations surrounding inference. These should only be used when implementing custom inference routines, such as real-time or performance-critical applications. They do not implement tracking (identity association).

  • High-level `Predictor`s which handle data loading, preprocessing, inference, tracking and postprocessing, including converting raw array results into SLEAP-specific data structures. These should be used for general-purpose prediction, including interactive inference and applications that require tracking (identity association).

For more information on tracking, see the sleap.nn.tracking module.

The recommended high-level API for loading saved models is the sleap.load_models function which provides a simplified interface for creating `Predictor`s.

class sleap.nn.inference.BottomUpInferenceLayer(*args, **kwargs)[source]

Keras layer that predicts instances from images using a trained model.

This layer encapsulates all of the inference operations required for generating predictions from a centered instance confidence map model. This includes preprocessing, model forward pass, peak finding and coordinate adjustment.

keras_model

A tf.keras.Model that accepts rank-4 images as input and predicts rank-4 confidence maps and part affinity fields as output.

paf_scorer

A sleap.nn.paf_grouping.PAFScorer instance configured to group instances based on peaks and PAFs produced by the model.

input_scale

Float indicating if the images should be resized before being passed to the model.

cm_output_stride

Output stride of the model, denoting the scale of the output confidence maps relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid. This will be inferred from the model shapes if not provided.

paf_output_stride

Output stride of the model, denoting the scale of the output part affinity fields relative to the images (after input scaling). This is used for adjusting the peak coordinates to the PAF grid. This will be inferred from the model shapes if not provided.

peak_threshold

Minimum confidence map value to consider a global peak as valid.

refinement

If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.

return_confmaps

If True, the confidence maps will be returned together with the predicted instances. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.

return_pafs

If True, the part affinity fields will be returned together with the predicted instances. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.

confmaps_ind

Index of the output tensor of the model corresponding to confidence maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "MultiInstanceConfmapsHead" in its name.

pafs_ind

Index of the output tensor of the model corresponding to part affinity fields. If None (the default), this will be detected automatically by searching for the first tensor that contains "PartAffinityFieldsHead" in its name.

offsets_ind

Index of the output tensor of the model corresponding to offset regression maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "OffsetRefinementHead" in its name. If the head is not present, the method specified in the refinement attribute will be used.

call(data)[source]

Predict instances for one batch of images.

Parameters
  • data – This may be either a single batch of images as a 4-D tensor of shape

  • ` (batch_size, height, width, channels) –

  • key. (image batch in the "images") –

Returns

"instance_peaks": (batch_size, n_instances, n_nodes, 2): Instance skeleton points.

"instance_peak_vals": (batch_size, n_instances, n_nodes): Confidence values for the instance skeleton points.

"instance_scores": (batch_size, n_instances): PAF matching score for each instance.

If BottomUpInferenceLayer.return_confmaps is True, the predicted confidence maps will be returned in the "confmaps" key.

If BottomUpInferenceLayer.return_pafs is True, the predicted PAFs will be returned in the "part_affinity_fields" key.

Return type

The predicted instances as a dictionary of tensors with keys

find_peaks(cms, offsets)[source]

Run peak finding on predicted confidence maps.

forward_pass(data)[source]

Run preprocessing and model inference on a batch.

class sleap.nn.inference.BottomUpInferenceModel(*args, **kwargs)[source]

Bottom-up instance prediction model.

This model encapsulates the bottom-up approach where points are first detected by local peak detection and then grouped into instances by connectivity scoring using part affinity fields.

bottomup_layer

A BottomUpInferenceLayer. This layer takes as input a full image and outputs the predicted instances.

call(example)[source]

Predict instances for one batch of images.

Parameters

example – This may be either a single batch of images as a 4-D tensor of shape (batch_size, height, width, channels), or a dictionary containing the image batch in the "images" key.

Returns

"instance_peaks": (batch_size, n_instances, n_nodes, 2): Instance skeleton

points.

"instance_peak_vals": (batch_size, n_instances, n_nodes): Confidence

values for the instance skeleton points.

"instance_scores": (batch_size, n_instances): PAF matching score for each

instance.

If BottomUpInferenceModel.bottomup_layer.return_confmaps is True, the predicted confidence maps will be returned in the "confmaps" key.

If BottomUpInferenceModel.bottomup_layer.return_pafs is True, the predicted PAFs will be returned in the "part_affinity_fields" key.

Return type

The predicted instances as a dictionary of tensors with keys

class sleap.nn.inference.BottomUpPredictor(bottomup_config: sleap.nn.config.training_job.TrainingJobConfig, bottomup_model: sleap.nn.model.Model, inference_model: Optional[sleap.nn.inference.BottomUpInferenceModel] = None, peak_threshold: float = 0.2, batch_size: int = 4, integral_refinement: bool = True, integral_patch_size: int = 5, max_edge_length_ratio: float = 0.5, paf_line_points: int = 10, min_line_scores: float = 0.25, *, verbosity: str = 'rich', report_rate: float = 2.0, model_paths: List[str] = NOTHING)[source]

Bottom-up multi-instance predictor.

This high-level class handles initialization, preprocessing and tracking using a trained bottom-up multi-instance SLEAP model.

This should be initialized using the from_trained_models() constructor or the high-level API (sleap.load_model).

bottomup_config

The sleap.nn.config.TrainingJobConfig containing the metadata for the trained bottomup model.

Type

sleap.nn.config.training_job.TrainingJobConfig

bottomup_model

A sleap.nn.model.Model instance created from the trained bottomup model. If None, ground truth centroids will be used if available from the data source.

Type

sleap.nn.model.Model

inference_model

A sleap.nn.inference.BottomUpInferenceModel that wraps a trained tf.keras.Model to implement preprocessing, centroid detection, cropping and peak finding.

Type

Optional[sleap.nn.inference.BottomUpInferenceModel]

pipeline

A sleap.nn.data.Pipeline that loads the data and batches input data. This will be updated dynamically if new data sources are used.

Type

Optional[sleap.nn.data.pipelines.Pipeline]

tracker

A sleap.nn.tracking.Tracker that will be called to associate detections over time. Predicted instances will not be assigned to tracks if if this is None.

Type

Optional[sleap.nn.tracking.Tracker]

batch_size

The default batch size to use when loading data for inference. Higher values increase inference speed at the cost of higher memory usage.

Type

int

peak_threshold

Minimum confidence map value to consider a local peak as valid.

Type

float

integral_refinement

If True, peaks will be refined with integral regression. If False, "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

Type

bool

integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.

Type

int

max_edge_length_ratio

The maximum expected length of a connected pair of points in relative image units. Candidate connections above this length will be penalized during matching.

Type

float

paf_line_points

Number of points to sample along the line integral.

Type

int

min_line_scores

Minimum line score (between -1 and 1) required to form a match between candidate point pairs. Useful for rejecting spurious detections when there are no better ones.

Type

float

classmethod from_trained_models(model_path: str, batch_size: int = 4, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, max_edge_length_ratio: float = 0.5, paf_line_points: int = 10, min_line_scores: float = 0.25) sleap.nn.inference.BottomUpPredictor[source]

Create predictor from a saved model.

Parameters
  • model_path – Path to a bottom-up model folder or training job JSON file inside a model folder. This folder should contain training_config.json and best_model.h5 files for a trained model.

  • batch_size – The default batch size to use when loading data for inference. Higher values increase inference speed at the cost of higher memory usage.

  • peak_threshold – Minimum confidence map value to consider a local peak as valid.

  • integral_refinement – If True, peaks will be refined with integral regression. If False, "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

  • integral_patch_size – Size of patches to crop around each rough peak for integral refinement as an integer scalar.

  • max_edge_length_ratio – The maximum expected length of a connected pair of points in relative image units. Candidate connections above this length will be penalized during matching.

  • paf_line_points – Number of points to sample along the line integral.

  • min_line_scores – Minimum line score (between -1 and 1) required to form a match between candidate point pairs. Useful for rejecting spurious detections when there are no better ones.

Returns

An instance of BottomUpPredictor with the loaded model.

class sleap.nn.inference.CentroidCrop(*args, **kwargs)[source]

Inference layer for applying centroid crop-based models.

This layer encapsulates all of the inference operations requires for generating predictions from a centroid confidence map model. This includes preprocessing, model forward pass, peak finding, coordinate adjustment and cropping.

keras_model

A tf.keras.Model that accepts rank-4 images as input and predicts rank-4 confidence maps as output. This should be a model that is trained on centroid/anchor confidence maps.

crop_size

Integer scalar specifying the height/width of the centered crops.

input_scale

Float indicating if the images should be resized before being passed to the model.

pad_to_stride

If not 1, input image will be paded to ensure that it is divisible by this value (after scaling). This should be set to the max stride of the model.

output_stride

Output stride of the model, denoting the scale of the output confidence maps relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid. This will be inferred from the model shapes if not provided.

peak_threshold

Minimum confidence map value to consider a global peak as valid.

refinement

If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.

return_confmaps

If True, the confidence maps will be returned together with the predicted peaks. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.

confmaps_ind

Index of the output tensor of the model corresponding to confidence maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "CentroidConfmapsHead" in its name.

offsets_ind

Index of the output tensor of the model corresponding to offset regression maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "OffsetRefinementHead" in its name. If the head is not present, the method specified in the refinement attribute will be used.

call(inputs)[source]

Predict centroid confidence maps and crop around peaks.

This layer can be chained with a FindInstancePeaks layer to create a top-down inference function from full images.

Parameters

inputs – Full frame images as a tf.Tensor of shape (samples, height, width, channels) or a dictionary with key: "image": Full frame images in the same format as above.

Returns

"crops": Cropped images of shape

(samples, ?, crop_size, crop_size, channels).

"crop_offsets": Coordinates of the top-left of the crops as (x, y)

offsets of shape (samples, ?, 2) for adjusting the predicted peak coordinates.

"centroids": The predicted centroids of shape (samples, ?, 2). "centroid_vals": The centroid confidence values of shape `(samples, ?).

If the return_confmaps attribute is set to True, the output will also contain a key named "centroid_confmaps" containing a tf.RaggedTensor of shape (samples, ?, output_height, output_width, 1) containing the confidence maps predicted by the model.

Return type

A dictionary of outputs grouped by sample with keys

class sleap.nn.inference.CentroidCropGroundTruth(*args, **kwargs)[source]

Keras layer that simulates a centroid cropping model using ground truth.

This layer is useful for testing and evaluating centered instance models.

crop_size

The length of the square box to extract around each centroid.

call(example_gt: Dict[str, tensorflow.python.framework.ops.Tensor]) Dict[str, tensorflow.python.framework.ops.Tensor][source]

Return the ground truth instance crops.

Parameters

example_gt

Dictionary generated from a labels pipeline with the keys: "image": (batch_size, height, width, channels) "centroids": (batch_size, n_centroids, 2): The input centroids.

Axis 1 is expected to be ragged.

These can be generated by the InstanceCentroidFinder transformer.

Returns

"crops": (batch_size, n_centroids, crop_size, crop_size, channels) "crop_offsets": (batch_size, n_centroids, crop_size, crop_size, channels)

These contain the top-left coordinates of each crop in the full images.

"centroids": (batch_size, n_centroids, 2) "centroid_vals": (batch_size, n_centroids)

Axis 1 of all keys are expected to be ragged.

"centroids" are from the input example and "centroid_vals" will be filled with ones.

Return type

Dictionary containing the output of the instance cropping layer with keys

class sleap.nn.inference.FindInstancePeaks(*args, **kwargs)[source]

Keras layer that predicts instance peaks from images using a trained model.

This layer encapsulates all of the inference operations required for generating predictions from a centered instance confidence map model. This includes preprocessing, model forward pass, peak finding and coordinate adjustment.

keras_model

A tf.keras.Model that accepts rank-4 images as input and predicts rank-4 confidence maps as output. This should be a model that is trained on centered instance confidence maps.

input_scale

Float indicating if the images should be resized before being passed to the model.

output_stride

Output stride of the model, denoting the scale of the output confidence maps relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid. This will be inferred from the model shapes if not provided.

peak_threshold

Minimum confidence map value to consider a global peak as valid.

refinement

If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.

return_confmaps

If True, the confidence maps will be returned together with the predicted peaks. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.

confmaps_ind

Index of the output tensor of the model corresponding to confidence maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "CenteredInstanceConfmapsHead" in its name.

offsets_ind

Index of the output tensor of the model corresponding to offset regression maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "OffsetRefinementHead" in its name. If the head is not present, the method specified in the refinement attribute will be used.

call(inputs: Union[Dict[str, tensorflow.python.framework.ops.Tensor], tensorflow.python.framework.ops.Tensor]) Dict[str, tensorflow.python.framework.ops.Tensor][source]

Predict confidence maps and infer peak coordinates.

This layer can be chained with a CentroidCrop layer to create a top-down inference function from full images.

Parameters

inputs

Instance-centered images as a tf.Tensor of shape (samples, height, width, channels) or tf.RaggedTensor of shape (samples, ?, height, width, channels) where images are grouped by sample and may contain a variable number of crops, or a dictionary with keys: "crops": Cropped images in either format above. "crop_offsets": (Optional) Coordinates of the top-left of the crops as

(x, y) offsets of shape (samples, ?, 2) for adjusting the predicted peak coordinates. No adjustment is performed if not provided.

"centroids": (Optional) If provided, will be passed through to the

output.

"centroid_vals": (Optional) If provided, will be passed through to the

output.

Returns

"instance_peaks": The predicted peaks for each instance in the batch as a

tf.RaggedTensor of shape (samples, ?, nodes, 2).

"instance_peak_vals": The value of the confidence maps at the predicted

peaks for each instance in the batch as a tf.RaggedTensor of shape (samples, ?, nodes).

If provided (e.g., from an input CentroidCrop layer), the centroids that generated the crops will also be included in the keys "centroids" and "centroid_vals".

If the return_confmaps attribute is set to True, the output will also contain a key named "instance_confmaps" containing a tf.RaggedTensor of shape (samples, ?, output_height, output_width, nodes) containing the confidence maps predicted by the model.

Return type

A dictionary of outputs with keys

class sleap.nn.inference.FindInstancePeaksGroundTruth(*args, **kwargs)[source]

Keras layer that simulates a centered instance peaks model.

This layer is useful for testing and evaluating centroid models.

call(example_gt: Dict[str, tensorflow.python.framework.ops.Tensor], crop_output: Dict[str, tensorflow.python.framework.ops.Tensor]) Dict[str, tensorflow.python.framework.ops.Tensor][source]

Return the ground truth instance peaks given a set of crops.

Parameters
  • example_gt

    Dictionary generated from a labels pipeline with the key: "instances": (batch_size, n_instances_gt, n_nodes, 2)

    Axes 1 and 2 are expected to be ragged dimensions.

  • crop_output

    Dictionary containing the output of the instance cropping layer with keys: "centroids": (batch_size, n_centroids, 2), "centroid_vals": (batch_size, n_centroids)

    Axis 1 of both keys are expected to be ragged.

Returns

A dictionary with the instance peaks for each frame. The peaks are just the ground truth instances matched to the crop output centroids via greedy matching of the closest node point to each centroid.

The output will have keys:

"centroids": (batch_size, n_centroids, 2): The input centroids. "centroid_vals": (batch_size, n_centroids): The input centroid

confidence values.

"instance_peaks": (batch_size, n_centroids, n_nodes, 2): The matched

instances.

"instance_peak_vals": (batch_size, n_centroids, n_nodes): Peak

confidence values (all 1.0).

class sleap.nn.inference.InferenceLayer(*args, **kwargs)[source]

Base layer for wrapping a Keras model into a layer with preprocessing.

This layer is useful for wrapping input preprocessing operations that would otherwise be handled by a separate pipeline.

This layer expects the same input as the model (rank-4 image) and automatically converts the input to a float if it is in integer form. This can help improve performance by enabling inference directly on uint8 inputs.

The call() method can be overloaded to create custom inference routines that take advantage of the preprocess() method.

keras_model

A tf.keras.Model that will be called on the input to this layer.

input_scale

If not 1.0, input image will be resized by this factor.

pad_to_stride

If not 1, input image will be paded to ensure that it is divisible by this value (after scaling).

ensure_grayscale

If True, converts inputs to grayscale if not already. If False, converts inputs to RGB if not already. If None (default), infer from the shape of the input layer of the model.

call(data: tensorflow.python.framework.ops.Tensor) tensorflow.python.framework.ops.Tensor[source]

Call the model with preprocessed data.

Parameters

data – Inputs to the model.

Returns

Output of the model after being called with preprocessing.

preprocess(imgs: tensorflow.python.framework.ops.Tensor) tensorflow.python.framework.ops.Tensor[source]

Apply all preprocessing operations configured for this layer.

Parameters

imgs – A batch of images as a tensor.

Returns

The input tensor after applying preprocessing operations. The tensor will always be a tf.float32, which will be adjusted to the range [0, 1] if it was previously an integer.

class sleap.nn.inference.InferenceModel(*args, **kwargs)[source]

SLEAP inference model base class.

This class wraps the tf.keras.Model class to provide SLEAP-specific inference utilities such as handling different input data types, preprocessing and variable output shapes.

predict(data: Union[numpy.ndarray, tensorflow.python.framework.ops.Tensor, Dict[str, tensorflow.python.framework.ops.Tensor], tensorflow.python.data.ops.dataset_ops.DatasetV2, sleap.nn.data.pipelines.Pipeline, sleap.io.video.Video], numpy: bool = True, batch_size: int = 4, **kwargs) Union[Dict[str, numpy.ndarray], Dict[str, Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.ops.ragged.ragged_tensor.RaggedTensor]]][source]

Predict instances in the data.

Parameters
  • data

    Input data in any form. Possible types: - np.ndarray, tf.Tensor: Images of shape

    (samples, height, width, channels)

    • dict with key "image" as a tensor

    • tf.data.Dataset that generates examples in one of the above formats.

    • sleap.Pipeline that generates examples in one of the above formats.

    • sleap.Video which will be converted into a pipeline that generates

      batches of batch_size frames.

  • numpy – If True (default), returned values will be converted to `np.ndarray`s or Python primitives if scalars.

  • batch_size – Batch size to use for inference. No effect if using a dataset or pipeline as input since those are expected to generate batches.

Returns

The model outputs as a dictionary of (potentially ragged) tensors or numpy arrays if numpy is True.

If numpy is False, values of the dictionary may be `tf.RaggedTensor`s with the same length for axis 0 (samples), but variable length axis 1 (instances).

If numpy is True and the output contained ragged tensors, they will be NaN-padded to the bounding shape and an additional key "n_valid" will be included to indicate the number of valid elements (before padding) in axis 1 of the tensors.

class sleap.nn.inference.Predictor(*, verbosity: str = 'rich', report_rate: float = 2.0, model_paths: List[str] = NOTHING)[source]

Base interface class for predictors.

classmethod from_model_paths(model_paths: Union[str, List[str]], peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, batch_size: int = 4) sleap.nn.inference.Predictor[source]

Create the appropriate Predictor subclass from a list of model paths.

Parameters
  • model_paths – A single or list of trained model paths.

  • peak_threshold – Minimum confidence map value to consider a peak as valid.

  • integral_refinement – If True, peaks will be refined with integral regression. If False, "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

  • integral_patch_size – Size of patches to crop around each rough peak for integral refinement as an integer scalar.

  • batch_size – The default batch size to use when loading data for inference. Higher values increase inference speed at the cost of higher memory usage.

Returns

A subclass of Predictor.

See also: SingleInstancePredictor, TopDownPredictor, BottomUpPredictor

make_pipeline(data_provider: Optional[Provider] = None) sleap.nn.data.pipelines.Pipeline[source]

Make a data loading pipeline.

Parameters

data_provider – If not None, the pipeline will be created with an instance of a sleap.pipelines.Provider.

Returns

The created sleap.pipelines.Pipeline with batching and prefetching.

Notes

This method also updates the class attribute for the pipeline and will be called automatically when predicting on data from a new source.

predict(data: Union[Provider, sleap.io.dataset.Labels, sleap.io.video.Video], make_labels: bool = True) Union[List[Dict[str, numpy.ndarray]], sleap.io.dataset.Labels][source]

Run inference on a data source.

Parameters
  • data – A sleap.pipelines.Provider, sleap.Labels or sleap.Video to run inference over.

  • make_labels – If True (the default), returns a sleap.Labels instance with sleap.PredictedInstance`s. If `False, just return a list of dictionaries containing the raw arrays returned by the inference model.

Returns

A sleap.Labels with sleap.PredictedInstance`s if `make_labels is True, otherwise a list of dictionaries containing batches of numpy arrays with the raw results.

property report_period: float

Time between progress reports in seconds.

class sleap.nn.inference.RateColumn[source]

Renders the progress rate.

render(task: Task) rich.text.Text[source]

Show progress rate.

class sleap.nn.inference.SingleInstanceInferenceLayer(*args, **kwargs)[source]

Inference layer for applying single instance models.

This layer encapsulates all of the inference operations requires for generating predictions from a single instance confidence map model. This includes preprocessing, model forward pass, peak finding and coordinate adjustment.

keras_model

A tf.keras.Model that accepts rank-4 images as input and predicts rank-4 confidence maps as output. This should be a model that is trained on single instance confidence maps.

input_scale

Float indicating if the images should be resized before being passed to the model.

pad_to_stride

If not 1, input image will be paded to ensure that it is divisible by this value (after scaling). This should be set to the max stride of the model.

output_stride

Output stride of the model, denoting the scale of the output confidence maps relative to the images (after input scaling). This is used for adjusting the peak coordinates to the image grid. This will be inferred from the model shapes if not provided.

peak_threshold

Minimum confidence map value to consider a global peak as valid.

refinement

If None, returns the grid-aligned peaks with no refinement. If "integral", peaks will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.

return_confmaps

If True, the confidence maps will be returned together with the predicted peaks. This will result in slower inference times since the data must be copied off of the GPU, but is useful for visualizing the raw output of the model.

confmaps_ind

Index of the output tensor of the model corresponding to confidence maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "SingleInstanceConfmapsHead" in its name.

offsets_ind

Index of the output tensor of the model corresponding to offset regression maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "OffsetRefinementHead" in its name. If the head is not present, the method specified in the refinement attribute will be used.

call(data)[source]

Predict instance confidence maps and find peaks.

Parameters

inputs – Full frame images as a tf.Tensor of shape (samples, height, width, channels) or a dictionary with key: "image": Full frame images in the same format as above.

Returns

"instance_peaks": The predicted peaks of shape (samples, 1, nodes, 2). "instance_peak_vals": The peak confidence values of shape `(samples, 1, nodes).

If the return_confmaps attribute is set to True, the output will also contain a key named "confmaps" containing a tf.Tensor of shape (samples, output_height, output_width, 1) containing the confidence maps predicted by the model.

Return type

A dictionary of outputs grouped by sample with keys

class sleap.nn.inference.SingleInstanceInferenceModel(*args, **kwargs)[source]

Single instance prediction model.

This model encapsulates the basic single instance approach where it is assumed that there is only one instance in the frame. The images are passed to a peak detector which is trained to detect all body parts for the instance assuming a single peak per body part.

single_instance_layer

A single instance instance peak detection layer. This layer takes as input full images and outputs the detected peaks.

call(example)[source]

Predict instances for one batch of images.

Parameters

example – This may be either a single batch of images as a 4-D tensor of shape (batch_size, height, width, channels), or a dictionary containing the image batch in the "images" key.

Returns

"instance_peaks": (batch_size, 1, n_nodes, 2): Instance skeleton points. "instance_peak_vals": (batch_size, 1, n_nodes): Confidence

values for the instance skeleton points.

Return type

The predicted instances as a dictionary of tensors with keys

class sleap.nn.inference.SingleInstancePredictor(confmap_config: sleap.nn.config.training_job.TrainingJobConfig, confmap_model: sleap.nn.model.Model, inference_model: Optional[sleap.nn.inference.SingleInstanceInferenceModel] = None, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, batch_size: int = 4, *, verbosity: str = 'rich', report_rate: float = 2.0, model_paths: List[str] = NOTHING)[source]

Single instance predictor.

This high-level class handles initialization, preprocessing and tracking using a trained single instance SLEAP model.

This should be initialized using the from_trained_models() constructor or the high-level API (sleap.load_model).

confmap_config

The sleap.nn.config.TrainingJobConfig containing the metadata for the trained model.

Type

sleap.nn.config.training_job.TrainingJobConfig

confmap_model

A sleap.nn.model.Model instance created from the trained model.

Type

sleap.nn.model.Model

inference_model

A sleap.nn.inference.SingleInstanceInferenceModel that wraps a trained tf.keras.Model to implement preprocessing and peak finding.

Type

Optional[sleap.nn.inference.SingleInstanceInferenceModel]

pipeline

A sleap.nn.data.Pipeline that loads the data and batches input data. This will be updated dynamically if new data sources are used.

Type

Optional[sleap.nn.data.pipelines.Pipeline]

peak_threshold

Minimum confidence map value to consider a global peak as valid.

Type

float

integral_refinement

If True, peaks will be refined with integral regression. If False, "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

Type

bool

integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.

Type

int

batch_size

The default batch size to use when loading data for inference. Higher values increase inference speed at the cost of higher memory usage.

Type

int

classmethod from_trained_models(model_path: str, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, batch_size: int = 4) sleap.nn.inference.SingleInstancePredictor[source]

Create the predictor from a saved model.

Parameters
  • model_path – Path to a model folder or training job JSON file inside a model folder. This folder should contain training_config.json and best_model.h5 files for a trained model.

  • peak_threshold – Minimum confidence map value to consider a global peak as valid.

  • integral_refinement – If True, peaks will be refined with integral regression. If False, "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

  • integral_patch_size – Size of patches to crop around each rough peak for integral refinement as an integer scalar.

  • batch_size – The default batch size to use when loading data for inference. Higher values increase inference speed at the cost of higher memory usage.

Returns

An instance of`SingleInstancePredictor` with the models loaded.

class sleap.nn.inference.TopDownInferenceModel(*args, **kwargs)[source]

Top-down instance prediction model.

This model encapsulates the top-down approach where instances are first detected by local peak detection of an anchor point and then cropped. These instance-centered crops are then passed to an instance peak detector which is trained to detect all remaining body parts for the instance that is centered within the crop.

centroid_crop

A centroid cropping layer. This can be either CentroidCrop or CentroidCropGroundTruth. This layer takes the full image as input and outputs a set of centroids and cropped boxes.

instance_peaks

A instance peak detection layer. This can be either FindInstancePeaks or FindInstancePeaksGroundTruth. This layer takes as input the output of the centroid cropper and outputs the detected peaks for the instances within each crop.

call(example: Union[Dict[str, tensorflow.python.framework.ops.Tensor], tensorflow.python.framework.ops.Tensor]) Dict[str, tensorflow.python.framework.ops.Tensor][source]

Predict instances for one batch of images.

Parameters

example – This may be either a single batch of images as a 4-D tensor of shape (batch_size, height, width, channels), or a dictionary containing the image batch in the "images" key. If using a ground truth model for either centroid cropping or instance peaks, the full example from a Pipeline is required for providing the metadata.

Returns

"centroids": (batch_size, n_instances, 2): Instance centroids. "centroid_vals": (batch_size, n_instances): Instance centroid confidence

values.

"instance_peaks": (batch_size, n_instances, n_nodes, 2): Instance skeleton

points.

"instance_peak_vals": (batch_size, n_instances, n_nodes): Confidence

values for the instance skeleton points.

Return type

The predicted instances as a dictionary of tensors with keys

class sleap.nn.inference.TopDownPredictor(centroid_config: Optional[sleap.nn.config.training_job.TrainingJobConfig] = None, centroid_model: Optional[sleap.nn.model.Model] = None, confmap_config: Optional[sleap.nn.config.training_job.TrainingJobConfig] = None, confmap_model: Optional[sleap.nn.model.Model] = None, inference_model: Optional[sleap.nn.inference.TopDownInferenceModel] = None, batch_size: int = 4, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, *, verbosity: str = 'rich', report_rate: float = 2.0, model_paths: List[str] = NOTHING)[source]

Top-down multi-instance predictor.

This high-level class handles initialization, preprocessing and tracking using a trained top-down multi-instance SLEAP model.

This should be initialized using the from_trained_models() constructor or the high-level API (sleap.load_model).

centroid_config

The sleap.nn.config.TrainingJobConfig containing the metadata for the trained centroid model. If None, ground truth centroids will be used if available from the data source.

Type

Optional[sleap.nn.config.training_job.TrainingJobConfig]

centroid_model

A sleap.nn.model.Model instance created from the trained centroid model. If None, ground truth centroids will be used if available from the data source.

Type

Optional[sleap.nn.model.Model]

confmap_config

The sleap.nn.config.TrainingJobConfig containing the metadata for the trained centered instance model. If None, ground truth instances will be used if available from the data source.

Type

Optional[sleap.nn.config.training_job.TrainingJobConfig]

confmap_model

A sleap.nn.model.Model instance created from the trained centered-instance model. If None, ground truth instances will be used if available from the data source.

Type

Optional[sleap.nn.model.Model]

inference_model

A sleap.nn.inference.TopDownInferenceModel that wraps a trained tf.keras.Model to implement preprocessing, centroid detection, cropping and peak finding.

Type

Optional[sleap.nn.inference.TopDownInferenceModel]

pipeline

A sleap.nn.data.Pipeline that loads the data and batches input data. This will be updated dynamically if new data sources are used.

Type

Optional[sleap.nn.data.pipelines.Pipeline]

tracker

A sleap.nn.tracking.Tracker that will be called to associate detections over time. Predicted instances will not be assigned to tracks if if this is None.

Type

Optional[sleap.nn.tracking.Tracker]

batch_size

The default batch size to use when loading data for inference. Higher values increase inference speed at the cost of higher memory usage.

Type

int

peak_threshold

Minimum confidence map value to consider a local peak as valid.

Type

float

integral_refinement

If True, peaks will be refined with integral regression. If False, "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

Type

bool

integral_patch_size

Size of patches to crop around each rough peak for integral refinement as an integer scalar.

Type

int

classmethod from_trained_models(centroid_model_path: Optional[str] = None, confmap_model_path: Optional[str] = None, batch_size: int = 4, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5) sleap.nn.inference.TopDownPredictor[source]

Create predictor from saved models.

Parameters
  • centroid_model_path – Path to a centroid model folder or training job JSON file inside a model folder. This folder should contain training_config.json and best_model.h5 files for a trained model.

  • confmap_model_path – Path to a centered instance model folder or training job JSON file inside a model folder. This folder should contain training_config.json and best_model.h5 files for a trained model.

  • batch_size – The default batch size to use when loading data for inference. Higher values increase inference speed at the cost of higher memory usage.

  • peak_threshold – Minimum confidence map value to consider a local peak as valid.

  • integral_refinement – If True, peaks will be refined with integral regression. If False, "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

  • integral_patch_size – Size of patches to crop around each rough peak for integral refinement as an integer scalar.

Returns

An instance of TopDownPredictor with the loaded models.

One of the two models can be left as None to perform inference with ground truth data. This will only work with LabelsReader as the provider.

make_pipeline(data_provider: Optional[Provider] = None) sleap.nn.data.pipelines.Pipeline[source]

Make a data loading pipeline.

Parameters

data_provider – If not None, the pipeline will be created with an instance of a sleap.pipelines.Provider.

Returns

The created sleap.pipelines.Pipeline with batching and prefetching.

Notes

This method also updates the class attribute for the pipeline and will be called automatically when predicting on data from a new source.

class sleap.nn.inference.VisualPredictor(config: sleap.nn.config.training_job.TrainingJobConfig, model: sleap.nn.model.Model, *, verbosity: str = 'rich', report_rate: float = 2.0, model_paths: List[str] = NOTHING)[source]

Predictor class for generating the visual output of model.

make_pipeline()[source]

Make a data loading pipeline.

Parameters

data_provider – If not None, the pipeline will be created with an instance of a sleap.pipelines.Provider.

Returns

The created sleap.pipelines.Pipeline with batching and prefetching.

Notes

This method also updates the class attribute for the pipeline and will be called automatically when predicting on data from a new source.

predict(data_provider: Provider)[source]

Run inference on a data source.

Parameters
  • data – A sleap.pipelines.Provider, sleap.Labels or sleap.Video to run inference over.

  • make_labels – If True (the default), returns a sleap.Labels instance with sleap.PredictedInstance`s. If `False, just return a list of dictionaries containing the raw arrays returned by the inference model.

Returns

A sleap.Labels with sleap.PredictedInstance`s if `make_labels is True, otherwise a list of dictionaries containing batches of numpy arrays with the raw results.

safely_generate(ds: tensorflow.python.data.ops.dataset_ops.DatasetV2, progress: bool = True)[source]

Yields examples from dataset, catching and logging exceptions.

sleap.nn.inference.find_head(model: tensorflow.python.keras.engine.training.Model, name: str) Optional[int][source]

Return the index of a head in a model’s outputs.

Parameters
  • model – A tf.keras.Model trained by SLEAP.

  • name – A string that is contained in the model output tensor name.

Returns

The index of the first output with a matched name or None if none were found.

Notes

SLEAP model heads are named: - "SingleInstanceConfmapsHead" - "CentroidConfmapsHead" - "CenteredInstanceConfmapsHead" - "MultiInstanceConfmapsHead" - "PartAffinityFieldsHead" - "OffsetRefinementHead"

sleap.nn.inference.get_keras_model_path(path: str) str[source]

Utility method for finding the path to a saved Keras model.

Parameters

path – Path to a model run folder or job file.

Returns

Path to best_model.h5 in the run folder.

sleap.nn.inference.get_model_output_stride(model: tensorflow.python.keras.engine.training.Model, input_ind: int = 0, output_ind: int = - 1) int[source]

Return the stride (1/scale) of the model outputs relative to the input.

Parameters
  • model – A tf.keras.Model.

  • input_ind – The index of the input to use as reference. Defaults to 0, indicating the first input for multi-output models.

  • output_ind – The index of the output to compute the stride for. Defaults to -1, indicating the last output for multi-output models.

Returns

The output stride of the model computed as the integer ratio of the input’s height relative to the output’s height, e.g., for a single input/output model:

model.input.shape[1] // model.output.shape[1]

Raises a warning if the shapes do not divide evenly.

sleap.nn.inference.load_model(model_path: Union[str, List[str]], batch_size: int = 4, peak_threshold: float = 0.2, refinement: str = 'integral', tracker: Optional[str] = None, tracker_window: int = 5, tracker_max_instances: Optional[int] = None, disable_gpu_preallocation: bool = True, progress_reporting: str = 'rich') sleap.nn.inference.Predictor[source]

Load a trained SLEAP model.

Parameters
  • model_path – Path to model or list of path to models that were trained by SLEAP. These should be the directories that contain training_job.json and best_model.h5.

  • batch_size – Number of frames to predict at a time. Larger values result in faster inference speeds, but require more memory.

  • peak_threshold – Minimum confidence map value to consider a peak as valid.

  • refinement – If "integral", peak locations will be refined with integral regression. If "local", peaks will be refined with quarter pixel local gradient offset. This has no effect if the model has an offset regression head.

  • tracker – Name of the tracker to use with the inference model. Must be one of "simple" or "flow". If None, no identity tracking across frames will be performed.

  • tracker_window – Number of frames of history to use when tracking. No effect when tracker is None.

  • tracker_max_instances – If not None, discard instances beyond this count when tracking. No effect when tracker is None.

  • disable_gpu_preallocation – If True (the default), initialize the GPU and disable preallocation of memory. This is necessary to prevent freezing on some systems with low GPU memory and has negligible impact on performance. If False, no GPU initialization is performed. No effect if running in CPU-only mode.

  • progress_reporting – Mode of inference progress reporting. If "rich" (the default), an updating progress bar is displayed in the console or notebook. If "json", a JSON-serialized message is printed out which can be captured for programmatic progress monitoring. If "none", nothing is displayed during inference – this is recommended when running on clusters or headless machines where the output is captured to a log file.

Returns

An instance of a Predictor based on which model type was detected.

If this is a top-down model, paths to the centroids model as well as the centered instance model must be provided. A TopDownPredictor instance will be returned.

If this is a bottom-up model, a BottomUpPredictor will be returned.

If this is a single-instance model, a SingleInstancePredictor will be returned.

If a tracker is specified, the predictor will also run identity tracking over time.

See also: TopDownPredictor, BottomUpPredictor, SingleInstancePredictor

sleap.nn.inference.main()[source]

Entrypoint for sleap-track CLI for running inference.