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

Optional[sleap.nn.inference.BottomUpMultiClassInferenceModel]

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]

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

Optional[List[sleap.instance.Track]]

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) sleap.nn.inference.BottomUpMultiClassPredictor[source]#

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.

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: 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.25, dist_penalty_weight: float = 1.0, 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 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

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

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.

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, 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.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: 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.

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.

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.

export_model(save_path: str, signatures: str = 'serving_default', save_traces: bool = True, model_name: Optional[str] = None, tensors: Optional[Dict[str, str]] = None)[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)

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

predict_on_batch(data: Union[numpy.ndarray, tensorflow.python.framework.ops.Tensor, Dict[str, tensorflow.python.framework.ops.Tensor]], numpy: bool = False, **kwargs) Union[Dict[str, numpy.ndarray], Dict[str, Union[tensorflow.python.framework.ops.Tensor, tensorflow.python.ops.ragged.ragged_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).

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.

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

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)

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

abstract property is_grayscale: bool#

Return whether the model expects grayscale inputs.

make_pipeline(data_provider: Optional[sleap.nn.data.pipelines.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[sleap.nn.data.pipelines.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(table_column: Optional[rich.table.Column] = None)[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

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

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)

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.

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: 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.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: 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.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_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.

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.

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 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: 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

export_model(save_path: str, signatures: str = 'serving_default', save_traces: bool = True, model_name: Optional[str] = None, tensors: Optional[Dict[str, str]] = None)[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)

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: 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.TopDownMultiClassInferenceModel] = None, batch_size: int = 4, peak_threshold: float = 0.2, integral_refinement: bool = True, integral_patch_size: int = 5, tracks: Optional[List[sleap.instance.Track]] = 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

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 TopDownMultiClassInferenceModel that wraps a trained tf.keras.Model to implement preprocessing, centroid detection, cropping, peak finding and classification.

Type

Optional[sleap.nn.inference.TopDownMultiClassInferenceModel]

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

tracks#

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

Type

Optional[List[sleap.instance.Track]]

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

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)

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

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: Optional[sleap.nn.data.pipelines.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.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

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

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)

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.

property is_grayscale: bool#

Return whether the model expects grayscale inputs.

make_pipeline(data_provider: Optional[sleap.nn.data.pipelines.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: sleap.nn.data.pipelines.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.export_cli()[source]#

CLI for sleap-export.

sleap.nn.inference.export_model(model_path: Union[str, List[str]], save_path: str = 'exported_model', signatures: str = 'serving_default', save_traces: bool = True, model_name: Optional[str] = None, tensors: Optional[Dict[str, str]] = 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).

sleap.nn.inference.find_head(model: 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: 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(args: Optional[list] = None)[source]#

Entrypoint for sleap-track CLI for running inference.

Parameters

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