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 therefinement
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
isTrue
, the predicted confidence maps will be returned in the"confmaps"
key.If
BottomUpInferenceLayer.return_pafs
isTrue
, 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.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 skeletonpoints.
"instance_peak_vals": (batch_size, n_instances, n_nodes)
: Confidencevalues for the instance skeleton points.
"instance_scores": (batch_size, n_instances)
: PAF matching score for eachinstance.
If
BottomUpInferenceModel.bottomup_layer.return_confmaps
isTrue
, the predicted confidence maps will be returned in the"confmaps"
key.If
BottomUpInferenceModel.bottomup_layer.return_pafs
isTrue
, 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 therefinement
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
isTrue
, the predicted confidence maps will be returned in the"confmaps"
key.If
inference_layer.return_class_maps
isTrue
, 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.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 skeletonpoints.
"instance_peak_vals": (batch_size, n_instances, n_nodes)
: Confidencevalues for the instance skeleton points.
"instance_scores": (batch_size, n_instances)
: PAF matching score for eachinstance.
If
inference_layer.return_confmaps
isTrue
, the predicted confidence maps will be returned in the"confmaps"
key.If
inference_layer.return_class_maps
isTrue
, 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.
- model#
A
sleap.nn.model.Model
instance created from the trained model. IfNone
, ground truth centroids will be used if available from the data source.- Type
- inference_model#
A
sleap.nn.inference.BottomUpMultiClassInferenceModel
that wraps a trainedtf.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. IfFalse
,"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, resize_input_layer: bool = True) 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
andbest_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. IfFalse
,"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: 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.
- bottomup_model#
A
sleap.nn.model.Model
instance created from the trained bottomup model. IfNone
, ground truth centroids will be used if available from the data source.- Type
- inference_model#
A
sleap.nn.inference.BottomUpInferenceModel
that wraps a trainedtf.keras.Model
to implement preprocessing, centroid detection, cropping and peak finding.- Type
- 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 isNone
.- 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. IfFalse
,"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, resize_input_layer: bool = True) 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
andbest_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. IfFalse
,"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).
- 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 therefinement
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 toFalse
.
- 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 toTrue
, the output will also contain a key named"centroid_confmaps"
containing atf.RaggedTensor
of shape(samples, ?, output_height, output_width, 1)
containing the confidence maps predicted by the model.If the
return_crops
attribute is set toTrue
, the output will also contain keys namedcrops
andcrop_offsets
. The former is atf.RaggedTensor
of cropped images of shape(samples, ?, crop_size, crop_size, channels)
. The latter is atf.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
orCentroidCropGroundTruth
. 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 aPipeline
is required for providing the metadata.- Returns
"centroids": (batch_size, n_instances, 2)
: Instance centroids."centroid_vals": (batch_size, n_instances)
: Instance centroid confidencevalues.
- 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 therefinement
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)
ortf.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 theoutput.
"centroid_vals"
: (Optional) If provided, will be passed through to theoutput.
- Returns
"instance_peaks"
: The predicted peaks for each instance in the batch as atf.RaggedTensor
of shape(samples, ?, nodes, 2)
."instance_peak_vals"
: The value of the confidence maps at the predictedpeaks 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 toTrue
, the output will also contain a key named"instance_confmaps"
containing atf.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 centroidconfidence values.
"instance_peaks": (batch_size, n_centroids, n_nodes, 2)
: The matchedinstances.
"instance_peak_vals": (batch_size, n_centroids, n_nodes)
: Peakconfidence 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 thepreprocess()
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. IfFalse
, converts inputs to RGB if not already. IfNone
(default), infer from the shape of the input layer of the model.
- ensure_float#
If
True
, converts inputs tofloat32
and scales the values to be between 0 and 1.
- 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, 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 layermodel_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(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 tensortf.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 generatesbatches 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
isTrue
.If
numpy
isFalse
, 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
isTrue
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
isTrue
.If
numpy
isFalse
, 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
isTrue
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.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.
- 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.
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 underlyingcall
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 tensor or None (no mask).
- 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: Optional[sleap.nn.inference.MoveNetInferenceModel] = 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 trainedtf.keras.Model
to implement preprocessing and peak finding.- Type
Optional[sleap.nn.inference.MoveNetInferenceModel]
- 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
- 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) sleap.nn.inference.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_model(save_path: str, signatures: str = 'serving_default', save_traces: bool = True, model_name: Optional[str] = None, tensors: Optional[Dict[str, str]] = None, unrag_outputs: bool = True, max_instances: Optional[int] = 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 layermodel_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 NaNsmax_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: Union[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) 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. 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. IfFalse
,"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
A subclass of
Predictor
.
- 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 asleap.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
orsleap.Video
to run inference over.make_labels – If
True
(the default), returns asleap.Labels
instance withsleap.PredictedInstance`s. If `False
, just return a list of dictionaries containing the raw arrays returned by the inference model.
- Returns
A
sleap.Labels
withsleap.PredictedInstance`s if `make_labels
isTrue
, 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.
- 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 therefinement
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 toTrue
, the output will also contain a key named"confmaps"
containing atf.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)
: Confidencevalues 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.
- confmap_model#
A
sleap.nn.model.Model
instance created from the trained model.- Type
- inference_model#
A
sleap.nn.inference.SingleInstanceInferenceModel
that wraps a trainedtf.keras.Model
to implement preprocessing and peak finding.- Type
- 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. IfFalse
,"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, unrag_outputs: bool = True, max_instances: Optional[int] = 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 layermodel_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 NaNsmax_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) 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
andbest_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. IfFalse
,"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
orCentroidCropGroundTruth
. 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
orFindInstancePeaksGroundTruth
. 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 aPipeline
is required for providing the metadata.- Returns
"centroids": (batch_size, n_instances, 2)
: Instance centroids."centroid_vals": (batch_size, n_instances)
: Instance centroid confidencevalues.
"instance_peaks": (batch_size, n_instances, n_nodes, 2)
: Instance skeletonpoints.
"instance_peak_vals": (batch_size, n_instances, n_nodes)
: Confidencevalues 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 therefinement
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 toFalse
.
- 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)
ortf.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 theoutput.
"centroid_vals"
: (Optional) If provided, will be passed through to theoutput.
- Returns
"instance_peaks"
: The predicted peaks for each instance in the batch as atf.Tensor
of shape(samples, n_classes, nodes, 2)
. Instances will be ordered by class and will be filled withNaN
where not found."instance_peak_vals"
: The value of the confidence maps at the predictedpeaks 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 toTrue
, the output will also contain a key named"instance_confmaps"
containing atf.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 toTrue
, 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 toTrue
, 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
orCentroidCropGroundTruth
. 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 aPipeline
is required for providing the metadata.- Returns
"centroids": (batch_size, n_instances, 2)
: Instance centroids."centroid_vals": (batch_size, n_instances)
: Instance centroid confidencevalues.
"instance_peaks": (batch_size, n_instances, n_nodes, 2)
: Instance skeletonpoints.
"instance_peak_vals": (batch_size, n_instances, n_nodes)
: Confidencevalues 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, 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 layermodel_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: 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. IfNone
, ground truth centroids will be used if available from the data source.- Type
- centroid_model#
A
sleap.nn.model.Model
instance created from the trained centroid model. IfNone
, 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. IfNone
, ground truth instances will be used if available from the data source.- Type
- confmap_model#
A
sleap.nn.model.Model
instance created from the trained centered-instance model. IfNone
, 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 trainedtf.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 isNone
.- 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. IfFalse
,"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, unrag_outputs: bool = True, max_instances: Optional[int] = 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 layermodel_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 NaNsmax_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: 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, resize_input_layer: bool = True) 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
andbest_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
andbest_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. IfFalse
,"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
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 withLabelsReader
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 asleap.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. IfNone
, ground truth centroids will be used if available from the data source.- Type
- centroid_model#
A
sleap.nn.model.Model
instance created from the trained centroid model. IfNone
, 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. IfNone
, ground truth instances will be used if available from the data source.- Type
- confmap_model#
A
sleap.nn.model.Model
instance created from the trained centered-instance model. IfNone
, 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 trainedtf.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 isNone
.- 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. IfFalse
,"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, unrag_outputs: bool = True, max_instances: Optional[int] = 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 layermodel_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 NaNsmax_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: 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, resize_input_layer: bool = True) 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
andbest_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
andbest_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. IfFalse
,"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
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 withLabelsReader
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 asleap.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 asleap.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
orsleap.Video
to run inference over.make_labels – If
True
(the default), returns asleap.Labels
instance withsleap.PredictedInstance`s. If `False
, just return a list of dictionaries containing the raw arrays returned by the inference model.
- Returns
A
sleap.Labels
withsleap.PredictedInstance`s if `make_labels
isTrue
, otherwise a list of dictionaries containing batches of numpy arrays with the raw results.
- 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, unrag_outputs: bool = True, max_instances: Optional[int] = 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
andbest_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 NaNsmax_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: 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', resize_input_layer: bool = True) 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
andbest_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"
. IfNone
, no identity tracking across frames will be performed.tracker_window – Number of frames of history to use when tracking. No effect when
tracker
isNone
.tracker_max_instances – If not
None
, discard instances beyond this count when tracking. No effect whentracker
isNone
.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. IfFalse
, 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).
- 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