sleap.nn.data.utils#

Miscellaneous utility functions for data processing.

sleap.nn.data.utils.describe_tensors(example: Dict[str, Tensor], return_description: bool = False) → Optional[str][source]#

Print the keys in a example.

Parameters:

  • example – Dictionary keyed by strings with tensors as values.

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

Returns:

String description if return_description is True, otherwise None.

sleap.nn.data.utils.ensure_list(x: Any) → List[Any][source]#

Convert the input into a list if it is not already.

sleap.nn.data.utils.expand_to_rank(x: Tensor, target_rank: int, prepend: bool = True) → Tensor[source]#

Expand a tensor to a target rank by adding singleton dimensions.

Parameters:

  • x – Any tf.Tensor with rank <= target_rank. If the rank is higher than target_rank, the tensor will be returned with the same shape.

  • target_rank – Rank to expand the input to.

  • prepend – If True, singleton dimensions are added before the first axis of the data. If False, singleton dimensions are added after the last axis.

Returns:

The expanded tensor of the same dtype as the input, but with rank target_rank.

The output has the same exact data as the input tensor and will be identical if they are both flattened.

sleap.nn.data.utils.gaussian_pdf(x: Tensor, sigma: float) → Tensor[source]#

Compute the PDF of an unnormalized 0-centered Gaussian distribution.

Parameters:

x – Any tensor of dtype tf.float32 with values to compute the PDF for.

Returns:

A tensor of the same shape as x, but with values of a PDF of an unnormalized Gaussian distribution. Values of 0 have an unnormalized PDF value of 1.0.

sleap.nn.data.utils.make_grid_vectors(image_height: int, image_width: int, output_stride: int = 1) → Tuple[Tensor, Tensor][source]#

Make sampling grid vectors from image dimensions.

This is a useful function for creating the x- and y-vectors that define a sampling grid over an image space. These vectors can be used to generate a full meshgrid or for equivalent broadcasting operations.

Parameters:

  • image_height – Height of the image grid that will be sampled, specified as a scalar integer.

  • image_width – width of the image grid that will be sampled, specified as a scalar integer.

  • output_stride – Sampling step size, specified as a scalar integer. This can be used to specify a sampling grid that has a smaller shape than the image grid but with values span the same range. This can be thought of as the reciprocal of the output scale, i.e., it will induce subsampling when set to values greater than 1.

Returns:

Tuple of grid vectors (xv, yv). These are tensors of dtype tf.float32 with shapes (grid_width,) and (grid_height,) respectively.

The grid dimensions are calculated as:

grid_width = image_width // output_stride grid_height = image_height // output_stride

sleap.nn.data.utils.unrag_example(example: Dict[str, Tensor], numpy: bool = False) → Dict[str, Tensor][source]#

Convert ragged tensors in an example into normal tensors with NaN padding.

Parameters:

  • example – Dictionary keyed by strings with tensors as values.

  • numpy – If True, convert values to numpy arrays or Python primitives.

Returns:

The same dictionary, but values of type tf.RaggedTensor will be converted to tensors of type tf.Tensor with NaN padding if the ragged dimensions are of variable length.

The output shapes will be the bounding shape of the ragged tensors.

If numpy is True, the values will be `numpy.ndarray`s or Python primitives depending on their data type and shape.

See also: keras.utils.sync_to_numpy_or_python_type

sleap.nn.data.utils.unrag_tensor(x: RaggedTensor, max_size: int, axis: int) → Tensor[source]#

Converts a ragged tensor to a full tensor by padding to a maximum size.

This function is useful for converting ragged tensors to a fixed size when one or more of the dimensions are of variable length.

Parameters:

  • x – Ragged tensor to convert.

  • max_size – Maximum size of the axis to pad.

  • axis – Axis of x to pad to max_size. This must specify ragged dimensions. If more than one axis is specified, max_size must be of the same length as axis.

Returns:

A padded version of x. Padding will use the equivalent of NaNs in the tensor’s native dtype.

This will replace the shape of the specified axis with max_size, leaving the remaining dimensions set to the bounding shape of the ragged tensor.