sleap.nn.paf_grouping

sleap.nn.paf_grouping#

This module provides a set of utilities for grouping peaks based on PAFs.

Part affinity fields (PAFs) are a representation used to resolve the peak grouping problem for multi-instance pose estimation [1].

They are a convenient way to represent directed graphs with support in image space. For each edge, a PAF can be represented by an image with two channels, corresponding to the x and y components of a unit vector pointing along the direction of the underlying directed graph formed by the connections of the landmarks belonging to an instance.

Given a pair of putatively connected landmarks, the agreement between the line segment that connects them and the PAF vectors found at the coordinates along the same line can be used as a measure of “connectedness”. These scores can then be used to guide the instance-wise grouping of landmarks.

This image space representation is particularly useful as it is amenable to neural network-based prediction from unlabeled images.

A high-level API for grouping based on PAFs is provided through the PAFScorer class.

References

class sleap.nn.paf_grouping.EdgeConnection(src_peak_ind: int, dst_peak_ind: int, score: float)[source]#

Indices to specify a matched connection between two peaks.

This is a convenience named tuple for use in the matching pipeline.

src_peak_ind#

Index of the source peak within all peaks.

Type:

int

dst_peak_ind#

Index of the destination peak within all peaks.

Type:

int

score#

Score of the match.

Type:

float

class sleap.nn.paf_grouping.EdgeType(src_node_ind: int, dst_node_ind: int)[source]#

Indices to uniquely identify a single edge type.

This is a convenience named tuple for use in the matching pipeline.

src_node_ind#

Index of the source node type within the skeleton edges.

Type:

int

dst_node_ind#

Index of the destination node type within the skeleton edges.

Type:

int

class sleap.nn.paf_grouping.PAFScorer(part_names: List[str], edges: List[Tuple[str, str]], pafs_stride: int, max_edge_length_ratio: float = 0.25, dist_penalty_weight: float = 1.0, n_points: int = 10, min_instance_peaks: int | float = 0, min_line_scores: float = 0.25)[source]#

Scoring pipeline based on part affinity fields.

This class facilitates grouping of predicted peaks based on PAFs. It holds a set of common parameters that are used across different steps of the pipeline.

part_names#

List of string node names in the skeleton.

Type:

List[str]

edges#

List of (src_node, dst_node) names in the skeleton.

Type:

List[Tuple[str, str]]

pafs_stride#

Output stride of the part affinity fields. This will be used to adjust the peak coordinates from full image to PAF subscripts.

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

n_points#

Number of points to sample along the line integral.

Type:

int

min_instance_peaks#

Minimum number of peaks the instance should have to be considered a real instance. Instances with fewer peaks than this will be discarded (useful for filtering spurious detections).

Type:

int | float

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

edge_inds#

The edges of the skeleton defined as a list of (source, destination) tuples of node indices. This is created automatically on initialization.

Type:

List[Tuple[int, int]]

edge_types#

A list of EdgeType instances representing the edges of the skeleton. This is created automatically on initialization.

Type:

List[sleap.nn.paf_grouping.EdgeType]

n_nodes#

The number of nodes in the skeleton as a scalar int. This is created automatically on initialization.

Type:

int

n_edges#

The number of edges in the skeleton as a scalar int. This is created automatically on initialization.

Type:

int

sorted_edge_inds#

A tuple of indices specifying the topological order that the edge types should be accessed in during instance assembly (assign_connections_to_instances).

Type:

Tuple[int]

Notes

This class provides high level APIs for grouping peaks into instances using PAFs.

The algorithm has three steps:

1. Find all candidate connections between peaks and compute their matching score based on the PAFs.

2. Match candidate connections using the connectivity score such that no peak is used in two connections of the same type.

  1. Group matched connections into complete instances.

In general, the output from a peak finder (such as multi-peak confidence map prediction network) can be passed into PAFScorer.predict() to get back complete instances.

For finer control over the grouping pipeline steps, use the instance methods in this class or the lower level functions in sleap.nn.paf_grouping.

classmethod from_config(config: MultiInstanceConfig, max_edge_length_ratio: float = 0.25, dist_penalty_weight: float = 1.0, n_points: int = 10, min_instance_peaks: int | float = 0, min_line_scores: float = 0.25) → PAFScorer[source]#

Initialize the PAF scorer from a MultiInstanceConfig head config.

Parameters:

  • config – MultiInstanceConfig from cfg.model.heads.multi_instance.

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

  • min_edge_score – Minimum score required to classify a connection as correct.

  • n_points – Number of points to sample along the line integral.

  • min_instance_peaks – Minimum number of peaks the instance should have to be considered a real instance. Instances with fewer peaks than this will be discarded (useful for filtering spurious detections).

  • min_line_scores – Minimum line score (between -1 and 1) required to form a match between candidate point pairs. Useful for rejecting spurious detections when there are no better ones.

Returns:

The initialized instance of PAFScorer.

group_instances(peaks: RaggedTensor, peak_vals: RaggedTensor, peak_channel_inds: RaggedTensor, match_edge_inds: RaggedTensor, match_src_peak_inds: RaggedTensor, match_dst_peak_inds: RaggedTensor, match_line_scores: RaggedTensor) → Tuple[RaggedTensor, RaggedTensor, RaggedTensor][source]#

Group matched connections into full instances for a batch.

Parameters:

  • peaks – The sample-grouped detected peaks in a batch as a tf.RaggedTensor of shape (n_samples, (n_peaks), 2) and dtype tf.float32. These should be (x, y) coordinates of each peak in the image scale.

  • peak_vals – The sample-grouped scores of the detected peaks in a batch as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.float32.

  • peak_channel_inds – The sample-grouped indices of the channel (node) that each detected peak is associated with as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.int32.

  • match_edge_inds – Sample-grouped indices of the skeleton edge that each connection corresponds to as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. This can be generated by PAFScorer.match_candidates().

  • match_src_peak_inds – Sample-grouped indices of the source peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in each sample. This can be generated by PAFScorer.match_candidates().

  • match_dst_peak_inds – Sample-grouped indices of the destination peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated by PAFScorer.match_candidates().

  • match_line_scores – Sample-grouped PAF line scores of the matched connections as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.float32. This can be generated by PAFScorer.match_candidates().

Returns:

A tuple of arrays with the grouped instances for the whole batch grouped by sample:

predicted_instances: The sample- and instance-grouped coordinates for each instance as tf.RaggedTensor of shape (n_samples, (n_instances), n_nodes, 2) and dtype tf.float32. Missing peaks are represented by `NaN`s.

predicted_peak_scores: The sample- and instance-grouped confidence map values for each peak as an array of (n_samples, (n_instances), n_nodes) and dtype tf.float32.

predicted_instance_scores: The sample-grouped instance grouping score for each instance as an array of shape (n_samples, (n_instances)) and dtype tf.float32.

Notes

This is a convenience wrapper for the standalone group_instances_batch().

See also: PAFScorer.match_candidates, group_instances_batch

match_candidates(edge_inds: RaggedTensor, edge_peak_inds: RaggedTensor, line_scores: RaggedTensor) → Tuple[RaggedTensor, RaggedTensor, RaggedTensor, RaggedTensor][source]#

Match candidate connections for a batch based on PAF scores.

Parameters:

  • edge_inds – Sample-grouped edge indices as a tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.int32 indicating the indices of the edge that each of the candidate connections belongs to. Can be generated using PAFScorer.score_paf_lines().

  • edge_peak_inds – Sample-grouped indices of the peaks that form the source and destination of each candidate connection as a tf.RaggedTensor of shape (n_samples, (n_candidates), 2) and dtype tf.int32. Can be generated using PAFScorer.score_paf_lines().

  • line_scores – Sample-grouped scores for each candidate connection as a tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.float32. Can be generated using PAFScorer.score_paf_lines().

Returns:

The connection peaks for each edge matched based on score as tuple of (match_edge_inds, match_src_peak_inds, match_dst_peak_inds, match_line_scores)

match_edge_inds: Sample-grouped indices of the skeleton edge for each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32.

match_src_peak_inds: Sample-grouped indices of the source peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample.

match_dst_peak_inds: Sample-grouped indices of the destination peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample.

match_line_scores: Sample-grouped PAF line scores of the matched connections as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.float32.

Notes

This is a convenience wrapper for the standalone match_candidates_batch().

See also: PAFScorer.score_paf_lines, match_candidates_batch

predict(pafs: Tensor, peaks: RaggedTensor, peak_vals: RaggedTensor, peak_channel_inds: RaggedTensor) → Tuple[RaggedTensor, RaggedTensor, RaggedTensor][source]#

Group a batch of predicted peaks into full instance predictions using PAFs.

Parameters:

  • pafs – The batch of part affinity fields as a tf.Tensor of shape (n_samples, height, width, 2 * n_edges) and type tf.float32.

  • peaks – The coordinates of the peaks grouped by sample as a tf.RaggedTensor of shape (n_samples, (n_peaks), 2).

  • peak_vals – The sample-grouped scores of the detected peaks in a batch as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.float32.

  • peak_channel_inds – The channel (node) that each peak in peaks corresponds to as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.int32.

Returns:

A tuple of arrays with the grouped instances for the whole batch grouped by sample:

predicted_instances: The sample- and instance-grouped coordinates for each instance as tf.RaggedTensor of shape (n_samples, (n_instances), n_nodes, 2) and dtype tf.float32. Missing peaks are represented by `NaN`s.

predicted_peak_scores: The sample- and instance-grouped confidence map values for each peak as an array of (n_samples, (n_instances), n_nodes) and dtype tf.float32.

predicted_instance_scores: The sample-grouped instance grouping score for each instance as an array of shape (n_samples, (n_instances)) and dtype tf.float32.

Notes

This is a high level API for grouping peaks into instances using PAFs.

See the PAFScorer class documentation for more details on the algorithm.

See also

PAFScorer.score_paf_lines, PAFScorer.match_candidates, PAFScorer.group_instances

score_paf_lines(pafs: Tensor, peaks: Tensor, peak_channel_inds: Tensor) → Tuple[RaggedTensor, RaggedTensor, RaggedTensor][source]#

Create and score PAF lines formed between connection candidates.

Parameters:

  • pafs – The batch of part affinity fields as a tf.Tensor of shape (n_samples, height, width, 2 * n_edges) and type tf.float32.

  • peaks – The coordinates of the peaks grouped by sample as a tf.RaggedTensor of shape (n_samples, (n_peaks), 2).

  • peak_channel_inds – The channel (node) that each peak in peaks corresponds to as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.int32.

Returns:

A tuple of (edge_inds, edge_peak_inds, line_scores) with the connections and their scores based on the PAFs.

edge_inds: Sample-grouped indices of the edge in the skeleton that each connection corresponds to as tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.int32.

edge_peak_inds: Sample-grouped indices of the peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_candidates), 2) and dtype tf.int32. The last axis corresponds to the [source, destination] peak indices. These index into the input peak_channel_inds.

line_scores: Sample-grouped scores for each candidate connection as a tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.float32.

Notes

This is a convenience wrapper for the standalone score_paf_lines_batch().

See also: score_paf_lines_batch

class sleap.nn.paf_grouping.PeakID(node_ind: int, peak_ind: int)[source]#

Indices to uniquely identify a single peak.

This is a convenience named tuple for use in the matching pipeline.

node_ind#

Index of the node type (channel) of the peak.

Type:

int

peak_ind#

Index of the peak within its node type.

Type:

int

sleap.nn.paf_grouping.assign_connections_to_instances(connections: Dict[EdgeType, List[EdgeConnection]], min_instance_peaks: int | float = 0, n_nodes: int | None = None) → Dict[PeakID, int][source]#

Assigns connected edges to instances via greedy graph partitioning.

Parameters:

  • connections – A dict that maps EdgeType to a list of EdgeConnections found through connection scoring. This can be generated by the filter_connection_candidates function.

  • min_instance_peaks – If this is greater than 0, grouped instances with fewer assigned peaks than this threshold will be excluded. If a float in the range (0., 1.] is provided, this is interpreted as a fraction of the total number of nodes in the skeleton. If an integer is provided, this is the absolute minimum number of peaks.

  • n_nodes – Total node type count. Used to convert min_instance_peaks to an absolute number when a fraction is specified. If not provided, the node count is inferred from the unique node inds in connections.

Returns:

A dict mapping PeakID to a unique instance ID specified as an integer.

A PeakID is a tuple of (node_type_ind, peak_ind), where the peak_ind is the index or identifier specified in a EdgeConnection as a src_peak_ind or dst_peak_ind.

Return type:

instance_assignments

Note

Instance IDs are not necessarily consecutive since some instances may be filtered out during the partitioning or filtering.

This function expects connections from a single sample/frame!

sleap.nn.paf_grouping.compute_distance_penalty(spatial_vec_lengths: Tensor, max_edge_length: float, dist_penalty_weight: float = 1.0) → Tensor[source]#

Compute the distance penalty component of the PAF line integral score.

Parameters:

  • spatial_vec_lengths – Euclidean distance between candidate source and destination points as a tf.float32 tensor of any shape (typically (n_candidates, 1)).

  • max_edge_length – Maximum length expected for any connection as a scalar float in units of pixels (corresponding to peaks_sample). Scores of lines longer than this will be penalized. Useful for ignoring spurious connections that are far apart in space.

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

Returns:

The distance penalty for each candidate as a tf.float32 tensor of the same shape as spatial_vec_lengths.

The penalty will be 0 (when below the threshold) and -1 as the distance approaches infinity. This is then scaled by the dist_penalty_weight.

Notes

The penalty is computed from the distances scaled by the max length:

``` if distance <= max_edge_length:

penalty = 0

else:

penalty = (max_edge_length / distance) - 1

```

For example, if the max length is 10 and the distance is 20, then the penalty will be: (10 / 20) - 1 == 0.5 - 1 == -0.5.

See also: score_paf_lines

sleap.nn.paf_grouping.get_connection_candidates(peak_channel_inds_sample: Tensor, skeleton_edges: Tensor, n_nodes: int) → Tuple[Tensor, Tensor][source]#

Find the indices of all the possible connections formed by the detected peaks.

Parameters:

  • peak_channel_inds_sample – The channel indices of the peaks found in a sample. This is a tf.Tensor of shape (n_peaks,) and dtype tf.int32 that is used to represent a detected peak by its channel/node index in the skeleton.

  • skeleton_edges – The indices of the nodes that form the skeleton graph as a tf.Tensor of shape (n_edges, 2) and dtype tf.int32 where each row corresponds to the source and destination node indices.

  • n_nodes – The total number of nodes in the skeleton as a scalar integer.

Returns:

A tuple of (edge_inds, edge_peak_inds).

edge_inds is a tf.Tensor of shape (n_candidates,) indicating the indices of the edge that each of the candidate connections belongs to.

edge_peak_inds is a tf.Tensor of shape (n_candidates, 2) with the indices of the peaks that form the source and destination of each candidate connection. This indexes into the input peak_channel_inds_sample.

sleap.nn.paf_grouping.get_paf_lines(pafs_sample: Tensor, peaks_sample: Tensor, edge_peak_inds: Tensor, edge_inds: Tensor, n_line_points: int, pafs_stride: int) → Tensor[source]#

Gets the PAF values at the lines formed between all detected peaks in a sample.

Parameters:

  • pafs_sample – The PAFs for the sample as a tf.Tensor of shape (height, width, 2 * n_edges).

  • peaks_sample – The detected peaks in a sample as a tf.Tensor of shape (n_peaks, 2) and dtype tf.float32. These should be (x, y) coordinates of each peak in the image scale (they will be scaled by the pafs_stride).

  • edge_peak_inds – A tf.Tensor of shape (n_candidates, 2) and dtype tf.int32 with the indices of the peaks that form the source and destination of each candidate connection. This indexes into the input peaks_sample. Can be generated using get_connection_candidates().

  • edge_inds – A tf.Tensor of shape (n_candidates,) and dtype tf.int32 indicating the indices of the edge that each of the candidate connections belongs to. Can be generated using get_connection_candidates().

  • n_line_points – The number of points to interpolate between source and destination peaks in each connection candidate as a scalar integer. Values ranging from 5 to 10 are pretty reasonable.

  • pafs_stride – The stride (1/scale) of the PAFs that these lines will need to index into relative to the image. Coordinates in peaks_sample will be divided by this value to adjust the indexing into the PAFs tensor.

Returns:

The PAF vectors at all of the line points as a tf.Tensor of shape (n_candidates, n_line_points, 2, 3) and dtype tf.int32. These subscripts can be used directly with tf.gather_nd to pull out the PAF values at the lines.

The last dimension of the line subscripts correspond to the full [row, col, channel] subscripts of each element of the lines. Axis -2 contains the same [row, col] for each line but channel is adjusted to match the channels in the PAFs tensor.

Notes

If only the subscripts are needed, use make_line_subs() to generate the lines without retrieving the PAF vector at the line points.

See also: get_connection_candidates, make_line_subs, score_paf_lines

sleap.nn.paf_grouping.group_instances_batch(peaks: RaggedTensor, peak_vals: RaggedTensor, peak_channel_inds: RaggedTensor, match_edge_inds: RaggedTensor, match_src_peak_inds: RaggedTensor, match_dst_peak_inds: RaggedTensor, match_line_scores: RaggedTensor, n_nodes: int, sorted_edge_inds: Tuple[int], edge_types: List[EdgeType], min_instance_peaks: int, min_line_scores: float = 0.25) → Tuple[RaggedTensor, RaggedTensor, RaggedTensor][source]#

Group matched connections into full instances for a batch.

Parameters:

  • peaks – The sample-grouped detected peaks in a batch as a tf.RaggedTensor of shape (n_samples, (n_peaks), 2) and dtype tf.float32. These should be (x, y) coordinates of each peak in the image scale.

  • peak_vals – The sample-grouped scores of the detected peaks in a batch as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.float32.

  • peak_channel_inds – The sample-grouped indices of the channel (node) that each detected peak is associated with as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.int32.

  • match_edge_inds – Sample-grouped indices of the skeleton edge that each connection corresponds to as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. This can be generated by match_candidates_batch().

  • match_src_peak_inds – Sample-grouped indices of the source peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in each sample. This can be generated by match_candidates_batch().

  • match_dst_peak_inds – Sample-grouped indices of the destination peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated by match_candidates_batch().

  • match_line_scores – Sample-grouped PAF line scores of the matched connections as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.float32. This can be generated by match_candidates_batch().

  • n_nodes – The total number of nodes in the skeleton as a scalar integer.

  • sorted_edge_inds – A tuple of indices specifying the topological order that the edge types should be accessed in during instance assembly (assign_connections_to_instances).

  • edge_types – A list of `EdgeType`s associated with the skeleton.

  • min_instance_peaks – If this is greater than 0, grouped instances with fewer assigned peaks than this threshold will be excluded. If a float in the range (0., 1.] is provided, this is interpreted as a fraction of the total number of nodes in the skeleton. If an int is provided, this is the absolute minimum number of peaks.

  • min_line_scores – Minimum line score (between -1 and 1) required to form a match between candidate point pairs.

Returns:

A tuple of arrays with the grouped instances for the whole batch grouped by sample:

predicted_instances: The sample- and instance-grouped coordinates for each instance as tf.RaggedTensor of shape (n_samples, (n_instances), n_nodes, 2) and dtype tf.float32. Missing peaks are represented by `NaN`s.

predicted_peak_scores: The sample- and instance-grouped confidence map values for each peak as an array of (n_samples, (n_instances), n_nodes) and dtype tf.float32.

predicted_instance_scores: The sample-grouped instance grouping score for each instance as an array of shape (n_samples, (n_instances)) and dtype tf.float32.

See also: match_candidates_batch, group_instances_sample

sleap.nn.paf_grouping.group_instances_sample(peaks_sample: Tensor, peak_scores_sample: Tensor, peak_channel_inds_sample: Tensor, match_edge_inds_sample: Tensor, match_src_peak_inds_sample: Tensor, match_dst_peak_inds_sample: Tensor, match_line_scores_sample: Tensor, n_nodes: int, sorted_edge_inds: Tuple[int], edge_types: List[EdgeType], min_instance_peaks: int, min_line_scores: float = 0.25) → Tuple[ndarray, ndarray, ndarray][source]#

Group matched connections into full instances for a single sample.

Parameters:

  • peaks_sample – The detected peaks in a sample as a tf.Tensor of shape (n_peaks, 2) and dtype tf.float32. These should be (x, y) coordinates of each peak in the image scale.

  • peak_scores_sample – The scores of the detected peaks in a sample as a tf.Tensor of shape (n_peaks,) and dtype tf.float32.

  • peak_channel_inds_sample – The indices of the channel (node) that each detected peak is associated with as a tf.Tensor of shape (n_peaks,) and dtype tf.int32.

  • match_edge_inds_sample – Indices of the skeleton edge that each connection corresponds to as a tf.Tensor of shape (n_connections,) and dtype tf.int32. This can be generated by match_candidates_sample().

  • match_src_peak_inds_sample – Indices of the source peaks that form each connection as a tf.Tensor of shape (n_connections,) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated by match_candidates_sample().

  • match_dst_peak_inds_sample – Indices of the destination peaks that form each connection as a tf.Tensor of shape (n_connections,) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated by match_candidates_sample().

  • match_line_scores_sample – PAF line scores of the matched connections as a tf.Tensor of shape (n_connections,) and dtype tf.float32. This can be generated by match_candidates_sample().

  • n_nodes – The total number of nodes in the skeleton as a scalar integer.

  • sorted_edge_inds – A tuple of indices specifying the topological order that the edge types should be accessed in during instance assembly (assign_connections_to_instances).

  • edge_types – A list of `EdgeType`s associated with the skeleton.

  • min_instance_peaks – If this is greater than 0, grouped instances with fewer assigned peaks than this threshold will be excluded. If a float in the range (0., 1.] is provided, this is interpreted as a fraction of the total number of nodes in the skeleton. If an int is provided, this is the absolute minimum number of peaks.

  • min_line_scores – Minimum line score (between -1 and 1) required to form a match between candidate point pairs.

Returns:

predicted_instances: The grouped coordinates for each instance as an array of shape (n_instances, n_nodes, 2) and dtype float32. Missing peaks are represented by `np.NaN`s.

predicted_peak_scores: The confidence map values for each peak as an array of (n_instances, n_nodes) and dtype float32.

predicted_instance_scores: The grouping score for each instance as an array of shape (n_instances,) and dtype float32.

Return type:

A tuple of arrays with the grouped instances

Notes

This function is meant to be run as a tf.py_function within a graph (see group_instances_batch()).

sleap.nn.paf_grouping.make_line_subs(peaks_sample: Tensor, edge_peak_inds: Tensor, edge_inds: Tensor, n_line_points: int, pafs_stride: int) → Tensor[source]#

Create the lines between candidate connections for evaluating the PAFs.

Parameters:

  • peaks_sample – The detected peaks in a sample as a tf.Tensor of shape (n_peaks, 2) and dtype tf.float32. These should be (x, y) coordinates of each peak in the image scale (they will be scaled by the pafs_stride).

  • edge_peak_inds – A tf.Tensor of shape (n_candidates, 2) and dtype tf.int32 with the indices of the peaks that form the source and destination of each candidate connection. This indexes into the input peaks_sample. Can be generated using get_connection_candidates().

  • edge_inds – A tf.Tensor of shape (n_candidates,) and dtype tf.int32 indicating the indices of the edge that each of the candidate connections belongs to. Can be generated using get_connection_candidates().

  • n_line_points – The number of points to interpolate between source and destination peaks in each connection candidate as a scalar integer. Values ranging from 5 to 10 are pretty reasonable.

  • pafs_stride – The stride (1/scale) of the PAFs that these lines will need to index into relative to the image. Coordinates in peaks_sample will be divided by this value to adjust the indexing into the PAFs tensor.

Returns:

The line subscripts as a tf.Tensor of shape (n_candidates, n_line_points, 2, 3) and dtype tf.int32. These subscripts can be used directly with tf.gather_nd to pull out the PAF values at the lines.

The last dimension of the line subscripts correspond to the full [row, col, channel] subscripts of each element of the lines. Axis -2 contains the same [row, col] for each line but channel is adjusted to match the channels in the PAFs tensor.

Notes

The subscripts are interpolated via nearest neighbor, so multiple fractional coordinates may map on to the same pixel if the line is short.

See also: get_connection_candidates

sleap.nn.paf_grouping.make_predicted_instances(peaks: array, peak_scores: array, connections: List[EdgeConnection], instance_assignments: Dict[PeakID, int]) → Tuple[array, array, array][source]#

Group peaks by assignments and accumulate scores.

Parameters:

  • peaks – Node-grouped peaks

  • peak_scores – Node-grouped peak scores

  • connections – `EdgeConnection`s grouped by edge type

  • instance_assignments – PeakID to instance ID mapping

Returns:

Tuple of (predicted_instances, predicted_peak_scores, predicted_instance_scores)

predicted_instances: (n_instances, n_nodes, 2) array predicted_peak_scores: (n_instances, n_nodes) array predicted_instance_scores: (n_instances,) array

sleap.nn.paf_grouping.match_candidates_batch(edge_inds: RaggedTensor, edge_peak_inds: RaggedTensor, line_scores: RaggedTensor, n_edges: int) → Tuple[RaggedTensor, RaggedTensor, RaggedTensor, RaggedTensor][source]#

Match candidate connections for a batch based on PAF scores.

Parameters:

  • edge_inds – Sample-grouped edge indices as a tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.int32 indicating the indices of the edge that each of the candidate connections belongs to. Can be generated using score_paf_lines_batch().

  • edge_peak_inds – Sample-grouped indices of the peaks that form the source and destination of each candidate connection as a tf.RaggedTensor of shape (n_samples, (n_candidates), 2) and dtype tf.int32. Can be generated using score_paf_lines_batch().

  • line_scores – Sample-grouped scores for each candidate connection as a tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.float32. Can be generated using score_paf_lines_batch().

  • n_edges – A scalar int denoting the number of edges in the skeleton.

Returns:

The connection peaks for each edge matched based on score as tuple of (match_edge_inds, match_src_peak_inds, match_dst_peak_inds, match_line_scores)

match_edge_inds: Sample-grouped indices of the skeleton edge for each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32.

match_src_peak_inds: Sample-grouped indices of the source peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample.

match_dst_peak_inds: Sample-grouped indices of the destination peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample.

match_line_scores: Sample-grouped PAF line scores of the matched connections as a tf.RaggedTensor of shape (n_samples, (n_connections)) and dtype tf.float32.

Notes

The matching is performed using the Munkres algorithm implemented in scipy.optimize.linear_sum_assignment() which is wrapped in tf_linear_sum_assignment() for execution within a graph.

See also: match_candidates_sample, score_paf_lines_batch, group_instances_batch

sleap.nn.paf_grouping.match_candidates_sample(edge_inds_sample: Tensor, edge_peak_inds_sample: Tensor, line_scores_sample: Tensor, n_edges: int) → Tuple[Tensor, Tensor, Tensor, Tensor][source]#

Match candidate connections for a sample based on PAF scores.

Parameters:

  • edge_inds_sample – A tf.Tensor of shape (n_candidates,) and dtype tf.int32 indicating the indices of the edge that each of the candidate connections belongs to for the sample. Can be generated using get_connection_candidates().

  • edge_peak_inds_sample – A tf.Tensor of shape (n_candidates, 2) and dtype tf.int32 with the indices of the peaks that form the source and destination of each candidate connection. Can be generated using get_connection_candidates().

  • line_scores_sample – Scores for each candidate connection in the sample as a tf.Tensor of shape (n_candidates,) and dtype tf.float32. Can be generated using score_paf_lines().

  • n_edges – A scalar int denoting the number of edges in the skeleton.

Returns:

The connection peaks for each edge matched based on score as tuple of (match_edge_inds, match_src_peak_inds, match_dst_peak_inds, match_line_scores)

match_edge_inds: Indices of the skeleton edge that each connection corresponds to as a tf.Tensor of shape (n_connections,) and dtype tf.int32.

match_src_peak_inds: Indices of the source peaks that form each connection as a tf.Tensor of shape (n_connections,) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample.

match_dst_peak_inds: Indices of the destination peaks that form each connection as a tf.Tensor of shape (n_connections,) and dtype tf.int32. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample.

match_line_scores: PAF line scores of the matched connections as a tf.Tensor of shape (n_connections,) and dtype tf.float32.

Notes

The matching is performed using the Munkres algorithm implemented in scipy.optimize.linear_sum_assignment() which is wrapped in tf_linear_sum_assignment() for execution within a graph.

See also: match_candidates_batch

sleap.nn.paf_grouping.score_paf_lines(paf_lines_sample: Tensor, peaks_sample: Tensor, edge_peak_inds_sample: Tensor, max_edge_length: float, dist_penalty_weight: float = 1.0) → Tensor[source]#

Compute the connectivity score for each PAF line in a sample.

Parameters:

  • paf_lines_sample – The PAF vectors evaluated at the lines formed between candidate conncetions as a tf.Tensor of shape (n_candidates, n_line_points, 2, 3) dtype tf.int32. This can be generated by get_paf_lines().

  • peaks_sample – The detected peaks in a sample as a tf.Tensor of shape (n_peaks, 2) and dtype tf.float32. These should be (x, y) coordinates of each peak in the image scale.

  • edge_peak_inds_sample – A tf.Tensor of shape (n_candidates, 2) and dtype tf.int32 with the indices of the peaks that form the source and destination of each candidate connection. This indexes into the input peaks_sample. Can be generated using get_connection_candidates().

  • max_edge_length – Maximum length expected for any connection as a scalar float in units of pixels (corresponding to peaks_sample). Scores of lines longer than this will be penalized. Useful for ignoring spurious connections that are far apart in space.

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

Returns:

The line scores as a tf.Tensor of shape (n_candidates,) and dtype tf.float32. Each score value is the average dot product between the PAFs and the normalized displacement vector between source and destination peaks.

Scores range from roughly -1.5 to 1.0, where larger values indicate a better connectivity score for the candidate. Values can be larger or smaller due to prediction error.

Notes

This function operates on a single sample (frame). For batches of multiple frames, use score_paf_lines_batch().

See also: get_paf_lines, score_paf_lines_batch, compute_distance_penalty

sleap.nn.paf_grouping.score_paf_lines_batch(pafs: Tensor, peaks: Tensor, peak_channel_inds: RaggedTensor, skeleton_edges: Tensor, n_line_points: int, pafs_stride: int, max_edge_length_ratio: float, dist_penalty_weight: float, n_nodes: int) → Tuple[RaggedTensor, RaggedTensor, RaggedTensor][source]#

Create and score PAF lines formed between connection candidates.

Parameters:

  • pafs – The batch of part affinity fields as a tf.Tensor of shape (n_samples, height, width, 2 * n_edges) and type tf.float32.

  • peaks – The coordinates of the peaks grouped by sample as a tf.RaggedTensor of shape (n_samples, (n_peaks), 2).

  • peak_channel_inds – The channel (node) that each peak in peaks corresponds to as a tf.RaggedTensor of shape (n_samples, (n_peaks)) and dtype tf.int32.

  • skeleton_edges – The indices of the nodes that form the skeleton graph as a tf.Tensor of shape (n_edges, 2) and dtype tf.int32 where each row corresponds to the source and destination node indices.

  • n_line_points – The number of points to interpolate between source and destination peaks in each connection candidate as a scalar integer. Values ranging from 5 to 10 are pretty reasonable.

  • pafs_stride – The stride (1/scale) of the PAFs that these lines will need to index into relative to the image. Coordinates in peaks will be divided by this value to adjust the indexing into the pafs tensor.

  • max_edge_length_ratio – The maximum expected length of a connected pair of points in relative image units. Candidate connections above 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.

  • n_nodes – The total number of nodes in the skeleton as a scalar integer.

Returns:

A tuple of (edge_inds, edge_peak_inds, line_scores) with the connections and their scores based on the PAFs.

edge_inds: Sample-grouped indices of the edge in the skeleton that each connection corresponds to as tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.int32.

edge_peak_inds: Sample-grouped indices of the peaks that form each connection as a tf.RaggedTensor of shape (n_samples, (n_candidates), 2) and dtype tf.int32. The last axis corresponds to the [source, destination] peak indices. These index into the input peak_channel_inds.

line_scores: Sample-grouped scores for each candidate connection as tf.RaggedTensor of shape (n_samples, (n_candidates)) and dtype tf.float32.

Notes

This function handles the looping over samples in the batch and applies:

  1. get_connection_candidates(): Find peaks that form connections.

  2. get_paf_lines(): Retrieve PAF vectors for each line.

  3. score_paf_lines(): Compute connectivity score for each candidate.

See also: get_connection_candidates, get_paf_lines, score_paf_lines

sleap.nn.paf_grouping.toposort_edges(edge_types: List[EdgeType]) → Tuple[int][source]#

Find a topological ordering for a list of edge types.

Parameters:

edge_types – A list of EdgeType instances describing a skeleton.

Returns:

A tuple of indices specifying the topological order that the edge types should be accessed in during instance assembly (assign_connections_to_instances).

This is important to ensure that instances are assembled starting at the root of the skeleton and moving down.

See also: assign_connections_to_instances

Contents