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: Union[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:
Union[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:
- 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.
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: Union[int, float] = 0, min_line_scores: float = 0.25) PAFScorer [source]#
Initialize the PAF scorer from a
MultiInstanceConfig
head config.- Parameters:
config –
MultiInstanceConfig
fromcfg.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 dtypetf.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 dtypetf.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 dtypetf.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 dtypetf.int32
. This can be generated byPAFScorer.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 dtypetf.int32
. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in each sample. This can be generated byPAFScorer.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 dtypetf.int32
. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated byPAFScorer.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 dtypetf.float32
. This can be generated byPAFScorer.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 astf.RaggedTensor
of shape(n_samples, (n_instances), n_nodes, 2)
and dtypetf.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 dtypetf.float32
.predicted_instance_scores
: The sample-grouped instance grouping score for each instance as an array of shape(n_samples, (n_instances))
and dtypetf.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 dtypetf.int32
indicating the indices of the edge that each of the candidate connections belongs to. Can be generated usingPAFScorer.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 dtypetf.int32
. Can be generated usingPAFScorer.score_paf_lines()
.line_scores – Sample-grouped scores for each candidate connection as a
tf.RaggedTensor
of shape(n_samples, (n_candidates))
and dtypetf.float32
. Can be generated usingPAFScorer.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 atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.int32
.match_src_peak_inds
: Sample-grouped indices of the source peaks that form each connection as atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.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 atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.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 atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.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 typetf.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 dtypetf.float32
.peak_channel_inds – The channel (node) that each peak in
peaks
corresponds to as atf.RaggedTensor
of shape(n_samples, (n_peaks))
and dtypetf.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 astf.RaggedTensor
of shape(n_samples, (n_instances), n_nodes, 2)
and dtypetf.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 dtypetf.float32
.predicted_instance_scores
: The sample-grouped instance grouping score for each instance as an array of shape(n_samples, (n_instances))
and dtypetf.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 typetf.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 atf.RaggedTensor
of shape(n_samples, (n_peaks))
and dtypetf.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 astf.RaggedTensor
of shape(n_samples, (n_candidates))
and dtypetf.int32
.edge_peak_inds
: Sample-grouped indices of the peaks that form each connection as atf.RaggedTensor
of shape(n_samples, (n_candidates), 2)
and dtypetf.int32
. The last axis corresponds to the[source, destination]
peak indices. These index into the inputpeak_channel_inds
.line_scores
: Sample-grouped scores for each candidate connection as atf.RaggedTensor
of shape(n_samples, (n_candidates))
and dtypetf.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: Union[int, float] = 0, n_nodes: Optional[int] = 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 topeaks_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 asspatial_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 dtypetf.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 dtypetf.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 atf.Tensor
of shape(n_candidates,)
indicating the indices of the edge that each of the candidate connections belongs to.edge_peak_inds
is atf.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 inputpeak_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 dtypetf.float32
. These should be(x, y)
coordinates of each peak in the image scale (they will be scaled by thepafs_stride
).edge_peak_inds – A
tf.Tensor
of shape(n_candidates, 2)
and dtypetf.int32
with the indices of the peaks that form the source and destination of each candidate connection. This indexes into the inputpeaks_sample
. Can be generated usingget_connection_candidates()
.edge_inds – A
tf.Tensor
of shape(n_candidates,)
and dtypetf.int32
indicating the indices of the edge that each of the candidate connections belongs to. Can be generated usingget_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 dtypetf.int32
. These subscripts can be used directly withtf.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 butchannel
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 dtypetf.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 dtypetf.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 dtypetf.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 dtypetf.int32
. This can be generated bymatch_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 dtypetf.int32
. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in each sample. This can be generated bymatch_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 dtypetf.int32
. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated bymatch_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 dtypetf.float32
. This can be generated bymatch_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 anint
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 astf.RaggedTensor
of shape(n_samples, (n_instances), n_nodes, 2)
and dtypetf.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 dtypetf.float32
.predicted_instance_scores
: The sample-grouped instance grouping score for each instance as an array of shape(n_samples, (n_instances))
and dtypetf.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 dtypetf.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 dtypetf.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 dtypetf.int32
.match_edge_inds_sample – Indices of the skeleton edge that each connection corresponds to as a
tf.Tensor
of shape(n_connections,)
and dtypetf.int32
. This can be generated bymatch_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 dtypetf.int32
. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated bymatch_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 dtypetf.int32
. Important: These indices correspond to the edge-grouped peaks, not the set of all peaks in the sample. This can be generated bymatch_candidates_sample()
.match_line_scores_sample – PAF line scores of the matched connections as a
tf.Tensor
of shape(n_connections,)
and dtypetf.float32
. This can be generated bymatch_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 anint
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 dtypefloat32
. 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 dtypefloat32
.predicted_instance_scores
: The grouping score for each instance as an array of shape(n_instances,)
and dtypefloat32
.- 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 (seegroup_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 dtypetf.float32
. These should be(x, y)
coordinates of each peak in the image scale (they will be scaled by thepafs_stride
).edge_peak_inds – A
tf.Tensor
of shape(n_candidates, 2)
and dtypetf.int32
with the indices of the peaks that form the source and destination of each candidate connection. This indexes into the inputpeaks_sample
. Can be generated usingget_connection_candidates()
.edge_inds – A
tf.Tensor
of shape(n_candidates,)
and dtypetf.int32
indicating the indices of the edge that each of the candidate connections belongs to. Can be generated usingget_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 dtypetf.int32
. These subscripts can be used directly withtf.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 butchannel
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:
- 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 dtypetf.int32
indicating the indices of the edge that each of the candidate connections belongs to. Can be generated usingscore_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 dtypetf.int32
. Can be generated usingscore_paf_lines_batch()
.line_scores – Sample-grouped scores for each candidate connection as a
tf.RaggedTensor
of shape(n_samples, (n_candidates))
and dtypetf.float32
. Can be generated usingscore_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 atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.int32
.match_src_peak_inds
: Sample-grouped indices of the source peaks that form each connection as atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.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 atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.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 atf.RaggedTensor
of shape(n_samples, (n_connections))
and dtypetf.float32
.
Notes
The matching is performed using the Munkres algorithm implemented in
scipy.optimize.linear_sum_assignment()
which is wrapped intf_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 dtypetf.int32
indicating the indices of the edge that each of the candidate connections belongs to for the sample. Can be generated usingget_connection_candidates()
.edge_peak_inds_sample – A
tf.Tensor
of shape(n_candidates, 2)
and dtypetf.int32
with the indices of the peaks that form the source and destination of each candidate connection. Can be generated usingget_connection_candidates()
.line_scores_sample – Scores for each candidate connection in the sample as a
tf.Tensor
of shape(n_candidates,)
and dtypetf.float32
. Can be generated usingscore_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 atf.Tensor
of shape(n_connections,)
and dtypetf.int32
.match_src_peak_inds
: Indices of the source peaks that form each connection as atf.Tensor
of shape(n_connections,)
and dtypetf.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 atf.Tensor
of shape(n_connections,)
and dtypetf.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 atf.Tensor
of shape(n_connections,)
and dtypetf.float32
.
Notes
The matching is performed using the Munkres algorithm implemented in
scipy.optimize.linear_sum_assignment()
which is wrapped intf_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)
dtypetf.int32
. This can be generated byget_paf_lines()
.peaks_sample – The detected peaks in a sample as a
tf.Tensor
of shape(n_peaks, 2)
and dtypetf.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 dtypetf.int32
with the indices of the peaks that form the source and destination of each candidate connection. This indexes into the inputpeaks_sample
. Can be generated usingget_connection_candidates()
.max_edge_length – Maximum length expected for any connection as a scalar
float
in units of pixels (corresponding topeaks_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 dtypetf.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 typetf.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 atf.RaggedTensor
of shape(n_samples, (n_peaks))
and dtypetf.int32
.skeleton_edges – The indices of the nodes that form the skeleton graph as a
tf.Tensor
of shape(n_edges, 2)
and dtypetf.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 thepafs
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 astf.RaggedTensor
of shape(n_samples, (n_candidates))
and dtypetf.int32
.edge_peak_inds
: Sample-grouped indices of the peaks that form each connection as atf.RaggedTensor
of shape(n_samples, (n_candidates), 2)
and dtypetf.int32
. The last axis corresponds to the[source, destination]
peak indices. These index into the inputpeak_channel_inds
.line_scores
: Sample-grouped scores for each candidate connection astf.RaggedTensor
of shape(n_samples, (n_candidates))
and dtypetf.float32
.
Notes
This function handles the looping over samples in the batch and applies:
get_connection_candidates()
: Find peaks that form connections.get_paf_lines()
: Retrieve PAF vectors for each line.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