utils

sleap_nn.tracking.utils

Helper functions for Tracker module.

Classes:

Name	Description
MaskFeature	Compact mask feature for fast IoU: a tight bbox crop + absolute offset.

Functions:

Name	Description
compute_cosine_sim	Return cosine similarity between appearance-embedding vectors a and b.
compute_euclidean_distance	Return the negative euclidean distance between vectors a and b.
compute_iou	Return the intersection over union for given a and b bounding boxes [xmin, ymin, xmax, ymax].
compute_mask_iou	Return the IoU of two segmentation masks (higher = more similar).
count_valid_points	Number of valid points used to gate track spawning/matching.
cull_frame_instances	Removes instances (for single frame) over instance per frame threshold.
cull_instances	Removes instances from frames over instance per frame threshold.
get_bbox	Return the bounding box coordinates for the PredictedInstance object.
get_centroid	Return the centroid of the PredictedInstance object.
get_embedding	Return the appearance re-ID embedding vector for a detection.
get_keypoints	Return keypoints as np.array from the PredictedInstance object.
get_mask	Return a compact :class:MaskFeature for a PredictedSegmentationMask.
greedy_matching	Match new instances to existing tracks using greedy bipartite matching.
hungarian_matching	Match new instances to existing tracks using Hungarian matching.
is_segmentation_mask	True if obj is a segmentation mask rather than a keypoint instance.
nms_fast	From: https://www.pyimagesearch.com/2015/02/16/faster-non-maximum-suppression-python/.
nms_instances	NMS for instances.

MaskFeature

Compact mask feature for fast IoU: a tight bbox crop + absolute offset.

A full-resolution segmentation mask (e.g. 1280x1024) is almost all background, so the per-pair mask-IoU only needs the foreground bbox: store the cropped boolean array (crop), its absolute top-left (y0, x0) in the original image frame, and the foreground area (px, precomputed). IoU between two features then AND-s only the overlap of their bboxes (usually tiny or empty) and reads area in O(1) — vs. allocating + AND/OR-ing a full-resolution canvas per pair. Coordinates are absolute, so this matches the full-canvas :func:sleap_nn.evaluation._mask_iou exactly.

Methods:

Name	Description
__init__	Store the bbox crop, its absolute top-left (y0, x0) and area.

Source code in sleap_nn/tracking/utils.py

class MaskFeature:
    """Compact mask feature for fast IoU: a tight bbox crop + absolute offset.

    A full-resolution segmentation mask (e.g. 1280x1024) is almost all
    background, so the per-pair mask-IoU only needs the foreground bbox: store
    the cropped boolean array (``crop``), its absolute top-left ``(y0, x0)`` in
    the original image frame, and the foreground ``area`` (px, precomputed). IoU
    between two features then AND-s only the *overlap* of their bboxes (usually
    tiny or empty) and reads ``area`` in O(1) — vs. allocating + AND/OR-ing a
    full-resolution canvas per pair. Coordinates are absolute, so this matches
    the full-canvas :func:`sleap_nn.evaluation._mask_iou` exactly.
    """

    __slots__ = ("crop", "y0", "x0", "area")

    def __init__(self, crop: np.ndarray, y0: int, x0: int, area: int):
        """Store the bbox crop, its absolute top-left ``(y0, x0)`` and area."""
        self.crop = crop
        self.y0 = int(y0)
        self.x0 = int(x0)
        self.area = int(area)

__init__(crop, y0, x0, area)

Store the bbox crop, its absolute top-left (y0, x0) and area.

Source code in sleap_nn/tracking/utils.py

def __init__(self, crop: np.ndarray, y0: int, x0: int, area: int):
    """Store the bbox crop, its absolute top-left ``(y0, x0)`` and area."""
    self.crop = crop
    self.y0 = int(y0)
    self.x0 = int(x0)
    self.area = int(area)

compute_cosine_sim(a, b)

Return cosine similarity between appearance-embedding vectors a and b.

The "cosine_sim" scoring method, paired with the "embeddings" feature (:func:get_embedding). Higher = more similar (a similarity like oks / iou; the cost negation cost = -score happens in :meth:Tracker.scores_to_cost_matrix, so it must NOT be negated here). Output range is [-1, 1].

Hardened against the degenerate inputs the embedding path can produce (utils.py:249, SPEC §7): a None feature (a detection with no appearance embedding -> :func:get_embedding returns None), a zero-norm vector, a shape mismatch, or non-finite values all reduce to NaN rather than a ZeroDivisionError / RuntimeWarning. NaN flows to inf cost in :meth:Tracker.scores_to_cost_matrix (no match), so a missing/garbage embedding spawns a fresh track instead of crashing the run.

Source code in sleap_nn/tracking/utils.py

def compute_cosine_sim(a, b):
    """Return cosine similarity between appearance-embedding vectors ``a`` and ``b``.

    The ``"cosine_sim"`` scoring method, paired with the ``"embeddings"`` feature
    (:func:`get_embedding`). Higher = more similar (a similarity like ``oks`` /
    ``iou``; the cost negation ``cost = -score`` happens in
    :meth:`Tracker.scores_to_cost_matrix`, so it must NOT be negated here).
    Output range is ``[-1, 1]``.

    Hardened against the degenerate inputs the embedding path can produce
    (utils.py:249, SPEC §7): a ``None`` feature (a detection with no appearance
    embedding -> :func:`get_embedding` returns ``None``), a zero-norm vector, a
    shape mismatch, or non-finite values all reduce to ``NaN`` rather than a
    ``ZeroDivisionError`` / ``RuntimeWarning``. ``NaN`` flows to ``inf`` cost in
    :meth:`Tracker.scores_to_cost_matrix` (no match), so a missing/garbage
    embedding spawns a fresh track instead of crashing the run.
    """
    if a is None or b is None:
        return np.nan
    a = np.asarray(a, dtype=np.float64).ravel()
    b = np.asarray(b, dtype=np.float64).ravel()
    if a.size == 0 or a.shape != b.shape:
        return np.nan
    denom = float(np.linalg.norm(a) * np.linalg.norm(b))
    if not np.isfinite(denom) or denom == 0.0:
        return np.nan
    return float(np.dot(a, b) / denom)

compute_euclidean_distance(a, b)

Return the negative euclidean distance between vectors a and b.

A similarity (higher = closer); the cost negation happens in :meth:Tracker.scores_to_cost_matrix. Hardened against the same degenerate inputs as :func:compute_cosine_sim (it is the one explicit alternative metric for features="embeddings", so a detection lacking an appearance vector -> :func:get_embedding returns None, and embeddings of differing dim can co-occur): a None operand, an empty/shape-mismatched pair, or a non-finite result reduces to NaN (-> inf cost -> no match) instead of raising. The real-array centroid path is unaffected (its inputs are never None).

Source code in sleap_nn/tracking/utils.py

def compute_euclidean_distance(a, b):
    """Return the negative euclidean distance between vectors ``a`` and ``b``.

    A similarity (higher = closer); the cost negation happens in
    :meth:`Tracker.scores_to_cost_matrix`. Hardened against the same degenerate
    inputs as :func:`compute_cosine_sim` (it is the one explicit alternative metric
    for ``features="embeddings"``, so a detection lacking an appearance vector ->
    :func:`get_embedding` returns ``None``, and embeddings of differing dim can
    co-occur): a ``None`` operand, an empty/shape-mismatched pair, or a non-finite
    result reduces to ``NaN`` (-> ``inf`` cost -> no match) instead of raising. The
    real-array centroid path is unaffected (its inputs are never ``None``).
    """
    if a is None or b is None:
        return np.nan
    # Fast path for the GEOMETRIC callers (centroids / keypoints), whose features
    # are always real 1-D float64 ndarrays of matching shape. `euclidean_dist` is
    # the metric `apply_tracking` auto-selects for single-node/centroid tracking, so
    # a default run pays the `asarray` + `ravel` normalization below on every pair
    # for no behavioral change; skipping it is ~29% (3.7 -> 2.6 us/pair measured).
    # The `float64` gate keeps this BIT-IDENTICAL to the general path -- a float32
    # embedding vector would otherwise accumulate in float32 and shift the score by
    # ~1e-7, which the appearance-matrix equivalence test catches.
    if (
        type(a) is np.ndarray
        and type(b) is np.ndarray
        and a.dtype == np.float64
        and b.dtype == np.float64
        and a.ndim == 1
        and a.shape == b.shape
        and a.size > 0
    ):
        # `math.sqrt(d.dot(d))` over `np.linalg.norm(d)`: same float64 result (checked
        # bit-for-bit over 250k random coordinate pairs at pixel magnitudes), ~35%
        # less dispatch overhead. The `isfinite` check below still catches an
        # overflow, which `norm`'s internal scaling would have avoided.
        d = a - b
        dist = math.sqrt(d.dot(d))
        return np.nan if not np.isfinite(dist) else -float(dist)
    a = np.asarray(a, dtype=np.float64).ravel()
    b = np.asarray(b, dtype=np.float64).ravel()
    if a.size == 0 or a.shape != b.shape:
        return np.nan
    dist = float(np.linalg.norm(a - b))
    return np.nan if not np.isfinite(dist) else -dist

compute_iou(a, b)

Return the intersection over union for given a and b bounding boxes [xmin, ymin, xmax, ymax].

Source code in sleap_nn/tracking/utils.py

def compute_iou(a, b):
    """Return the intersection over union for given a and b bounding boxes [xmin, ymin, xmax, ymax]."""
    (xmin1, ymin1, xmax1, ymax1), (xmin2, ymin2, xmax2, ymax2) = a, b

    xmin_intersection = max(xmin1, xmin2)
    ymin_intersection = max(ymin1, ymin2)
    xmax_intersection = min(xmax1, xmax2)
    ymax_intersection = min(ymax1, ymax2)

    intersection_area = max(0, xmax_intersection - xmin_intersection + 1) * max(
        0, ymax_intersection - ymin_intersection + 1
    )
    bbox1_area = (xmax1 - xmin1 + 1) * (ymax1 - ymin1 + 1)
    bbox2_area = (xmax2 - xmin2 + 1) * (ymax2 - ymin2 + 1)
    union_area = bbox1_area + bbox2_area - intersection_area

    iou = intersection_area / union_area
    return iou

compute_mask_iou(a, b)

Return the IoU of two segmentation masks (higher = more similar).

The "mask_iou" scoring method. Operates on :class:MaskFeature (the "masks" feature); raw dense bool arrays are accepted too (coerced via :func:get_mask). The intersection is computed only over the overlap of the two foreground bboxes — see :class:MaskFeature — which is numerically identical to the full-canvas :func:sleap_nn.evaluation._mask_iou (top-left aligned, shape-mismatch safe, empty/empty -> 1.0) but avoids touching the background. This is a similarity, like oks/iou; the cost negation (cost = -score) happens in :meth:Tracker.scores_to_cost_matrix, so it must NOT be negated here. Pixel IoU (not bbox-IoU on mask.bbox) sidesteps the XYWH-vs-XYXY bbox-format issue.

Source code in sleap_nn/tracking/utils.py

def compute_mask_iou(a, b) -> float:
    """Return the IoU of two segmentation masks (higher = more similar).

    The ``"mask_iou"`` scoring method. Operates on :class:`MaskFeature` (the
    ``"masks"`` feature); raw dense bool arrays are accepted too (coerced via
    :func:`get_mask`). The intersection is computed only over the *overlap* of
    the two foreground bboxes — see :class:`MaskFeature` — which is numerically
    identical to the full-canvas :func:`sleap_nn.evaluation._mask_iou` (top-left
    aligned, shape-mismatch safe, empty/empty -> 1.0) but avoids touching the
    background. This is a similarity, like ``oks``/``iou``; the cost negation
    (``cost = -score``) happens in :meth:`Tracker.scores_to_cost_matrix`, so it
    must NOT be negated here. Pixel IoU (not bbox-IoU on ``mask.bbox``) sidesteps
    the XYWH-vs-XYXY bbox-format issue.
    """
    fa = a if isinstance(a, MaskFeature) else get_mask(a)
    fb = b if isinstance(b, MaskFeature) else get_mask(b)
    inter = _mask_feature_intersection(fa, fb)
    union = fa.area + fb.area - inter
    # union == 0 only when both masks are empty -> identical -> 1.0 (matches the
    # _mask_iou degenerate contract).
    return 1.0 if union == 0 else float(inter / union)

count_valid_points(obj)

Number of valid points used to gate track spawning/matching.

For a keypoint instance this is the count of non-NaN nodes; for a segmentation mask there are no keypoints, so the mask area (foreground px) is the analogous "support" measure (min_new_track_points / min_match_points then read as a pixel-area threshold; default 0 keeps every non-empty mask and drops empty ones).

Source code in sleap_nn/tracking/utils.py

def count_valid_points(obj) -> int:
    """Number of valid points used to gate track spawning/matching.

    For a keypoint instance this is the count of non-NaN nodes; for a
    segmentation mask there are no keypoints, so the mask area (foreground px)
    is the analogous "support" measure (``min_new_track_points`` /
    ``min_match_points`` then read as a pixel-area threshold; default 0 keeps
    every non-empty mask and drops empty ones).
    """
    if is_segmentation_mask(obj):
        return int(obj.area)
    points = obj if isinstance(obj, np.ndarray) else obj.numpy()
    return int((~np.isnan(points).any(axis=1)).sum())

cull_frame_instances(instances_list, instance_count, iou_threshold=None)

Removes instances (for single frame) over instance per frame threshold.

Parameters:

Name	Type	Description	Default
instances_list	List[PredictedInstance]	The list of instances for a single frame.	required
instance_count	int	The maximum number of instances we want per frame.	required
iou_threshold	Optional[float]	Intersection over Union (IOU) threshold to use when removing overlapping instances over target count; if None, then only use score to determine which instances to remove.	None

Returns:

Type	Description
List[PredictedInstance]	Updated list of frames, also modifies frames in place.

Source code in sleap_nn/tracking/utils.py

def cull_frame_instances(
    instances_list: List[sio.PredictedInstance],
    instance_count: int,
    iou_threshold: Optional[float] = None,
) -> List[sio.PredictedInstance]:
    """Removes instances (for single frame) over instance per frame threshold.

    Args:
        instances_list: The list of instances for a single frame.
        instance_count: The maximum number of instances we want per frame.
        iou_threshold: Intersection over Union (IOU) threshold to use when
            removing overlapping instances over target count; if None, then
            only use score to determine which instances to remove.

    Returns:
        Updated list of frames, also modifies frames in place.
    """
    if not instances_list:
        return

    if len(instances_list) > instance_count:
        # List of instances which we'll pare down
        keep_instances = instances_list

        # Use NMS to remove overlapping instances over target count
        if iou_threshold:
            keep_instances, extra_instances = nms_instances(
                keep_instances,
                iou_threshold=iou_threshold,
                target_count=instance_count,
            )
            updated_instances_list = []
            # Remove the extra instances
            for inst in extra_instances:
                for instance in instances_list:
                    if not instance.same_pose_as(inst):
                        updated_instances_list.append(instance)
            instances_list = updated_instances_list

        # Use lower score to remove instances over target count
        if len(keep_instances) > instance_count:
            # Sort by ascending score, get target number of instances
            # from the end of list (i.e., with highest score)
            extra_instances = sorted(keep_instances, key=operator.attrgetter("score"))[
                :-instance_count
            ]

            # Remove the extra instances
            updated_instances_list = []
            for inst in extra_instances:
                for instance in instances_list:
                    if instance.same_pose_as(inst):
                        updated_instances_list.append(instance)
            instances_list = updated_instances_list

    return instances_list

cull_instances(frames, instance_count, iou_threshold=None)

Removes instances from frames over instance per frame threshold.

Parameters:

Name	Type	Description	Default
frames	List[LabeledFrame]	The list of LabeledFrame objects with predictions.	required
instance_count	int	The maximum number of instances we want per frame.	required
iou_threshold	Optional[float]	Intersection over Union (IOU) threshold to use when removing overlapping instances over target count; if None, then only use score to determine which instances to remove.	None

Returns:

Type	Description
	The (possibly empty) input frames, also modified in place.

Source code in sleap_nn/tracking/utils.py

def cull_instances(
    frames: List[sio.LabeledFrame],
    instance_count: int,
    iou_threshold: Optional[float] = None,
):
    """Removes instances from frames over instance per frame threshold.

    Args:
        frames: The list of `LabeledFrame` objects with predictions.
        instance_count: The maximum number of instances we want per frame.
        iou_threshold: Intersection over Union (IOU) threshold to use when
            removing overlapping instances over target count; if None, then
            only use score to determine which instances to remove.

    Returns:
        The (possibly empty) input `frames`, also modified in place.
    """
    if not frames:
        return frames

    frames.sort(key=lambda lf: lf.frame_idx)

    lf_inst_list = []
    # Find all frames with more instances than the desired threshold
    for lf in frames:
        if len(lf.predicted_instances) > instance_count:
            # List of instances which we'll pare down
            keep_instances = lf.predicted_instances

            # Use NMS to remove overlapping instances over target count
            if iou_threshold:
                keep_instances, extra_instances = nms_instances(
                    keep_instances,
                    iou_threshold=iou_threshold,
                    target_count=instance_count,
                )
                # Mark for removal
                lf_inst_list.extend([(lf, inst) for inst in extra_instances])

            # Use lower score to remove instances over target count
            if len(keep_instances) > instance_count:
                # Sort by ascending score, get target number of instances
                # from the end of list (i.e., with highest score)
                extra_instances = sorted(
                    keep_instances, key=operator.attrgetter("score")
                )[:-instance_count]

                # Mark for removal
                lf_inst_list.extend([(lf, inst) for inst in extra_instances])

    # Remove instances over per frame threshold
    for lf, inst in lf_inst_list:
        filtered_instances = []
        for instance in lf.instances:
            if not instance.same_pose_as(inst):
                filtered_instances.append(instance)
        lf.instances = filtered_instances

    return frames

get_bbox(pred_instance)

Return the bounding box coordinates for the PredictedInstance object.

Source code in sleap_nn/tracking/utils.py

def get_bbox(pred_instance: Union[sio.PredictedInstance, np.ndarray]):
    """Return the bounding box coordinates for the `PredictedInstance` object."""
    points = (
        pred_instance.numpy()
        if not isinstance(pred_instance, np.ndarray)
        else pred_instance
    )
    bbox = np.concatenate(
        [
            np.nanmin(points, axis=0),
            np.nanmax(points, axis=0),
        ]  # [xmin, ymin, xmax, ymax]
    )
    return bbox

get_centroid(pred_instance)

Return the centroid of the PredictedInstance object.

Source code in sleap_nn/tracking/utils.py

def get_centroid(pred_instance: Union[sio.PredictedInstance, np.ndarray]):
    """Return the centroid of the `PredictedInstance` object."""
    pts = pred_instance
    if not isinstance(pred_instance, np.ndarray):
        pts = pred_instance.numpy()
    centroid = np.nanmedian(pts, axis=0)
    return centroid

get_embedding(pred_instance)

Return the appearance re-ID embedding vector for a detection.

The "embeddings" feature extractor (mirrors :func:get_keypoints / :func:get_mask). Reads the appearance vector attached by the embedding (re-ID) model into the single identity_embedding slot (sleap-io #535); works on any embedding-carrying detection (PredictedInstance or PredictedSegmentationMask — both carry identity_embedding), so embedding tracking is keypoint/mask agnostic. Scored by :func:compute_cosine_sim.

Returns the vector as a float32 np.ndarray of shape (D,), or None when the detection carries no embedding. None is the "no feature" sentinel: :func:compute_cosine_sim maps it to NaN (-> inf cost in :meth:Tracker.scores_to_cost_matrix), so a detection that is missing its embedding simply never matches and spawns a fresh track instead of crashing the run. An np.ndarray is passed through unchanged (a precomputed feature).

Source code in sleap_nn/tracking/utils.py

def get_embedding(pred_instance):
    """Return the appearance re-ID embedding vector for a detection.

    The ``"embeddings"`` feature extractor (mirrors :func:`get_keypoints` /
    :func:`get_mask`). Reads the appearance vector attached by the ``embedding``
    (re-ID) model into the single ``identity_embedding`` slot (sleap-io #535); works
    on any embedding-carrying detection (``PredictedInstance`` *or*
    ``PredictedSegmentationMask`` — both carry ``identity_embedding``), so embedding
    tracking is keypoint/mask agnostic. Scored by :func:`compute_cosine_sim`.

    Returns the vector as a ``float32`` ``np.ndarray`` of shape ``(D,)``, or
    ``None`` when the detection carries no embedding. ``None`` is the "no feature"
    sentinel: :func:`compute_cosine_sim` maps it to ``NaN`` (-> ``inf`` cost in
    :meth:`Tracker.scores_to_cost_matrix`), so a detection that is missing its
    embedding simply never matches and spawns a fresh track instead of crashing the
    run. An ``np.ndarray`` is passed through unchanged (a precomputed feature).
    """
    if isinstance(pred_instance, np.ndarray):
        return pred_instance
    emb = getattr(pred_instance, "identity_embedding", None)
    if emb is None:
        return None
    return np.asarray(emb.vector, dtype=np.float32)

get_keypoints(pred_instance)

Return keypoints as np.array from the PredictedInstance object.

Source code in sleap_nn/tracking/utils.py

def get_keypoints(pred_instance: Union[sio.PredictedInstance, np.ndarray]):
    """Return keypoints as np.array from the `PredictedInstance` object."""
    if isinstance(pred_instance, np.ndarray):
        return pred_instance
    return pred_instance.numpy()

get_mask(pred_mask)

Return a compact :class:MaskFeature for a PredictedSegmentationMask.

Mirrors :func:get_keypoints/:func:get_bbox as the "masks" feature extractor. The mask is decoded onto the image-pixel grid first, then cropped to its foreground bbox (using the mask's own .bbox when available, avoiding a scan); the crop is cached as the candidate feature so scoring reuses it without re-decoding and only touches the foreground region.

The image-grid decode is essential: sio .data decodes at the mask's stored resolution, which for the default inference path (full_res_masks=False, masks encoded at output-stride, scale~=0.5) is NOT the image grid, while .bbox is always in IMAGE space. Cropping the stride-res .data with image-space bbox indices would slice the wrong region (often entirely out of bounds -> empty -> compute_mask_iou 1.0 for every pair, scrambling identity). :func:decode_mask_to_image_res is a zero-copy passthrough for scale==1 masks (legacy full-res), so that path is unchanged.

Source code in sleap_nn/tracking/utils.py

def get_mask(
    pred_mask: Union["sio.PredictedSegmentationMask", np.ndarray, MaskFeature],
) -> MaskFeature:
    """Return a compact :class:`MaskFeature` for a `PredictedSegmentationMask`.

    Mirrors :func:`get_keypoints`/:func:`get_bbox` as the ``"masks"`` feature
    extractor. The mask is decoded onto the **image-pixel grid** first, then
    cropped to its foreground bbox (using the mask's own ``.bbox`` when available,
    avoiding a scan); the crop is cached as the candidate feature so scoring
    reuses it without re-decoding and only touches the foreground region.

    The image-grid decode is essential: sio ``.data`` decodes at the mask's
    *stored* resolution, which for the default inference path
    (``full_res_masks=False``, masks encoded at output-stride, ``scale~=0.5``) is
    NOT the image grid, while ``.bbox`` is always in IMAGE space. Cropping the
    stride-res ``.data`` with image-space bbox indices would slice the wrong
    region (often entirely out of bounds -> empty -> ``compute_mask_iou`` 1.0 for
    every pair, scrambling identity). :func:`decode_mask_to_image_res` is a
    zero-copy passthrough for ``scale==1`` masks (legacy full-res), so that path
    is unchanged.
    """
    if isinstance(pred_mask, MaskFeature):
        return pred_mask
    if isinstance(pred_mask, np.ndarray):
        return _mask_feature_from_dense(pred_mask)
    from sleap_nn.inference.segmentation_convert import decode_mask_to_image_res

    data = decode_mask_to_image_res(pred_mask)
    bbox = getattr(pred_mask, "bbox", None)
    if bbox is not None:
        # PredictedSegmentationMask.bbox is (x, y, width, height) (XYWH), in IMAGE
        # space -- consistent with the image-grid `data` decoded above.
        x, y, w, h = (int(round(float(v))) for v in bbox)
        height, width = data.shape
        y0, x0 = max(0, y), max(0, x)
        y1, x1 = min(height, y + h), min(width, x + w)
        if y1 > y0 and x1 > x0:
            crop = data[y0:y1, x0:x1]
            return MaskFeature(crop, y0, x0, int(np.count_nonzero(crop)))
    return _mask_feature_from_dense(data)

greedy_matching(cost_matrix)

Match new instances to existing tracks using greedy bipartite matching.

Source code in sleap_nn/tracking/utils.py

def greedy_matching(cost_matrix: np.ndarray) -> List[Tuple[int, int]]:
    """Match new instances to existing tracks using greedy bipartite matching."""
    # Sort edges by ascending cost.
    rows, cols = np.unravel_index(np.argsort(cost_matrix, axis=None), cost_matrix.shape)
    unassigned_edges = list(zip(rows, cols))

    # Greedily assign edges.
    row_inds, col_inds = [], []
    while len(unassigned_edges) > 0:
        # Assign the lowest cost edge.
        row_ind, col_ind = unassigned_edges.pop(0)
        row_inds.append(row_ind)
        col_inds.append(col_ind)

        # Remove all other edges that contain either node (in reverse order).
        for i in range(len(unassigned_edges) - 1, -1, -1):
            if unassigned_edges[i][0] == row_ind or unassigned_edges[i][1] == col_ind:
                del unassigned_edges[i]

    return row_inds, col_inds

hungarian_matching(cost_matrix)

Match new instances to existing tracks using Hungarian matching.

Source code in sleap_nn/tracking/utils.py

def hungarian_matching(cost_matrix: np.ndarray) -> List[Tuple[int, int]]:
    """Match new instances to existing tracks using Hungarian matching."""
    # Replace inf/nan with a large finite value so linear_sum_assignment doesn't
    # raise "cost matrix is infeasible".
    invalid = ~np.isfinite(cost_matrix)
    if invalid.any():
        cost_matrix = np.copy(cost_matrix)
        finite_vals = cost_matrix[~invalid]
        fill = (np.abs(finite_vals).max() * 10 + 1) if finite_vals.size > 0 else 1e6
        cost_matrix[invalid] = fill

    row_ids, col_ids = linear_sum_assignment(cost_matrix)
    return row_ids, col_ids

is_segmentation_mask(obj)

True if obj is a segmentation mask rather than a keypoint instance.

Mask tracking flows sio.PredictedSegmentationMask objects through the same code paths as sio.PredictedInstance; this is the single predicate used to dispatch the pose-vs-mask differences (no .numpy() keypoints).

Source code in sleap_nn/tracking/utils.py

def is_segmentation_mask(obj) -> bool:
    """True if ``obj`` is a segmentation mask rather than a keypoint instance.

    Mask tracking flows ``sio.PredictedSegmentationMask`` objects through the
    same code paths as ``sio.PredictedInstance``; this is the single predicate
    used to dispatch the pose-vs-mask differences (no ``.numpy()`` keypoints).
    """
    return isinstance(obj, (sio.PredictedSegmentationMask, sio.SegmentationMask))

nms_fast(boxes, scores, iou_threshold, target_count=None)

From: https://www.pyimagesearch.com/2015/02/16/faster-non-maximum-suppression-python/.

Source code in sleap_nn/tracking/utils.py

def nms_fast(boxes, scores, iou_threshold, target_count=None) -> List[int]:
    """From: https://www.pyimagesearch.com/2015/02/16/faster-non-maximum-suppression-python/."""
    # if there are no boxes, return an empty list
    if len(boxes) == 0:
        return []

    # if we already have fewer boxes than the target count, return all boxes
    if target_count and len(boxes) < target_count:
        return list(range(len(boxes)))

    # if the bounding boxes coordinates are integers, convert them to floats --
    # this is important since we'll be doing a bunch of divisions
    if boxes.dtype.kind == "i":
        boxes = boxes.astype("float")

    # initialize the list of picked indexes
    picked_idxs = []

    # init list of boxes removed by nms
    nms_idxs = []
    # grab the coordinates of the bounding boxes
    x1 = boxes[:, 0]
    y1 = boxes[:, 1]
    x2 = boxes[:, 2]
    y2 = boxes[:, 3]

    # compute the area of the bounding boxes and sort the bounding
    # boxes by their scores
    area = (x2 - x1 + 1) * (y2 - y1 + 1)
    idxs = np.argsort(scores)

    # keep looping while some indexes still remain in the indexes list
    while len(idxs) > 0:
        # we want to add the best box which is the last box in sorted list
        picked_box_idx = idxs[-1]

        # last = len(idxs) - 1
        # i = idxs[last]
        picked_idxs.append(picked_box_idx)

        # find the largest (x, y) coordinates for the start of
        # the bounding box and the smallest (x, y) coordinates
        # for the end of the bounding box
        xx1 = np.maximum(x1[picked_box_idx], x1[idxs[:-1]])
        yy1 = np.maximum(y1[picked_box_idx], y1[idxs[:-1]])
        xx2 = np.minimum(x2[picked_box_idx], x2[idxs[:-1]])
        yy2 = np.minimum(y2[picked_box_idx], y2[idxs[:-1]])

        # compute the width and height of the bounding box
        w = np.maximum(0, xx2 - xx1 + 1)
        h = np.maximum(0, yy2 - yy1 + 1)

        # compute the ratio of overlap
        overlap = (w * h) / area[idxs[:-1]]

        # find boxes with iou over threshold
        nms_for_new_box = np.where(overlap > iou_threshold)[0]
        nms_idxs.extend(list(idxs[nms_for_new_box]))

        # delete new box (last in list) plus nms boxes
        idxs = np.delete(idxs, nms_for_new_box)[:-1]

    # if we're below the target number of boxes, add some back
    if target_count and nms_idxs and len(picked_idxs) < target_count:
        # sort by descending score
        nms_idxs.sort(key=lambda idx: -scores[idx])

        add_back_count = min(len(nms_idxs), len(picked_idxs) - target_count)
        picked_idxs.extend(nms_idxs[:add_back_count])

    # return the list of picked boxes
    return picked_idxs

nms_instances(instances, iou_threshold, target_count=None)

NMS for instances.

Source code in sleap_nn/tracking/utils.py

def nms_instances(
    instances, iou_threshold, target_count=None
) -> Tuple[List[sio.PredictedInstance], List[sio.PredictedInstance]]:
    """NMS for instances."""
    # get_bbox: # [xmin, ymin, xmax, ymax]
    boxes = np.array([get_bbox(inst) for inst in instances])
    scores = np.array([inst.score for inst in instances])
    picks = nms_fast(boxes, scores, iou_threshold, target_count)

    to_keep = [inst for i, inst in enumerate(instances) if i in picks]
    to_remove = [inst for i, inst in enumerate(instances) if i not in picks]

    return to_keep, to_remove