association

sleap_nn.tracking.mot.association

Similarity matrices, gates and assignment helpers for the MOT trackers.

Every function here is vectorized and NaN-free on output unless documented. Similarities are in [0, 1]; costs are 1 - similarity. Rows are tracks (or the first argument), columns are detections (or the second argument). Every function accepts empty inputs and returns correctly shaped (n, 0) / (0, m) results.

Functions:

Name	Description
adaptive_weight	Deep OC-SORT adaptive weighting of the appearance matrix E.
appearance_margins	Top-1 minus top-2 cosine per row and per column over finite entries.
centroid_similarity	Gaussian similarity exp(-d^2 / (2 scale^2)) of (n, 2) vs (m, 2) points.
cosine_matrix	Cosine similarity of L2-normalized vectors; NaN where either is missing.
decisive_cells	Cells where appearance decides a pair (per cell, not per row/column).
fuse_score	ByteTrack score fusion: sim * det_scores[None, :].
ioa_min_matrix	Intersection over the smaller area (duplicate detection); NaN boxes score 0.
iou_matrix	IoU between (n, 4) and (m, 4) xyxy boxes; NaN boxes score 0.
linear_assignment	Minimum-cost assignment with a cost cap.
linear_assignment_masked	Maximum-score assignment, keeping only valid pairs.
ocm_matrix	OC-SORT observation-centric momentum term.
oks_matrix	Object keypoint similarity between two sets of poses.
one_to_one_pattern	OC-SORT/SORT shortcut: the valid pairs if they form a 1:1 pattern.
speed_direction	Unit (dy, dx) direction between box centers, row-wise for (n, 4) inputs.

adaptive_weight(E, aw_param=0.5, base=0.75)

Deep OC-SORT adaptive weighting of the appearance matrix E.

W = base + 0.5 * min(row_top1 - row_top2, aw_param) + 0.5 * min(col_top1 - col_top2, aw_param), computed only when E has at least 2 rows and 2 columns; otherwise every entry is base.

Non-finite entries of E are ignored when ranking; a row or column with fewer than two finite entries gets no bonus.

Parameters:

Name	Type	Description	Default
E	ndarray	(n, m) appearance similarities.	required
aw_param	float	Cap on each margin bonus (Deep OC-SORT aw_param).	0.5
base	float	Constant weight (Deep OC-SORT w_assoc_emb).	0.75

Returns:

Type	Description
ndarray	(n, m) weights.

Source code in sleap_nn/tracking/mot/association.py

def adaptive_weight(
    E: np.ndarray, aw_param: float = 0.5, base: float = 0.75
) -> np.ndarray:
    """Deep OC-SORT adaptive weighting of the appearance matrix ``E``.

    ``W = base + 0.5 * min(row_top1 - row_top2, aw_param) + 0.5 * min(col_top1 -
    col_top2, aw_param)``, computed only when ``E`` has at least 2 rows and 2
    columns; otherwise every entry is ``base``.

    Non-finite entries of ``E`` are ignored when ranking; a row or column with
    fewer than two finite entries gets no bonus.

    Args:
        E: ``(n, m)`` appearance similarities.
        aw_param: Cap on each margin bonus (Deep OC-SORT ``aw_param``).
        base: Constant weight (Deep OC-SORT ``w_assoc_emb``).

    Returns:
        ``(n, m)`` weights.
    """
    E = np.asarray(E, dtype=np.float64)
    W = np.full(E.shape, float(base), dtype=np.float64)
    if E.ndim != 2 or E.shape[0] < 2 or E.shape[1] < 2:
        return W
    row_margin, col_margin = appearance_margins(E)
    row_bonus = np.where(np.isfinite(row_margin), row_margin, 0.0)
    col_bonus = np.where(np.isfinite(col_margin), col_margin, 0.0)
    row_bonus = 0.5 * np.minimum(row_bonus, aw_param)
    col_bonus = 0.5 * np.minimum(col_bonus, aw_param)
    return W + row_bonus[:, None] + col_bonus[None, :]

appearance_margins(cos)

Top-1 minus top-2 cosine per row and per column over finite entries.

A row/column with a single finite entry has margin inf; with none, nan.

Parameters:

Name	Type	Description	Default
cos	ndarray	(n, m) cosines (NaN where unavailable).	required

Returns:

Type	Description
Tuple[ndarray, ndarray]	(row_margin (n,), col_margin (m,)).

Source code in sleap_nn/tracking/mot/association.py

def appearance_margins(cos: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
    """Top-1 minus top-2 cosine per row and per column over finite entries.

    A row/column with a single finite entry has margin ``inf``; with none, ``nan``.

    Args:
        cos: ``(n, m)`` cosines (NaN where unavailable).

    Returns:
        ``(row_margin (n,), col_margin (m,))``.
    """
    cos = np.asarray(cos, dtype=np.float64)
    if cos.ndim != 2:
        raise ValueError(f"cos must be a 2-D matrix, got shape {cos.shape}.")
    return _top2_margin(cos), _top2_margin(cos.T)

centroid_similarity(c_a, c_b, scale)

Gaussian similarity exp(-d^2 / (2 scale^2)) of (n, 2) vs (m, 2) points.

NaN points score 0.

Parameters:

Name	Type	Description	Default
c_a	ndarray	(n, 2) points.	required
c_b	ndarray	(m, 2) points.	required
scale	float	Gaussian scale in px (body_size / 2 for the trackers).	required

Returns:

Type	Description
ndarray	(n, m) similarities in [0, 1].

Raises:

Type	Description
ValueError	If scale is not positive and finite.

Source code in sleap_nn/tracking/mot/association.py

def centroid_similarity(c_a: np.ndarray, c_b: np.ndarray, scale: float) -> np.ndarray:
    """Gaussian similarity ``exp(-d^2 / (2 scale^2))`` of ``(n, 2)`` vs ``(m, 2)`` points.

    NaN points score 0.

    Args:
        c_a: ``(n, 2)`` points.
        c_b: ``(m, 2)`` points.
        scale: Gaussian scale in px (``body_size / 2`` for the trackers).

    Returns:
        ``(n, m)`` similarities in ``[0, 1]``.

    Raises:
        ValueError: If ``scale`` is not positive and finite.
    """
    scale = float(scale)
    if not (np.isfinite(scale) and scale > 0):
        raise ValueError(f"Centroid similarity scale must be > 0, got {scale}.")
    c_a = np.asarray(c_a, dtype=np.float64).reshape(-1, 2)
    c_b = np.asarray(c_b, dtype=np.float64).reshape(-1, 2)
    if len(c_a) == 0 or len(c_b) == 0:
        return np.zeros((len(c_a), len(c_b)), dtype=np.float64)
    d2 = ((c_a[:, None] - c_b[None]) ** 2).sum(axis=-1)
    ok = np.isfinite(d2)
    return np.where(ok, np.exp(-np.where(ok, d2, 0.0) / (2.0 * scale**2)), 0.0)

cosine_matrix(e_a, ok_a, e_b, ok_b)

Cosine similarity of L2-normalized vectors; NaN where either is missing.

Parameters:

Name	Type	Description	Default
e_a	ndarray	(n, D) vectors (rows where ok_a is False are ignored), or None (all missing).	required
ok_a	ndarray	(n,) bool availability.	required
e_b	ndarray	(m, D) vectors, or None.	required
ok_b	ndarray	(m,) bool availability.	required

Returns:

Type	Description
ndarray	(n, m) cosines in [-1, 1], NaN where unavailable (also everywhere when the two vector dimensions differ).

Source code in sleap_nn/tracking/mot/association.py

def cosine_matrix(
    e_a: np.ndarray, ok_a: np.ndarray, e_b: np.ndarray, ok_b: np.ndarray
) -> np.ndarray:
    """Cosine similarity of L2-normalized vectors; NaN where either is missing.

    Args:
        e_a: ``(n, D)`` vectors (rows where ``ok_a`` is False are ignored), or
            ``None`` (all missing).
        ok_a: ``(n,)`` bool availability.
        e_b: ``(m, D)`` vectors, or ``None``.
        ok_b: ``(m,)`` bool availability.

    Returns:
        ``(n, m)`` cosines in ``[-1, 1]``, NaN where unavailable (also everywhere
        when the two vector dimensions differ).
    """
    ok_a = np.asarray(ok_a, dtype=bool).reshape(-1)
    ok_b = np.asarray(ok_b, dtype=bool).reshape(-1)
    out = np.full((len(ok_a), len(ok_b)), np.nan, dtype=np.float64)
    if out.size == 0 or e_a is None or e_b is None:
        return out
    e_a = np.asarray(e_a, dtype=np.float64)
    e_b = np.asarray(e_b, dtype=np.float64)
    if e_a.ndim != 2 or e_b.ndim != 2 or e_a.shape[1] != e_b.shape[1]:
        return out
    ok_a = ok_a & np.isfinite(e_a).all(axis=1)
    ok_b = ok_b & np.isfinite(e_b).all(axis=1)
    ia, ib = np.flatnonzero(ok_a), np.flatnonzero(ok_b)
    if len(ia) and len(ib):
        out[np.ix_(ia, ib)] = np.clip(e_a[ia] @ e_b[ib].T, -1.0, 1.0)
    return out

decisive_cells(cos, admissible, margin=0.0)

Cells where appearance decides a pair (per cell, not per row/column).

A cell is decisive when it is admissible and its cosine is the top admissible cosine of its row AND of its column, with each top-1 minus top-2 margin strictly positive and at least margin (a lone admissible cell has an infinite margin). A row's runner-up is never decisive, even when the row's margin is large: appearance ranks it below another candidate. Shared by BoT-SORT's appearance discount (admissible: inside the proximity and appearance gates) and Deep OC-SORT's rescue rule (admissible: geometrically plausible).

Parameters:

Name	Type	Description	Default
cos	ndarray	(n, m) cosines (NaN where unavailable).	required
admissible	ndarray	(n, m) bool cells appearance may decide at all.	required
margin	float	Minimum row and column margin (the calibrated m_assoc).	0.0

Returns:

Type	Description
ndarray	(n, m) bool mask.

Source code in sleap_nn/tracking/mot/association.py

def decisive_cells(
    cos: np.ndarray, admissible: np.ndarray, margin: float = 0.0
) -> np.ndarray:
    """Cells where appearance decides a pair (per cell, not per row/column).

    A cell is decisive when it is admissible and its cosine is the top admissible
    cosine of its row AND of its column, with each top-1 minus top-2 margin
    strictly positive and at least ``margin`` (a lone admissible cell has an
    infinite margin). A row's runner-up is never decisive, even when the row's
    margin is large: appearance ranks it below another candidate. Shared by
    BoT-SORT's appearance discount (admissible: inside the proximity and
    appearance gates) and Deep OC-SORT's rescue rule (admissible: geometrically
    plausible).

    Args:
        cos: ``(n, m)`` cosines (NaN where unavailable).
        admissible: ``(n, m)`` bool cells appearance may decide at all.
        margin: Minimum row and column margin (the calibrated ``m_assoc``).

    Returns:
        ``(n, m)`` bool mask.
    """
    cos = np.asarray(cos, dtype=np.float64)
    admissible = np.asarray(admissible, dtype=bool) & np.isfinite(cos)
    if cos.size == 0 or not admissible.any():
        return np.zeros(cos.shape, dtype=bool)
    row_margin, col_margin = appearance_margins(np.where(admissible, cos, np.nan))
    filled = np.where(admissible, cos, -np.inf)
    is_top = (filled >= filled.max(axis=1, keepdims=True)) & (
        filled >= filled.max(axis=0, keepdims=True)
    )
    row_margin = np.nan_to_num(row_margin, nan=0.0)
    col_margin = np.nan_to_num(col_margin, nan=0.0)
    margin = max(float(margin), 0.0)
    row_ok = (row_margin > 0.0) & (row_margin >= margin)
    col_ok = (col_margin > 0.0) & (col_margin >= margin)
    return admissible & is_top & row_ok[:, None] & col_ok[None, :]

fuse_score(sim, det_scores)

ByteTrack score fusion: sim * det_scores[None, :].

Parameters:

Name	Type	Description	Default
sim	ndarray	(n, m) similarities.	required
det_scores	ndarray	(m,) detection scores in [0, 1].	required

Returns:

Type	Description
ndarray	(n, m) fused similarities.

Source code in sleap_nn/tracking/mot/association.py

def fuse_score(sim: np.ndarray, det_scores: np.ndarray) -> np.ndarray:
    """ByteTrack score fusion: ``sim * det_scores[None, :]``.

    Args:
        sim: ``(n, m)`` similarities.
        det_scores: ``(m,)`` detection scores in ``[0, 1]``.

    Returns:
        ``(n, m)`` fused similarities.
    """
    sim = np.asarray(sim, dtype=np.float64)
    det_scores = np.asarray(det_scores, dtype=np.float64).reshape(-1)
    return sim * det_scores[None, :]

ioa_min_matrix(a, b)

Intersection over the smaller area (duplicate detection); NaN boxes score 0.

Parameters:

Name	Type	Description	Default
a	ndarray	(n, 4) xyxy boxes.	required
b	ndarray	(m, 4) xyxy boxes.	required

Returns:

Type	Description
ndarray	(n, m) values in [0, 1]: 1 when one box contains the other. A pair involving a zero-area or non-finite box scores 0.

Source code in sleap_nn/tracking/mot/association.py

def ioa_min_matrix(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """Intersection over the smaller area (duplicate detection); NaN boxes score 0.

    Args:
        a: ``(n, 4)`` ``xyxy`` boxes.
        b: ``(m, 4)`` ``xyxy`` boxes.

    Returns:
        ``(n, m)`` values in ``[0, 1]``: 1 when one box contains the other. A
        pair involving a zero-area or non-finite box scores 0.
    """
    a, b = _as_boxes(a), _as_boxes(b)
    if len(a) == 0 or len(b) == 0:
        return np.zeros((len(a), len(b)), dtype=np.float64)
    inter, area_a, area_b, _ = _intersection_and_areas(a, b)
    return _safe_ratio(inter, np.minimum(area_a[:, None], area_b[None, :]))

iou_matrix(a, b)

IoU between (n, 4) and (m, 4) xyxy boxes; NaN boxes score 0.

Parameters:

Name	Type	Description	Default
a	ndarray	(n, 4) boxes.	required
b	ndarray	(m, 4) boxes.	required

Returns:

Type	Description
ndarray	(n, m) IoU in [0, 1]. Boxes with a non-finite coordinate, and pairs whose union is empty (two zero-area boxes), score 0.

Source code in sleap_nn/tracking/mot/association.py

def iou_matrix(a: np.ndarray, b: np.ndarray) -> np.ndarray:
    """IoU between ``(n, 4)`` and ``(m, 4)`` ``xyxy`` boxes; NaN boxes score 0.

    Args:
        a: ``(n, 4)`` boxes.
        b: ``(m, 4)`` boxes.

    Returns:
        ``(n, m)`` IoU in ``[0, 1]``. Boxes with a non-finite coordinate, and
        pairs whose union is empty (two zero-area boxes), score 0.
    """
    a, b = _as_boxes(a), _as_boxes(b)
    if len(a) == 0 or len(b) == 0:
        return np.zeros((len(a), len(b)), dtype=np.float64)
    inter, area_a, area_b, _ = _intersection_and_areas(a, b)
    union = area_a[:, None] + area_b[None, :] - inter
    return _safe_ratio(inter, union)

linear_assignment(cost, max_cost)

Minimum-cost assignment with a cost cap.

Entries above max_cost or non-finite are infeasible: they are replaced by max_cost + 1e-5 before solving with :func:scipy.optimize.linear_sum_assignment, and matched pairs that were infeasible are then dropped. Pricing a forbidden pair just above the cap (DeepSORT's max_distance + 1e-5; equivalently ByteTrack/BoT-SORT's lapjv(cost_limit=max_cost)) means leaving a row and a column unmatched costs about max_cost: the solver never pulls a track off a confident match just to squeeze in more matches. With an infinite max_cost only non-finite entries are infeasible (priced INFEASIBLE_COST above the largest finite cost). Empty matrices are handled; the result is deterministic.

Parameters:

Name	Type	Description	Default
cost	ndarray	(n, m) costs (lower is better).	required
max_cost	float	Largest admissible cost of a matched pair (inclusive).	required

Returns:

Type	Description
Tuple[ndarray, ndarray, ndarray]	(matches, unmatched_rows, unmatched_cols) where matches is an (k, 2) int array of (row, col) sorted by row, and the unmatched arrays are sorted int arrays.

Raises:

Type	Description
ValueError	If max_cost is NaN or cost is not 2-D.

Source code in sleap_nn/tracking/mot/association.py

def linear_assignment(
    cost: np.ndarray, max_cost: float
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Minimum-cost assignment with a cost cap.

    Entries above ``max_cost`` or non-finite are infeasible: they are replaced by
    ``max_cost + 1e-5`` before solving with
    :func:`scipy.optimize.linear_sum_assignment`, and matched pairs that were
    infeasible are then dropped. Pricing a forbidden pair just above the cap
    (DeepSORT's ``max_distance + 1e-5``; equivalently ByteTrack/BoT-SORT's
    ``lapjv(cost_limit=max_cost)``) means leaving a row and a column unmatched
    costs about ``max_cost``: the solver never pulls a track off a confident
    match just to squeeze in more matches. With an infinite ``max_cost`` only
    non-finite entries are infeasible (priced ``INFEASIBLE_COST`` above the
    largest finite cost). Empty matrices are handled; the result is
    deterministic.

    Args:
        cost: ``(n, m)`` costs (lower is better).
        max_cost: Largest admissible cost of a matched pair (inclusive).

    Returns:
        ``(matches, unmatched_rows, unmatched_cols)`` where ``matches`` is an
        ``(k, 2)`` int array of ``(row, col)`` sorted by row, and the unmatched
        arrays are sorted int arrays.

    Raises:
        ValueError: If ``max_cost`` is NaN or ``cost`` is not 2-D.
    """
    cost = np.asarray(cost, dtype=np.float64)
    if cost.ndim != 2:
        raise ValueError(f"cost must be a 2-D matrix, got shape {cost.shape}.")
    max_cost = float(max_cost)
    if np.isnan(max_cost):
        raise ValueError("max_cost must not be NaN.")
    n, m = cost.shape
    if n == 0 or m == 0:
        return _assignment_result([], [], n, m)

    feasible = np.isfinite(cost) & (cost <= max_cost)
    if not feasible.any():
        return _assignment_result([], [], n, m)
    if np.isfinite(max_cost):
        fill = max_cost + _CAP_EPS
    else:
        fill = float(cost[feasible].max()) + INFEASIBLE_COST
    solve = np.where(feasible, cost, fill)
    rows, cols = linear_sum_assignment(solve)
    keep = feasible[rows, cols]
    return _assignment_result(rows[keep], cols[keep], n, m)

linear_assignment_masked(score, valid)

Maximum-score assignment, keeping only valid pairs.

OC-SORT semantics: validity is decided separately (on geometry alone) and the score (geometry + momentum + appearance) only orders the pairs. As in the OC-SORT / Deep OC-SORT reference, the assignment maximizes the score over the WHOLE matrix and pairs failing the validity test are then rejected: an invalid pair's own (low) score acts as the value of leaving that row and column unmatched, so a confident valid match is not traded for two weaker ones. A pair with a non-finite score is never matched.

Parameters:

Name	Type	Description	Default
score	ndarray	(n, m) scores (higher is better).	required
valid	ndarray	(n, m) bool admissibility.	required

Returns:

Type	Description
Tuple[ndarray, ndarray, ndarray]	(matches, unmatched_rows, unmatched_cols) as in :func:linear_assignment; no invalid pair is ever matched.

Raises:

Type	Description
ValueError	If score and valid shapes differ or are not 2-D.

Source code in sleap_nn/tracking/mot/association.py

def linear_assignment_masked(
    score: np.ndarray, valid: np.ndarray
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """Maximum-score assignment, keeping only ``valid`` pairs.

    OC-SORT semantics: validity is decided separately (on geometry alone) and the
    score (geometry + momentum + appearance) only orders the pairs. As in the
    OC-SORT / Deep OC-SORT reference, the assignment maximizes the score over the
    WHOLE matrix and pairs failing the validity test are then rejected: an
    invalid pair's own (low) score acts as the value of leaving that row and
    column unmatched, so a confident valid match is not traded for two weaker
    ones. A pair with a non-finite score is never matched.

    Args:
        score: ``(n, m)`` scores (higher is better).
        valid: ``(n, m)`` bool admissibility.

    Returns:
        ``(matches, unmatched_rows, unmatched_cols)`` as in
        :func:`linear_assignment`; no invalid pair is ever matched.

    Raises:
        ValueError: If ``score`` and ``valid`` shapes differ or are not 2-D.
    """
    score = np.asarray(score, dtype=np.float64)
    valid = np.asarray(valid, dtype=bool)
    if score.ndim != 2 or score.shape != valid.shape:
        raise ValueError(
            f"score and valid must be 2-D with equal shapes, got {score.shape} "
            f"and {valid.shape}."
        )
    n, m = score.shape
    if n == 0 or m == 0:
        return _assignment_result([], [], n, m)
    valid = valid & np.isfinite(score)
    if not valid.any():
        return _assignment_result([], [], n, m)
    finite = np.isfinite(score)
    floor = float(score[finite].min()) - INFEASIBLE_COST
    rows, cols = linear_sum_assignment(-np.where(finite, score, floor))
    keep = valid[rows, cols]
    return _assignment_result(rows[keep], cols[keep], n, m)

ocm_matrix(track_dirs, track_prev_obs, track_has_dir, det_boxes, det_scores, inertia)

OC-SORT observation-centric momentum term.

For track i with unit velocity direction track_dirs[i] ((dy, dx), from an observation delta_t steps back to its last observation) and the direction from track_prev_obs[i] (that earlier observation) to each detection: theta = |arccos(clip(dot, -1, 1))| and ocm = inertia * ((pi/2 - theta) / pi) * has_dir[:, None] * det_scores[None, :].

A detection centered exactly on the earlier observation (no direction) has dot = 0 and contributes 0, as in the reference.

Parameters:

Name	Type	Description	Default
track_dirs	ndarray	(n, 2) unit (dy, dx) track directions.	required
track_prev_obs	ndarray	(n, 4) xyxy earlier observations.	required
track_has_dir	ndarray	(n,) bool: the track has a direction (and an earlier observation).	required
det_boxes	ndarray	(m, 4) xyxy detection boxes.	required
det_scores	ndarray	(m,) detection scores.	required
inertia	float	Weight of the term (OC-SORT: 0.2).	required

Returns:

Type	Description
ndarray	(n, m) additive similarity term (may be negative), in [-inertia/2, inertia/2] for scores in [0, 1].

Source code in sleap_nn/tracking/mot/association.py

def ocm_matrix(
    track_dirs: np.ndarray,
    track_prev_obs: np.ndarray,
    track_has_dir: np.ndarray,
    det_boxes: np.ndarray,
    det_scores: np.ndarray,
    inertia: float,
) -> np.ndarray:
    """OC-SORT observation-centric momentum term.

    For track ``i`` with unit velocity direction ``track_dirs[i]`` (``(dy, dx)``,
    from an observation ``delta_t`` steps back to its last observation) and the
    direction from ``track_prev_obs[i]`` (that earlier observation) to each
    detection: ``theta = |arccos(clip(dot, -1, 1))|`` and
    ``ocm = inertia * ((pi/2 - theta) / pi) * has_dir[:, None] * det_scores[None, :]``.

    A detection centered exactly on the earlier observation (no direction) has
    ``dot = 0`` and contributes 0, as in the reference.

    Args:
        track_dirs: ``(n, 2)`` unit ``(dy, dx)`` track directions.
        track_prev_obs: ``(n, 4)`` ``xyxy`` earlier observations.
        track_has_dir: ``(n,)`` bool: the track has a direction (and an earlier
            observation).
        det_boxes: ``(m, 4)`` ``xyxy`` detection boxes.
        det_scores: ``(m,)`` detection scores.
        inertia: Weight of the term (OC-SORT: 0.2).

    Returns:
        ``(n, m)`` additive similarity term (may be negative), in
        ``[-inertia/2, inertia/2]`` for scores in ``[0, 1]``.
    """
    track_dirs = np.asarray(track_dirs, dtype=np.float64).reshape(-1, 2)
    prev = _as_boxes(track_prev_obs)
    dets = _as_boxes(det_boxes)
    n, m = len(track_dirs), len(dets)
    if n == 0 or m == 0:
        return np.zeros((n, m), dtype=np.float64)
    has_dir = np.asarray(track_has_dir, dtype=bool).reshape(-1)
    has_dir = has_dir & np.isfinite(track_dirs).all(axis=1)
    track_dirs = np.where(has_dir[:, None], track_dirs, 0.0)
    scores = np.asarray(det_scores, dtype=np.float64).reshape(-1)
    scores = np.where(np.isfinite(scores), scores, 0.0)

    # Direction from each track's earlier observation to each detection (dy, dx).
    prev_cy = 0.5 * (prev[:, 1] + prev[:, 3])
    prev_cx = 0.5 * (prev[:, 0] + prev[:, 2])
    det_cy = 0.5 * (dets[:, 1] + dets[:, 3])
    det_cx = 0.5 * (dets[:, 0] + dets[:, 2])
    to_det = _unit_directions(
        det_cy[None, :] - prev_cy[:, None], det_cx[None, :] - prev_cx[:, None]
    )  # (n, m, 2)

    cos = np.clip((track_dirs[:, None, :] * to_det).sum(axis=-1), -1.0, 1.0)
    theta = np.abs(np.arccos(cos))
    term = (np.pi / 2.0 - theta) / np.pi
    return float(inertia) * term * has_dir[:, None] * scores[None, :]

oks_matrix(kp_a, kp_b, area_b, sigma)

Object keypoint similarity between two sets of poses.

Uses the COCO/cocoeval form, identical in constants to :func:sleap_nn.evaluation.compute_oks with use_cocoeval=True: e = d^2 / (2 sigma)^2 / (area + eps) / 2 and oks = mean(exp(-e)) over the nodes visible in BOTH poses (symmetric visibility: tracking compares two observations). Pairs sharing no visible node score 0.

Parameters:

Name	Type	Description	Default
kp_a	ndarray	(n, K, 2) keypoints with NaN for missing nodes.	required
kp_b	ndarray	(m, K, 2) keypoints with NaN for missing nodes.	required
area_b	ndarray	(m,) scale areas (px^2) of the b poses, floored at 1.	required
sigma		Scalar or (K,) per-node standard deviation.	required

Returns:

Type	Description
ndarray	(n, m) similarities in [0, 1].

Raises:

Type	Description
ValueError	If the two pose sets have different node counts or sigma is not positive and finite.

Source code in sleap_nn/tracking/mot/association.py

def oks_matrix(
    kp_a: np.ndarray, kp_b: np.ndarray, area_b: np.ndarray, sigma
) -> np.ndarray:
    """Object keypoint similarity between two sets of poses.

    Uses the COCO/cocoeval form, identical in constants to
    :func:`sleap_nn.evaluation.compute_oks` with ``use_cocoeval=True``:
    ``e = d^2 / (2 sigma)^2 / (area + eps) / 2`` and ``oks = mean(exp(-e))`` over
    the nodes visible in BOTH poses (symmetric visibility: tracking compares two
    observations). Pairs sharing no visible node score 0.

    Args:
        kp_a: ``(n, K, 2)`` keypoints with NaN for missing nodes.
        kp_b: ``(m, K, 2)`` keypoints with NaN for missing nodes.
        area_b: ``(m,)`` scale areas (px^2) of the ``b`` poses, floored at 1.
        sigma: Scalar or ``(K,)`` per-node standard deviation.

    Returns:
        ``(n, m)`` similarities in ``[0, 1]``.

    Raises:
        ValueError: If the two pose sets have different node counts or
            ``sigma`` is not positive and finite.
    """
    kp_a = np.asarray(kp_a, dtype=np.float64)
    kp_b = np.asarray(kp_b, dtype=np.float64)
    n, m = kp_a.shape[0], kp_b.shape[0]
    if n == 0 or m == 0:
        return np.zeros((n, m), dtype=np.float64)
    if kp_a.ndim != 3 or kp_b.ndim != 3 or kp_a.shape[1:] != kp_b.shape[1:]:
        raise ValueError(
            f"Keypoint arrays must be (n, K, 2) and (m, K, 2), got {kp_a.shape} "
            f"and {kp_b.shape}."
        )
    n_nodes = kp_a.shape[1]
    sigma = np.broadcast_to(np.asarray(sigma, dtype=np.float64), (n_nodes,))
    if not (np.all(np.isfinite(sigma)) and np.all(sigma > 0)):
        raise ValueError(f"OKS sigma must be positive and finite, got {sigma}.")

    # Scale: non-finite areas (an invalid pose) and tiny ones are floored at 1.
    area = np.asarray(area_b, dtype=np.float64).reshape(-1)
    area = np.where(np.isfinite(area), np.maximum(area, 1.0), 1.0)

    d2 = ((kp_a[:, None] - kp_b[None]) ** 2).sum(axis=-1)  # (n, m, K)
    shared = np.isfinite(d2)  # node visible in both poses
    spread = (2.0 * sigma) ** 2  # (K,)
    e = np.where(shared, d2, 0.0) / spread / (area[None, :, None] + _OKS_EPS) / 2.0
    ks = np.where(shared, np.exp(-e), 0.0)
    n_shared = shared.sum(axis=-1)
    oks = ks.sum(axis=-1) / np.maximum(n_shared, 1)
    return np.where(n_shared > 0, np.clip(oks, 0.0, 1.0), 0.0)

one_to_one_pattern(valid)

OC-SORT/SORT shortcut: the valid pairs if they form a 1:1 pattern.

Parameters:

Name	Type	Description	Default
valid	ndarray	(n, m) bool admissibility.	required

Returns:

Type	Description
Optional[ndarray]	(k, 2) matches (sorted by row) when every row and column of valid has at most one True entry (and at least one True overall), else None.

Source code in sleap_nn/tracking/mot/association.py

def one_to_one_pattern(valid: np.ndarray) -> Optional[np.ndarray]:
    """OC-SORT/SORT shortcut: the valid pairs if they form a 1:1 pattern.

    Args:
        valid: ``(n, m)`` bool admissibility.

    Returns:
        ``(k, 2)`` matches (sorted by row) when every row and column of ``valid``
        has at most one True entry (and at least one True overall), else
        ``None``.
    """
    valid = np.asarray(valid, dtype=bool)
    if valid.ndim != 2 or valid.size == 0 or not valid.any():
        return None
    if valid.sum(axis=1).max() > 1 or valid.sum(axis=0).max() > 1:
        return None
    return np.argwhere(valid).astype(np.int64)

speed_direction(prev_xyxy, cur_xyxy)

Unit (dy, dx) direction between box centers, row-wise for (n, 4) inputs.

Zero-length motion gives (0, 0).

Parameters:

Name	Type	Description	Default
prev_xyxy	ndarray	(n, 4) (or (4,)) earlier boxes.	required
cur_xyxy	ndarray	(n, 4) (or (4,)) later boxes.	required

Returns:

Type	Description
ndarray	(n, 2) (or (2,) for single boxes) unit directions from the earlier box center to the later one; (0, 0) for zero-length motion or non-finite boxes.

Source code in sleap_nn/tracking/mot/association.py

def speed_direction(prev_xyxy: np.ndarray, cur_xyxy: np.ndarray) -> np.ndarray:
    """Unit ``(dy, dx)`` direction between box centers, row-wise for ``(n, 4)`` inputs.

    Zero-length motion gives ``(0, 0)``.

    Args:
        prev_xyxy: ``(n, 4)`` (or ``(4,)``) earlier boxes.
        cur_xyxy: ``(n, 4)`` (or ``(4,)``) later boxes.

    Returns:
        ``(n, 2)`` (or ``(2,)`` for single boxes) unit directions from the
        earlier box center to the later one; ``(0, 0)`` for zero-length motion
        or non-finite boxes.
    """
    single = np.ndim(prev_xyxy) == 1 and np.ndim(cur_xyxy) == 1
    prev = _as_boxes(prev_xyxy)
    cur = _as_boxes(cur_xyxy)
    dy = 0.5 * (cur[:, 1] + cur[:, 3]) - 0.5 * (prev[:, 1] + prev[:, 3])
    dx = 0.5 * (cur[:, 0] + cur[:, 2]) - 0.5 * (prev[:, 0] + prev[:, 2])
    out = _unit_directions(dy, dx)
    return out[0] if single else out