splitting

sleap_nn.data.splitting

Group-aware train/val splitter, decided before training.

For the embedding (crop -> vector re-ID) model type, the train/val partition is the generalization axis: the model must only ever see the training partition, with val/test held out by a group key so there is no leakage of the held-out group into training.

Three group keys are supported (mirroring the standalone reference scratch/.../embedding/splits.py):

• frame : stratified-random over LabeledFrame units (frames mixed, identity-balanced via :class:~sklearn.model_selection.StratifiedGroupKFold grouped by frame). Both train and val contain all identities -- the in-distribution headline.
• video : hold out whole videos (by sio video index) via :class:~sklearn.model_selection.GroupKFold -- the honest cross-session number.
• identity : hold out whole track names (true open-vocab) via :class:~sklearn.model_selection.GroupKFold grouped by identity. Train and val identity sets are disjoint (degenerate for few identities; supported for completeness / verification-only).

The splitter operates at the LabeledFrame / detection level and returns new sio.Labels objects containing the selected detections, so it composes directly with the existing dataset builders (which iterate sio.Labels -> detections). For frame/video each source frame stays whole on one side; for identity a source frame's detections are filtered by identity, so a frame may contribute (disjoint) detections to both sides.

Both carriers are split. A detection is a pose instance (lf.instances) or a mask (lf.masks); a mask LINKED to one of its frame's instances (mask.instance, as a top-down segmentation or SAM output has) is not a detection of its own: it goes wherever its instance goes, link intact. Keeping only one carrier per frame used to drop every mask of a pose + mask frame, which silently flipped an embedding model from mask to pose mode. Mask-only labels (e.g. the gerbil instance-segmentation data) carry all their detections on lf.masks.

Functions:

Name	Description
split_labels_list_train_val	Apply :func:split_labels_train_val to each sio.Labels in a list.
split_labels_train_val	Partition a single sio.Labels into (train, val) by a group key.

split_labels_list_train_val(labels_list, split_config)

Apply :func:split_labels_train_val to each sio.Labels in a list.

Splits each input sio.Labels independently (video indices and track names are scoped per file) and returns (train_labels_list, val_labels_list) aligned to the input list, so the result drops straight into ModelTrainer.train_labels / .val_labels.

Parameters:

Name	Type	Description	Default
labels_list	List[Labels]	List of source sio.Labels (typically the loaded train files).	required
split_config		A SplitConfig (or any object exposing split_by / n_folds / fold / seed).	required

Returns:

Type	Description
Tuple[List[Labels], List[Labels]]	Tuple (train_labels_list, val_labels_list) of equal length to labels_list.

Source code in sleap_nn/data/splitting.py

def split_labels_list_train_val(
    labels_list: List[sio.Labels],
    split_config,
) -> Tuple[List[sio.Labels], List[sio.Labels]]:
    """Apply :func:`split_labels_train_val` to each ``sio.Labels`` in a list.

    Splits each input ``sio.Labels`` independently (video indices and track names are scoped
    per file) and returns ``(train_labels_list, val_labels_list)`` aligned to the input list,
    so the result drops straight into ``ModelTrainer.train_labels`` / ``.val_labels``.

    Args:
        labels_list: List of source ``sio.Labels`` (typically the loaded train files).
        split_config: A ``SplitConfig`` (or any object exposing ``split_by`` / ``n_folds`` /
            ``fold`` / ``seed``).

    Returns:
        Tuple ``(train_labels_list, val_labels_list)`` of equal length to ``labels_list``.
    """
    split_by = getattr(split_config, "split_by", "frame")
    n_folds = int(getattr(split_config, "n_folds", 5))
    fold = int(getattr(split_config, "fold", 0))
    seed = int(getattr(split_config, "seed", 0))

    # Each file is split independently, so a group-aware split is only group-aware
    # WITHIN a file: with `split_by="identity"` and several files, one animal can
    # land in file A's train side and file B's val side, which is the leakage the
    # mode exists to prevent. Say so rather than reporting a clean-looking split.
    if split_by == "identity" and len(labels_list) > 1:
        logger.warning(
            f"split_by='identity' with {len(labels_list)} labels files: each file is "
            "split independently, so an identity present in more than one file can "
            "still appear on both sides of the split. Pass explicit `val_labels_path` "
            "for a guaranteed identity-disjoint validation set, or merge the files "
            "first."
        )

    train_list, val_list = [], []
    for labels in labels_list:
        train_labels, val_labels = split_labels_train_val(
            labels,
            split_by=split_by,
            n_folds=n_folds,
            fold=fold,
            seed=seed,
        )
        train_list.append(train_labels)
        val_list.append(val_labels)

    logger.info(
        f"Group-aware split (split_by={split_by!r}, n_folds={n_folds}, fold={fold}, "
        f"seed={seed}): "
        f"train frames={sum(len(t) for t in train_list)}, "
        f"val frames={sum(len(v) for v in val_list)}."
    )
    return train_list, val_list

split_labels_train_val(source, *, split_by, n_folds, fold, seed)

Partition a single sio.Labels into (train, val) by a group key.

Parameters:

Name	Type	Description	Default
source	Labels	The sio.Labels to partition.	required
split_by	str	One of 'frame' | 'video' | 'identity' (see module docstring).	required
n_folds	int	Number of CV folds; the val partition is 1 / n_folds of the data.	required
fold	int	Which fold (0-based) to hold out as validation.	required
seed	int	Random seed for the (shuffled) splitter.	required

Returns:

Type	Description
Tuple[Labels, Labels]	Tuple of new sio.Labels (train_labels, val_labels) with no group leakage.

Source code in sleap_nn/data/splitting.py

def split_labels_train_val(
    source: sio.Labels,
    *,
    split_by: str,
    n_folds: int,
    fold: int,
    seed: int,
) -> Tuple[sio.Labels, sio.Labels]:
    """Partition a single ``sio.Labels`` into (train, val) by a group key.

    Args:
        source: The ``sio.Labels`` to partition.
        split_by: One of ``'frame'`` | ``'video'`` | ``'identity'`` (see module docstring).
        n_folds: Number of CV folds; the val partition is ``1 / n_folds`` of the data.
        fold: Which fold (0-based) to hold out as validation.
        seed: Random seed for the (shuffled) splitter.

    Returns:
        Tuple of new ``sio.Labels`` ``(train_labels, val_labels)`` with no group leakage.
    """
    lf_idx, det_idx, track_names, video_idx, identity_names = _build_pool(source)
    n = len(lf_idx)

    # Frames with zero detections (e.g. user-confirmed negatives) have no identity to
    # fold, so they would otherwise be silently dropped from BOTH sides. Carry them on
    # the train side (deterministic) — mirroring the non-split path's negative handling.
    frames_with_dets = set(int(i) for i in lf_idx)
    negative_lf_idx = [
        li for li in range(len(source.labeled_frames)) if li not in frames_with_dets
    ]

    if n == 0:
        # No detections at all (e.g. an all-negative or empty .slp). Don't abort the
        # run — keep any frames on the train side, leave val empty.
        if negative_lf_idx:
            logger.warning(
                "Group-aware split: a labels file has no detections; keeping its "
                f"{len(negative_lf_idx)} frame(s) on the train side (val empty)."
            )
        empty = np.zeros(0, dtype=bool)
        train_labels = _rebuild_labels(
            source, lf_idx, det_idx, empty, extra_lf_idx=negative_lf_idx
        )
        val_labels = _rebuild_labels(source, lf_idx, det_idx, empty)
        return train_labels, val_labels

    # Identity label per detection (used as stratification target and as identity group).
    y = track_names
    if split_by == "video":
        groups = video_idx
    elif split_by == "identity":
        # Group AND stratify on the global identity, not the per-video track name
        # (see `_identity_name`): the two differ exactly when an animal carries a
        # `sio.Identity` plus per-video tracks, which is the leak this split is
        # meant to prevent.
        groups = identity_names
        y = identity_names
    else:  # frame
        groups = lf_idx

    val_mask = _select_val_fold(
        n,
        y,
        groups,
        split_by=split_by,
        n_folds=n_folds,
        fold=fold,
        seed=seed,
    )

    if negative_lf_idx:
        logger.info(
            f"Group-aware split: carried {len(negative_lf_idx)} zero-detection "
            "frame(s) to the train side."
        )
    train_labels = _rebuild_labels(
        source, lf_idx, det_idx, ~val_mask, extra_lf_idx=negative_lf_idx
    )
    val_labels = _rebuild_labels(source, lf_idx, det_idx, val_mask)
    return train_labels, val_labels