sleap.nn.inference

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.

return_paf_graph#: If True, the part affinity field graph will be returned together with the predicted instances. The graph is obtained by parsing the part affinity fields with the paf_scorer instance and is an intermediate representation used during instance grouping.

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.

If BottomUpInferenceLayer.return_paf_graph is True, the predicted PAF graph will be returned in the "peaks", "peak_vals", "peak_channel_inds", "edge_inds", "edge_peak_inds" and "line_scores" keys.

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.BottomUpMultiClassInferenceLayer(*args, **kwargs)[source]#

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

keras_model#: A tf.keras.Model that accepts rank-4 images as input and predicts rank-4 confidence maps and class maps as output.

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.

class_maps_output_stride#: Output stride of the model, denoting the scale of the output class maps relative to the images (after input scaling). This is used for adjusting the peak coordinates to the class maps 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_class_maps#: If True, the class 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.

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.

class_maps_ind#: Index of the output tensor of the model corresponding to class maps. If None (the default), this will be detected automatically by searching for the first tensor that contains "ClassMapsHead" 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 inference_layer.return_confmaps is True, the predicted confidence maps will be returned in the "confmaps" key.

If inference_layer.return_class_maps is True, the predicted class maps will be returned in the "class_maps" 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.BottomUpMultiClassInferenceModel(*args, **kwargs)[source]#

Bottom-up multi-class instance prediction model.

This model encapsulates the bottom-up multi-class approach where points are first detected by local peak finding and then grouped into instances by their identity classifications.

inference_layer#: A BottomUpMultiClassInferenceLayer. 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 inference_layer.return_confmaps is True, the predicted confidence maps will be returned in the "confmaps" key.

If inference_layer.return_class_maps is True, the predicted class maps will be returned in the "class_maps" key.

Return type:

The predicted instances as a dictionary of tensors with keys

class sleap.nn.inference.BottomUpMultiClassPredictor(config: TrainingJobConfig, model: Model, inference_model: BottomUpMultiClassInferenceModel | None = None, peak_threshold: float = 0.2, batch_size: int = 4, integral_refinement: bool = True, integral_patch_size: int = 5, tracks: List[Track] | None = None, *, 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).

config#

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

Type:: sleap.nn.config.training_job.TrainingJobConfig

model#

A sleap.nn.model.Model instance created from the trained 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.BottomUpMultiClassInferenceModel that wraps a trained tf.keras.Model to implement preprocessing, peak finding and classification.

Type:: sleap.nn.inference.BottomUpMultiClassInferenceModel | None

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:: sleap.nn.data.pipelines.Pipeline | None

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

tracks#

If provided, instances will be created using these track instances. If not, instances will be assigned tracks from the provider if possible.

Type:: List[sleap.instance.Track] | None

Create 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.
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.
resize_input_layer – If True, the the input layer of the tf.Keras.model is resized to (None, None, None, num_color_channels).

Returns:

An instance of BottomUpPredictor with the loaded model.

property is_grayscale: bool#: Return whether the model expects grayscale inputs.

class sleap.nn.inference.BottomUpPredictor(bottomup_config: TrainingJobConfig, bottomup_model: Model, inference_model: BottomUpInferenceModel | None = 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.25, dist_penalty_weight: float = 1.0, paf_line_points: int = 10, min_line_scores: float = 0.25, max_instances: int | None = None, *, 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:: sleap.nn.inference.BottomUpInferenceModel | None

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:: sleap.nn.data.pipelines.Pipeline | None

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:: sleap.nn.tracking.Tracker | None

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#

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 as a fraction of the image size. Candidate connections longer than this length will be penalized during matching.

Type:: float

dist_penalty_weight#

A coefficient to scale weight of the distance penalty as a scalar float. Set to values greater than 1.0 to enforce the distance penalty more strictly.

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

max_instances#

If not None, discard instances beyond this count when predicting, regardless of whether filtering is done at the tracking stage. This is useful for preventing extraneous instances from being created when tracking is not being applied.

Type:: int | None

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.25, dist_penalty_weight: float = 1.0, paf_line_points: int = 10, min_line_scores: float = 0.25, resize_input_layer: bool = True, max_instances: int | None = None) → 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 as a fraction of the image size. Candidate connections longer than this length will be penalized during matching.
dist_penalty_weight – A coefficient to scale weight of the distance penalty as a scalar float. Set to values greater than 1.0 to enforce the distance penalty more strictly.
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.
resize_input_layer – If True, the the input layer of the tf.Keras.model is resized to (None, None, None, num_color_channels).
max_instances – If not None, discard instances beyond this count when predicting, regardless of whether filtering is done at the tracking stage. This is useful for preventing extraneous instances from being created when tracking is not being applied.

Returns:

An instance of BottomUpPredictor with the loaded model.

property is_grayscale: bool#: Return whether the model expects grayscale inputs.

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.

return_crops#: If True, the crops and offsets will be returned together with the predicted peaks. This is true by default since crops are used for finding instance peaks in a top down model. If using a centroid only inference model, this should be set to False.

max_instances#: If set, determines the max number of instances that a multi-instance model returns.

call(inputs)#

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:

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

If the return_crops attribute is set to True, the output will also contain keys named crops and crop_offsets. The former is a tf.RaggedTensor of cropped images of shape (samples, ?, crop_size, crop_size, channels). The latter is a tf.RaggedTensor of Coordinates of the top-left of the crops as (x, y) offsets of shape (samples, ?, 2) for adjusting the predicted peak coordinates.

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, Tensor]) → Dict[str, 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.CentroidInferenceModel(*args, **kwargs)[source]#

Centroid only instance prediction model.

This model encapsulates the first step in a top-down approach where instances are detected by local peak detection of an anchor point and then cropped.

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.

call(example: Dict[str, Tensor] | Tensor) → Dict[str, 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.

Return type:

The predicted instances as a dictionary of tensors with keys

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

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

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: Dict[str, Tensor] | Tensor) → Dict[str, 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, Tensor], crop_output: Dict[str, Tensor]) → Dict[str, 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 padded 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.

ensure_float#: If True, converts inputs to float32 and scales the values to be between 0 and 1.

call(data: Tensor) → 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: Tensor) → 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.

export_model(save_path: str, signatures: str = 'serving_default', save_traces: bool = True, model_name: str | None = None, tensors: Dict[str, str] | None = None, unrag_outputs: bool = True)[source]#

Save the frozen graph of a model.

Parameters:

save_path – Path to output directory to store the frozen graph
signatures – String defining the input and output types for computation.
save_traces – If True (default) the SavedModel will store the function traces for each layer
model_name – (Optional) Name to give the model. If given, will be added to the output json file containing meta information about the model
tensors – (Optional) Dictionary describing the predicted tensors (see sleap.nn.data.utils.describe_tensors as an example)
unrag_outputs – If True (default), any ragged tensors will be converted to normal tensors and padded with NaNs

Notes

This function call writes relevant meta data to an info.json file in the given save_path in addition to the frozen_graph.pb file

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.

predict_on_batch(data: ndarray | Tensor | Dict[str, Tensor], numpy: bool = False, **kwargs) → Dict[str, ndarray] | Dict[str, Tensor | RaggedTensor][source]#

Predict a single batch of samples.

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
numpy – If True (default), returned values will be converted to `np.ndarray`s or Python primitives if scalars.

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

class sleap.nn.inference.MoveNetInferenceLayer(*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.

model_name#: Variable indicating which model of MoveNet to use, either “lightning” or “thunder.”

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.

ensure_grayscale#: Boolean indicating whether the type of input data is grayscale or not.

ensure_float#: Boolean indicating whether the type of data is float or not.

call(ex)[source]#

Call the model with preprocessed data.

Parameters:: data – Inputs to the model.
Returns:: Output of the model after being called with preprocessing.

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

MoveNet prediction model.

This model encapsulates the basic MoveNet approach. The images are passed to a model which is trained to detect all body parts (17 joints in total).

inference_layer#: A MoveNet layer. This layer takes as input full images/videos and

outputs the detected peaks.

call(x)[source]#

Calls the model on new inputs and returns the outputs as tensors.

In this case call() just reapplies all ops in the graph to the new inputs (e.g. build a new computational graph from the provided inputs).

Note: This method should not be called directly. It is only meant to be overridden when subclassing tf.keras.Model. To call a model on an input, always use the __call__() method, i.e. model(inputs), which relies on the underlying call() method.

Parameters:

inputs – Input tensor, or dict/list/tuple of input tensors.
training – Boolean or boolean scalar tensor, indicating whether to run the Network in training mode or inference mode.
mask –
A mask or list of masks. A mask can be either a boolean tensor or None (no mask). For more details, check the guide

[here](https://www.tensorflow.org/guide/keras/masking_and_padding).

Returns:

A tensor if there is a single output, or a list of tensors if there are more than one outputs.

class sleap.nn.inference.MoveNetPredictor(inference_model: MoveNetInferenceModel | None = None, peak_threshold: float = 0.2, batch_size: int = 1, model_name: str = 'lightning', *, verbosity: str = 'rich', report_rate: float = 2.0, model_paths: List[str] = NOTHING)[source]#

MoveNet predictor.

This high-level class handles initialization, preprocessing and tracking using a trained MoveNet model. This should be initialized using the from_trained_models() constructor or the high-level API (sleap.load_model).

inference_model#

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

Type:: sleap.nn.inference.MoveNetInferenceModel | None

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:: sleap.nn.data.pipelines.Pipeline | None

peak_threshold#

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

Type:: float

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

model_name#

Variable indicating which model of MoveNet to use, either “lightning” or “thunder.”

Type:: str

classmethod from_trained_models(model_name: str, peak_threshold: float = 0.2) → MoveNetPredictor[source]#

Create the predictor from a saved model.

Parameters:

model_name – Variable indicating which model of MoveNet to use, either “lightning” or “thunder.”
peak_threshold – Minimum confidence map value to consider a global peak as valid.

Returns:

An instance of`MoveNetPredictor` with the models loaded.

property is_grayscale: bool#: Return whether the model expects grayscale inputs.

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

Base interface class for predictors.

Export a trained SLEAP model as a frozen graph. Initializes model, creates a dummy tracing batch and passes it through the model. The frozen graph is saved along with training meta info.

Parameters:

save_path – Path to output directory to store the frozen graph
signatures – String defining the input and output types for computation.
save_traces – If True (default) the SavedModel will store the function traces for each layer
model_name – (Optional) Name to give the model. If given, will be added to the output json file containing meta information about the model
tensors – (Optional) Dictionary describing the predicted tensors (see sleap.nn.data.utils.describe_tensors as an example)
unrag_outputs – If True (default), any ragged tensors will be converted to normal tensors and padded with NaNs
max_instances – If set, determines the max number of instances that a multi-instance model returns. This is enforced during centroid cropping and therefore only compatible with TopDown models.

classmethod from_model_paths(model_paths: str | List[str], peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, batch_size: int = 4, resize_input_layer: bool = True, max_instances: int | None = None) → Predictor[source]#

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

Parameters:

model_paths – A single or list of trained model paths. Special cases of non-SLEAP models include “movenet-thunder” and “movenet-lightning”.
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.
resize_input_layer – If True, the the input layer of the tf.Keras.model is resized to (None, None, None, num_color_channels).
max_instances – If not None, discard instances beyond this count when predicting, regardless of whether filtering is done at the tracking stage. This is useful for preventing extraneous instances from being created when tracking is not being applied.

Returns:

A subclass of Predictor.

See also: SingleInstancePredictor, TopDownPredictor, BottomUpPredictor,: MoveNetPredictor, TopDownMultiClassPredictor, BottomUpMultiClassPredictor.

abstract property is_grayscale: bool#: Return whether the model expects grayscale inputs.

make_pipeline(data_provider: Provider | None = None) → 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 | Labels | Video, make_labels: bool = True) → List[Dict[str, ndarray]] | 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(table_column: Column | None = None)[source]#

Renders the progress rate.

render(task: Task) → Text[source]#: Show progress rate.

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

Inference layer for applying single instance models.

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: TrainingJobConfig, confmap_model: Model, inference_model: SingleInstanceInferenceModel | None = 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:: sleap.nn.inference.SingleInstanceInferenceModel | None

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:: sleap.nn.data.pipelines.Pipeline | None

peak_threshold#

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

Type:: float

integral_refinement#

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

Export a trained SLEAP model as a frozen graph. Initializes model, creates a dummy tracing batch and passes it through the model. The frozen graph is saved along with training meta info.

Parameters:

save_path – Path to output directory to store the frozen graph
signatures – String defining the input and output types for computation.
save_traces – If True (default) the SavedModel will store the function traces for each layer
model_name – (Optional) Name to give the model. If given, will be added to the output json file containing meta information about the model
tensors – (Optional) Dictionary describing the predicted tensors (see sleap.nn.data.utils.describe_tensors as an example)
unrag_outputs – If True (default), any ragged tensors will be converted to normal tensors and padded with NaNs
max_instances – If set, determines the max number of instances that a multi-instance model returns. This is enforced during centroid cropping and therefore only compatible with TopDown models.

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, resize_input_layer: bool = True) → 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.
resize_input_layer – If True, the the input layer of the tf.Keras.model is resized to (None, None, None, num_color_channels).

Returns:

An instance of`SingleInstancePredictor` with the models loaded.

property is_grayscale: bool#: Return whether the model expects grayscale inputs.

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: Dict[str, Tensor] | Tensor) → Dict[str, Tensor][source]#

Predict instances for one batch of images.

Parameters:

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.TopDownMultiClassFindPeaks(*args, **kwargs)[source]#

Keras layer that predicts and classifies 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 and multi-class model. This includes preprocessing, model forward pass, peak finding, coordinate adjustment, and classification.

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 and classification.

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.

return_class_vectors#: If True, the classification probabilities will be returned together with the predicted peaks. This will not line up with the grouped instances, for which the associtated class probabilities will always be returned in "instance_scores".

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.

class_vectors_ind#: Index of the output tensor of the model corresponding to the classification vectors. If None (the default), this will be detected automatically by searching for the first tensor that contains "ClassVectorsHead" in its name.

optimal_grouping#: If True (the default), group peaks from classification probabilities. If saving a frozen graph of the model, this will be overridden to False.

call(inputs: Dict[str, Tensor] | Tensor) → Dict[str, 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 –

(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.Tensor of shape (samples, n_classes, nodes, 2). Instances will be ordered by class and will be filled with NaN where not found.
"instance_peak_vals": The value of the confidence maps at the predicted: peaks for each instance in the batch as a tf.Tensor of shape (samples, n_classes, 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_class_vectors attribute is set to True, the output will also contain a key named "class_vectors" containing the full classification probabilities for all crops.

If the optimal_grouping attribute is set to True, peaks are grouped from classification properties. This is overridden to False if exporting a frozen graph to allow for tracing. Note: If set to False this will change the output dict keys and shapes.

Return type:

A dictionary of outputs with keys

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

Top-down instance prediction model.

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 and classification layer, an instance of TopDownMultiClassFindPeaks. This layer takes as input the output of the centroid cropper and outputs the detected peaks and classes for the instances within each crop.

call(example: Dict[str, Tensor] | Tensor) → Dict[str, Tensor][source]#

Predict instances for one batch of images.

Parameters:

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

Save the frozen graph of a model.

Parameters:

save_path – Path to output directory to store the frozen graph
signatures – String defining the input and output types for computation.
save_traces – If True (default) the SavedModel will store the function traces for each layer
model_name – (Optional) Name to give the model. If given, will be added to the output json file containing meta information about the model
tensors – (Optional) Dictionary describing the predicted tensors (see sleap.nn.data.utils.describe_tensors as an example)
unrag_outputs – If True (default), any ragged tensors will be converted to normal tensors and padded with NaNs

Notes

This function call writes relevant meta data to an info.json file in the given save_path in addition to the frozen_graph.pb file

class sleap.nn.inference.TopDownMultiClassPredictor(centroid_config: TrainingJobConfig | None = None, centroid_model: Model | None = None, confmap_config: TrainingJobConfig | None = None, confmap_model: Model | None = None, inference_model: TopDownMultiClassInferenceModel | None = None, batch_size: int = 4, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, tracks: List[Track] | None = None, *, verbosity: str = 'rich', report_rate: float = 2.0, model_paths: List[str] = NOTHING)[source]#

Top-down multi-instance predictor with classification.

This high-level class handles initialization, preprocessing and tracking using a trained top-down multi-instance classification 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:: sleap.nn.config.training_job.TrainingJobConfig | None

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:: sleap.nn.model.Model | None

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:: sleap.nn.config.training_job.TrainingJobConfig | None

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:: sleap.nn.model.Model | None

inference_model#

A TopDownMultiClassInferenceModel that wraps a trained tf.keras.Model to implement preprocessing, centroid detection, cropping, peak finding and classification.

Type:: sleap.nn.inference.TopDownMultiClassInferenceModel | None

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:: sleap.nn.data.pipelines.Pipeline | None

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:: sleap.nn.tracking.Tracker | None

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#

Type:: bool

integral_patch_size#

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

Type:: int

tracks#

If provided, instances will be created using these track instances. If not, instances will be assigned tracks from the provider if possible.

Type:: List[sleap.instance.Track] | None

Export a trained SLEAP model as a frozen graph. Initializes model, creates a dummy tracing batch and passes it through the model. The frozen graph is saved along with training meta info.

Parameters:

save_path – Path to output directory to store the frozen graph
signatures – String defining the input and output types for computation.
save_traces – If True (default) the SavedModel will store the function traces for each layer
model_name – (Optional) Name to give the model. If given, will be added to the output json file containing meta information about the model
tensors – (Optional) Dictionary describing the predicted tensors (see sleap.nn.data.utils.describe_tensors as an example)
unrag_outputs – If True (default), any ragged tensors will be converted to normal tensors and padded with NaNs
max_instances – If set, determines the max number of instances that a multi-instance model returns. This is enforced during centroid cropping and therefore only compatible with TopDown models.

classmethod from_trained_models(centroid_model_path: str | None = None, confmap_model_path: str | None = None, batch_size: int = 4, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, resize_input_layer: bool = True) → TopDownMultiClassPredictor[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.
resize_input_layer – If True, the the input layer of the tf.Keras.model is resized to (None, None, None, num_color_channels).

Returns:

An instance of TopDownMultiClassPredictor 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.

property is_grayscale: bool#: Return whether the model expects grayscale inputs.

make_pipeline(data_provider: Provider | None = None) → 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.TopDownPredictor(centroid_config: TrainingJobConfig | None = None, centroid_model: Model | None = None, confmap_config: TrainingJobConfig | None = None, confmap_model: Model | None = None, inference_model: TopDownInferenceModel | None = None, batch_size: int = 4, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, max_instances: int | None = None, *, 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:: sleap.nn.config.training_job.TrainingJobConfig | None

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:: sleap.nn.model.Model | None

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:: sleap.nn.config.training_job.TrainingJobConfig | None

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:: sleap.nn.model.Model | None

inference_model#

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

Type:: sleap.nn.inference.TopDownInferenceModel | None

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:: sleap.nn.data.pipelines.Pipeline | None

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:: sleap.nn.tracking.Tracker | None

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#

Type:: bool

integral_patch_size#

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

Type:: int

max_instances#

Type:: int | None

Export a trained SLEAP model as a frozen graph. Initializes model, creates a dummy tracing batch and passes it through the model. The frozen graph is saved along with training meta info.

Parameters:

save_path – Path to output directory to store the frozen graph
signatures – String defining the input and output types for computation.
save_traces – If True (default) the SavedModel will store the function traces for each layer
model_name – (Optional) Name to give the model. If given, will be added to the output json file containing meta information about the model
tensors – (Optional) Dictionary describing the predicted tensors (see sleap.nn.data.utils.describe_tensors as an example)
unrag_outputs – If True (default), any ragged tensors will be converted to normal tensors and padded with NaNs
max_instances – If set, determines the max number of instances that a multi-instance model returns. This is enforced during centroid cropping and therefore only compatible with TopDown models.

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.
resize_input_layer – If True, the the input layer of the tf.Keras.model is resized to (None, None, None, num_color_channels).
max_instances – If not None, discard instances beyond this count when predicting, regardless of whether filtering is done at the tracking stage. This is useful for preventing extraneous instances from being created when tracking is not being applied.

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.

property is_grayscale: bool#: Return whether the model expects grayscale inputs.

make_pipeline(data_provider: Provider | None = None) → 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: TrainingJobConfig, 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: DatasetV2, progress: bool = True)[source]#: Yields examples from dataset, catching and logging exceptions.

sleap.nn.inference.export_cli(args: list | None = None)[source]#: CLI for sleap-export.

sleap.nn.inference.export_model(model_path: str | List[str], save_path: str = 'exported_model', signatures: str = 'serving_default', save_traces: bool = True, model_name: str | None = None, tensors: Dict[str, str] | None = None, unrag_outputs: bool = True, max_instances: int | None = None)[source]#

High level export of a trained SLEAP model as a frozen graph.

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.
save_path – Path to output directory to store the frozen graph.
signatures – String defining the input and output types for computation.
save_traces – If True (default) the SavedModel will store the function traces for each layer.
model_name – (Optional) Name to give the model. If given, will be added to the output json file containing meta information about the model.
tensors – (Optional) Dictionary describing the predicted tensors (see sleap.nn.data.utils.describe_tensors as an example).
unrag_outputs – If True (default), any ragged tensors will be converted to normal tensors and padded with NaNs
max_instances – If set, determines the max number of instances that a multi-instance model returns. This is enforced during centroid cropping and therefore only compatible with TopDown models.

sleap.nn.inference.find_head(model: Model, name: str) → int | None[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: 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: str | List[str], batch_size: int = 4, peak_threshold: float = 0.2, refinement: str = 'integral', tracker: str | None = None, tracker_window: int = 5, tracker_max_instances: int | None = None, disable_gpu_preallocation: bool = True, progress_reporting: str = 'rich', resize_input_layer: bool = True, max_instances: int | None = None) → 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. Special cases of non-SLEAP models include “movenet-thunder” and “movenet-lightning”.
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, create at most this many tracks.
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.
resize_input_layer – If True, the the input layer of the tf.Keras.model is resized to (None, None, None, num_color_channels).
max_instances – If not None, discard instances beyond this count when predicting, regardless of whether filtering is done at the tracking stage. This is useful for preventing extraneous instances from being created when tracking is not being applied.

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(args: list | None = None)[source]#

Entrypoint for sleap-track CLI for running inference.

Parameters:: args – A list of arguments to be passed into sleap-track.

sleap.nn.inference.make_model_movenet(model_name: str) → Model[source]#

Load a MoveNet model by name.

Parameters:: model_name – Name of the model (“lightning” or “thunder”)
Returns:: A tf.keras.Model ready for inference.

sleap.nn.inference

Contents

sleap.nn.inference#