sleap.nn.data.pipelines

Contents

sleap.nn.data.pipelines#

This module defines high level pipeline configurations from providers/transformers.

The Pipeline class has the capability to create sequences of data I/O and processing operations wrapped in a tf.data-based pipeline.

This allows for convenient ways to configure individual variants of common pipelines, as well as to define training vs inference versions based on the same configurations.

class sleap.nn.data.pipelines.BottomUpMultiClassPipeline(data_config: DataConfig, optimization_config: OptimizationConfig, confmaps_head: MultiInstanceConfmapsHead, class_maps_head: ClassMapsHead, offsets_head: OffsetRefinementHead | None = None)[source]#

Pipeline builder for confidence maps and class maps models.

data_config#

Data-related configuration.

Type:

sleap.nn.config.data.DataConfig

optimization_config#

Optimization-related configuration.

Type:

sleap.nn.config.optimization.OptimizationConfig

confmaps_head#

Instantiated head describing the output confidence maps tensor.

Type:

sleap.nn.heads.MultiInstanceConfmapsHead

class_maps_head#

Instantiated head describing the output class maps tensor.

Type:

sleap.nn.heads.ClassMapsHead

offsets_head#

Optional head describing the offset refinement maps.

Type:

sleap.nn.heads.OffsetRefinementHead | None

make_base_pipeline(data_provider: Provider) Pipeline[source]#

Create base pipeline with input data only.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce input examples.

make_training_pipeline(data_provider: Provider) Pipeline[source]#

Create full training pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce all data keys required for training.

Notes

This does not remap keys to model outputs. Use KeyMapper to pull out keys with the appropriate format for the instantiated tf.keras.Model.

make_viz_pipeline(data_provider: Provider, keras_model: Model) Pipeline[source]#

Create visualization pipeline.

Parameters:
  • data_provider – A Provider that generates data examples, typically a LabelsReader instance.

  • keras_model – A tf.keras.Model that can be used for inference.

Returns:

A Pipeline instance configured to fetch data and run inference to generate predictions useful for visualization during training.

class sleap.nn.data.pipelines.BottomUpPipeline(data_config: DataConfig, optimization_config: OptimizationConfig, confmaps_head: MultiInstanceConfmapsHead, pafs_head: PartAffinityFieldsHead, offsets_head: OffsetRefinementHead | None = None)[source]#

Pipeline builder for confidence maps + part affinity fields models.

data_config#

Data-related configuration.

Type:

sleap.nn.config.data.DataConfig

optimization_config#

Optimization-related configuration.

Type:

sleap.nn.config.optimization.OptimizationConfig

confmaps_head#

Instantiated head describing the output confidence maps tensor.

Type:

sleap.nn.heads.MultiInstanceConfmapsHead

pafs_head#

Instantiated head describing the output PAFs tensor.

Type:

sleap.nn.heads.PartAffinityFieldsHead

offsets_head#

Optional head describing the offset refinement maps.

Type:

sleap.nn.heads.OffsetRefinementHead | None

make_base_pipeline(data_provider: Provider) Pipeline[source]#

Create base pipeline with input data only.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce input examples.

make_training_pipeline(data_provider: Provider) Pipeline[source]#

Create full training pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce all data keys required for training.

Notes

This does not remap keys to model outputs. Use KeyMapper to pull out keys with the appropriate format for the instantiated tf.keras.Model.

make_viz_pipeline(data_provider: Provider, keras_model: Model) Pipeline[source]#

Create visualization pipeline.

Parameters:
  • data_provider – A Provider that generates data examples, typically a LabelsReader instance.

  • keras_model – A tf.keras.Model that can be used for inference.

Returns:

A Pipeline instance configured to fetch data and run inference to generate predictions useful for visualization during training.

class sleap.nn.data.pipelines.CentroidConfmapsPipeline(data_config: DataConfig, optimization_config: OptimizationConfig, centroid_confmap_head: CentroidConfmapsHead, offsets_head: OffsetRefinementHead | None = None)[source]#

Pipeline builder for centroid confidence map models.

data_config#

Data-related configuration.

Type:

sleap.nn.config.data.DataConfig

optimization_config#

Optimization-related configuration.

Type:

sleap.nn.config.optimization.OptimizationConfig

centroid_confmap_head#

Instantiated head describing the output centroid confidence maps tensor.

Type:

sleap.nn.heads.CentroidConfmapsHead

offsets_head#

Optional head describing the offset refinement maps.

Type:

sleap.nn.heads.OffsetRefinementHead | None

make_base_pipeline(data_provider: Provider) Pipeline[source]#

Create base pipeline with input data only.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce input examples.

make_training_pipeline(data_provider: Provider) Pipeline[source]#

Create full training pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce all data keys required for training.

Notes

This does not remap keys to model outputs. Use KeyMapper to pull out keys with the appropriate format for the instantiated tf.keras.Model.

make_viz_pipeline(data_provider: Provider, keras_model: Model) Pipeline[source]#

Create visualization pipeline.

Parameters:
  • data_provider – A Provider that generates data examples, typically a LabelsReader instance.

  • keras_model – A tf.keras.Model that can be used for inference.

Returns:

A Pipeline instance configured to fetch data and run inference to generate predictions useful for visualization during training.

class sleap.nn.data.pipelines.Pipeline(providers: Any = NOTHING, transformers: Any = NOTHING)[source]#

Pipeline composed of providers and transformers.

providers#

A single or a list of data providers.

Type:

List[sleap.nn.data.pipelines.Provider]

transformers#

A single or a list of transformers.

Type:

List[sleap.nn.data.pipelines.Transformer]

append(other: Pipeline | Transformer | List[Transformer])[source]#

Append one or more blocks to this pipeline instance.

Parameters:

other – A single Pipeline, Transformer or list of `Transformer`s to append to the end of this pipeline.

Raises:

ValueError – If blocks provided are not a Pipeline, Transformer or list of `Transformer`s.

describe(return_description: bool = False) str | None[source]#

Prints the keys in the examples generated by the pipeline.

Parameters:

return_description – If True, returns the string description instead of printing it.

Returns:

String description if return_description is True, otherwise None.

classmethod from_blocks(blocks: Provider | Transformer | Sequence[Provider | Transformer]) Pipeline[source]#

Create a pipeline from a sequence of providers and transformers.

Parameters:

blocks – List or tuple of providers and transformer instances.

Returns:

An instantiated pipeline with all blocks chained.

classmethod from_pipelines(pipelines: Sequence[Pipeline]) Pipeline[source]#

Create a new pipeline instance by chaining together multiple pipelines.

Parameters:

pipelines – A sequence of Pipeline instances.

Returns:

A new Pipeline instance formed by concatenating the individual pipelines.

make_dataset() DatasetV2[source]#

Create a dataset instance that generates examples from the pipeline.

Returns:

The instantiated tf.data.Dataset pipeline that generates examples with the keys in the output_keys attribute.

property output_keys: List[str]#

Return the keys in examples from a dataset generated from this pipeline.

peek(n: int = 1) Dict[str, Tensor] | List[Dict[str, Tensor]][source]#

Build and return the first n examples from the pipeline.

This function is useful for quickly inspecting the output of a pipeline.

Parameters:

n – Number of examples to get from the pipeline.

Returns:

A dictionary with tensors if n = 1, or a list of dictionaries if n > 1.

run() List[Dict[str, Tensor]][source]#

Build and evaluate the pipeline.

Returns:

List of example dictionaries after processing the pipeline.

validate_pipeline() List[str][source]#

Check that all pipeline blocks meet the data requirements.

Returns:

The final keys that will be present in each example.

Raises:

ValueError – If keys required for a block are dropped at some point in the pipeline.

class sleap.nn.data.pipelines.SingleInstanceConfmapsPipeline(data_config: DataConfig, optimization_config: OptimizationConfig, single_instance_confmap_head: SingleInstanceConfmapsHead, offsets_head: OffsetRefinementHead | None = None)[source]#

Pipeline builder for single-instance confidence map models.

data_config#

Data-related configuration.

Type:

sleap.nn.config.data.DataConfig

optimization_config#

Optimization-related configuration.

Type:

sleap.nn.config.optimization.OptimizationConfig

single_instance_confmap_head#

Instantiated head describing the output confidence maps tensor.

Type:

sleap.nn.heads.SingleInstanceConfmapsHead

offsets_head#

Optional head describing the offset refinement maps.

Type:

sleap.nn.heads.OffsetRefinementHead | None

make_base_pipeline(data_provider: Provider) Pipeline[source]#

Create base pipeline with input data only.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce input examples.

make_training_pipeline(data_provider: Provider) Pipeline[source]#

Create full training pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce all data keys required for training.

Notes

This does not remap keys to model outputs. Use KeyMapper to pull out keys with the appropriate format for the instantiated tf.keras.Model.

make_viz_pipeline(data_provider: Provider, keras_model: Model) Pipeline[source]#

Create visualization pipeline.

Parameters:
  • data_provider – A Provider that generates data examples, typically a LabelsReader instance.

  • keras_model – A tf.keras.Model that can be used for inference.

Returns:

A Pipeline instance configured to fetch data and run inference to generate predictions useful for visualization during training.

class sleap.nn.data.pipelines.TopDownMultiClassPipeline(data_config: DataConfig, optimization_config: OptimizationConfig, instance_confmap_head: CenteredInstanceConfmapsHead, class_vectors_head: ClassVectorsHead, offsets_head: OffsetRefinementHead | None = None)[source]#

Pipeline builder for confidence maps and class maps models.

data_config#

Data-related configuration.

Type:

sleap.nn.config.data.DataConfig

optimization_config#

Optimization-related configuration.

Type:

sleap.nn.config.optimization.OptimizationConfig

confmaps_head#

Instantiated head describing the output confidence maps tensor.

class_vectors_head#

Instantiated head describing the output class vectors tensor.

Type:

sleap.nn.heads.ClassVectorsHead

offsets_head#

Optional head describing the offset refinement maps.

Type:

sleap.nn.heads.OffsetRefinementHead | None

make_base_pipeline(data_provider: Provider) Pipeline[source]#

Create base pipeline with input data only.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce input examples.

make_training_pipeline(data_provider: Provider) Pipeline[source]#

Create full training pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce all data keys required for training.

Notes

This does not remap keys to model outputs. Use KeyMapper to pull out keys with the appropriate format for the instantiated tf.keras.Model.

make_viz_pipeline(data_provider: Provider) Pipeline[source]#

Create visualization pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to fetch data and run inference to generate predictions useful for visualization during training.

class sleap.nn.data.pipelines.TopdownConfmapsPipeline(data_config: DataConfig, optimization_config: OptimizationConfig, instance_confmap_head: CenteredInstanceConfmapsHead, offsets_head: OffsetRefinementHead | None = None)[source]#

Pipeline builder for instance-centered confidence map models.

data_config#

Data-related configuration.

Type:

sleap.nn.config.data.DataConfig

optimization_config#

Optimization-related configuration.

Type:

sleap.nn.config.optimization.OptimizationConfig

instance_confmap_head#

Instantiated head describing the output centered confidence maps tensor.

Type:

sleap.nn.heads.CenteredInstanceConfmapsHead

offsets_head#

Optional head describing the offset refinement maps.

Type:

sleap.nn.heads.OffsetRefinementHead | None

make_base_pipeline(data_provider: Provider) Pipeline[source]#

Create base pipeline with input data only.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce input examples.

make_training_pipeline(data_provider: Provider) Pipeline[source]#

Create full training pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to produce all data keys required for training.

Notes

This does not remap keys to model outputs. Use KeyMapper to pull out keys with the appropriate format for the instantiated tf.keras.Model.

make_viz_pipeline(data_provider: Provider) Pipeline[source]#

Create visualization pipeline.

Parameters:

data_provider – A Provider that generates data examples, typically a LabelsReader instance.

Returns:

A Pipeline instance configured to fetch data and for running inference to generate predictions useful for visualization during training.