sleap.nn.data.edge_maps

Transformers for generating edge confidence maps and part affinity fields.

class sleap.nn.data.edge_maps.PartAffinityFieldsGenerator(sigma=1.0, output_stride=1, skeletons: Any | None = None, flatten_channels: bool = False)

Transformer to generate part affinity fields.

sigma

Standard deviation of the 2D Gaussian distribution sampled to weight the part affinity fields by their distance to the edges. This defines the spread in units of the input image’s grid, i.e., it does not take scaling in previous steps into account.

Type:

float

output_stride

Relative stride of the generated confidence maps. This is effectively the reciprocal of the output scale, i.e., increase this to generate confidence maps that are smaller than the input images.

Type:

int

skeletons

List of `sleap.Skeleton`s to use for looking up the index of the edges.

Type:

List[sleap.skeleton.Skeleton] | None

flatten_channels

If False, the generated tensors are of shape [height, width, n_edges, 2]. If True, generated tensors are of shape [height, width, n_edges * 2] by flattening the last 2 axes.

Type:

bool

property input_keys: List[str]

Return the keys that incoming elements are expected to have.

property output_keys: List[str]

Return the keys that outgoing elements will have.

transform_dataset(input_ds: DatasetV2) → DatasetV2

Create a dataset that contains the generated confidence maps.

Parameters:

input_ds – A dataset with elements that contain the keys “image”, “instances” and “skeleton_inds”.

Returns:

A tf.data.Dataset with the same keys as the input, as well as “part_affinity_fields”.

The “part_affinity_fields” key will be a tensor of shape (grid_height, grid_width, n_edges, 2) containing the combined part affinity fields of all instances in the frame.

If the flatten_channels attribute is set to True, the last 2 axes of the “part_affinity_fields” are flattened to produce a tensor of shape (grid_height, grid_width, n_edges * 2). This is a convenient form when training models as a rank-4 (batched) tensor will generally be expected.

Notes

The output stride is relative to the current scale of the image. To map points on the part affinity fields to the raw image, first multiply them by the output stride, and then scale the x- and y-coordinates by the “scale” key.

Importantly, the sigma will be proportional to the current image grid, not the original grid prior to scaling operations.

sleap.nn.data.edge_maps.distance_to_edge(points: Tensor, edge_source: Tensor, edge_destination: Tensor) → Tensor

Compute pairwise distance between points and undirected edges.

Parameters:

• points – Tensor of dtype tf.float32 of shape (d_0, …, d_n, 2) where the last axis corresponds to x- and y-coordinates. Distances will be broadcast across all point dimensions.

• edge_source – Tensor of dtype tf.float32 of shape (n_edges, 2) where the last axis corresponds to x- and y-coordinates of the source points of each edge.

• edge_destination – Tensor of dtype tf.float32 of shape (n_edges, 2) where the last axis corresponds to x- and y-coordinates of the source points of each edge.

Returns:

A tensor of dtype tf.float32 of shape (d_0, …, d_n, n_edges) where the first axes correspond to the initial dimensions of points, and the last indicates the distance of each point to each edge.

sleap.nn.data.edge_maps.get_edge_points(instances: Tensor, edge_inds: Tensor) → Tuple[Tensor, Tensor]

Return the points in each instance that form a directed graph.

Parameters:

• instances – A tensor of shape (n_instances, n_nodes, 2) and dtype tf.float32 containing instance points where the last axis corresponds to (x, y) pixel coordinates on the image. This must be rank-3 even if a single instance is present.

• edge_inds – A tensor of shape (n_edges, 2) and dtype tf.int32 containing the node indices that define a directed graph, where the last axis corresponds to the source and destination node indices.

Returns:

Tuple of (edge_sources, edge_destinations) containing the edge and destination points respectively. Both will be tensors of shape (n_instances, n_edges, 2), where the last axis corresponds to (x, y) pixel coordinates on the image.

sleap.nn.data.edge_maps.make_edge_maps(xv: Tensor, yv: Tensor, edge_source: Tensor, edge_destination: Tensor, sigma: float) → Tensor

Generate confidence maps for a set of undirected edges.

Parameters:

• xv – Sampling grid vector for x-coordinates of shape (grid_width,) and dtype tf.float32. This can be generated by sleap.nn.data.utils.make_grid_vectors.

• yv – Sampling grid vector for y-coordinates of shape (grid_height,) and dtype tf.float32. This can be generated by sleap.nn.data.utils.make_grid_vectors.

• edge_source – Tensor of dtype tf.float32 of shape (n_edges, 2) where the last axis corresponds to x- and y-coordinates of the source points of each edge.

• edge_destination – Tensor of dtype tf.float32 of shape (n_edges, 2) where the last axis corresponds to x- and y-coordinates of the destination points of each edge.

• sigma – Standard deviation of the 2D Gaussian distribution sampled to generate confidence maps.

Returns:

A set of confidence maps corresponding to the probability of each point on a sampling grid being on each edge. These will be in a tensor of shape (grid_height, grid_width, n_edges) of dtype tf.float32.

sleap.nn.data.edge_maps.make_multi_pafs(xv: Tensor, yv: Tensor, edge_sources: Tensor, edge_destinations: Tensor, sigma: float) → Tensor

Make multiple instance PAFs with max reduction.

Parameters:

• xv – Sampling grid vector for x-coordinates of shape (grid_width,) and dtype tf.float32. This can be generated by sleap.nn.data.utils.make_grid_vectors.

• yv – Sampling grid vector for y-coordinates of shape (grid_height,) and dtype tf.float32. This can be generated by sleap.nn.data.utils.make_grid_vectors.

• edge_sources – Tensor of dtype tf.float32 of shape (n_instances, n_edges, 2) where the last axis corresponds to x- and y-coordinates of the source points of each edge.

• edge_destinations – Tensor of dtype tf.float32 of shape (n_instances, n_edges, 2) where the last axis corresponds to x- and y-coordinates of the destination points of each edge.

• sigma – Standard deviation of the 2D Gaussian distribution sampled to generate the edge maps for masking the PAFs.

Returns:

A set of part affinity fields generated for each instance. These will be in a tensor of shape (grid_height, grid_width, n_edges, 2). If multiple instance PAFs are defined on the same pixel, they will be summed.

sleap.nn.data.edge_maps.make_pafs(xv: Tensor, yv: Tensor, edge_source: Tensor, edge_destination: Tensor, sigma: float) → Tensor

Generate part affinity fields for a set of directed edges.

Parameters:

• xv – Sampling grid vector for x-coordinates of shape (grid_width,) and dtype tf.float32. This can be generated by sleap.nn.data.utils.make_grid_vectors.

• yv – Sampling grid vector for y-coordinates of shape (grid_height,) and dtype tf.float32. This can be generated by sleap.nn.data.utils.make_grid_vectors.

• edge_source – Tensor of dtype tf.float32 of shape (n_edges, 2) where the last axis corresponds to x- and y-coordinates of the source points of each edge.

• edge_destination – Tensor of dtype tf.float32 of shape (n_edges, 2) where the last axis corresponds to x- and y-coordinates of the destination points of each edge.

• sigma – Standard deviation of the 2D Gaussian distribution sampled to generate the edge maps for masking the PAFs.

Returns:

A set of part affinity fields corresponding to the unit vector pointing along the direction of each edge weighted by the probability of each point on a sampling grid being on each edge. These will be in a tensor of shape (grid_height, grid_width, n_edges, 2) of dtype tf.float32. The last axis corresponds to the x- and y-coordinates of the unit vectors.