detections

sleap_nn.tracking.mot.detections

Per-frame detection arrays for the MOT trackers.

:func:extract_detections turns one frame's sleap-io detections (pose Instance / PredictedInstance or PredictedSegmentationMask) into the vectorized :class:FrameDetections every tracker, the calibration and the identity linker operate on. The source objects are kept (sources) so tracks can be written back onto them.

Geometry conventions (shared by every consumer -- never mix them):

• raw_boxes is the tight extent of the visible keypoints (or the mask's bounding box), xyxy in image pixels.
• boxes is raw_boxes padded by box_padding * max(w, h) on each side.
• centroids is the CENTER OF raw_boxes -- the same reference point the Kalman filters track (box center). Do not use the node median here: mixing the two reference points injects a fake per-frame displacement.

Classes:

Name	Description
FrameDetections	Vectorized detections of one frame.

Functions:

Name	Description
extract_detections	Build :class:FrameDetections from one frame's sleap-io detections.
has_usable_embedding	Whether obj carries a usable appearance vector (see :func:_usable_vector).
normalize_scores	Divide dets.scores by normalizer and clip to [0, 1], in place.
widen_boxes	Floor dets.boxes (padded) at min_box_size per axis, in place.

FrameDetections

Vectorized detections of one frame.

Attributes:

Name	Type	Description
frame_idx	int	Frame index of the detections.
sources	list	The original sleap-io objects, in input order.
boxes	ndarray	(n, 4) float64 padded xyxy boxes; NaN rows where invalid.
raw_boxes	ndarray	(n, 4) float64 unpadded extents; NaN rows where invalid.
centroids	ndarray	(n, 2) float64 raw-box centers; NaN rows where invalid.
keypoints	Optional[ndarray]	(n, K, 2) float64 keypoints with NaN for missing nodes, or None for masks.
scores	ndarray	(n,) float64 confidences in [0, 1].
embeddings	Optional[ndarray]	(n, D) float32 L2-normalized appearance vectors with zero rows where missing, or None when no detection has one.
has_embedding	ndarray	(n,) bool.
valid	ndarray	(n,) bool: the detection has usable geometry.
n_visible	ndarray	(n,) int visible nodes (masks: 1 if non-empty else 0).
n_nodes	int	Nodes per detection K (masks: 1).

Methods:

Name	Description
__len__	Number of detections.
empty	Return an empty frame (zero detections) with correctly shaped arrays.
subset	Return the detections at idx (index array or bool mask), copied.

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

@attrs.define
class FrameDetections:
    """Vectorized detections of one frame.

    Attributes:
        frame_idx: Frame index of the detections.
        sources: The original sleap-io objects, in input order.
        boxes: ``(n, 4)`` float64 padded ``xyxy`` boxes; NaN rows where invalid.
        raw_boxes: ``(n, 4)`` float64 unpadded extents; NaN rows where invalid.
        centroids: ``(n, 2)`` float64 raw-box centers; NaN rows where invalid.
        keypoints: ``(n, K, 2)`` float64 keypoints with NaN for missing nodes, or
            ``None`` for masks.
        scores: ``(n,)`` float64 confidences in ``[0, 1]``.
        embeddings: ``(n, D)`` float32 L2-normalized appearance vectors with zero
            rows where missing, or ``None`` when no detection has one.
        has_embedding: ``(n,)`` bool.
        valid: ``(n,)`` bool: the detection has usable geometry.
        n_visible: ``(n,)`` int visible nodes (masks: 1 if non-empty else 0).
        n_nodes: Nodes per detection ``K`` (masks: 1).
    """

    frame_idx: int
    sources: list
    boxes: np.ndarray
    raw_boxes: np.ndarray
    centroids: np.ndarray
    keypoints: Optional[np.ndarray]
    scores: np.ndarray
    embeddings: Optional[np.ndarray]
    has_embedding: np.ndarray
    valid: np.ndarray
    n_visible: np.ndarray
    n_nodes: int

    def __len__(self) -> int:
        """Number of detections."""
        return len(self.sources)

    def subset(self, idx: np.ndarray) -> "FrameDetections":
        """Return the detections at ``idx`` (index array or bool mask), copied.

        Args:
            idx: Integer indices (any order) or a boolean mask of length
                ``len(self)``.

        Returns:
            A new :class:`FrameDetections` holding copies of the selected rows.
            ``embeddings`` is ``None`` when none of the selected detections has
            a vector (same invariant as :func:`extract_detections`).

        Raises:
            ValueError: If a boolean mask has the wrong length.
        """
        idx = np.asarray(idx)
        if idx.dtype == bool:
            if idx.shape != (len(self),):
                raise ValueError(
                    f"Boolean subset mask has shape {idx.shape}, expected "
                    f"({len(self)},)."
                )
            idx = np.flatnonzero(idx)
        idx = idx.astype(np.intp, copy=False).reshape(-1)

        has_embedding = self.has_embedding[idx].copy()
        embeddings = None
        if self.embeddings is not None and has_embedding.any():
            embeddings = self.embeddings[idx].copy()

        return FrameDetections(
            frame_idx=self.frame_idx,
            sources=[self.sources[i] for i in idx],
            boxes=self.boxes[idx].copy(),
            raw_boxes=self.raw_boxes[idx].copy(),
            centroids=self.centroids[idx].copy(),
            keypoints=None if self.keypoints is None else self.keypoints[idx].copy(),
            scores=self.scores[idx].copy(),
            embeddings=embeddings,
            has_embedding=has_embedding,
            valid=self.valid[idx].copy(),
            n_visible=self.n_visible[idx].copy(),
            n_nodes=self.n_nodes,
        )

    @classmethod
    def empty(cls, frame_idx: int, n_nodes: int = 1) -> "FrameDetections":
        """Return an empty frame (zero detections) with correctly shaped arrays.

        Args:
            frame_idx: Frame index.
            n_nodes: Nodes per detection; sets the ``(0, n_nodes, 2)`` keypoint
                shape.

        Returns:
            A :class:`FrameDetections` with ``len() == 0`` and no embeddings.
        """
        n_nodes = int(n_nodes)
        return cls(
            frame_idx=int(frame_idx),
            sources=[],
            boxes=np.zeros((0, 4), dtype=np.float64),
            raw_boxes=np.zeros((0, 4), dtype=np.float64),
            centroids=np.zeros((0, 2), dtype=np.float64),
            keypoints=np.zeros((0, n_nodes, 2), dtype=np.float64),
            scores=np.zeros((0,), dtype=np.float64),
            embeddings=None,
            has_embedding=np.zeros((0,), dtype=bool),
            valid=np.zeros((0,), dtype=bool),
            n_visible=np.zeros((0,), dtype=np.int64),
            n_nodes=n_nodes,
        )

__len__()

Number of detections.

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

def __len__(self) -> int:
    """Number of detections."""
    return len(self.sources)

empty(frame_idx, n_nodes=1) classmethod

Return an empty frame (zero detections) with correctly shaped arrays.

Parameters:

Name	Type	Description	Default
frame_idx	int	Frame index.	required
n_nodes	int	Nodes per detection; sets the (0, n_nodes, 2) keypoint shape.	1

Returns:

Name	Type	Description
A	FrameDetections	class:FrameDetections with len() == 0 and no embeddings.

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

@classmethod
def empty(cls, frame_idx: int, n_nodes: int = 1) -> "FrameDetections":
    """Return an empty frame (zero detections) with correctly shaped arrays.

    Args:
        frame_idx: Frame index.
        n_nodes: Nodes per detection; sets the ``(0, n_nodes, 2)`` keypoint
            shape.

    Returns:
        A :class:`FrameDetections` with ``len() == 0`` and no embeddings.
    """
    n_nodes = int(n_nodes)
    return cls(
        frame_idx=int(frame_idx),
        sources=[],
        boxes=np.zeros((0, 4), dtype=np.float64),
        raw_boxes=np.zeros((0, 4), dtype=np.float64),
        centroids=np.zeros((0, 2), dtype=np.float64),
        keypoints=np.zeros((0, n_nodes, 2), dtype=np.float64),
        scores=np.zeros((0,), dtype=np.float64),
        embeddings=None,
        has_embedding=np.zeros((0,), dtype=bool),
        valid=np.zeros((0,), dtype=bool),
        n_visible=np.zeros((0,), dtype=np.int64),
        n_nodes=n_nodes,
    )

subset(idx)

Return the detections at idx (index array or bool mask), copied.

Parameters:

Name	Type	Description	Default
idx	ndarray	Integer indices (any order) or a boolean mask of length len(self).	required

Returns:

Type	Description
FrameDetections	A new :class:FrameDetections holding copies of the selected rows. embeddings is None when none of the selected detections has a vector (same invariant as :func:extract_detections).

Raises:

Type	Description
ValueError	If a boolean mask has the wrong length.

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

def subset(self, idx: np.ndarray) -> "FrameDetections":
    """Return the detections at ``idx`` (index array or bool mask), copied.

    Args:
        idx: Integer indices (any order) or a boolean mask of length
            ``len(self)``.

    Returns:
        A new :class:`FrameDetections` holding copies of the selected rows.
        ``embeddings`` is ``None`` when none of the selected detections has
        a vector (same invariant as :func:`extract_detections`).

    Raises:
        ValueError: If a boolean mask has the wrong length.
    """
    idx = np.asarray(idx)
    if idx.dtype == bool:
        if idx.shape != (len(self),):
            raise ValueError(
                f"Boolean subset mask has shape {idx.shape}, expected "
                f"({len(self)},)."
            )
        idx = np.flatnonzero(idx)
    idx = idx.astype(np.intp, copy=False).reshape(-1)

    has_embedding = self.has_embedding[idx].copy()
    embeddings = None
    if self.embeddings is not None and has_embedding.any():
        embeddings = self.embeddings[idx].copy()

    return FrameDetections(
        frame_idx=self.frame_idx,
        sources=[self.sources[i] for i in idx],
        boxes=self.boxes[idx].copy(),
        raw_boxes=self.raw_boxes[idx].copy(),
        centroids=self.centroids[idx].copy(),
        keypoints=None if self.keypoints is None else self.keypoints[idx].copy(),
        scores=self.scores[idx].copy(),
        embeddings=embeddings,
        has_embedding=has_embedding,
        valid=self.valid[idx].copy(),
        n_visible=self.n_visible[idx].copy(),
        n_nodes=self.n_nodes,
    )

extract_detections(objs, frame_idx, *, carrier='pose', box_padding=0.1, pseudo_box_size=None, score_mode='auto', embedding_dim=None, min_box_size=None)

Build :class:FrameDetections from one frame's sleap-io detections.

Geometry: a pose's raw box is the extent of its visible (finite) nodes; a mask's is mask.bbox (image space), and an empty mask is invalid. Each object is dispatched with :func:~sleap_nn.tracking.utils.is_segmentation_mask. A valid detection is DEGENERATE when it is a pose with a single visible node or its raw box is narrower than 1 px on either axis: its raw and padded boxes then both become the square of side pseudo_box_size centered on its centroid (with pseudo_box_size=None, the raw box is instead floored at 1 px per axis around its center and padded as usual). Centroids are raw-box centers.

Scores (score_mode), then non-finite -> 1.0 and clipped to [0, 1]:

• "auto": predicted pose -> mean point score over its visible nodes (obj.score if none is finite); predicted mask -> obj.score; user Instance / user mask (no score) -> 1.0.
• "instance": obj.score (1.0 if absent).
• "none": 1.0 everywhere.

Parameters:

Name	Type	Description	Default
objs	Sequence	The frame's detections: pose instances (carrier="pose") or segmentation masks (carrier="mask").	required
frame_idx	int	Frame index of the detections.	required
carrier	str	"pose" or "mask".	'pose'
box_padding	float	Fraction of max(w, h) added to each side of the box.	0.1
pseudo_box_size	Optional[float]	Side (px) of the square box used for degenerate geometry (a single visible node, or a raw box narrower than 1 px). None floors the box at 1 px instead.	None
score_mode	str	"auto", "instance" or "none" (see :class:sleap_nn.tracking.mot.params.MOTParams).	'auto'
embedding_dim	Optional[int]	Expected appearance-vector dimension; vectors of another dimension are treated as missing. None accepts the first seen.	None
min_box_size	Optional[float]	Minimum side (px) of the PADDED box on each axis; narrower boxes are widened symmetrically around their center. Keeps IoU meaningful for thin poses (nodes nearly collinear along one axis). Raw boxes (the OKS scale) are not changed. None disables it.	None

Returns:

Type	Description
FrameDetections	The frame's :class:FrameDetections. keypoints is None for the mask carrier; on the pose carrier a (stray) mask gets an all-NaN row.

Raises:

Type	Description
ValueError	On an unknown carrier or score_mode.

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

def extract_detections(
    objs: Sequence,
    frame_idx: int,
    *,
    carrier: str = "pose",
    box_padding: float = 0.1,
    pseudo_box_size: Optional[float] = None,
    score_mode: str = "auto",
    embedding_dim: Optional[int] = None,
    min_box_size: Optional[float] = None,
) -> FrameDetections:
    """Build :class:`FrameDetections` from one frame's sleap-io detections.

    Geometry: a pose's raw box is the extent of its visible (finite) nodes; a
    mask's is ``mask.bbox`` (image space), and an empty mask is invalid. Each
    object is dispatched with
    :func:`~sleap_nn.tracking.utils.is_segmentation_mask`. A valid detection is
    DEGENERATE when it is a pose with a single visible node or its raw box is
    narrower than 1 px on either axis: its raw and padded boxes then both become
    the square of side ``pseudo_box_size`` centered on its centroid (with
    ``pseudo_box_size=None``, the raw box is instead floored at 1 px per axis
    around its center and padded as usual). Centroids are raw-box centers.

    Scores (``score_mode``), then non-finite -> 1.0 and clipped to ``[0, 1]``:

    - ``"auto"``: predicted pose -> mean point score over its visible nodes
      (``obj.score`` if none is finite); predicted mask -> ``obj.score``; user
      ``Instance`` / user mask (no ``score``) -> 1.0.
    - ``"instance"``: ``obj.score`` (1.0 if absent).
    - ``"none"``: 1.0 everywhere.

    Args:
        objs: The frame's detections: pose instances (``carrier="pose"``) or
            segmentation masks (``carrier="mask"``).
        frame_idx: Frame index of the detections.
        carrier: ``"pose"`` or ``"mask"``.
        box_padding: Fraction of ``max(w, h)`` added to each side of the box.
        pseudo_box_size: Side (px) of the square box used for degenerate
            geometry (a single visible node, or a raw box narrower than 1 px).
            ``None`` floors the box at 1 px instead.
        score_mode: ``"auto"``, ``"instance"`` or ``"none"`` (see
            :class:`sleap_nn.tracking.mot.params.MOTParams`).
        embedding_dim: Expected appearance-vector dimension; vectors of another
            dimension are treated as missing. ``None`` accepts the first seen.
        min_box_size: Minimum side (px) of the PADDED box on each axis; narrower
            boxes are widened symmetrically around their center. Keeps IoU
            meaningful for thin poses (nodes nearly collinear along one axis).
            Raw boxes (the OKS scale) are not changed. ``None`` disables it.

    Returns:
        The frame's :class:`FrameDetections`. ``keypoints`` is ``None`` for the
        mask carrier; on the pose carrier a (stray) mask gets an all-NaN row.

    Raises:
        ValueError: On an unknown ``carrier`` or ``score_mode``.
    """
    if carrier not in CARRIERS:
        raise ValueError(f"carrier must be one of {CARRIERS}, got {carrier!r}.")
    if score_mode not in SCORE_MODES:
        raise ValueError(
            f"score_mode must be one of {SCORE_MODES}, got {score_mode!r}."
        )
    objs = list(objs)
    n = len(objs)
    is_mask = np.array([is_segmentation_mask(o) for o in objs], dtype=bool)
    pose_idx = np.flatnonzero(~is_mask)
    mask_idx = np.flatnonzero(is_mask)

    raw = np.full((n, 4), np.nan, dtype=np.float64)
    valid = np.zeros(n, dtype=bool)
    n_visible = np.zeros(n, dtype=np.int64)
    scores = np.ones(n, dtype=np.float64)
    single_node = np.zeros(n, dtype=bool)
    xy = None

    # Pose detections: one numpy() call per object, everything else vectorized.
    if len(pose_idx) > 0:
        pose_objs = [objs[i] for i in pose_idx]
        xy, point_scores, is_predicted = _pose_arrays(pose_objs)
        pose_raw, pose_n_visible, visible = _pose_geometry(xy)
        raw[pose_idx] = pose_raw
        valid[pose_idx] = pose_n_visible > 0
        n_visible[pose_idx] = pose_n_visible
        single_node[pose_idx] = pose_n_visible == 1
        scores[pose_idx] = _pose_scores(
            pose_objs, point_scores, visible, is_predicted, score_mode
        )

    # Mask detections: geometry from the image-space bbox; score is the mask's
    # own (a user mask has no `score` -> 1.0) in both "auto" and "instance".
    if len(mask_idx) > 0:
        mask_objs = [objs[i] for i in mask_idx]
        mask_raw, nonempty = _mask_geometry(mask_objs)
        raw[mask_idx] = mask_raw
        valid[mask_idx] = nonempty
        n_visible[mask_idx] = nonempty.astype(np.int64)
        if score_mode != "none":
            scores[mask_idx] = [_instance_score(o) for o in mask_objs]

    scores = np.clip(np.where(np.isfinite(scores), scores, 1.0), 0.0, 1.0)
    valid &= np.isfinite(raw).all(axis=1)
    raw[~valid] = np.nan

    # Keypoints (pose carrier only), NaN-padded to the longest skeleton.
    keypoints = None
    n_nodes = 1
    if carrier == "pose":
        n_nodes = max(1, 0 if xy is None else xy.shape[1])
        keypoints = np.full((n, n_nodes, 2), np.nan, dtype=np.float64)
        if xy is not None:
            keypoints[pose_idx, : xy.shape[1]] = xy

    if n == 0:
        out = FrameDetections.empty(frame_idx, n_nodes=n_nodes)
        out.keypoints = keypoints
        return out

    # Padding, then degenerate geometry (overrides the padded box).
    wh = raw[:, 2:] - raw[:, :2]
    centroids = 0.5 * (raw[:, :2] + raw[:, 2:])
    pad = float(box_padding) * wh.max(axis=1)
    boxes = raw + pad[:, None] * np.array([-1.0, -1.0, 1.0, 1.0])
    degenerate = valid & (single_node | (wh < _MIN_BOX_SIDE).any(axis=1))
    if degenerate.any():
        c = centroids[degenerate]
        if pseudo_box_size is not None and np.isfinite(pseudo_box_size):
            # Square of the calibrated body size: raw (OKS scale) and padded.
            half = np.full_like(c, 0.5 * max(float(pseudo_box_size), _MIN_BOX_SIDE))
            raw[degenerate] = np.concatenate([c - half, c + half], axis=1)
            boxes[degenerate] = raw[degenerate]
        else:
            # Standalone callers: 1-px floor per axis, then the usual padding.
            half = 0.5 * np.maximum(wh[degenerate], _MIN_BOX_SIDE)
            raw[degenerate] = np.concatenate([c - half, c + half], axis=1)
            dpad = float(box_padding) * 2.0 * half.max(axis=1)
            boxes[degenerate] = raw[degenerate] + dpad[:, None] * np.array(
                [-1.0, -1.0, 1.0, 1.0]
            )

    boxes = _widened(boxes, valid, min_box_size)

    embeddings, has_embedding = _embedding_arrays(objs, embedding_dim)

    return FrameDetections(
        frame_idx=int(frame_idx),
        sources=objs,
        boxes=boxes,
        raw_boxes=raw,
        centroids=centroids,
        keypoints=keypoints,
        scores=scores,
        embeddings=embeddings,
        has_embedding=has_embedding,
        valid=valid,
        n_visible=n_visible,
        n_nodes=int(n_nodes),
    )

has_usable_embedding(obj)

Whether obj carries a usable appearance vector (see :func:_usable_vector).

A frame-level check can additionally drop a usable vector whose dimension differs from the frame's (:func:extract_detections); per detection, only presence, finiteness and a non-zero norm can be judged.

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

def has_usable_embedding(obj) -> bool:
    """Whether ``obj`` carries a usable appearance vector (see :func:`_usable_vector`).

    A frame-level check can additionally drop a usable vector whose dimension
    differs from the frame's (:func:`extract_detections`); per detection, only
    presence, finiteness and a non-zero norm can be judged.
    """
    return _usable_vector(obj) is not None

normalize_scores(dets, normalizer)

Divide dets.scores by normalizer and clip to [0, 1], in place.

The per-video score normalization of MOTParams.normalize_scores (the normalizer is :func:~sleap_nn.tracking.mot.calibration.score_normalizer: the video's score q95). Apply it exactly once per extraction.

Parameters:

Name	Type	Description	Default
dets	FrameDetections	Detections to modify.	required
normalizer	Optional[float]	Positive divisor; None, non-finite or non-positive is a no-op.	required

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

def normalize_scores(dets: FrameDetections, normalizer: Optional[float]) -> None:
    """Divide ``dets.scores`` by ``normalizer`` and clip to ``[0, 1]``, in place.

    The per-video score normalization of ``MOTParams.normalize_scores`` (the
    normalizer is :func:`~sleap_nn.tracking.mot.calibration.score_normalizer`:
    the video's score q95). Apply it exactly once per extraction.

    Args:
        dets: Detections to modify.
        normalizer: Positive divisor; ``None``, non-finite or non-positive is a
            no-op.
    """
    if normalizer is None or not np.isfinite(normalizer) or normalizer <= 0:
        return
    if len(dets):
        dets.scores = np.clip(
            np.asarray(dets.scores, dtype=np.float64) / float(normalizer), 0.0, 1.0
        )

widen_boxes(dets, min_box_size)

Floor dets.boxes (padded) at min_box_size per axis, in place.

Same rule as extract_detections(min_box_size=...), for detections that were already extracted. Raw boxes are not changed.

Parameters:

Name	Type	Description	Default
dets	FrameDetections	Detections to modify.	required
min_box_size	Optional[float]	Minimum padded side (px); None or non-positive is a no-op.	required

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

def widen_boxes(dets: FrameDetections, min_box_size: Optional[float]) -> None:
    """Floor ``dets.boxes`` (padded) at ``min_box_size`` per axis, in place.

    Same rule as ``extract_detections(min_box_size=...)``, for detections that
    were already extracted. Raw boxes are not changed.

    Args:
        dets: Detections to modify.
        min_box_size: Minimum padded side (px); ``None`` or non-positive is a no-op.
    """
    if len(dets):
        dets.boxes = _widened(dets.boxes, dets.valid, min_box_size)