custom_datasets

sleap_nn.data.custom_datasets

Custom torch.utils.data.Datasets for different model types.

Classes:

Name	Description
BBoxDataset	Whole-frame dataset for the bounding-box detectors (bbox_* model types).
BaseDataset	Base class for custom torch Datasets.
BottomUpDataset	Dataset class for bottom-up models.
BottomUpMultiClassDataset	Dataset class for bottom-up ID models.
BottomUpSegmentationDataset	Dataset class for bottom-up instance segmentation models.
BottomUpSegmentationTiledDataset	Bottom-up segmentation dataset that emits fixed-size tiles (Phase C).
CenteredInstanceDataset	Dataset class for instance-centered confidence map models.
CenteredInstanceSegmentationDataset	Dataset for top-down (crop-centered) instance segmentation (#622).
CentroidDataset	Dataset class for centroid models.
EmbeddingDataset	Dataset for the embedding (crop -> vector, re-ID) model type.
EmbeddingMembership	Which detections an embedding dataset uses, and the groups they fall in.
GroupAwareBatchSampler	Group-aware batch sampler for contrastive embedding training.
InfiniteDataLoader	Dataloader that reuses workers for infinite iteration.
ParallelCacheFiller	Parallel implementation of image caching using thread-local video copies.
SemanticSegmentationDataset	Dataset class for whole-frame semantic (foreground/background) segmentation.
SemanticSegmentationTiledDataset	Whole-frame semantic (fg/bg) segmentation dataset emitting fixed-size tiles.
SingleInstanceDataset	Dataset class for single-instance models.
SingleInstanceTiledDataset	Single-instance dataset that emits fixed-size tiles instead of whole frames.
TopDownCenteredInstanceMultiClassDataset	Dataset class for instance-centered confidence map ID models.

Functions:

Name	Description
embedding_membership_from_config	The training :class:EmbeddingMembership of an embedding config.
embedding_preferred_carrier	The carrier an embedding model was trained on (pose / mask), if known.
embedding_training_groups	(group_ids, video_ids, frame_ids) of the crops an embedding dataset would index.
get_embedding_batch_sampler	Build the group-aware batch sampler an embedding training run uses.
get_steps_per_epoch	Compute the number of steps (iterations) per epoch for the given dataset.
get_train_samples_per_step	Number of training samples one optimizer step consumes (per device).
get_train_val_dataloaders	Return the train and val dataloaders.
get_train_val_datasets	Return the train and val datasets.
labels_have_user_centroids	Return True if any labeled frame carries a usable UserCentroid.
predict_embedding_carrier	The carrier :func:resolve_embedding_detection_mode will pick, before the labels exist.
resolve_centroid_source	Resolve the centroid target source to one dataset-wide mode.
resolve_embedding_class_names	Collect the sorted global-identity vocabulary (the global_id / eval grouping).
resolve_embedding_detection_mode	Pick the carrier (pose / mask) an embedding dataset crops.
resolve_embedding_grouping	Return (track_names_are_global, id_scope) for an embedding training config.

BBoxDataset

Bases: BaseDataset

Whole-frame dataset for the bounding-box detectors (bbox_* model types).

Indexes frames by their user boxes (LabeledFrame.bboxes; derive them from poses/masks beforehand with :func:sleap_nn.data.bboxes.add_bboxes_to_labels) and, optionally, the user-confirmed negative frames. Each sample carries the preprocessed image, the frame's boxes in model-input pixels (bboxes, NaN-padded to (max_instances, 4), with bbox_classes -1-padded) and the training targets of the model's detection style, produced by target_encoder (see :mod:sleap_nn.data.bboxes).

Boxes are pushed through the shared preprocessing/augmentation chain as 4-corner pseudo-instances, so rotation/scale/translate/flip stay registered with the image; afterwards each box is the extent of its corners, clipped to the frame, and dropped when too little of it remains visible.

Attributes:

Name	Type	Description
target_encoder		Callable (boxes (n, 4), classes (n,), img_hw) -> dict of target tensors (CenterNetTargetEncoder / SSDTargetEncoder).
categories		Ordered class names; box class ids index into it.
min_visible_fraction		A box keeping less of its area than this inside the frame after augmentation is dropped from the targets.

Methods:

Name	Description
__getitem__	Return the image, padded GT boxes and detection targets for index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class BBoxDataset(BaseDataset):
    """Whole-frame dataset for the bounding-box detectors (``bbox_*`` model types).

    Indexes frames by their user boxes (``LabeledFrame.bboxes``; derive them from
    poses/masks beforehand with :func:`sleap_nn.data.bboxes.add_bboxes_to_labels`)
    and, optionally, the user-confirmed negative frames. Each sample carries the
    preprocessed image, the frame's boxes in model-input pixels (``bboxes``,
    NaN-padded to ``(max_instances, 4)``, with ``bbox_classes`` ``-1``-padded) and
    the training targets of the model's detection style, produced by
    ``target_encoder`` (see :mod:`sleap_nn.data.bboxes`).

    Boxes are pushed through the shared preprocessing/augmentation chain as
    4-corner pseudo-instances, so rotation/scale/translate/flip stay registered with
    the image; afterwards each box is the extent of its corners, clipped to the
    frame, and dropped when too little of it remains visible.

    Attributes:
        target_encoder: Callable ``(boxes (n, 4), classes (n,), img_hw) -> dict`` of
            target tensors (``CenterNetTargetEncoder`` / ``SSDTargetEncoder``).
        categories: Ordered class names; box class ids index into it.
        min_visible_fraction: A box keeping less of its area than this inside the
            frame after augmentation is dropped from the targets.
    """

    _flip_needs_symmetries = False

    def __init__(
        self,
        labels: List[sio.Labels],
        target_encoder: Any,
        categories: List[str],
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
        min_visible_fraction: float = 0.25,
    ) -> None:
        """Initialize class attributes."""
        # Read by `_get_lf_idx_list`, which runs inside `BaseDataset.__init__`.
        self.target_encoder = target_encoder
        self.categories = list(categories)
        self.min_visible_fraction = min_visible_fraction
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        # A box's corners are not skeleton nodes: a horizontal flip mirrors them and
        # re-boxing takes their extent, so no index swap is needed (or wanted).
        self.symmetric_inds = []

    def _get_lf_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return samples for frames with user boxes (+ negative frames if enabled).

        The boxes are captured into each sample as plain arrays so ``__getitem__``
        needs no live ``Labels`` handle (the image-caching paths drop it).
        """
        lf_idx_list = []
        max_boxes = 0
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                xyxy, classes = frame_bbox_arrays(lf, self.categories)
                if xyxy.shape[0] == 0:
                    continue
                max_boxes = max(max_boxes, xyxy.shape[0])
                lf_idx_list.append(
                    {
                        "labels_idx": labels_idx,
                        "lf_idx": lf_idx,
                        "video_idx": label.videos.index(lf.video),
                        "frame_idx": lf.frame_idx,
                        "is_negative": False,
                        "instances": None,
                        "bboxes": xyxy,
                        "bbox_classes": classes,
                    }
                )
        if self.use_negative_frames and lf_idx_list:
            neg_samples = self._collect_negative_frames(labels)
            for neg in neg_samples:
                neg["bboxes"] = np.zeros((0, 4), dtype=np.float32)
                neg["bbox_classes"] = np.zeros((0,), dtype=np.int64)
            if neg_samples:
                logger.info(
                    f"Added {len(neg_samples)} negative samples "
                    f"to {len(lf_idx_list)} positive samples."
                )
                lf_idx_list.extend(neg_samples)
        # GT boxes are padded to this many slots per sample (collation needs one
        # shape). Overrides the pose-instance count `BaseDataset` computed.
        self.max_instances = max(max_boxes, 1)
        return lf_idx_list

    def _read_image(self, sample: Dict) -> np.ndarray:
        """The sample's frame as ``(H, W, C)``, from the cache or the video."""
        labels_idx, lf_idx = sample["labels_idx"], sample["lf_idx"]
        if self.cache_img == "disk":
            img = np.array(
                Image.open(f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg")
            )
        elif self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)].copy()
        elif sample["is_negative"]:
            video = self.labels_list[labels_idx].videos[sample["video_idx"]]
            img = video[sample["frame_idx"]]
        else:
            img = self.labels_list[labels_idx][lf_idx].image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)
        return img

    def __getitem__(self, index) -> Dict:
        """Return the image, padded GT boxes and detection targets for ``index``."""
        sample = self.lf_idx_list[index]
        img = self._read_image(sample)
        image = torch.from_numpy(np.transpose(img, (2, 0, 1))[None].copy())

        xyxy = sample["bboxes"]
        classes = torch.from_numpy(sample["bbox_classes"]).long()
        n_boxes = xyxy.shape[0]
        # Boxes ride the instance-coordinate path of the shared preprocessing /
        # augmentation as (1, n, 4 corners, 2) pseudo-instances.
        corners = (
            torch.from_numpy(bboxes_to_corners(xyxy))[None]
            if n_boxes
            else torch.full((1, 1, 4, 2), float("nan"))
        )
        sample_dict = {
            "image": image,
            "instances": corners,
            "video_idx": torch.tensor(sample["video_idx"], dtype=torch.int32),
            "frame_idx": torch.tensor(sample["frame_idx"], dtype=torch.int32),
            "orig_size": torch.Tensor([image.shape[-2], image.shape[-1]]).unsqueeze(0),
        }
        sample_dict = self._apply_common_preprocessing(sample_dict)
        img_hw = tuple(sample_dict["image"].shape[-2:])

        corners = sample_dict.pop("instances")[0, :n_boxes]
        boxes, keep = clip_bboxes(
            corners_to_bboxes(corners),
            img_hw,
            min_visible_fraction=self.min_visible_fraction,
        )
        boxes, classes = boxes[keep], classes[keep]
        sample_dict.update(self.target_encoder(boxes, classes, img_hw))

        n_kept = boxes.shape[0]
        padded = torch.full((self.max_instances, 4), float("nan"))
        padded[:n_kept] = boxes
        padded_cls = torch.full((self.max_instances,), -1, dtype=torch.long)
        padded_cls[:n_kept] = classes
        sample_dict["bboxes"] = padded
        sample_dict["bbox_classes"] = padded_cls
        sample_dict["num_instances"] = n_kept
        sample_dict["is_negative"] = bool(sample["is_negative"])
        sample_dict["labels_idx"] = sample["labels_idx"]
        return sample_dict

__getitem__(index)

Return the image, padded GT boxes and detection targets for index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return the image, padded GT boxes and detection targets for ``index``."""
    sample = self.lf_idx_list[index]
    img = self._read_image(sample)
    image = torch.from_numpy(np.transpose(img, (2, 0, 1))[None].copy())

    xyxy = sample["bboxes"]
    classes = torch.from_numpy(sample["bbox_classes"]).long()
    n_boxes = xyxy.shape[0]
    # Boxes ride the instance-coordinate path of the shared preprocessing /
    # augmentation as (1, n, 4 corners, 2) pseudo-instances.
    corners = (
        torch.from_numpy(bboxes_to_corners(xyxy))[None]
        if n_boxes
        else torch.full((1, 1, 4, 2), float("nan"))
    )
    sample_dict = {
        "image": image,
        "instances": corners,
        "video_idx": torch.tensor(sample["video_idx"], dtype=torch.int32),
        "frame_idx": torch.tensor(sample["frame_idx"], dtype=torch.int32),
        "orig_size": torch.Tensor([image.shape[-2], image.shape[-1]]).unsqueeze(0),
    }
    sample_dict = self._apply_common_preprocessing(sample_dict)
    img_hw = tuple(sample_dict["image"].shape[-2:])

    corners = sample_dict.pop("instances")[0, :n_boxes]
    boxes, keep = clip_bboxes(
        corners_to_bboxes(corners),
        img_hw,
        min_visible_fraction=self.min_visible_fraction,
    )
    boxes, classes = boxes[keep], classes[keep]
    sample_dict.update(self.target_encoder(boxes, classes, img_hw))

    n_kept = boxes.shape[0]
    padded = torch.full((self.max_instances, 4), float("nan"))
    padded[:n_kept] = boxes
    padded_cls = torch.full((self.max_instances,), -1, dtype=torch.long)
    padded_cls[:n_kept] = classes
    sample_dict["bboxes"] = padded
    sample_dict["bbox_classes"] = padded_cls
    sample_dict["num_instances"] = n_kept
    sample_dict["is_negative"] = bool(sample["is_negative"])
    sample_dict["labels_idx"] = sample["labels_idx"]
    return sample_dict

__init__(labels, target_encoder, categories, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False, min_visible_fraction=0.25)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    target_encoder: Any,
    categories: List[str],
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
    min_visible_fraction: float = 0.25,
) -> None:
    """Initialize class attributes."""
    # Read by `_get_lf_idx_list`, which runs inside `BaseDataset.__init__`.
    self.target_encoder = target_encoder
    self.categories = list(categories)
    self.min_visible_fraction = min_visible_fraction
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=use_negative_frames,
    )
    # A box's corners are not skeleton nodes: a horizontal flip mirrors them and
    # re-boxing takes their extent, so no index swap is needed (or wanted).
    self.symmetric_inds = []

BaseDataset

Bases: Dataset

Base class for custom torch Datasets.

Attributes:

Name	Type	Description
max_stride		Scalar integer specifying the maximum stride that the image must be divisible by.
user_instances_only		True if only user labeled instances should be used for training. If False, both user labeled and predicted instances would be used.
ensure_rgb		(bool) True if the input image should have 3 channels (RGB image). If input has only one
is	replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default	False.
ensure_grayscale		(bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
image.	If the source image has only one channel and this is set to False, then we retain the single channel input. Default	False.
intensity_aug		Intensity augmentation configuration. Can be: - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness'] - List of strings: Multiple intensity augmentations from the allowed values - Dictionary: Custom intensity configuration - None: No intensity augmentation applied
geometric_aug		Geometric augmentation configuration. Can be: - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup'] - List of strings: Multiple geometric augmentations from the allowed values - Dictionary: Custom geometric configuration - None: No geometric augmentation applied
scale		Factor to resize the image dimensions by, specified as a float. Default: 1.0.
apply_aug		True if augmentations should be applied to the data pipeline, else False. Default: False.
max_hw		Maximum height and width of images across the labels file. If max_height and max_width in the config is None, then max_hw is used (computed with sleap_nn.data.providers.get_max_height_width). Else the values in the config are used.
cache_img		String to indicate which caching to use: memory or disk. If None, the images aren't cached and loaded from the .slp file on each access.
cache_img_path		Path to save the .jpg files. If None, current working dir is used.
use_existing_imgs		Use existing imgs/ chunks in the cache_img_path.
rank		Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.
parallel_caching		If True, use parallel processing for caching (faster for large datasets). Default: True.
cache_workers		Number of worker threads for parallel caching. If 0, uses min(4, cpu_count). Default: 0.
labels_list		List of sio.Labels objects. Used to store the labels in the cache. (only used if cache_img is None)

Methods:

Name	Description
__getitem__	Returns the sample dict for given index.
__getstate__	Drop the per-process frame LRU so it is never pickled to workers.
__init__	Initialize class attributes.
__iter__	Returns an iterator.
__len__	Return the number of samples in the dataset.
__next__	Get the next sample from the dataset.

Source code in sleap_nn/data/custom_datasets.py

class BaseDataset(Dataset):
    """Base class for custom torch Datasets.

    Attributes:
        max_stride: Scalar integer specifying the maximum stride that the image must be
            divisible by.
        user_instances_only: `True` if only user labeled instances should be used for training. If `False`,
            both user labeled and predicted instances would be used.
        ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one
        channel when this is set to `True`, then the images from single-channel
        is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default: `False`.
        ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
        is set to True, then we convert the image to grayscale (single-channel)
        image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
        intensity_aug: Intensity augmentation configuration. Can be:
            - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness']
            - List of strings: Multiple intensity augmentations from the allowed values
            - Dictionary: Custom intensity configuration
            - None: No intensity augmentation applied
        geometric_aug: Geometric augmentation configuration. Can be:
            - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup']
            - List of strings: Multiple geometric augmentations from the allowed values
            - Dictionary: Custom geometric configuration
            - None: No geometric augmentation applied
        scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
        apply_aug: `True` if augmentations should be applied to the data pipeline,
            else `False`. Default: `False`.
        max_hw: Maximum height and width of images across the labels file. If `max_height` and
           `max_width` in the config is None, then `max_hw` is used (computed with
            `sleap_nn.data.providers.get_max_height_width`). Else the values in the config
            are used.
        cache_img: String to indicate which caching to use: `memory` or `disk`. If `None`,
            the images aren't cached and loaded from the `.slp` file on each access.
        cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
        use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        parallel_caching: If True, use parallel processing for caching (faster for large datasets). Default: True.
        cache_workers: Number of worker threads for parallel caching. If 0, uses min(4, cpu_count). Default: 0.
        labels_list: List of `sio.Labels` objects. Used to store the labels in the cache. (only used if `cache_img` is `None`)
    """

    # Subclasses set this True to keep frames that carry user centroid
    # annotations but no pose instances (pure-centroid seeding). Class-level so
    # it is resolved during `__init__` -> `_get_lf_idx_list`, before subclass
    # instance attributes are assigned. Only `CentroidDataset` opts in; every
    # other dataset keeps requiring a pose instance per frame.
    _include_centroid_only_frames: bool = False

    # Whether flip augmentation needs the skeleton's left/right symmetries to keep
    # labels correct. Datasets whose targets are mirror-invariant (bounding boxes)
    # set this False so a skeleton-less dataset does not warn about flipping.
    _flip_needs_symmetries: bool = True

    def __init__(
        self,
        labels: List[sio.Labels],
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
        tiling: Optional[Union[DictConfig, Any]] = None,
        output_stride: Optional[int] = None,
        base_seed: int = 0,
    ) -> None:
        """Initialize class attributes."""
        super().__init__()
        self.user_instances_only = user_instances_only
        self.use_negative_frames = use_negative_frames
        self.ensure_rgb = ensure_rgb
        self.ensure_grayscale = ensure_grayscale

        # --- Tiling awareness (Phase A) -------------------------------------
        # When a non-None, enabled tiling config is passed, the dataset switches
        # into tiled mode: the sizematcher step in `_apply_common_preprocessing`
        # is replaced by a per-tile slice (`extract_tile`) and subclasses emit
        # one (frame, tile-slot) sample per tile. When tiling is None/disabled
        # everything below stays inert and the dataset is byte-identical to
        # before (the regression guarantee).
        self.tiling = tiling
        self.tiling_enabled = tiling is not None and bool(
            getattr(tiling, "enabled", False)
        )
        self.base_seed = base_seed
        if self.tiling_enabled:
            self.tile_size = int(tiling.tile_size)
            self.overlap = int(tiling.overlap)
            # BaseDataset has no head; the head output stride is passed in by the
            # subclass so the tile grid can snap to the prediction grid.
            self.output_stride = int(output_stride) if output_stride is not None else 1
            self.tile_sampling = tiling.sampling
            self.min_overlap_fraction = float(tiling.min_overlap_fraction)
            self.tile_fg_fraction = float(tiling.tile_fg_fraction)
            self.center_jitter = float(tiling.center_jitter)
            self.min_visible_keypoints = int(tiling.min_visible_keypoints)
            self.samples_per_frame = int(tiling.samples_per_frame or 1)
            # Epoch counter in shared memory so persistent/spawned workers see the
            # main-process epoch updates (they never observe plain attribute writes).
            self._epoch = torch.zeros((), dtype=torch.long).share_memory_()

        # Handle intensity augmentation
        if intensity_aug is not None:
            if not isinstance(intensity_aug, DictConfig):
                intensity_aug = get_aug_config(intensity_aug=intensity_aug)
                config = OmegaConf.structured(intensity_aug)
                OmegaConf.to_container(config, resolve=True, throw_on_missing=True)
                intensity_aug = DictConfig(config.intensity)
        self.intensity_aug = intensity_aug

        # Handle geometric augmentation
        if geometric_aug is not None:
            if not isinstance(geometric_aug, DictConfig):
                geometric_aug = get_aug_config(geometric_aug=geometric_aug)
                config = OmegaConf.structured(geometric_aug)
                OmegaConf.to_container(config, resolve=True, throw_on_missing=True)
                geometric_aug = DictConfig(config.geometric)
        self.geometric_aug = geometric_aug
        self.curr_idx = 0
        self.max_stride = max_stride
        self.scale = scale
        self.apply_aug = apply_aug
        self.max_hw = max_hw
        self.rank = rank
        self.max_instances = 0
        for x in labels:
            max_instances = get_max_instances(x) if x else None

            if max_instances > self.max_instances:
                self.max_instances = max_instances

        # Store num_nodes for negative frame generation. Guard against mask-only
        # labels (e.g. instance segmentation) that may carry no skeleton.
        self.num_nodes = (
            len(labels[0].skeletons[0].nodes) if labels and labels[0].skeletons else 0
        )

        # Resolve symmetric node-index pairs once for flip augmentation. After
        # mirroring an image, left/right symmetric parts must be swapped to keep
        # labels correct. Read from the raw skeleton (version-proof; see
        # `get_symmetric_inds`).
        self.symmetric_inds = (
            get_symmetric_inds(labels[0].skeletons[0])
            if labels and labels[0].skeletons
            else []
        )
        # Warn about the silent correctness footgun: flipping a left/right
        # asymmetric skeleton without symmetries teaches the model wrong labels.
        flip_p = (
            self.geometric_aug.get("flip_p", 0.0)
            if self.geometric_aug is not None
            else 0.0
        )
        if (
            self.apply_aug
            and flip_p
            and flip_p > 0
            and not self.symmetric_inds
            and self._flip_needs_symmetries
        ):
            logger.warning(
                "Flip augmentation is enabled (flip_p > 0) but the skeleton has no "
                "symmetries. Flipping will not swap any nodes, which is only correct "
                "if the labeled animal is truly left/right symmetric. Add symmetry "
                "pairs to the skeleton to fix this."
            )

        self.cache_img = cache_img
        self.cache_img_path = cache_img_path
        self.use_existing_imgs = use_existing_imgs
        self.parallel_caching = parallel_caching
        self.cache_workers = cache_workers
        if self.cache_img is not None and "disk" in self.cache_img:
            if self.cache_img_path is None:
                self.cache_img_path = "."
            path = (
                Path(self.cache_img_path)
                if isinstance(self.cache_img_path, str)
                else self.cache_img_path
            )
            if not path.is_dir():
                path.mkdir(parents=True, exist_ok=True)

        self.lf_idx_list = self._get_lf_idx_list(labels)

        # Fail fast with an actionable message when no frame yields a training
        # sample. Otherwise the empty dataset surfaces much later as a cryptic
        # ``IndexError: list index out of range`` the first time ``dataset[0]``
        # is accessed (e.g. the trainer's "Input image shape" log line).
        if not self.lf_idx_list:
            n_frames = sum(len(label) for label in labels)
            raise ValueError(
                f"{type(self).__name__} has no training samples: none of the "
                f"{n_frames} labeled frame(s) in the provided labels contain "
                "user-labeled data usable by this model. Predicted instances "
                "and suggestion frames are not used as training targets (nor "
                "are standalone centroid annotations, except by centroid "
                "models). Verify that the .slp file passed for training "
                "contains user-labeled instances."
            )

        self.labels_list = None
        # this is to ensure that the labels are not passed to the multiprocessing pool when caching is enabled
        # (h5py objects can't be pickled error with num_workers > 0) in mac and windows
        if self.cache_img is None:
            self.labels_list = labels

        self.transform_to_pil = T.ToPILImage()
        self.transform_pil_to_tensor = T.ToTensor()
        self.cache = {}

        if self.cache_img is not None:
            if self.cache_img == "memory":
                # Every rank fills its own in-memory cache independently; sync
                # so one rank's failure aborts all ranks instead of leaving
                # survivors to hang at a later collective op.
                _run_cache_fill_with_dist_sync(
                    lambda: self._fill_cache(
                        labels,
                        parallel=self.parallel_caching,
                        num_workers=self.cache_workers,
                    )
                )
            elif self.cache_img == "disk" and not self.use_existing_imgs:
                is_cache_writer = self.rank is None or self.rank == -1 or self.rank == 0
                _run_cache_fill_with_dist_sync(
                    (
                        lambda: self._fill_cache(
                            labels,
                            parallel=self.parallel_caching,
                            num_workers=self.cache_workers,
                        )
                    )
                    if is_cache_writer
                    else (lambda: None)
                )
                # Synchronize all ranks after cache creation
                if is_distributed_initialized():
                    dist.barrier()

    @staticmethod
    def _extract_user_centroid_xy(
        lf: sio.LabeledFrame,
    ) -> Optional[List[List[float]]]:
        """Return a frame's user-annotated centroids as ``[[x, y], ...]`` or None.

        First-class centroid annotations (``sio.UserCentroid`` on
        ``LabeledFrame.centroids``) are the preferred confmap target for
        centroid-model training. Predicted centroids (``is_predicted=True``)
        and entries with NaN coordinates are skipped. Returns None when the
        frame carries no usable user centroids, or when the installed sleap-io
        predates first-class centroids (no ``.centroids`` attribute).

        Coordinates are pulled as plain Python floats here (at index-build time)
        so they pickle cheaply to caching workers, mirroring why ``instances``
        is only carried when caching (h5py-backed objects don't pickle).
        """
        centroids = getattr(lf, "centroids", None)
        if not centroids:
            return None
        xy: List[List[float]] = []
        for c in centroids:
            if getattr(c, "is_predicted", False):
                continue
            x = float(getattr(c, "x", float("nan")))
            y = float(getattr(c, "y", float("nan")))
            if math.isnan(x) or math.isnan(y):
                continue
            xy.append([x, y])
        return xy if xy else None

    def _get_lf_idx_list(self, labels: List[sio.Labels]) -> List[Tuple[int]]:
        """Return list of indices of labelled frames (and optionally negative frames).

        If ``self.use_negative_frames`` is True, all user-confirmed negative
        frames (``labels.negative_frames``) are appended to the sample list.
        """
        lf_idx_list = []
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                # User-annotated centroids as plain floats (picklable). Computed
                # before instance filtering (centroids are independent of pose
                # instances).
                user_centroids = self._extract_user_centroid_xy(lf)
                # A frame with user centroids but no pose instances is a valid
                # sample for the centroid model only (pure-centroid seeding).
                centroid_only_ok = self._include_centroid_only_frames and bool(
                    user_centroids
                )
                # Filter to user instances
                if self.user_instances_only:
                    if lf.user_instances is not None and len(lf.user_instances) > 0:
                        lf.instances = lf.user_instances
                    elif centroid_only_ok:
                        lf.instances = []
                    else:
                        # Skip frames without user instances
                        continue
                is_empty = True
                for _, inst in enumerate(lf.instances):
                    if not inst.is_empty:  # filter all NaN instances.
                        is_empty = False
                if (not is_empty) or centroid_only_ok:
                    video_idx = labels[labels_idx].videos.index(lf.video)
                    sample = {
                        "labels_idx": labels_idx,
                        "lf_idx": lf_idx,
                        "video_idx": video_idx,
                        "frame_idx": lf.frame_idx,
                        "is_negative": False,
                        "instances": (
                            lf.instances if self.cache_img is not None else None
                        ),
                        # Carried unconditionally: the confmap target prefers
                        # these over instance-keypoint-derived centroids. None
                        # when the frame has no user centroids (fallback path).
                        "user_centroids": user_centroids,
                        # Whether the frame has at least one non-empty pose
                        # instance. Lets CentroidDataset drop pose-less
                        # (centroid-only) frames when training on computed
                        # centroids, the mirror of dropping centroid-less frames
                        # when training on user centroids.
                        "has_pose_instances": (not is_empty),
                    }
                    lf_idx_list.append(sample)
                    # This is to ensure that the labels are not passed to the multiprocessing pool (h5py objects can't be pickled)

        # Add negative frames if requested
        if self.use_negative_frames and len(lf_idx_list) > 0:
            neg_samples = self._collect_negative_frames(labels)
            if neg_samples:
                n_positive = len(lf_idx_list)
                lf_idx_list.extend(neg_samples)
                logger.info(
                    f"Added {len(neg_samples)} negative samples "
                    f"to {n_positive} positive samples."
                )

        return lf_idx_list

    def _collect_negative_frames(
        self,
        labels: List[sio.Labels],
    ) -> List[Dict]:
        """Collect all user-confirmed negative frames from labels.

        Only frames explicitly marked by the user as negative are used
        (``LabeledFrame`` objects with ``is_negative=True``, accessed via
        ``labels.negative_frames``).  Unlabeled frames are **not** included
        because they may contain animals that simply haven't been annotated yet.

        Args:
            labels: List of sio.Labels objects.

        Returns:
            List of sample dicts with ``is_negative=True`` (one per unique
            negative frame).
        """
        neg_samples: List[Dict] = []

        for labels_idx, label in enumerate(labels):
            if not hasattr(label, "negative_frames"):
                continue
            for idx, lf in enumerate(label.negative_frames):
                video_idx = label.videos.index(lf.video)
                neg_samples.append(
                    {
                        "labels_idx": labels_idx,
                        "lf_idx": f"neg_{labels_idx}_{idx}",
                        "video_idx": video_idx,
                        "frame_idx": lf.frame_idx,
                        "is_negative": True,
                        "instances": None,
                    }
                )

        return neg_samples

    def _frame_lru(self) -> _FrameLRU:
        """Return this process's decoded-frame LRU, building it lazily.

        The cache is stored keyed by ``os.getpid()`` so forked / persistent
        dataloader workers each get their own instance and never share (or
        pickle) decoded-frame tensors. It is excluded from pickling via
        ``__getstate__`` and rebuilt on first access in each process.
        """
        pid = os.getpid()
        store = self.__dict__.get("_frame_lru_store")
        if store is None or store.get("pid") != pid:
            store = {"pid": pid, "lru": _FrameLRU(_FRAME_LRU_CAPACITY)}
            self.__dict__["_frame_lru_store"] = store
        return store["lru"]

    def __getstate__(self):
        """Drop the per-process frame LRU so it is never pickled to workers."""
        state = self.__dict__.copy()
        state.pop("_frame_lru_store", None)
        return state

    def _frame_sized_hw(self, lf: sio.LabeledFrame) -> Tuple[int, int]:
        """Return a labeled frame's ``(H, W)`` after ``scale`` (sizematcher bypassed).

        Uses video metadata when available to avoid decoding the frame. Matches
        the ``int(dim * scale)`` truncation used by :func:`resize_image` so grid
        origins align with the sized frame produced by ``apply_resizer``.
        """
        shape = getattr(lf.video, "shape", None)
        if shape is not None and len(shape) >= 3:
            height, width = int(shape[1]), int(shape[2])
        else:
            img = lf.image
            height, width = int(img.shape[0]), int(img.shape[1])
        if self.scale != 1.0:
            height = int(height * self.scale)
            width = int(width * self.scale)
        return height, width

    def _get_tile_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return per-(frame, tile-slot) sample descriptors for tiled datasets.

        Mirrors ``_get_lf_idx_list`` (same user-instance filtering, empty-frame
        skip, and cache-aware ``instances`` storage), but explodes each frame
        into multiple tile descriptors: one per grid tile (``sampling="grid"``,
        pinned origins) or ``samples_per_frame`` runtime-drawn tiles
        (``sampling="foreground"``, ``tile_origin=None``). A frame's descriptors
        form a contiguous run (the block the sampler groups on).
        """
        tile_idx_list: List[Dict] = []
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                if self.user_instances_only:
                    if lf.user_instances is not None and len(lf.user_instances) > 0:
                        lf.instances = lf.user_instances
                    else:
                        continue
                if all(inst.is_empty for inst in lf.instances):
                    continue
                video_idx = labels[labels_idx].videos.index(lf.video)

                if self.tile_sampling == "grid":
                    origins = generate_tile_grid(
                        self._frame_sized_hw(lf),
                        tile_size=self.tile_size,
                        overlap=self.overlap,
                        output_stride=self.output_stride,
                        max_stride=self.max_stride,
                        min_overlap_fraction=self.min_overlap_fraction,
                    )
                else:
                    origins = [None] * self.samples_per_frame

                for sample_k, origin in enumerate(origins):
                    tile_idx_list.append(
                        {
                            "labels_idx": labels_idx,
                            "lf_idx": lf_idx,
                            "video_idx": video_idx,
                            "frame_idx": lf.frame_idx,
                            "instances": (
                                lf.instances if self.cache_img is not None else None
                            ),
                            "sample_k": sample_k,
                            "tile_origin": origin,
                            "is_grid": self.tile_sampling == "grid",
                            "is_negative": False,
                        }
                    )

        if self.use_negative_frames:
            for neg in self._collect_negative_frames(labels):
                for sample_k in range(self.samples_per_frame):
                    d = dict(neg)
                    d.update(
                        sample_k=sample_k,
                        tile_origin=None,
                        is_grid=False,
                        is_negative=True,
                    )
                    tile_idx_list.append(d)

        return tile_idx_list

    @staticmethod
    def _build_frame_blocks(tile_idx_list: List[Dict]) -> List[List[int]]:
        """Group contiguous per-frame runs of ``tile_idx_list`` into index blocks."""
        blocks: List[List[int]] = []
        prev_key = object()
        for idx, d in enumerate(tile_idx_list):
            key = (d["labels_idx"], d["lf_idx"])
            if key != prev_key:
                blocks.append([])
                prev_key = key
            blocks[-1].append(idx)
        return blocks

    def __next__(self):
        """Get the next sample from the dataset."""
        if self.curr_idx >= len(self):
            raise StopIteration

        sample = self.__getitem__(self.curr_idx)
        self.curr_idx += 1
        return sample

    def __iter__(self):
        """Returns an iterator."""
        return self

    def _fill_cache(
        self,
        labels: List[sio.Labels],
        parallel: bool = True,
        num_workers: int = 0,
    ):
        """Load all samples to cache.

        Args:
            labels: List of sio.Labels objects containing the data.
            parallel: If True, use parallel processing for caching (faster for large
                datasets). Default: True.
            num_workers: Number of worker threads for parallel caching. If 0, uses
                min(4, cpu_count). Default: 0.
        """
        total_samples = len(self.lf_idx_list)
        cache_type = "disk" if self.cache_img == "disk" else "memory"

        # Check for NO_COLOR env var to disable progress bar
        no_color = (
            os.environ.get("NO_COLOR") is not None
            or os.environ.get("FORCE_COLOR") == "0"
        )
        use_progress = not no_color

        # Use parallel caching for larger datasets
        use_parallel = parallel and total_samples >= MIN_SAMPLES_FOR_PARALLEL_CACHING

        logger.info(f"Caching {total_samples} images to {cache_type}...")

        if use_parallel:
            self._fill_cache_parallel(
                labels, total_samples, cache_type, use_progress, num_workers
            )
        else:
            self._fill_cache_sequential(labels, total_samples, cache_type, use_progress)

        logger.info(f"Caching complete.")

    def _fill_cache_sequential(
        self,
        labels: List[sio.Labels],
        total_samples: int,
        cache_type: str,
        use_progress: bool,
    ):
        """Sequential implementation of cache filling.

        Args:
            labels: List of sio.Labels objects.
            total_samples: Total number of samples to cache.
            cache_type: Either "disk" or "memory".
            use_progress: Whether to show a progress bar.
        """

        def process_samples(progress=None, task=None):
            for sample in self.lf_idx_list:
                labels_idx = sample["labels_idx"]
                lf_idx = sample["lf_idx"]
                try:
                    if sample.get("is_negative", False):
                        video_idx = sample["video_idx"]
                        frame_idx = sample["frame_idx"]
                        img = labels[labels_idx].videos[video_idx][frame_idx]
                    else:
                        img = labels[labels_idx][lf_idx].image
                    if img.shape[-1] == 1:
                        img = np.squeeze(img)
                    if self.cache_img == "disk":
                        f_name = (
                            f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                        )
                        Image.fromarray(img).save(f_name, format="JPEG")
                    if self.cache_img == "memory":
                        self.cache[(labels_idx, lf_idx)] = img
                except Exception as e:
                    _raise_cache_fill_error(
                        [(labels_idx, lf_idx, f"{type(e).__name__}: {e}")],
                        total_samples,
                        cache_type,
                    )
                if progress is not None:
                    progress.update(task, advance=1)

        if use_progress:
            with _cache_progress() as progress:
                task = progress.add_task(
                    f"Caching images to {cache_type}", total=total_samples
                )
                process_samples(progress, task)
        else:
            process_samples()

    def _fill_cache_parallel(
        self,
        labels: List[sio.Labels],
        total_samples: int,
        cache_type: str,
        use_progress: bool,
        num_workers: int = 0,
    ):
        """Parallel implementation of cache filling using thread-local video copies.

        Args:
            labels: List of sio.Labels objects.
            total_samples: Total number of samples to cache.
            cache_type: Either "disk" or "memory".
            use_progress: Whether to show a progress bar.
            num_workers: Number of worker threads. If 0, uses min(4, cpu_count).
        """
        # Determine number of workers
        if num_workers <= 0:
            num_workers = min(4, os.cpu_count() or 1)

        cache_path = Path(self.cache_img_path) if self.cache_img_path else None

        filler = ParallelCacheFiller(
            labels=labels,
            lf_idx_list=self.lf_idx_list,
            cache_type=cache_type,
            cache_path=cache_path,
            num_workers=num_workers,
        )

        if use_progress:
            with _cache_progress() as progress:
                task = progress.add_task(
                    f"Caching images to {cache_type} (parallel, {num_workers} workers)",
                    total=total_samples,
                )

                def progress_callback(completed):
                    progress.update(task, completed=completed)

                cache, errors = filler.fill_cache(progress_callback)
        else:
            logger.info(
                f"Caching {total_samples} images to {cache_type} "
                f"(parallel, {num_workers} workers)..."
            )
            cache, errors = filler.fill_cache()

        # Update instance cache
        if cache_type == "memory":
            self.cache.update(cache)

        # A frame that failed to cache is silently missing from `cache`/disk;
        # letting training proceed means it surfaces as a confusing
        # FileNotFoundError/KeyError in a random later DataLoader batch
        # instead of here, where we know exactly which frame failed and why.
        if errors:
            _raise_cache_fill_error(errors, total_samples, cache_type)

    def _apply_common_preprocessing(self, sample: Dict) -> Dict:
        """Apply common preprocessing steps shared across all dataset types.

        Handles: RGB/grayscale conversion, size matching, scaling, padding,
        and augmentation.

        Args:
            sample: Sample dict with at least ``image`` and ``instances`` keys.

        Returns:
            The sample dict with preprocessing applied in-place.
        """
        if self.ensure_rgb:
            sample["image"] = convert_to_rgb(sample["image"])
        elif self.ensure_grayscale:
            sample["image"] = convert_to_grayscale(sample["image"])

        if not self.tiling_enabled:
            # size matcher
            sample["image"], eff_scale = apply_sizematcher(
                sample["image"],
                max_height=self.max_hw[0],
                max_width=self.max_hw[1],
            )
            sample["instances"] = sample["instances"] * eff_scale
            sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

            # resize image
            sample["image"], sample["instances"] = apply_resizer(
                sample["image"],
                sample["instances"],
                scale=self.scale,
            )

            # Co-transform whole-frame segmentation masks (a (1, K, H, W) float
            # tensor the seg datasets place into the sample) through the IDENTICAL
            # size-match + scale as the image — the SAME helpers, so the mask
            # target stays pixel-aligned with the image by construction (bilinear
            # via ``tvf.resize``, matching the image). Non-seg samples have no
            # ``masks`` key, so this is a no-op for them.
            if sample.get("masks") is not None:
                sample["masks"], _ = apply_sizematcher(
                    sample["masks"],
                    max_height=self.max_hw[0],
                    max_width=self.max_hw[1],
                )
                sample["masks"], _ = apply_resizer(
                    sample["masks"], torch.zeros(1), scale=self.scale
                )
        else:
            # TILING: slice a tile IN PLACE OF the sizematcher (constant-zero pad
            # only). The incoming frame is already scaled + channel-coerced (via
            # `_to_sized_frame`), so `scale` is not re-applied here (that would
            # double-scale); tiles are extracted in the model's input space where
            # `tile_size` is divisible by the network strides. Geometric aug is
            # folded into `extract_tile` (halo path) when enabled.
            sample["image"], sample["instances"] = extract_tile(
                image=sample["image"],
                instances=sample["instances"],
                tile_origin=sample["tile_origin"],
                tile_size=self.tile_size,
                apply_geometric=(self.apply_aug and self.geometric_aug is not None),
                geometric_kwargs=(
                    dict(self.geometric_aug) if self.geometric_aug is not None else None
                ),
                symmetric_inds=self.symmetric_inds,
                rng_seed=sample.get("aug_seed"),
            )
            sample["eff_scale"] = torch.tensor(1.0, dtype=torch.float32)
            sample["tile_origin"] = torch.tensor(
                sample["tile_origin"], dtype=torch.int32
            )

        # Pad the image (if needed) according max stride. Per-tile under tiling;
        # a no-op when `tile_size % max_stride == 0` (guaranteed by write-back).
        sample["image"] = apply_pad_to_stride(
            sample["image"], max_stride=self.max_stride
        )
        # Pad segmentation masks to the SAME bottom-right stride multiple as the
        # image so the whole-frame target stays registered to the padded image.
        if sample.get("masks") is not None:
            sample["masks"] = apply_pad_to_stride(
                sample["masks"], max_stride=self.max_stride
            )

        # apply augmentation
        if self.apply_aug:
            if self.intensity_aug is not None:
                sample["image"], sample["instances"] = apply_intensity_augmentation(
                    sample["image"],
                    sample["instances"],
                    **self.intensity_aug,
                )

            # Under tiling, geometric augmentation is folded into `extract_tile`
            # (halo path); do NOT run it again here.
            if not self.tiling_enabled and self.geometric_aug is not None:
                sample["image"], sample["instances"] = apply_geometric_augmentation(
                    sample["image"],
                    sample["instances"],
                    symmetric_inds=self.symmetric_inds,
                    **self.geometric_aug,
                )

        return sample

    def _load_negative_sample(self, sample: Dict) -> Dict:
        """Load and preprocess a negative frame (no instances).

        Reads the image from the cache (if available) or video and returns a
        sample dict with all-NaN instances.  Downstream code will generate
        all-zero confidence maps from these NaN instances.

        Args:
            sample: Sample dict from lf_idx_list with ``is_negative=True``.

        Returns:
            Preprocessed sample dict.
        """
        labels_idx = sample["labels_idx"]
        video_idx = sample["video_idx"]
        frame_idx = sample["frame_idx"]
        lf_idx = sample["lf_idx"]

        if self.cache_img == "disk":
            img = np.array(
                Image.open(f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg")
            )
        elif self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            video = self.labels_list[labels_idx].videos[video_idx]
            img = video[frame_idx]
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        return process_negative_lf(
            img=img,
            frame_idx=frame_idx,
            video_idx=video_idx,
            max_instances=self.max_instances,
            num_nodes=self.num_nodes,
        )

    def __len__(self) -> int:
        """Return the number of samples in the dataset."""
        return len(self.lf_idx_list)

    def __getitem__(self, index) -> Dict:
        """Returns the sample dict for given index."""
        message = "Subclasses must implement __getitem__"
        logger.error(message)
        raise NotImplementedError(message)

__getitem__(index)

Returns the sample dict for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Returns the sample dict for given index."""
    message = "Subclasses must implement __getitem__"
    logger.error(message)
    raise NotImplementedError(message)

__getstate__()

Drop the per-process frame LRU so it is never pickled to workers.

Source code in sleap_nn/data/custom_datasets.py

def __getstate__(self):
    """Drop the per-process frame LRU so it is never pickled to workers."""
    state = self.__dict__.copy()
    state.pop("_frame_lru_store", None)
    return state

__init__(labels, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False, tiling=None, output_stride=None, base_seed=0)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
    tiling: Optional[Union[DictConfig, Any]] = None,
    output_stride: Optional[int] = None,
    base_seed: int = 0,
) -> None:
    """Initialize class attributes."""
    super().__init__()
    self.user_instances_only = user_instances_only
    self.use_negative_frames = use_negative_frames
    self.ensure_rgb = ensure_rgb
    self.ensure_grayscale = ensure_grayscale

    # --- Tiling awareness (Phase A) -------------------------------------
    # When a non-None, enabled tiling config is passed, the dataset switches
    # into tiled mode: the sizematcher step in `_apply_common_preprocessing`
    # is replaced by a per-tile slice (`extract_tile`) and subclasses emit
    # one (frame, tile-slot) sample per tile. When tiling is None/disabled
    # everything below stays inert and the dataset is byte-identical to
    # before (the regression guarantee).
    self.tiling = tiling
    self.tiling_enabled = tiling is not None and bool(
        getattr(tiling, "enabled", False)
    )
    self.base_seed = base_seed
    if self.tiling_enabled:
        self.tile_size = int(tiling.tile_size)
        self.overlap = int(tiling.overlap)
        # BaseDataset has no head; the head output stride is passed in by the
        # subclass so the tile grid can snap to the prediction grid.
        self.output_stride = int(output_stride) if output_stride is not None else 1
        self.tile_sampling = tiling.sampling
        self.min_overlap_fraction = float(tiling.min_overlap_fraction)
        self.tile_fg_fraction = float(tiling.tile_fg_fraction)
        self.center_jitter = float(tiling.center_jitter)
        self.min_visible_keypoints = int(tiling.min_visible_keypoints)
        self.samples_per_frame = int(tiling.samples_per_frame or 1)
        # Epoch counter in shared memory so persistent/spawned workers see the
        # main-process epoch updates (they never observe plain attribute writes).
        self._epoch = torch.zeros((), dtype=torch.long).share_memory_()

    # Handle intensity augmentation
    if intensity_aug is not None:
        if not isinstance(intensity_aug, DictConfig):
            intensity_aug = get_aug_config(intensity_aug=intensity_aug)
            config = OmegaConf.structured(intensity_aug)
            OmegaConf.to_container(config, resolve=True, throw_on_missing=True)
            intensity_aug = DictConfig(config.intensity)
    self.intensity_aug = intensity_aug

    # Handle geometric augmentation
    if geometric_aug is not None:
        if not isinstance(geometric_aug, DictConfig):
            geometric_aug = get_aug_config(geometric_aug=geometric_aug)
            config = OmegaConf.structured(geometric_aug)
            OmegaConf.to_container(config, resolve=True, throw_on_missing=True)
            geometric_aug = DictConfig(config.geometric)
    self.geometric_aug = geometric_aug
    self.curr_idx = 0
    self.max_stride = max_stride
    self.scale = scale
    self.apply_aug = apply_aug
    self.max_hw = max_hw
    self.rank = rank
    self.max_instances = 0
    for x in labels:
        max_instances = get_max_instances(x) if x else None

        if max_instances > self.max_instances:
            self.max_instances = max_instances

    # Store num_nodes for negative frame generation. Guard against mask-only
    # labels (e.g. instance segmentation) that may carry no skeleton.
    self.num_nodes = (
        len(labels[0].skeletons[0].nodes) if labels and labels[0].skeletons else 0
    )

    # Resolve symmetric node-index pairs once for flip augmentation. After
    # mirroring an image, left/right symmetric parts must be swapped to keep
    # labels correct. Read from the raw skeleton (version-proof; see
    # `get_symmetric_inds`).
    self.symmetric_inds = (
        get_symmetric_inds(labels[0].skeletons[0])
        if labels and labels[0].skeletons
        else []
    )
    # Warn about the silent correctness footgun: flipping a left/right
    # asymmetric skeleton without symmetries teaches the model wrong labels.
    flip_p = (
        self.geometric_aug.get("flip_p", 0.0)
        if self.geometric_aug is not None
        else 0.0
    )
    if (
        self.apply_aug
        and flip_p
        and flip_p > 0
        and not self.symmetric_inds
        and self._flip_needs_symmetries
    ):
        logger.warning(
            "Flip augmentation is enabled (flip_p > 0) but the skeleton has no "
            "symmetries. Flipping will not swap any nodes, which is only correct "
            "if the labeled animal is truly left/right symmetric. Add symmetry "
            "pairs to the skeleton to fix this."
        )

    self.cache_img = cache_img
    self.cache_img_path = cache_img_path
    self.use_existing_imgs = use_existing_imgs
    self.parallel_caching = parallel_caching
    self.cache_workers = cache_workers
    if self.cache_img is not None and "disk" in self.cache_img:
        if self.cache_img_path is None:
            self.cache_img_path = "."
        path = (
            Path(self.cache_img_path)
            if isinstance(self.cache_img_path, str)
            else self.cache_img_path
        )
        if not path.is_dir():
            path.mkdir(parents=True, exist_ok=True)

    self.lf_idx_list = self._get_lf_idx_list(labels)

    # Fail fast with an actionable message when no frame yields a training
    # sample. Otherwise the empty dataset surfaces much later as a cryptic
    # ``IndexError: list index out of range`` the first time ``dataset[0]``
    # is accessed (e.g. the trainer's "Input image shape" log line).
    if not self.lf_idx_list:
        n_frames = sum(len(label) for label in labels)
        raise ValueError(
            f"{type(self).__name__} has no training samples: none of the "
            f"{n_frames} labeled frame(s) in the provided labels contain "
            "user-labeled data usable by this model. Predicted instances "
            "and suggestion frames are not used as training targets (nor "
            "are standalone centroid annotations, except by centroid "
            "models). Verify that the .slp file passed for training "
            "contains user-labeled instances."
        )

    self.labels_list = None
    # this is to ensure that the labels are not passed to the multiprocessing pool when caching is enabled
    # (h5py objects can't be pickled error with num_workers > 0) in mac and windows
    if self.cache_img is None:
        self.labels_list = labels

    self.transform_to_pil = T.ToPILImage()
    self.transform_pil_to_tensor = T.ToTensor()
    self.cache = {}

    if self.cache_img is not None:
        if self.cache_img == "memory":
            # Every rank fills its own in-memory cache independently; sync
            # so one rank's failure aborts all ranks instead of leaving
            # survivors to hang at a later collective op.
            _run_cache_fill_with_dist_sync(
                lambda: self._fill_cache(
                    labels,
                    parallel=self.parallel_caching,
                    num_workers=self.cache_workers,
                )
            )
        elif self.cache_img == "disk" and not self.use_existing_imgs:
            is_cache_writer = self.rank is None or self.rank == -1 or self.rank == 0
            _run_cache_fill_with_dist_sync(
                (
                    lambda: self._fill_cache(
                        labels,
                        parallel=self.parallel_caching,
                        num_workers=self.cache_workers,
                    )
                )
                if is_cache_writer
                else (lambda: None)
            )
            # Synchronize all ranks after cache creation
            if is_distributed_initialized():
                dist.barrier()

__iter__()

Returns an iterator.

Source code in sleap_nn/data/custom_datasets.py

def __iter__(self):
    """Returns an iterator."""
    return self

__len__()

Return the number of samples in the dataset.

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Return the number of samples in the dataset."""
    return len(self.lf_idx_list)

__next__()

Get the next sample from the dataset.

Source code in sleap_nn/data/custom_datasets.py

def __next__(self):
    """Get the next sample from the dataset."""
    if self.curr_idx >= len(self):
        raise StopIteration

    sample = self.__getitem__(self.curr_idx)
    self.curr_idx += 1
    return sample

BottomUpDataset

Bases: BaseDataset

Dataset class for bottom-up models.

Attributes:

Name	Type	Description
max_stride		Scalar integer specifying the maximum stride that the image must be divisible by.
user_instances_only		True if only user labeled instances should be used for training. If False, both user labeled and predicted instances would be used.
ensure_rgb		(bool) True if the input image should have 3 channels (RGB image). If input has only one
is	replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default	False.
ensure_grayscale		(bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
image.	If the source image has only one channel and this is set to False, then we retain the single channel input. Default	False.
intensity_aug		Intensity augmentation configuration. Can be: - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness'] - List of strings: Multiple intensity augmentations from the allowed values - Dictionary: Custom intensity configuration - None: No intensity augmentation applied
geometric_aug		Geometric augmentation configuration. Can be: - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup'] - List of strings: Multiple geometric augmentations from the allowed values - Dictionary: Custom geometric configuration - None: No geometric augmentation applied
scale		Factor to resize the image dimensions by, specified as a float. Default: 1.0.
apply_aug		True if augmentations should be applied to the data pipeline, else False. Default: False.
max_hw		Maximum height and width of images across the labels file. If max_height and max_width in the config is None, then max_hw is used (computed with sleap_nn.data.providers.get_max_height_width). Else the values in the config are used.
confmap_head_config		DictConfig object with all the keys in the head_config section. (required keys: sigma, output_stride and anchor_part depending on the model type ).
pafs_head_config		DictConfig object with all the keys in the head_config section (required keys: sigma, output_stride and anchor_part depending on the model type ) for PAFs.
cache_img		String to indicate which caching to use: memory or disk. If None, the images aren't cached and loaded from the .slp file on each access.
cache_img_path		Path to save the .jpg files. If None, current working dir is used.
use_existing_imgs		Use existing imgs/ chunks in the cache_img_path.
rank		Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.
labels_list		List of sio.Labels objects. Used to store the labels in the cache. (only used if cache_img is None)

Methods:

Name	Description
__getitem__	Return dict with image, confmaps and pafs for given index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class BottomUpDataset(BaseDataset):
    """Dataset class for bottom-up models.

    Attributes:
        max_stride: Scalar integer specifying the maximum stride that the image must be
            divisible by.
        user_instances_only: `True` if only user labeled instances should be used for training. If `False`,
            both user labeled and predicted instances would be used.
        ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one
        channel when this is set to `True`, then the images from single-channel
        is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default: `False`.
        ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
        is set to True, then we convert the image to grayscale (single-channel)
        image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
        intensity_aug: Intensity augmentation configuration. Can be:
            - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness']
            - List of strings: Multiple intensity augmentations from the allowed values
            - Dictionary: Custom intensity configuration
            - None: No intensity augmentation applied
        geometric_aug: Geometric augmentation configuration. Can be:
            - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup']
            - List of strings: Multiple geometric augmentations from the allowed values
            - Dictionary: Custom geometric configuration
            - None: No geometric augmentation applied
        scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
        apply_aug: `True` if augmentations should be applied to the data pipeline,
            else `False`. Default: `False`.
        max_hw: Maximum height and width of images across the labels file. If `max_height` and
           `max_width` in the config is None, then `max_hw` is used (computed with
            `sleap_nn.data.providers.get_max_height_width`). Else the values in the config
            are used.
        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
            (required keys: `sigma`, `output_stride` and `anchor_part` depending on the model type ).
        pafs_head_config: DictConfig object with all the keys in the `head_config` section
            (required keys: `sigma`, `output_stride` and `anchor_part` depending on the model type )
            for PAFs.
        cache_img: String to indicate which caching to use: `memory` or `disk`. If `None`,
            the images aren't cached and loaded from the `.slp` file on each access.
        cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
        use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        labels_list: List of `sio.Labels` objects. Used to store the labels in the cache. (only used if `cache_img` is `None`)
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        confmap_head_config: DictConfig,
        pafs_head_config: DictConfig,
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
    ) -> None:
        """Initialize class attributes."""
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        self.confmap_head_config = confmap_head_config
        self.pafs_head_config = pafs_head_config

        self.edge_inds = labels[0].skeletons[0].edge_inds

    def __getitem__(self, index) -> Dict:
        """Return dict with image, confmaps and pafs for given index."""
        sample = self.lf_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        video_idx = sample["video_idx"]
        frame_idx = sample["frame_idx"]

        if sample.get("is_negative", False):
            sample = self._load_negative_sample(sample)
        else:
            if self.cache_img is not None:
                instances = sample["instances"]
                if self.cache_img == "disk":
                    img = np.array(
                        Image.open(
                            f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                        )
                    )
                elif self.cache_img == "memory":
                    img = self.cache[(labels_idx, lf_idx)].copy()
            else:
                lf = self.labels_list[labels_idx][lf_idx]
                instances = lf.instances
                img = lf.image

            if img.ndim == 2:
                img = np.expand_dims(img, axis=2)

            # get dict
            sample = process_lf(
                instances_list=instances,
                img=img,
                frame_idx=frame_idx,
                video_idx=video_idx,
                max_instances=self.max_instances,
                user_instances_only=self.user_instances_only,
            )

        sample = self._apply_common_preprocessing(sample)

        img_hw = sample["image"].shape[-2:]

        # Generate confidence maps
        confidence_maps = generate_multiconfmaps(
            sample["instances"],
            img_hw=img_hw,
            num_instances=sample["num_instances"],
            sigma=self.confmap_head_config.sigma,
            output_stride=self.confmap_head_config.output_stride,
            is_centroids=False,
        )

        # pafs
        pafs = generate_pafs(
            sample["instances"],
            img_hw=img_hw,
            sigma=self.pafs_head_config.sigma,
            output_stride=self.pafs_head_config.output_stride,
            edge_inds=torch.Tensor(self.edge_inds),
            flatten_channels=True,
        )

        sample["confidence_maps"] = confidence_maps
        sample["part_affinity_fields"] = pafs
        sample["labels_idx"] = labels_idx
        if self.use_negative_frames:
            sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

        return sample

__getitem__(index)

Return dict with image, confmaps and pafs for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image, confmaps and pafs for given index."""
    sample = self.lf_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    video_idx = sample["video_idx"]
    frame_idx = sample["frame_idx"]

    if sample.get("is_negative", False):
        sample = self._load_negative_sample(sample)
    else:
        if self.cache_img is not None:
            instances = sample["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances = lf.instances
            img = lf.image

        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        # get dict
        sample = process_lf(
            instances_list=instances,
            img=img,
            frame_idx=frame_idx,
            video_idx=video_idx,
            max_instances=self.max_instances,
            user_instances_only=self.user_instances_only,
        )

    sample = self._apply_common_preprocessing(sample)

    img_hw = sample["image"].shape[-2:]

    # Generate confidence maps
    confidence_maps = generate_multiconfmaps(
        sample["instances"],
        img_hw=img_hw,
        num_instances=sample["num_instances"],
        sigma=self.confmap_head_config.sigma,
        output_stride=self.confmap_head_config.output_stride,
        is_centroids=False,
    )

    # pafs
    pafs = generate_pafs(
        sample["instances"],
        img_hw=img_hw,
        sigma=self.pafs_head_config.sigma,
        output_stride=self.pafs_head_config.output_stride,
        edge_inds=torch.Tensor(self.edge_inds),
        flatten_channels=True,
    )

    sample["confidence_maps"] = confidence_maps
    sample["part_affinity_fields"] = pafs
    sample["labels_idx"] = labels_idx
    if self.use_negative_frames:
        sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

    return sample

__init__(labels, confmap_head_config, pafs_head_config, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    confmap_head_config: DictConfig,
    pafs_head_config: DictConfig,
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
) -> None:
    """Initialize class attributes."""
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=use_negative_frames,
    )
    self.confmap_head_config = confmap_head_config
    self.pafs_head_config = pafs_head_config

    self.edge_inds = labels[0].skeletons[0].edge_inds

BottomUpMultiClassDataset

Bases: BaseDataset

Dataset class for bottom-up ID models.

Attributes:

Name	Type	Description
max_stride		Scalar integer specifying the maximum stride that the image must be divisible by.
class_map_threshold		Minimum confidence map value below which map values will be replaced with zeros.
user_instances_only		True if only user labeled instances should be used for training. If False, both user labeled and predicted instances would be used.
ensure_rgb		(bool) True if the input image should have 3 channels (RGB image). If input has only one
is	replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default	False.
ensure_grayscale		(bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
image.	If the source image has only one channel and this is set to False, then we retain the single channel input. Default	False.
intensity_aug		Intensity augmentation configuration. Can be: - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness'] - List of strings: Multiple intensity augmentations from the allowed values - Dictionary: Custom intensity configuration - None: No intensity augmentation applied
geometric_aug		Geometric augmentation configuration. Can be: - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup'] - List of strings: Multiple geometric augmentations from the allowed values - Dictionary: Custom geometric configuration - None: No geometric augmentation applied
scale		Factor to resize the image dimensions by, specified as a float. Default: 1.0.
apply_aug		True if augmentations should be applied to the data pipeline, else False. Default: False.
max_hw		Maximum height and width of images across the labels file. If max_height and max_width in the config is None, then max_hw is used (computed with sleap_nn.data.providers.get_max_height_width). Else the values in the config are used.
confmap_head_config		DictConfig object with all the keys in the head_config section. (required keys: sigma, output_stride and anchor_part depending on the model type ).
class_maps_head_config		DictConfig object with all the keys in the head_config section (required keys: sigma, output_stride and classes) for class maps.
cache_img		String to indicate which caching to use: memory or disk. If None, the images aren't cached and loaded from the .slp file on each access.
cache_img_path		Path to save the .jpg files. If None, current working dir is used.
use_existing_imgs		Use existing imgs/ chunks in the cache_img_path.
rank		Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.
labels_list		List of sio.Labels objects. Used to store the labels in the cache. (only used if cache_img is None)

Methods:

Name	Description
__getitem__	Return dict with image, confmaps and class maps for given index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class BottomUpMultiClassDataset(BaseDataset):
    """Dataset class for bottom-up ID models.

    Attributes:
        max_stride: Scalar integer specifying the maximum stride that the image must be
            divisible by.
        class_map_threshold: Minimum confidence map value below which map values will be
            replaced with zeros.
        user_instances_only: `True` if only user labeled instances should be used for training. If `False`,
            both user labeled and predicted instances would be used.
        ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one
        channel when this is set to `True`, then the images from single-channel
        is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default: `False`.
        ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
        is set to True, then we convert the image to grayscale (single-channel)
        image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
        intensity_aug: Intensity augmentation configuration. Can be:
            - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness']
            - List of strings: Multiple intensity augmentations from the allowed values
            - Dictionary: Custom intensity configuration
            - None: No intensity augmentation applied
        geometric_aug: Geometric augmentation configuration. Can be:
            - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup']
            - List of strings: Multiple geometric augmentations from the allowed values
            - Dictionary: Custom geometric configuration
            - None: No geometric augmentation applied
        scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
        apply_aug: `True` if augmentations should be applied to the data pipeline,
            else `False`. Default: `False`.
        max_hw: Maximum height and width of images across the labels file. If `max_height` and
           `max_width` in the config is None, then `max_hw` is used (computed with
            `sleap_nn.data.providers.get_max_height_width`). Else the values in the config
            are used.
        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
            (required keys: `sigma`, `output_stride` and `anchor_part` depending on the model type ).
        class_maps_head_config: DictConfig object with all the keys in the `head_config` section
            (required keys: `sigma`, `output_stride` and `classes`)
            for class maps.
        cache_img: String to indicate which caching to use: `memory` or `disk`. If `None`,
            the images aren't cached and loaded from the `.slp` file on each access.
        cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
        use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        labels_list: List of `sio.Labels` objects. Used to store the labels in the cache. (only used if `cache_img` is `None`)
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        confmap_head_config: DictConfig,
        class_maps_head_config: DictConfig,
        max_stride: int,
        class_map_threshold: float = 0.2,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
    ) -> None:
        """Initialize class attributes."""
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        self.confmap_head_config = confmap_head_config
        self.class_maps_head_config = class_maps_head_config

        self.class_names = self.class_maps_head_config.classes
        self.class_map_threshold = class_map_threshold

    def __getitem__(self, index) -> Dict:
        """Return dict with image, confmaps and class maps for given index."""
        sample = self.lf_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        video_idx = sample["video_idx"]
        frame_idx = sample["frame_idx"]

        if sample.get("is_negative", False):
            sample = self._load_negative_sample(sample)
            track_ids = torch.zeros(0, dtype=torch.int32)
        else:
            if self.cache_img is not None:
                instances = sample["instances"]
                if self.cache_img == "disk":
                    img = np.array(
                        Image.open(
                            f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                        )
                    )
                elif self.cache_img == "memory":
                    img = self.cache[(labels_idx, lf_idx)].copy()
            else:
                lf = self.labels_list[labels_idx][lf_idx]
                instances = lf.instances
                img = lf.image

            if img.ndim == 2:
                img = np.expand_dims(img, axis=2)

            # get dict
            sample = process_lf(
                instances_list=instances,
                img=img,
                frame_idx=frame_idx,
                video_idx=video_idx,
                max_instances=self.max_instances,
                user_instances_only=self.user_instances_only,
            )

            track_ids = torch.Tensor(
                [
                    (
                        self.class_names.index(instances[idx].track.name)
                        if instances[idx].track is not None
                        else -1
                    )
                    for idx in range(sample["num_instances"])
                ]
            ).to(torch.int32)

        sample["num_tracks"] = torch.tensor(len(self.class_names), dtype=torch.int32)

        sample = self._apply_common_preprocessing(sample)

        img_hw = sample["image"].shape[-2:]

        # Generate confidence maps
        confidence_maps = generate_multiconfmaps(
            sample["instances"],
            img_hw=img_hw,
            num_instances=sample["num_instances"],
            sigma=self.confmap_head_config.sigma,
            output_stride=self.confmap_head_config.output_stride,
            is_centroids=False,
        )

        # class maps
        class_maps = generate_class_maps(
            instances=sample["instances"],
            img_hw=img_hw,
            num_instances=sample["num_instances"],
            class_inds=track_ids,
            num_tracks=sample["num_tracks"],
            class_map_threshold=self.class_map_threshold,
            sigma=self.class_maps_head_config.sigma,
            output_stride=self.class_maps_head_config.output_stride,
            is_centroids=False,
        )

        sample["confidence_maps"] = confidence_maps
        sample["class_maps"] = class_maps
        sample["labels_idx"] = labels_idx
        if self.use_negative_frames:
            sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

        return sample

__getitem__(index)

Return dict with image, confmaps and class maps for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image, confmaps and class maps for given index."""
    sample = self.lf_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    video_idx = sample["video_idx"]
    frame_idx = sample["frame_idx"]

    if sample.get("is_negative", False):
        sample = self._load_negative_sample(sample)
        track_ids = torch.zeros(0, dtype=torch.int32)
    else:
        if self.cache_img is not None:
            instances = sample["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances = lf.instances
            img = lf.image

        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        # get dict
        sample = process_lf(
            instances_list=instances,
            img=img,
            frame_idx=frame_idx,
            video_idx=video_idx,
            max_instances=self.max_instances,
            user_instances_only=self.user_instances_only,
        )

        track_ids = torch.Tensor(
            [
                (
                    self.class_names.index(instances[idx].track.name)
                    if instances[idx].track is not None
                    else -1
                )
                for idx in range(sample["num_instances"])
            ]
        ).to(torch.int32)

    sample["num_tracks"] = torch.tensor(len(self.class_names), dtype=torch.int32)

    sample = self._apply_common_preprocessing(sample)

    img_hw = sample["image"].shape[-2:]

    # Generate confidence maps
    confidence_maps = generate_multiconfmaps(
        sample["instances"],
        img_hw=img_hw,
        num_instances=sample["num_instances"],
        sigma=self.confmap_head_config.sigma,
        output_stride=self.confmap_head_config.output_stride,
        is_centroids=False,
    )

    # class maps
    class_maps = generate_class_maps(
        instances=sample["instances"],
        img_hw=img_hw,
        num_instances=sample["num_instances"],
        class_inds=track_ids,
        num_tracks=sample["num_tracks"],
        class_map_threshold=self.class_map_threshold,
        sigma=self.class_maps_head_config.sigma,
        output_stride=self.class_maps_head_config.output_stride,
        is_centroids=False,
    )

    sample["confidence_maps"] = confidence_maps
    sample["class_maps"] = class_maps
    sample["labels_idx"] = labels_idx
    if self.use_negative_frames:
        sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

    return sample

__init__(labels, confmap_head_config, class_maps_head_config, max_stride, class_map_threshold=0.2, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    confmap_head_config: DictConfig,
    class_maps_head_config: DictConfig,
    max_stride: int,
    class_map_threshold: float = 0.2,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
) -> None:
    """Initialize class attributes."""
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=use_negative_frames,
    )
    self.confmap_head_config = confmap_head_config
    self.class_maps_head_config = class_maps_head_config

    self.class_names = self.class_maps_head_config.classes
    self.class_map_threshold = class_map_threshold

BottomUpSegmentationDataset

Bases: BaseDataset

Dataset class for bottom-up instance segmentation models.

Loads per-instance segmentation masks from LabeledFrame.masks and generates ground truth tensors for a center-offset instance segmentation pipeline (foreground mask, instance-center heatmap, and per-pixel offsets).

Masks are captured into the sample index at construction time (mirroring how keypoint instances are captured for caching), so __getitem__ never needs a live Labels handle — this keeps it correct under the memory/disk image caching paths (where self.labels_list is None).

Note

Augmentation: intensity aug is applied to the image; geometric aug (rotation/scale/translate/flip) co-transforms every per-instance mask with the SAME affine matrix as the image (nearest-neighbor, re-binarized) at the preprocessed resolution, before center/offset targets are derived. Erase/mixup stay image-only. Mask resizing to the preprocessed image size handles scaling but is not pad-aware; for v1 train with scale=1.0 and input dims divisible by max_stride (and prefer small rotation ranges, since a full-frame rotation can clip instances at the frame edge, as it does for bottom-up pose).

Attributes:

Name	Type	Description
seg_head_config		Configuration for the segmentation head.
center_head_config		Configuration for the instance center heatmap head.
offset_head_config		Configuration for the center offset head.

Methods:

Name	Description
__getitem__	Return dict with image and segmentation GT for given index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class BottomUpSegmentationDataset(BaseDataset):
    """Dataset class for bottom-up instance segmentation models.

    Loads per-instance segmentation masks from ``LabeledFrame.masks`` and
    generates ground truth tensors for a center-offset instance segmentation
    pipeline (foreground mask, instance-center heatmap, and per-pixel offsets).

    Masks are captured into the sample index at construction time (mirroring how
    keypoint instances are captured for caching), so ``__getitem__`` never needs
    a live ``Labels`` handle — this keeps it correct under the memory/disk image
    caching paths (where ``self.labels_list`` is ``None``).

    Note:
        Augmentation: intensity aug is applied to the image; geometric aug
        (rotation/scale/translate/flip) co-transforms every per-instance mask with the
        SAME affine matrix as the image (nearest-neighbor, re-binarized) at the
        preprocessed resolution, before center/offset targets are derived. Erase/mixup
        stay image-only. Mask resizing to the preprocessed image size handles scaling
        but is not pad-aware; for v1 train with ``scale=1.0`` and input dims divisible
        by ``max_stride`` (and prefer small rotation ranges, since a full-frame rotation
        can clip instances at the frame edge, as it does for bottom-up pose).

    Attributes:
        seg_head_config: Configuration for the segmentation head.
        center_head_config: Configuration for the instance center heatmap head.
        offset_head_config: Configuration for the center offset head.
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        seg_head_config: DictConfig,
        center_head_config: DictConfig,
        offset_head_config: DictConfig,
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
    ) -> None:
        """Initialize class attributes."""
        self.seg_head_config = seg_head_config
        self.center_head_config = center_head_config
        self.offset_head_config = offset_head_config
        # Segmentation never uses negative frames (degenerate num_nodes/instances).
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=False,
        )

    def _apply_common_preprocessing(self, sample: Dict) -> Dict:
        """Apply common preprocessing with geometric augmentation deferred.

        Geometric augmentation must co-transform the segmentation masks, but those
        masks are not present in ``sample`` here (they are loaded separately), so the
        base method would warp only the image. We disable geometric aug for the base
        call (intensity aug still applies) and re-apply it in ``__getitem__`` once the
        masks have been resized to the preprocessed resolution, co-transforming image
        and masks with the same matrix.

        Args:
            sample: Sample dict with at least ``image`` and ``instances`` keys.

        Returns:
            The sample dict with preprocessing applied in-place.
        """
        saved_geometric_aug = self.geometric_aug
        self.geometric_aug = None
        try:
            sample = super()._apply_common_preprocessing(sample)
        finally:
            self.geometric_aug = saved_geometric_aug
        return sample

    def _get_lf_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return samples for frames that have segmentation masks.

        Overrides the base class to index frames by their masks rather than
        keypoint instances. The decoded mask arrays are captured into each
        sample so ``__getitem__`` does not depend on a live ``Labels`` handle.
        """
        lf_idx_list = []
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                lf_masks = getattr(lf, "masks", None)
                if not lf_masks:
                    continue
                # Scale-aware decode: masks written by the segmentation inference
                # layer are encoded at output-stride (non-identity scale); decode
                # them up to the IMAGE-pixel grid so self-training / pseudo-label
                # ``.slp`` files yield correctly-scaled targets (a stride-res
                # ``m.data`` would silently mis-scale the GT, since __getitem__'s
                # resize branch only fires on a preprocessing size change). Scale-1
                # GT masks take the zero-copy fast path.
                from sleap_nn.inference.segmentation_convert import (
                    decode_mask_to_image_res,
                )

                mask_arrays = [decode_mask_to_image_res(m) for m in lf_masks]
                if len(mask_arrays) == 0:
                    continue
                video_idx = label.videos.index(lf.video)
                sample = {
                    "labels_idx": labels_idx,
                    "lf_idx": lf_idx,
                    "video_idx": video_idx,
                    "frame_idx": lf.frame_idx,
                    "is_negative": False,
                    "instances": None,
                    "masks": mask_arrays,
                }
                lf_idx_list.append(sample)

        return lf_idx_list

    def __getitem__(self, index) -> Dict:
        """Return dict with image and segmentation GT for given index."""
        sample = self.lf_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        video_idx = sample["video_idx"]
        frame_idx = sample["frame_idx"]

        # Load image
        if self.cache_img is not None:
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            img = lf.image

        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
        image = np.expand_dims(image, axis=0)  # (1, C, H, W)
        image = torch.from_numpy(image.copy())

        # Dummy instances tensor (needed for preprocessing compatibility)
        instances = torch.zeros((1, 1, 1, 2), dtype=torch.float32)

        sample_dict = {
            "image": image,
            "instances": instances,
            "video_idx": torch.tensor(video_idx, dtype=torch.int32),
            "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
            "orig_size": torch.Tensor([image.shape[-2], image.shape[-1]]).unsqueeze(0),
            "num_instances": 0,
        }

        # Masks captured at index-build time (decoded bool arrays at orig res).
        # Place them on the full-frame image grid as one float (1, K, H, W) tensor
        # so they ride the IDENTICAL size-match / scale / stride-pad chain as the
        # image inside ``_apply_common_preprocessing`` (kept as a float tensor
        # end-to-end and binarized ONCE just before target generation). This is
        # what keeps the whole-frame target registered to the padded image and
        # makes ragged/offset-carrying decoded masks well-defined (see
        # ``_masks_to_frame_canvas``).
        mask_arrays = [np.asarray(m, dtype=bool) for m in sample["masks"]]
        orig_img_hw = (image.shape[-2], image.shape[-1])
        sample_dict["masks"] = _masks_to_frame_canvas(mask_arrays, orig_img_hw)

        # Apply common preprocessing (RGB/grayscale, size matching, scaling, padding).
        # Co-transforms ``sample_dict["masks"]`` with the same geometry as the image.
        sample_dict = self._apply_common_preprocessing(sample_dict)

        img_hw = sample_dict["image"].shape[-2:]

        # Geometric augmentation: co-transform the per-instance masks with the SAME
        # flip/affine matrix as the image (nearest-neighbor). Applied here (post
        # size-match / resize / pad) so image and masks share a resolution, and
        # BEFORE centroid/heatmap/offset generation so those targets are derived
        # from the augmented masks. Bottom-up has no keypoints, so a dummy instances
        # tensor rides along; erase/mixup stay image-only.
        if (
            self.apply_aug
            and self.geometric_aug is not None
            and sample_dict["masks"].shape[1] > 0
        ):
            (
                sample_dict["image"],
                _,
                sample_dict["masks"],
            ) = apply_geometric_augmentation(
                sample_dict["image"],
                torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                masks=sample_dict["masks"],
                **self.geometric_aug,
            )

        # Single re-binarization to bool arrays right before target generation.
        masks_t = sample_dict.pop("masks")
        mask_arrays = [masks_t[0, k].numpy() > 0.5 for k in range(masks_t.shape[1])]

        # Pre-compute mask centroids once for both center heatmap and offset heads
        centers = _compute_mask_centroids(mask_arrays) if len(mask_arrays) > 0 else []

        # Generate GT tensors
        foreground_mask = generate_foreground_mask(
            mask_arrays,
            img_hw=img_hw,
            output_stride=self.seg_head_config.output_stride,
            maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
        )

        center_heatmap = generate_center_heatmap(
            mask_arrays,
            img_hw=img_hw,
            output_stride=self.center_head_config.output_stride,
            sigma=self.center_head_config.sigma,
            centers=centers,
        )

        center_offsets, foreground_weight = generate_center_offsets(
            mask_arrays,
            img_hw=img_hw,
            output_stride=self.offset_head_config.output_stride,
            centers=centers,
        )

        sample_dict["foreground_mask"] = foreground_mask
        sample_dict["center_heatmap"] = center_heatmap
        sample_dict["center_offsets"] = center_offsets
        sample_dict["foreground_weight"] = foreground_weight
        sample_dict["labels_idx"] = labels_idx

        return sample_dict

__getitem__(index)

Return dict with image and segmentation GT for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image and segmentation GT for given index."""
    sample = self.lf_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    video_idx = sample["video_idx"]
    frame_idx = sample["frame_idx"]

    # Load image
    if self.cache_img is not None:
        if self.cache_img == "disk":
            img = np.array(
                Image.open(
                    f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                )
            )
        elif self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)].copy()
    else:
        lf = self.labels_list[labels_idx][lf_idx]
        img = lf.image

    if img.ndim == 2:
        img = np.expand_dims(img, axis=2)

    image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    image = np.expand_dims(image, axis=0)  # (1, C, H, W)
    image = torch.from_numpy(image.copy())

    # Dummy instances tensor (needed for preprocessing compatibility)
    instances = torch.zeros((1, 1, 1, 2), dtype=torch.float32)

    sample_dict = {
        "image": image,
        "instances": instances,
        "video_idx": torch.tensor(video_idx, dtype=torch.int32),
        "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
        "orig_size": torch.Tensor([image.shape[-2], image.shape[-1]]).unsqueeze(0),
        "num_instances": 0,
    }

    # Masks captured at index-build time (decoded bool arrays at orig res).
    # Place them on the full-frame image grid as one float (1, K, H, W) tensor
    # so they ride the IDENTICAL size-match / scale / stride-pad chain as the
    # image inside ``_apply_common_preprocessing`` (kept as a float tensor
    # end-to-end and binarized ONCE just before target generation). This is
    # what keeps the whole-frame target registered to the padded image and
    # makes ragged/offset-carrying decoded masks well-defined (see
    # ``_masks_to_frame_canvas``).
    mask_arrays = [np.asarray(m, dtype=bool) for m in sample["masks"]]
    orig_img_hw = (image.shape[-2], image.shape[-1])
    sample_dict["masks"] = _masks_to_frame_canvas(mask_arrays, orig_img_hw)

    # Apply common preprocessing (RGB/grayscale, size matching, scaling, padding).
    # Co-transforms ``sample_dict["masks"]`` with the same geometry as the image.
    sample_dict = self._apply_common_preprocessing(sample_dict)

    img_hw = sample_dict["image"].shape[-2:]

    # Geometric augmentation: co-transform the per-instance masks with the SAME
    # flip/affine matrix as the image (nearest-neighbor). Applied here (post
    # size-match / resize / pad) so image and masks share a resolution, and
    # BEFORE centroid/heatmap/offset generation so those targets are derived
    # from the augmented masks. Bottom-up has no keypoints, so a dummy instances
    # tensor rides along; erase/mixup stay image-only.
    if (
        self.apply_aug
        and self.geometric_aug is not None
        and sample_dict["masks"].shape[1] > 0
    ):
        (
            sample_dict["image"],
            _,
            sample_dict["masks"],
        ) = apply_geometric_augmentation(
            sample_dict["image"],
            torch.zeros((1, 1, 1, 2), dtype=torch.float32),
            masks=sample_dict["masks"],
            **self.geometric_aug,
        )

    # Single re-binarization to bool arrays right before target generation.
    masks_t = sample_dict.pop("masks")
    mask_arrays = [masks_t[0, k].numpy() > 0.5 for k in range(masks_t.shape[1])]

    # Pre-compute mask centroids once for both center heatmap and offset heads
    centers = _compute_mask_centroids(mask_arrays) if len(mask_arrays) > 0 else []

    # Generate GT tensors
    foreground_mask = generate_foreground_mask(
        mask_arrays,
        img_hw=img_hw,
        output_stride=self.seg_head_config.output_stride,
        maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
    )

    center_heatmap = generate_center_heatmap(
        mask_arrays,
        img_hw=img_hw,
        output_stride=self.center_head_config.output_stride,
        sigma=self.center_head_config.sigma,
        centers=centers,
    )

    center_offsets, foreground_weight = generate_center_offsets(
        mask_arrays,
        img_hw=img_hw,
        output_stride=self.offset_head_config.output_stride,
        centers=centers,
    )

    sample_dict["foreground_mask"] = foreground_mask
    sample_dict["center_heatmap"] = center_heatmap
    sample_dict["center_offsets"] = center_offsets
    sample_dict["foreground_weight"] = foreground_weight
    sample_dict["labels_idx"] = labels_idx

    return sample_dict

__init__(labels, seg_head_config, center_head_config, offset_head_config, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    seg_head_config: DictConfig,
    center_head_config: DictConfig,
    offset_head_config: DictConfig,
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
) -> None:
    """Initialize class attributes."""
    self.seg_head_config = seg_head_config
    self.center_head_config = center_head_config
    self.offset_head_config = offset_head_config
    # Segmentation never uses negative frames (degenerate num_nodes/instances).
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=False,
    )

BottomUpSegmentationTiledDataset

Bases: BaseDataset

Bottom-up segmentation dataset that emits fixed-size tiles (Phase C).

Structural analogue of :class:SingleInstanceTiledDataset for the center-offset instance-segmentation pipeline. Each frame is decomposed into overlapping square tiles (foreground-aware random draws for training, a deterministic grid for validation). A frame is decoded / channel-coerced / scaled once (cached in a per-worker LRU together with its decoded per-instance masks) and reused across all of its tiles. Each tile is then cut out with a constant-zero pad; when geometric augmentation is enabled the tile is taken via a sqrt(2) halo so a rotation has valid context, co-transforming every per-instance mask with the SAME affine matrix as the image (nearest-neighbor, re-binarized). An ownership filter keeps only the masks whose (tile-local) centroid lands inside the tile so off-tile instances do not seed spurious center-offset regression, and the segmentation GT tensors (foreground mask, instance-center heatmap, per-pixel offsets) are generated on tile-local coordinates.

Emits one sample per (frame, tile-slot); __len__ is the total number of tile slots. Returned samples match the BottomUpSegmentationDataset key contract (plus an int32 tile_origin of shape (2,)), so the default collate and the BottomUpSegmentationLightningModule apply with no changes.

Attributes:

Name	Type	Description
seg_head_config		Configuration for the segmentation (foreground) head.
center_head_config		Configuration for the instance center heatmap head.
offset_head_config		Configuration for the center offset head.

Methods:

Name	Description
__getitem__	Return dict with image + segmentation GT for one tile of one frame.
__init__	Initialize class attributes.
__len__	Return the number of tile samples (frames x tiles-per-frame).

Source code in sleap_nn/data/custom_datasets.py

class BottomUpSegmentationTiledDataset(BaseDataset):
    """Bottom-up segmentation dataset that emits fixed-size tiles (Phase C).

    Structural analogue of :class:`SingleInstanceTiledDataset` for the center-offset
    instance-segmentation pipeline. Each frame is decomposed into overlapping square
    tiles (foreground-aware random draws for training, a deterministic grid for
    validation). A frame is decoded / channel-coerced / scaled once (cached in a
    per-worker LRU together with its decoded per-instance masks) and reused across all
    of its tiles. Each tile is then cut out with a constant-zero pad; when geometric
    augmentation is enabled the tile is taken via a ``sqrt(2)`` halo so a rotation has
    valid context, co-transforming every per-instance mask with the SAME affine matrix
    as the image (nearest-neighbor, re-binarized). An ownership filter keeps only the
    masks whose (tile-local) centroid lands inside the tile so off-tile instances do
    not seed spurious center-offset regression, and the segmentation GT tensors
    (foreground mask, instance-center heatmap, per-pixel offsets) are generated on
    tile-local coordinates.

    Emits one sample per ``(frame, tile-slot)``; ``__len__`` is the total number of
    tile slots. Returned samples match the ``BottomUpSegmentationDataset`` key contract
    (plus an ``int32`` ``tile_origin`` of shape ``(2,)``), so the default collate and
    the ``BottomUpSegmentationLightningModule`` apply with no changes.

    Attributes:
        seg_head_config: Configuration for the segmentation (foreground) head.
        center_head_config: Configuration for the instance center heatmap head.
        offset_head_config: Configuration for the center offset head.
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        seg_head_config: DictConfig,
        center_head_config: DictConfig,
        offset_head_config: DictConfig,
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
        tiling: Optional[Union[DictConfig, Any]] = None,
        base_seed: int = 0,
    ) -> None:
        """Initialize class attributes."""
        self.seg_head_config = seg_head_config
        self.center_head_config = center_head_config
        self.offset_head_config = offset_head_config
        # Segmentation never uses negative frames (degenerate num_nodes/instances).
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=False,
            tiling=tiling,
            output_stride=seg_head_config.output_stride,
            base_seed=base_seed,
        )

        # Per-(frame, tile-slot) descriptors + contiguous per-frame index blocks.
        self.tile_idx_list = self._get_tile_idx_list(labels)
        self.frame_blocks = self._build_frame_blocks(self.tile_idx_list)

    def _get_lf_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return per-frame samples for frames that have segmentation masks.

        Mirrors ``BottomUpSegmentationDataset._get_lf_idx_list``: indexes frames by
        their masks (not keypoint instances) and captures the decoded mask arrays so
        ``__getitem__`` never needs a live ``Labels`` handle. This is what the base
        ``__init__`` stores as ``self.lf_idx_list`` (used for image caching);
        ``_get_tile_idx_list`` explodes it into per-tile descriptors.
        """
        from sleap_nn.inference.segmentation_convert import decode_mask_to_image_res

        lf_idx_list = []
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                lf_masks = getattr(lf, "masks", None)
                if not lf_masks:
                    continue
                # Scale-aware decode up to the IMAGE-pixel grid (scale-1 GT masks
                # take the zero-copy fast path); see BottomUpSegmentationDataset.
                mask_arrays = [decode_mask_to_image_res(m) for m in lf_masks]
                if len(mask_arrays) == 0:
                    continue
                video_idx = label.videos.index(lf.video)
                lf_idx_list.append(
                    {
                        "labels_idx": labels_idx,
                        "lf_idx": lf_idx,
                        "video_idx": video_idx,
                        "frame_idx": lf.frame_idx,
                        "is_negative": False,
                        "instances": None,
                        "masks": mask_arrays,
                    }
                )
        return lf_idx_list

    def _get_tile_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return per-(frame, tile-slot) descriptors for the tiled seg dataset.

        Mirrors :meth:`BaseDataset._get_tile_idx_list` but keys off the mask-indexed
        per-frame list (``self.lf_idx_list``, built by :meth:`_get_lf_idx_list`) so the
        decoded ``mask_arrays`` ride along on every descriptor. Grid (val) pins one
        descriptor per :func:`generate_tile_grid` origin (sized-frame ``H, W`` taken
        from the decoded mask shape); foreground (train) emits ``samples_per_frame``
        slots with ``tile_origin=None`` (drawn at runtime). A frame's descriptors form
        a contiguous run (the block the sampler groups on).
        """
        tile_idx_list: List[Dict] = []
        for f in self.lf_idx_list:
            mask_arrays = f["masks"]
            # Masks are at image resolution; the sized (post-scale) frame H, W match
            # `apply_resizer`'s int(dim * scale) truncation used by `_frame_sized_hw`.
            mh, mw = mask_arrays[0].shape[:2]
            if self.scale != 1.0:
                sized_hw = (int(mh * self.scale), int(mw * self.scale))
            else:
                sized_hw = (int(mh), int(mw))

            if self.tile_sampling == "grid":
                origins = generate_tile_grid(
                    sized_hw,
                    tile_size=self.tile_size,
                    overlap=self.overlap,
                    output_stride=self.output_stride,
                    max_stride=self.max_stride,
                    min_overlap_fraction=self.min_overlap_fraction,
                )
            else:
                origins = [None] * self.samples_per_frame

            for sample_k, origin in enumerate(origins):
                tile_idx_list.append(
                    {
                        "labels_idx": f["labels_idx"],
                        "lf_idx": f["lf_idx"],
                        "video_idx": f["video_idx"],
                        "frame_idx": f["frame_idx"],
                        "masks": mask_arrays,
                        "sample_k": sample_k,
                        "tile_origin": origin,
                        "is_grid": self.tile_sampling == "grid",
                        "is_negative": False,
                    }
                )
        return tile_idx_list

    def __len__(self) -> int:
        """Return the number of tile samples (frames x tiles-per-frame)."""
        return len(self.tile_idx_list)

    def _read_image(self, d: Dict) -> np.ndarray:
        """Read a frame's raw HWC image (cache/disk/labels), restoring 2D -> 3D."""
        labels_idx = d["labels_idx"]
        lf_idx = d["lf_idx"]
        if self.cache_img is not None:
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            img = lf.image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)
        return img

    def _load_sized_frame(
        self, d: Dict
    ) -> Tuple[torch.Tensor, List[np.ndarray], tuple]:
        """Decode a frame once: channel-coerce + scale the image, size-match masks.

        Returns ``(image, mask_arrays, orig_hw)`` where ``image`` is a sized
        ``(1, C, H, W)`` tensor, ``mask_arrays`` are bool arrays at the sized image
        resolution, and ``orig_hw`` is the raw full-frame ``(H, W)`` before scaling.
        """
        img = self._read_image(d)  # HWC
        orig_hw = (int(img.shape[0]), int(img.shape[1]))

        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
        image = np.expand_dims(image, axis=0)  # (1, C, H, W)
        image = torch.from_numpy(image.copy())

        if self.ensure_rgb:
            image = convert_to_rgb(image)
        elif self.ensure_grayscale:
            image = convert_to_grayscale(image)

        dummy = torch.zeros((1, 1, 1, 2), dtype=torch.float32)
        image, _ = apply_resizer(image, dummy, scale=self.scale)

        mask_arrays = [np.asarray(m, dtype=bool) for m in d["masks"]]
        sized_hw = (image.shape[-2], image.shape[-1])
        if (sized_hw != orig_hw) and len(mask_arrays) > 0:
            # Resize masks to the sized image resolution with antialiased bilinear,
            # matching the image's ``apply_resizer`` / ``tvf.resize`` above
            # (standardized across every seg mask-geometry resize).
            target_h, target_w = sized_hw
            resized = []
            for m in mask_arrays:
                m_tensor = (
                    torch.from_numpy(m.astype(np.float32)).unsqueeze(0).unsqueeze(0)
                )
                m_resized = F.interpolate(
                    m_tensor,
                    size=(target_h, target_w),
                    mode="bilinear",
                    align_corners=False,
                    antialias=True,
                )
                resized.append(m_resized.squeeze().numpy() > 0.5)
            mask_arrays = resized

        return image, mask_arrays, orig_hw

    def _slice_halo(
        self,
        image: torch.Tensor,
        mask_arrays: List[np.ndarray],
        hy0: int,
        hx0: int,
        side: int,
    ) -> Tuple[torch.Tensor, List[np.ndarray]]:
        """Slice a ``side x side`` window (top-left ``hy0, hx0``) with constant-zero pad.

        Slices both the image and every per-instance mask at the SAME offsets so they
        stay pixel-aligned. Out-of-bounds regions are zero.
        """
        _, C, H, W = image.shape
        ys, xs = max(0, hy0), max(0, hx0)
        ye, xe = min(H, hy0 + side), min(W, hx0 + side)
        win_img = image.new_zeros((1, C, side, side))
        win_masks = [np.zeros((side, side), dtype=bool) for _ in mask_arrays]
        if ye > ys and xe > xs:
            win_img[:, :, ys - hy0 : ye - hy0, xs - hx0 : xe - hx0] = image[
                :, :, ys:ye, xs:xe
            ]
            for k, m in enumerate(mask_arrays):
                win_masks[k][ys - hy0 : ye - hy0, xs - hx0 : xe - hx0] = m[ys:ye, xs:xe]
        return win_img, win_masks

    def __getitem__(self, index) -> Dict:
        """Return dict with image + segmentation GT for one tile of one frame."""
        d = self.tile_idx_list[index]
        labels_idx = d["labels_idx"]
        video_idx = d["video_idx"]
        frame_idx = d["frame_idx"]
        epoch = int(self._epoch)
        ts = self.tile_size

        # 1. Decode the full frame once per (labels_idx, lf_idx), via per-worker LRU.
        cached = self._frame_lru().get((labels_idx, d["lf_idx"]))
        if cached is None:
            cached = self._load_sized_frame(d)
            self._frame_lru().put((labels_idx, d["lf_idx"]), cached)
        image, mask_arrays, orig_hw = cached
        sized_hw = (image.shape[-2], image.shape[-1])

        # 2. Resolve the tile origin: pinned for grid/val, drawn for train.
        if d["is_grid"]:
            tile_origin = tuple(int(v) for v in d["tile_origin"])
            aug_seed = None
        else:
            rng = np.random.default_rng(
                tile_sample_seed(
                    self.base_seed, epoch, video_idx, frame_idx, d["sample_k"]
                )
            )
            cents = _compute_mask_centroids(mask_arrays)  # list of (x, y)
            centers = (
                torch.tensor(cents, dtype=torch.float32).reshape(-1, 2)
                if len(cents) > 0
                else torch.zeros((0, 2), dtype=torch.float32)
            )
            tile_origin = draw_tile_origin(
                centers,
                sized_hw,
                ts,
                d["sample_k"],
                self.samples_per_frame,
                self.tile_fg_fraction,
                self.center_jitter,
                rng,
            )
            aug_seed = tile_sample_seed(
                self.base_seed, epoch, video_idx, frame_idx, d["sample_k"], salt=1
            )

        y0, x0 = tile_origin
        apply_geo = self.apply_aug and self.geometric_aug is not None

        # 3. Cut the tile out of the frame (image + co-transformed masks). Under
        #    geometric aug, take a sqrt(2) halo centered on the tile center so the
        #    rotation has valid context, augment image + masks with the SAME matrix
        #    (nearest-neighbor, re-binarized), then trim the center tile back out.
        if apply_geo:
            halo = int(math.ceil(ts * math.sqrt(2)))
            hy0 = y0 - (halo - ts) // 2
            hx0 = x0 - (halo - ts) // 2
            halo_img, halo_masks = self._slice_halo(image, mask_arrays, hy0, hx0, halo)

            if len(halo_masks) > 0:
                halo_masks_t = torch.from_numpy(
                    np.stack([hm.astype(np.float32) for hm in halo_masks])
                ).unsqueeze(
                    0
                )  # (1, K, halo, halo)
                # The skia geometric backend samples its transform from the GLOBAL
                # numpy RNG (and torch); seed both so the halo path is reproducible.
                np.random.seed(aug_seed & 0xFFFFFFFF)
                torch.manual_seed(aug_seed)
                halo_img, _, halo_masks_t = apply_geometric_augmentation(
                    halo_img,
                    torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                    masks=halo_masks_t,
                    **dict(self.geometric_aug),
                )
                halo_masks = [
                    halo_masks_t[0, k].numpy() > 0.5
                    for k in range(halo_masks_t.shape[1])
                ]

            # Trim the augmented halo back to `ts`, centered on the halo center:
            # crop_and_resize for the image (codebase convention), an equivalent
            # integer center-slice for the (axis-aligned) mask arrays.
            c = halo / 2.0
            bbox = make_centered_bboxes(
                torch.tensor([[c, c]], dtype=torch.float32), ts, ts
            )
            tile_image = crop_and_resize(halo_img, boxes=bbox, size=(ts, ts))
            off = (halo - ts) // 2
            tile_masks = [hm[off : off + ts, off : off + ts] for hm in halo_masks]
        else:
            # Fast path (no aug): direct slice + constant-zero pad, byte-identical.
            tile_image, tile_masks = self._slice_halo(image, mask_arrays, y0, x0, ts)

        # 4. Intensity aug (image only) + pad to stride (no-op when ts % max_stride==0).
        if self.apply_aug and self.intensity_aug is not None:
            tile_image, _ = apply_intensity_augmentation(
                tile_image,
                torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                **self.intensity_aug,
            )
        tile_image = apply_pad_to_stride(tile_image, max_stride=self.max_stride)

        # 5. Ownership filter: keep only masks whose tile-local centroid lands inside
        #    [0, ts) x [0, ts) AND that have >= a few foreground px. Instances owned by
        #    a neighbor tile (centroid off-tile) are dropped so they do not seed
        #    off-tile center-offset regression.
        tile_centers = _compute_mask_centroids(tile_masks) if tile_masks else []
        owned_masks: List[np.ndarray] = []
        owned_centers: List[Tuple[float, float]] = []
        for m, (cx, cy) in zip(tile_masks, tile_centers):
            if int(m.sum()) < _MIN_OWNED_FG_PX:
                continue
            if 0.0 <= cx < ts and 0.0 <= cy < ts:
                owned_masks.append(m)
                owned_centers.append((cx, cy))

        # 6. Generate GT tensors on tile-local coordinates from the OWNED masks.
        img_hw = tile_image.shape[-2:]
        foreground_mask = generate_foreground_mask(
            owned_masks,
            img_hw=img_hw,
            output_stride=self.seg_head_config.output_stride,
            maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
        )
        center_heatmap = generate_center_heatmap(
            owned_masks,
            img_hw=img_hw,
            output_stride=self.center_head_config.output_stride,
            sigma=self.center_head_config.sigma,
            centers=owned_centers,
        )
        center_offsets, foreground_weight = generate_center_offsets(
            owned_masks,
            img_hw=img_hw,
            output_stride=self.offset_head_config.output_stride,
            centers=owned_centers,
        )

        return {
            "image": tile_image,
            "instances": torch.zeros((1, 1, 1, 2), dtype=torch.float32),
            "video_idx": torch.tensor(video_idx, dtype=torch.int32),
            "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
            "orig_size": torch.Tensor([orig_hw[0], orig_hw[1]]).unsqueeze(0),
            "num_instances": len(owned_masks),
            # Tiles are extracted in the model's input space (sizematcher bypassed),
            # so the effective scale is 1.0 (matches SingleInstanceTiledDataset).
            "eff_scale": torch.tensor(1.0, dtype=torch.float32),
            "foreground_mask": foreground_mask,
            "center_heatmap": center_heatmap,
            "center_offsets": center_offsets,
            "foreground_weight": foreground_weight,
            "labels_idx": labels_idx,
            "tile_origin": torch.tensor(tile_origin, dtype=torch.int32),
        }

__getitem__(index)

Return dict with image + segmentation GT for one tile of one frame.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image + segmentation GT for one tile of one frame."""
    d = self.tile_idx_list[index]
    labels_idx = d["labels_idx"]
    video_idx = d["video_idx"]
    frame_idx = d["frame_idx"]
    epoch = int(self._epoch)
    ts = self.tile_size

    # 1. Decode the full frame once per (labels_idx, lf_idx), via per-worker LRU.
    cached = self._frame_lru().get((labels_idx, d["lf_idx"]))
    if cached is None:
        cached = self._load_sized_frame(d)
        self._frame_lru().put((labels_idx, d["lf_idx"]), cached)
    image, mask_arrays, orig_hw = cached
    sized_hw = (image.shape[-2], image.shape[-1])

    # 2. Resolve the tile origin: pinned for grid/val, drawn for train.
    if d["is_grid"]:
        tile_origin = tuple(int(v) for v in d["tile_origin"])
        aug_seed = None
    else:
        rng = np.random.default_rng(
            tile_sample_seed(
                self.base_seed, epoch, video_idx, frame_idx, d["sample_k"]
            )
        )
        cents = _compute_mask_centroids(mask_arrays)  # list of (x, y)
        centers = (
            torch.tensor(cents, dtype=torch.float32).reshape(-1, 2)
            if len(cents) > 0
            else torch.zeros((0, 2), dtype=torch.float32)
        )
        tile_origin = draw_tile_origin(
            centers,
            sized_hw,
            ts,
            d["sample_k"],
            self.samples_per_frame,
            self.tile_fg_fraction,
            self.center_jitter,
            rng,
        )
        aug_seed = tile_sample_seed(
            self.base_seed, epoch, video_idx, frame_idx, d["sample_k"], salt=1
        )

    y0, x0 = tile_origin
    apply_geo = self.apply_aug and self.geometric_aug is not None

    # 3. Cut the tile out of the frame (image + co-transformed masks). Under
    #    geometric aug, take a sqrt(2) halo centered on the tile center so the
    #    rotation has valid context, augment image + masks with the SAME matrix
    #    (nearest-neighbor, re-binarized), then trim the center tile back out.
    if apply_geo:
        halo = int(math.ceil(ts * math.sqrt(2)))
        hy0 = y0 - (halo - ts) // 2
        hx0 = x0 - (halo - ts) // 2
        halo_img, halo_masks = self._slice_halo(image, mask_arrays, hy0, hx0, halo)

        if len(halo_masks) > 0:
            halo_masks_t = torch.from_numpy(
                np.stack([hm.astype(np.float32) for hm in halo_masks])
            ).unsqueeze(
                0
            )  # (1, K, halo, halo)
            # The skia geometric backend samples its transform from the GLOBAL
            # numpy RNG (and torch); seed both so the halo path is reproducible.
            np.random.seed(aug_seed & 0xFFFFFFFF)
            torch.manual_seed(aug_seed)
            halo_img, _, halo_masks_t = apply_geometric_augmentation(
                halo_img,
                torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                masks=halo_masks_t,
                **dict(self.geometric_aug),
            )
            halo_masks = [
                halo_masks_t[0, k].numpy() > 0.5
                for k in range(halo_masks_t.shape[1])
            ]

        # Trim the augmented halo back to `ts`, centered on the halo center:
        # crop_and_resize for the image (codebase convention), an equivalent
        # integer center-slice for the (axis-aligned) mask arrays.
        c = halo / 2.0
        bbox = make_centered_bboxes(
            torch.tensor([[c, c]], dtype=torch.float32), ts, ts
        )
        tile_image = crop_and_resize(halo_img, boxes=bbox, size=(ts, ts))
        off = (halo - ts) // 2
        tile_masks = [hm[off : off + ts, off : off + ts] for hm in halo_masks]
    else:
        # Fast path (no aug): direct slice + constant-zero pad, byte-identical.
        tile_image, tile_masks = self._slice_halo(image, mask_arrays, y0, x0, ts)

    # 4. Intensity aug (image only) + pad to stride (no-op when ts % max_stride==0).
    if self.apply_aug and self.intensity_aug is not None:
        tile_image, _ = apply_intensity_augmentation(
            tile_image,
            torch.zeros((1, 1, 1, 2), dtype=torch.float32),
            **self.intensity_aug,
        )
    tile_image = apply_pad_to_stride(tile_image, max_stride=self.max_stride)

    # 5. Ownership filter: keep only masks whose tile-local centroid lands inside
    #    [0, ts) x [0, ts) AND that have >= a few foreground px. Instances owned by
    #    a neighbor tile (centroid off-tile) are dropped so they do not seed
    #    off-tile center-offset regression.
    tile_centers = _compute_mask_centroids(tile_masks) if tile_masks else []
    owned_masks: List[np.ndarray] = []
    owned_centers: List[Tuple[float, float]] = []
    for m, (cx, cy) in zip(tile_masks, tile_centers):
        if int(m.sum()) < _MIN_OWNED_FG_PX:
            continue
        if 0.0 <= cx < ts and 0.0 <= cy < ts:
            owned_masks.append(m)
            owned_centers.append((cx, cy))

    # 6. Generate GT tensors on tile-local coordinates from the OWNED masks.
    img_hw = tile_image.shape[-2:]
    foreground_mask = generate_foreground_mask(
        owned_masks,
        img_hw=img_hw,
        output_stride=self.seg_head_config.output_stride,
        maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
    )
    center_heatmap = generate_center_heatmap(
        owned_masks,
        img_hw=img_hw,
        output_stride=self.center_head_config.output_stride,
        sigma=self.center_head_config.sigma,
        centers=owned_centers,
    )
    center_offsets, foreground_weight = generate_center_offsets(
        owned_masks,
        img_hw=img_hw,
        output_stride=self.offset_head_config.output_stride,
        centers=owned_centers,
    )

    return {
        "image": tile_image,
        "instances": torch.zeros((1, 1, 1, 2), dtype=torch.float32),
        "video_idx": torch.tensor(video_idx, dtype=torch.int32),
        "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
        "orig_size": torch.Tensor([orig_hw[0], orig_hw[1]]).unsqueeze(0),
        "num_instances": len(owned_masks),
        # Tiles are extracted in the model's input space (sizematcher bypassed),
        # so the effective scale is 1.0 (matches SingleInstanceTiledDataset).
        "eff_scale": torch.tensor(1.0, dtype=torch.float32),
        "foreground_mask": foreground_mask,
        "center_heatmap": center_heatmap,
        "center_offsets": center_offsets,
        "foreground_weight": foreground_weight,
        "labels_idx": labels_idx,
        "tile_origin": torch.tensor(tile_origin, dtype=torch.int32),
    }

__init__(labels, seg_head_config, center_head_config, offset_head_config, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False, tiling=None, base_seed=0)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    seg_head_config: DictConfig,
    center_head_config: DictConfig,
    offset_head_config: DictConfig,
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
    tiling: Optional[Union[DictConfig, Any]] = None,
    base_seed: int = 0,
) -> None:
    """Initialize class attributes."""
    self.seg_head_config = seg_head_config
    self.center_head_config = center_head_config
    self.offset_head_config = offset_head_config
    # Segmentation never uses negative frames (degenerate num_nodes/instances).
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=False,
        tiling=tiling,
        output_stride=seg_head_config.output_stride,
        base_seed=base_seed,
    )

    # Per-(frame, tile-slot) descriptors + contiguous per-frame index blocks.
    self.tile_idx_list = self._get_tile_idx_list(labels)
    self.frame_blocks = self._build_frame_blocks(self.tile_idx_list)

__len__()

Return the number of tile samples (frames x tiles-per-frame).

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Return the number of tile samples (frames x tiles-per-frame)."""
    return len(self.tile_idx_list)

CenteredInstanceDataset

Bases: BaseDataset

Dataset class for instance-centered confidence map models.

Attributes:

Name	Type	Description
max_stride		Scalar integer specifying the maximum stride that the image must be divisible by.
anchor_ind		Index of the node to use as the anchor point, based on its index in the ordered list of skeleton nodes.
user_instances_only		True if only user labeled instances should be used for training. If False, both user labeled and predicted instances would be used.
ensure_rgb		(bool) True if the input image should have 3 channels (RGB image). If input has only one
is	replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default	False.
ensure_grayscale		(bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
image.	If the source image has only one channel and this is set to False, then we retain the single channel input. Default	False.
intensity_aug		Intensity augmentation configuration. Can be: - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness'] - List of strings: Multiple intensity augmentations from the allowed values - Dictionary: Custom intensity configuration - None: No intensity augmentation applied
geometric_aug		Geometric augmentation configuration. Can be: - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup'] - List of strings: Multiple geometric augmentations from the allowed values - Dictionary: Custom geometric configuration - None: No geometric augmentation applied
scale		Factor to resize the image dimensions by, specified as a float. Default: 1.0.
apply_aug		True if augmentations should be applied to the data pipeline, else False. Default: False.
max_hw		Maximum height and width of images across the labels file. If max_height and max_width in the config is None, then max_hw is used (computed with sleap_nn.data.providers.get_max_height_width). Else the values in the config are used.
cache_img		String to indicate which caching to use: memory or disk. If None, the images aren't cached and loaded from the .slp file on each access.
cache_img_path		Path to save the .jpg files. If None, current working dir is used.
use_existing_imgs		Use existing imgs/ chunks in the cache_img_path.
crop_size		Crop size of each instance for centered-instance model. If scale is provided, then the cropped image will be resized according to scale.
rank		Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.
confmap_head_config		DictConfig object with all the keys in the head_config section. (required keys: sigma, output_stride, part_names and anchor_part depending on the model type ).
labels_list		List of sio.Labels objects. Used to store the labels in the cache. (only used if cache_img is None)
bbox_crop		How box-scaled crops are cut (preprocessing.crop_mode="bbox"), or None for fixed-size crops centered on the instance centroid. With a spec, each instance's box (derived like a detector's ground truth) is resized to crop_size x crop_size, exactly as the top-down inference layer crops a detected box.

Methods:

Name	Description
__getitem__	Return dict with cropped image and confmaps of instance for given index.
__init__	Initialize class attributes.
__len__	Return number of instances in the labels object.

Source code in sleap_nn/data/custom_datasets.py

class CenteredInstanceDataset(BaseDataset):
    """Dataset class for instance-centered confidence map models.

    Attributes:
        max_stride: Scalar integer specifying the maximum stride that the image must be
            divisible by.
        anchor_ind: Index of the node to use as the anchor point, based on its index in the
            ordered list of skeleton nodes.
        user_instances_only: `True` if only user labeled instances should be used for training. If `False`,
            both user labeled and predicted instances would be used.
        ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one
        channel when this is set to `True`, then the images from single-channel
        is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default: `False`.
        ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
        is set to True, then we convert the image to grayscale (single-channel)
        image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
        intensity_aug: Intensity augmentation configuration. Can be:
            - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness']
            - List of strings: Multiple intensity augmentations from the allowed values
            - Dictionary: Custom intensity configuration
            - None: No intensity augmentation applied
        geometric_aug: Geometric augmentation configuration. Can be:
            - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup']
            - List of strings: Multiple geometric augmentations from the allowed values
            - Dictionary: Custom geometric configuration
            - None: No geometric augmentation applied
        scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
        apply_aug: `True` if augmentations should be applied to the data pipeline,
            else `False`. Default: `False`.
        max_hw: Maximum height and width of images across the labels file. If `max_height` and
           `max_width` in the config is None, then `max_hw` is used (computed with
            `sleap_nn.data.providers.get_max_height_width`). Else the values in the config
            are used.
        cache_img: String to indicate which caching to use: `memory` or `disk`. If `None`,
            the images aren't cached and loaded from the `.slp` file on each access.
        cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
        use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
        crop_size: Crop size of each instance for centered-instance model. If `scale` is provided, then the cropped image will be resized according to `scale`.
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
            (required keys: `sigma`, `output_stride`, `part_names` and `anchor_part` depending on the model type ).
        labels_list: List of `sio.Labels` objects. Used to store the labels in the cache. (only used if `cache_img` is `None`)
        bbox_crop: How box-scaled crops are cut (`preprocessing.crop_mode="bbox"`),
            or `None` for fixed-size crops centered on the instance centroid. With a
            spec, each instance's box (derived like a detector's ground truth) is
            resized to `crop_size` x `crop_size`, exactly as the top-down inference
            layer crops a detected box.
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        crop_size: int,
        confmap_head_config: DictConfig,
        max_stride: int,
        anchor_ind: Optional[int] = None,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        bbox_crop: Optional[BBoxCropSpec] = None,
    ) -> None:
        """Initialize class attributes."""
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
        )
        self.labels = None
        self.crop_size = crop_size
        self.anchor_ind = anchor_ind
        self.confmap_head_config = confmap_head_config
        # (method, fallback) for the crop center, resolved once from the head
        # config so training crops and inference crops agree (#586).
        self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
            *centroid_method_from_config(confmap_head_config), anchor_ind
        )
        self.bbox_crop = bbox_crop
        self.instance_idx_list = self._get_instance_idx_list(labels)
        self.cache_lf = [None, None]

    def _get_instance_idx_list(self, labels: List[sio.Labels]) -> List[Tuple[int]]:
        """Return list of tuples with indices of labelled frames and instances."""
        instance_idx_list = []
        unlinked = 0
        for labels_idx, label in enumerate(labels):
            box_source = (
                resolve_bbox_source(label, self.bbox_crop.bbox_source)
                if self.bbox_crop is not None
                else None
            )
            for lf_idx, lf in enumerate(label):
                # Filter to user instances
                if self.user_instances_only:
                    if lf.user_instances is not None and len(lf.user_instances) > 0:
                        lf.instances = lf.user_instances
                    else:
                        # Skip frames without user instances
                        continue
                for inst_idx, inst in enumerate(lf.instances):
                    if not inst.is_empty:  # filter all NaN instances.
                        video_idx = labels[labels_idx].videos.index(lf.video)
                        sample = {
                            "labels_idx": labels_idx,
                            "lf_idx": lf_idx,
                            "inst_idx": inst_idx,
                            "video_idx": video_idx,
                            "instances": (
                                lf.instances if self.cache_img is not None else None
                            ),
                            "frame_idx": lf.frame_idx,
                        }
                        if box_source is not None:
                            # Box-scaled crops: the instance's box, in original
                            # pixels, derived as for a detector's ground truth.
                            box = instance_crop_bbox(
                                lf, inst, box_source, self.bbox_crop.bbox_padding
                            )
                            if box is None:
                                continue
                            sample["bbox"] = box
                            if box_source == "bboxes" and not any(
                                getattr(b, "instance", None) is inst
                                for b in frame_user_bboxes(lf)
                            ):
                                unlinked += 1
                        instance_idx_list.append(sample)
                        # This is to ensure that the labels are not passed to the multiprocessing pool (h5py objects can't be pickled)
        if unlinked:
            logger.warning(
                f"Box-scaled crops: {unlinked} of {len(instance_idx_list)} instances "
                "have no user box linked to them (box.instance), so their training "
                "boxes come from their keypoints, while a detector trained with the "
                "same bbox_source learns the drawn boxes. Link the boxes, or set "
                "bbox_source: instances on both models."
            )
        return instance_idx_list

    def __len__(self) -> int:
        """Return number of instances in the labels object."""
        return len(self.instance_idx_list)

    def _crop_for_augmentation(
        self,
        image: torch.Tensor,
        instance: torch.Tensor,
        centroid: torch.Tensor,
        box: Optional[np.ndarray],
        eff_scale: float,
    ) -> Dict:
        """Cut the oversized crop that augmentation runs on.

        Fixed crops: a ``crop_size * sqrt(2)`` window around ``centroid`` at the
        frame's pixel size. Box-scaled crops: the window
        :meth:`BBoxCropSpec.windows` gives the instance's ``box``, grown about its
        center to :func:`pre_crop_size` at the same scale, so that without
        augmentation the centered ``crop_size`` re-crop (:meth:`_final_crop`) is
        exactly the crop the top-down inference layer cuts for that box.

        Args:
            image: Size-matched frame ``(1, C, H, W)``.
            instance: Size-matched keypoints ``(n_nodes, 2)``.
            centroid: Size-matched centroid ``(2,)`` (fixed crops only).
            box: Instance box ``(4,)`` in original pixels (box-scaled crops only).
            eff_scale: The size matcher's scale.
        """
        if self.bbox_crop is None:
            crop_size = np.array([self.crop_size, self.crop_size]) * np.sqrt(
                2
            )  # crop extra for rotation augmentation
            crop_size = crop_size.astype(np.int32).tolist()
            return generate_crops(image, instance, centroid, crop_size)

        box = torch.as_tensor(box, dtype=torch.float32).reshape(1, 4) * float(eff_scale)
        windows, scales, keep_half = self.bbox_crop.windows(box, self.crop_size)
        m = pre_crop_size(self.crop_size)
        pre_window = scale_windows(windows, m / self.crop_size)
        corners = window_corners(pre_window)  # (1, 4, 2)
        crop = crop_and_resize(image, boxes=corners, size=(m, m))
        if self.bbox_crop.aspect == "letterbox":
            crop = letterbox_crops(crop, keep_half)
        return {
            "instance_image": crop,
            "instance_bbox": corners,
            "instance": points_to_crop(instance[None], pre_window, scales),
            "centroid": torch.full((1, 2), m / 2 - 0.5),
            # Frame -> final-crop geometry (size-matched pixels), for mapping
            # predictions on un-augmented samples back to the frame.
            "crop_window": windows,
            "crop_scale": scales,
            "input_scale": torch.tensor(float(self.scale)),
        }

    def _final_crop(self, sample: Dict) -> Dict:
        """Re-crop the (augmented) oversized crop to ``crop_size`` around its center."""
        if self.bbox_crop is None:
            sample["instance_bbox"] = torch.unsqueeze(
                make_centered_bboxes(
                    sample["centroid"][0], self.crop_size, self.crop_size
                ),
                0,
            )  # (n_samples=1, 4, 2)

            sample["instance_image"] = crop_and_resize(
                sample["instance_image"],
                boxes=sample["instance_bbox"],
                size=(self.crop_size, self.crop_size),
            )
            point = sample["instance_bbox"][0][0]
            center_instance = sample["instance"] - point
            centered_centroid = sample["centroid"] - point

            sample["instance"] = center_instance  # (n_samples=1, n_nodes, 2)
            sample["centroid"] = centered_centroid  # (n_samples=1, 2)
            return sample

        # Box-scaled crops: a window of exactly `crop_size` pixels on whole-pixel
        # edges (`pre_crop_size` keeps the parity), i.e. a copy, not a resample.
        origin = sample["centroid"][0] + 0.5 - self.crop_size / 2
        window = torch.cat([origin, origin + self.crop_size])[None]
        sample["instance_bbox"] = window_corners(window)  # (n_samples=1, 4, 2)
        sample["instance_image"] = crop_and_resize(
            sample["instance_image"],
            boxes=sample["instance_bbox"],
            size=(self.crop_size, self.crop_size),
        )
        sample["instance"] = sample["instance"] - origin
        sample["centroid"] = sample["centroid"] - origin
        return sample

    def __getitem__(self, index) -> Dict:
        """Return dict with cropped image and confmaps of instance for given index."""
        sample = self.instance_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        inst_idx = sample["inst_idx"]
        video_idx = sample["video_idx"]
        lf_frame_idx = sample["frame_idx"]
        crop_box = sample.get("bbox")

        if self.cache_img is not None:
            instances_list = sample["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances_list = lf.instances
            img = lf.image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW

        instances = []
        for inst in instances_list:
            instances.append(
                inst.numpy()
            )  # no need to filter empty instances; handled while creating instance_idx_list
        instances = np.stack(instances, axis=0)

        # Add singleton time dimension for single frames.
        image = np.expand_dims(image, axis=0)  # (n_samples=1, C, H, W)
        instances = np.expand_dims(
            instances, axis=0
        )  # (n_samples=1, num_instances, num_nodes, 2)

        instances = torch.from_numpy(instances.astype("float32"))
        image = torch.from_numpy(image.copy())

        num_instances, _ = instances.shape[1:3]
        orig_img_height, orig_img_width = image.shape[-2:]

        instances = instances[:, inst_idx]

        if self.ensure_rgb:
            image = convert_to_rgb(image)
        elif self.ensure_grayscale:
            image = convert_to_grayscale(image)

        # size matcher
        image, eff_scale = apply_sizematcher(
            image,
            max_height=self.max_hw[0],
            max_width=self.max_hw[1],
        )
        instances = instances * eff_scale

        # get the centroids based on the anchor idx
        centroids = generate_centroids(
            instances,
            anchor_ind=self.anchor_ind,
            method=self.centroid_method,
            fallback=self.centroid_fallback,
        )

        instance, centroid = instances[0], centroids[0]  # (n_samples=1)

        sample = self._crop_for_augmentation(
            image, instance, centroid, crop_box, eff_scale
        )

        sample["frame_idx"] = torch.tensor(lf_frame_idx, dtype=torch.int32)
        sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
        sample["num_instances"] = num_instances
        sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
            0
        )
        sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

        # apply augmentation
        if self.apply_aug:
            if self.intensity_aug is not None:
                (
                    sample["instance_image"],
                    sample["instance"],
                ) = apply_intensity_augmentation(
                    sample["instance_image"],
                    sample["instance"],
                    **self.intensity_aug,
                )

            if self.geometric_aug is not None:
                (
                    sample["instance_image"],
                    sample["instance"],
                ) = apply_geometric_augmentation(
                    sample["instance_image"],
                    sample["instance"],
                    symmetric_inds=self.symmetric_inds,
                    **self.geometric_aug,
                )

        # re-crop to original crop size
        sample = self._final_crop(sample)

        # resize the cropped image
        sample["instance_image"], sample["instance"] = apply_resizer(
            sample["instance_image"],
            sample["instance"],
            scale=self.scale,
        )

        # Pad the image (if needed) according max stride
        sample["instance_image"] = apply_pad_to_stride(
            sample["instance_image"], max_stride=self.max_stride
        )

        img_hw = sample["instance_image"].shape[-2:]

        # Drop keypoints pushed outside the crop by augmentation so their target
        # confidence map is empty rather than a partial blob at the crop edge.
        sample["instance"] = filter_oob_points(sample["instance"], img_hw[0], img_hw[1])

        # Generate confidence maps
        confidence_maps = generate_confmaps(
            sample["instance"],
            img_hw=img_hw,
            sigma=self.confmap_head_config.sigma,
            output_stride=self.confmap_head_config.output_stride,
        )

        sample["confidence_maps"] = confidence_maps
        sample["labels_idx"] = labels_idx

        return sample

__getitem__(index)

Return dict with cropped image and confmaps of instance for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with cropped image and confmaps of instance for given index."""
    sample = self.instance_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    inst_idx = sample["inst_idx"]
    video_idx = sample["video_idx"]
    lf_frame_idx = sample["frame_idx"]
    crop_box = sample.get("bbox")

    if self.cache_img is not None:
        instances_list = sample["instances"]
        if self.cache_img == "disk":
            img = np.array(
                Image.open(
                    f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                )
            )
        elif self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)].copy()
    else:
        lf = self.labels_list[labels_idx][lf_idx]
        instances_list = lf.instances
        img = lf.image
    if img.ndim == 2:
        img = np.expand_dims(img, axis=2)

    image = np.transpose(img, (2, 0, 1))  # HWC -> CHW

    instances = []
    for inst in instances_list:
        instances.append(
            inst.numpy()
        )  # no need to filter empty instances; handled while creating instance_idx_list
    instances = np.stack(instances, axis=0)

    # Add singleton time dimension for single frames.
    image = np.expand_dims(image, axis=0)  # (n_samples=1, C, H, W)
    instances = np.expand_dims(
        instances, axis=0
    )  # (n_samples=1, num_instances, num_nodes, 2)

    instances = torch.from_numpy(instances.astype("float32"))
    image = torch.from_numpy(image.copy())

    num_instances, _ = instances.shape[1:3]
    orig_img_height, orig_img_width = image.shape[-2:]

    instances = instances[:, inst_idx]

    if self.ensure_rgb:
        image = convert_to_rgb(image)
    elif self.ensure_grayscale:
        image = convert_to_grayscale(image)

    # size matcher
    image, eff_scale = apply_sizematcher(
        image,
        max_height=self.max_hw[0],
        max_width=self.max_hw[1],
    )
    instances = instances * eff_scale

    # get the centroids based on the anchor idx
    centroids = generate_centroids(
        instances,
        anchor_ind=self.anchor_ind,
        method=self.centroid_method,
        fallback=self.centroid_fallback,
    )

    instance, centroid = instances[0], centroids[0]  # (n_samples=1)

    sample = self._crop_for_augmentation(
        image, instance, centroid, crop_box, eff_scale
    )

    sample["frame_idx"] = torch.tensor(lf_frame_idx, dtype=torch.int32)
    sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
    sample["num_instances"] = num_instances
    sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
        0
    )
    sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

    # apply augmentation
    if self.apply_aug:
        if self.intensity_aug is not None:
            (
                sample["instance_image"],
                sample["instance"],
            ) = apply_intensity_augmentation(
                sample["instance_image"],
                sample["instance"],
                **self.intensity_aug,
            )

        if self.geometric_aug is not None:
            (
                sample["instance_image"],
                sample["instance"],
            ) = apply_geometric_augmentation(
                sample["instance_image"],
                sample["instance"],
                symmetric_inds=self.symmetric_inds,
                **self.geometric_aug,
            )

    # re-crop to original crop size
    sample = self._final_crop(sample)

    # resize the cropped image
    sample["instance_image"], sample["instance"] = apply_resizer(
        sample["instance_image"],
        sample["instance"],
        scale=self.scale,
    )

    # Pad the image (if needed) according max stride
    sample["instance_image"] = apply_pad_to_stride(
        sample["instance_image"], max_stride=self.max_stride
    )

    img_hw = sample["instance_image"].shape[-2:]

    # Drop keypoints pushed outside the crop by augmentation so their target
    # confidence map is empty rather than a partial blob at the crop edge.
    sample["instance"] = filter_oob_points(sample["instance"], img_hw[0], img_hw[1])

    # Generate confidence maps
    confidence_maps = generate_confmaps(
        sample["instance"],
        img_hw=img_hw,
        sigma=self.confmap_head_config.sigma,
        output_stride=self.confmap_head_config.output_stride,
    )

    sample["confidence_maps"] = confidence_maps
    sample["labels_idx"] = labels_idx

    return sample

__init__(labels, crop_size, confmap_head_config, max_stride, anchor_ind=None, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, bbox_crop=None)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    crop_size: int,
    confmap_head_config: DictConfig,
    max_stride: int,
    anchor_ind: Optional[int] = None,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    bbox_crop: Optional[BBoxCropSpec] = None,
) -> None:
    """Initialize class attributes."""
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
    )
    self.labels = None
    self.crop_size = crop_size
    self.anchor_ind = anchor_ind
    self.confmap_head_config = confmap_head_config
    # (method, fallback) for the crop center, resolved once from the head
    # config so training crops and inference crops agree (#586).
    self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
        *centroid_method_from_config(confmap_head_config), anchor_ind
    )
    self.bbox_crop = bbox_crop
    self.instance_idx_list = self._get_instance_idx_list(labels)
    self.cache_lf = [None, None]

__len__()

Return number of instances in the labels object.

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Return number of instances in the labels object."""
    return len(self.instance_idx_list)

CenteredInstanceSegmentationDataset

Bases: CenteredInstanceDataset

Dataset for top-down (crop-centered) instance segmentation (#622).

Subclasses :class:CenteredInstanceDataset: reuses the centroid-crop pipeline but replaces the keypoint confidence-map GT with a single binary foreground mask of ONLY the centered instance (other instances' foreground inside the crop is background). The centered instance's full-frame mask is captured into the sample index at construction time (via the one-way mask.instance link, with a bbox-IoU fallback), decoded on access, and carried through the SAME size-matcher / crop / resize / pad operations as the image so it stays pixel-aligned with instance_image; it is then downsampled to the segmentation head's output stride.

Note

Augmentation: intensity aug is applied to the image; geometric aug (rotation/scale/translate/flip) co-transforms the centered-instance mask with the SAME affine matrix as the image+keypoints (nearest-neighbor, re-binarized) on the oversized sqrt(2) crop, which provides rotation headroom before the re-crop to crop_size. Erase/mixup stay image-only. Train at scale=1.0 with crop dims divisible by max_stride (the mask resize/pad mirror the image but are not sub-pixel pad-aware).

Attributes:

Name	Type	Description
seg_head_config		Configuration for the segmentation head (output_stride).

Methods:

Name	Description
__getitem__	Return dict with cropped image and the centered-instance mask GT.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class CenteredInstanceSegmentationDataset(CenteredInstanceDataset):
    """Dataset for top-down (crop-centered) instance segmentation (#622).

    Subclasses :class:`CenteredInstanceDataset`: reuses the centroid-crop
    pipeline but replaces the keypoint confidence-map GT with a single binary
    foreground mask of ONLY the centered instance (other instances' foreground
    inside the crop is background). The centered instance's full-frame mask is
    captured into the sample index at construction time (via the one-way
    ``mask.instance`` link, with a bbox-IoU fallback), decoded on access, and
    carried through the SAME size-matcher / crop / resize / pad operations as the
    image so it stays pixel-aligned with ``instance_image``; it is then
    downsampled to the segmentation head's output stride.

    Note:
        Augmentation: intensity aug is applied to the image; geometric aug
        (rotation/scale/translate/flip) co-transforms the centered-instance mask with
        the SAME affine matrix as the image+keypoints (nearest-neighbor, re-binarized)
        on the oversized ``sqrt(2)`` crop, which provides rotation headroom before the
        re-crop to ``crop_size``. Erase/mixup stay image-only. Train at ``scale=1.0``
        with crop dims divisible by ``max_stride`` (the mask resize/pad mirror the
        image but are not sub-pixel pad-aware).

    Attributes:
        seg_head_config: Configuration for the segmentation head (``output_stride``).
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        crop_size: int,
        seg_head_config: DictConfig,
        max_stride: int,
        anchor_ind: Optional[int] = None,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
    ) -> None:
        """Initialize class attributes."""
        self.seg_head_config = seg_head_config
        super().__init__(
            labels=labels,
            crop_size=crop_size,
            # The base class only uses confmap_head_config in its (overridden)
            # __getitem__; reuse the seg head config (it carries `output_stride`).
            confmap_head_config=seg_head_config,
            max_stride=max_stride,
            anchor_ind=anchor_ind,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
        )

    def _get_instance_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Index per (frame, instance), capturing each instance's mask object.

        The associated ``sio.SegmentationMask`` (RLE-compact, picklable) is stored
        in each sample so ``__getitem__`` never needs a live ``Labels`` handle
        (correct under the memory/disk image caching paths). Instances with no
        associated mask are skipped.
        """
        instance_idx_list = []
        n_missing = 0
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                if self.user_instances_only:
                    if lf.user_instances is not None and len(lf.user_instances) > 0:
                        lf.instances = lf.user_instances
                    else:
                        continue
                lf_masks = getattr(lf, "masks", None) or []
                # One-to-one mask<->instance assignment per frame (no mask reused
                # across overlapping instances).
                assigned = _associate_masks(lf.instances, lf_masks)
                for inst_idx, inst in enumerate(lf.instances):
                    if inst.is_empty:
                        continue
                    mask_obj = assigned.get(inst_idx)
                    if mask_obj is None:
                        n_missing += 1
                        continue
                    video_idx = labels[labels_idx].videos.index(lf.video)
                    instance_idx_list.append(
                        {
                            "labels_idx": labels_idx,
                            "lf_idx": lf_idx,
                            "inst_idx": inst_idx,
                            "video_idx": video_idx,
                            "instances": (
                                lf.instances if self.cache_img is not None else None
                            ),
                            "frame_idx": lf.frame_idx,
                            "mask_obj": mask_obj,
                        }
                    )
        if n_missing:
            logger.warning(
                f"CenteredInstanceSegmentationDataset: skipped {n_missing} "
                f"instance(s) with no associated segmentation mask."
            )
        return instance_idx_list

    def __getitem__(self, index) -> Dict:
        """Return dict with cropped image and the centered-instance mask GT."""
        from sleap_nn.inference.segmentation_convert import decode_mask_to_image_res

        meta = self.instance_idx_list[index]
        labels_idx = meta["labels_idx"]
        lf_idx = meta["lf_idx"]
        inst_idx = meta["inst_idx"]
        video_idx = meta["video_idx"]
        lf_frame_idx = meta["frame_idx"]

        if self.cache_img is not None:
            instances_list = meta["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances_list = lf.instances
            img = lf.image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
        instances = np.stack([inst.numpy() for inst in instances_list], axis=0)
        image = np.expand_dims(image, axis=0)  # (1, C, H, W)
        instances = np.expand_dims(instances, axis=0)  # (1, num, nodes, 2)
        instances = torch.from_numpy(instances.astype("float32"))
        image = torch.from_numpy(image.copy())

        num_instances = instances.shape[1]
        orig_img_height, orig_img_width = image.shape[-2:]
        instances = instances[:, inst_idx]

        # Decode the centered instance's mask and place it on the EXACT full-frame
        # canvas so it shares the image's pixel grid. `decode_mask_to_image_res`
        # returns a partial-frame array for masks carrying a non-identity
        # scale/offset (e.g. a pseudo-label `.slp` written by a top-down seg
        # model); without this placement the mask desyncs from the (full-frame)
        # image under the shared size-matcher/crop and the GT is silently
        # corrupted. Clamp to frame bounds (decoded extent can be +/-1 px).
        mask_np = decode_mask_to_image_res(meta["mask_obj"])
        if mask_np.shape[:2] != (orig_img_height, orig_img_width):
            full = np.zeros((orig_img_height, orig_img_width), dtype=bool)
            h0 = min(mask_np.shape[0], orig_img_height)
            w0 = min(mask_np.shape[1], orig_img_width)
            full[:h0, :w0] = mask_np[:h0, :w0]
            mask_np = full
        mask_t = torch.from_numpy(np.ascontiguousarray(mask_np, dtype=np.float32))[
            None, None
        ]

        if self.ensure_rgb:
            image = convert_to_rgb(image)
        elif self.ensure_grayscale:
            image = convert_to_grayscale(image)

        # Size matcher (apply the SAME geometry to image and mask).
        image, eff_scale = apply_sizematcher(
            image, max_height=self.max_hw[0], max_width=self.max_hw[1]
        )
        mask_t, _ = apply_sizematcher(
            mask_t, max_height=self.max_hw[0], max_width=self.max_hw[1]
        )
        instances = instances * eff_scale

        centroids = generate_centroids(
            instances,
            anchor_ind=self.anchor_ind,
            method=self.centroid_method,
            fallback=self.centroid_fallback,
        )
        instance, centroid = instances[0], centroids[0]

        # Oversized (sqrt(2)) crop for rotation headroom — kept for parity with
        # CenteredInstanceDataset even though geometric aug is skipped here.
        crop_size_aug = (
            (np.array([self.crop_size, self.crop_size]) * np.sqrt(2))
            .astype(np.int32)
            .tolist()
        )
        sample = generate_crops(image, instance, centroid, crop_size_aug)
        # Carry the mask through the IDENTICAL crop bbox.
        mask_t = crop_and_resize(
            mask_t, boxes=sample["instance_bbox"], size=crop_size_aug
        )

        sample["frame_idx"] = torch.tensor(lf_frame_idx, dtype=torch.int32)
        sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
        sample["num_instances"] = num_instances
        sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
            0
        )
        sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

        # Intensity augmentation on the image only (mask is intensity-invariant).
        if self.apply_aug and self.intensity_aug is not None:
            (
                sample["instance_image"],
                sample["instance"],
            ) = apply_intensity_augmentation(
                sample["instance_image"], sample["instance"], **self.intensity_aug
            )
        # Geometric augmentation: co-transform the centered-instance mask with the
        # SAME flip/affine matrix as the image + keypoints (nearest-neighbor, then
        # re-binarized). The oversized sqrt(2) crop above provides the rotation
        # headroom; the re-crop below trims back to crop_size so no out-of-frame
        # corners reach the model.
        if self.apply_aug and self.geometric_aug is not None:
            (
                sample["instance_image"],
                sample["instance"],
                mask_t,
            ) = apply_geometric_augmentation(
                sample["instance_image"],
                sample["instance"],
                symmetric_inds=self.symmetric_inds,
                masks=mask_t,
                **self.geometric_aug,
            )

        # Re-crop to the exact crop size (same bbox for image and mask).
        sample["instance_bbox"] = torch.unsqueeze(
            make_centered_bboxes(sample["centroid"][0], self.crop_size, self.crop_size),
            0,
        )
        sample["instance_image"] = crop_and_resize(
            sample["instance_image"],
            boxes=sample["instance_bbox"],
            size=(self.crop_size, self.crop_size),
        )
        mask_t = crop_and_resize(
            mask_t, boxes=sample["instance_bbox"], size=(self.crop_size, self.crop_size)
        )
        point = sample["instance_bbox"][0][0]
        sample["instance"] = sample["instance"] - point
        sample["centroid"] = sample["centroid"] - point

        # Resize image + keypoints by scale; match the mask to the image size.
        sample["instance_image"], sample["instance"] = apply_resizer(
            sample["instance_image"], sample["instance"], scale=self.scale
        )
        tgt_hw = sample["instance_image"].shape[-2:]
        if tuple(mask_t.shape[-2:]) != tuple(tgt_hw):
            # Antialiased bilinear, matching the image's ``apply_resizer`` /
            # ``tvf.resize`` above so the mask stays registered to the image
            # (standardized across every seg mask-geometry resize).
            mask_t = F.interpolate(
                mask_t,
                size=(int(tgt_hw[0]), int(tgt_hw[1])),
                mode="bilinear",
                align_corners=False,
                antialias=True,
            )

        # Pad both to the model's max stride (bottom-right).
        sample["instance_image"] = apply_pad_to_stride(
            sample["instance_image"], max_stride=self.max_stride
        )
        mask_t = apply_pad_to_stride(mask_t, max_stride=self.max_stride)

        img_hw = sample["instance_image"].shape[-2:]
        mask_bool = mask_t[0, 0].numpy() > 0.5
        # Single-element list -> the centered instance's mask only (no union),
        # downsampled + thresholded to the seg head's output stride.
        sample["foreground_mask"] = generate_foreground_mask(
            [mask_bool],
            img_hw=img_hw,
            output_stride=self.seg_head_config.output_stride,
        )
        sample["labels_idx"] = labels_idx

        return sample

__getitem__(index)

Return dict with cropped image and the centered-instance mask GT.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with cropped image and the centered-instance mask GT."""
    from sleap_nn.inference.segmentation_convert import decode_mask_to_image_res

    meta = self.instance_idx_list[index]
    labels_idx = meta["labels_idx"]
    lf_idx = meta["lf_idx"]
    inst_idx = meta["inst_idx"]
    video_idx = meta["video_idx"]
    lf_frame_idx = meta["frame_idx"]

    if self.cache_img is not None:
        instances_list = meta["instances"]
        if self.cache_img == "disk":
            img = np.array(
                Image.open(
                    f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                )
            )
        elif self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)].copy()
    else:
        lf = self.labels_list[labels_idx][lf_idx]
        instances_list = lf.instances
        img = lf.image
    if img.ndim == 2:
        img = np.expand_dims(img, axis=2)

    image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    instances = np.stack([inst.numpy() for inst in instances_list], axis=0)
    image = np.expand_dims(image, axis=0)  # (1, C, H, W)
    instances = np.expand_dims(instances, axis=0)  # (1, num, nodes, 2)
    instances = torch.from_numpy(instances.astype("float32"))
    image = torch.from_numpy(image.copy())

    num_instances = instances.shape[1]
    orig_img_height, orig_img_width = image.shape[-2:]
    instances = instances[:, inst_idx]

    # Decode the centered instance's mask and place it on the EXACT full-frame
    # canvas so it shares the image's pixel grid. `decode_mask_to_image_res`
    # returns a partial-frame array for masks carrying a non-identity
    # scale/offset (e.g. a pseudo-label `.slp` written by a top-down seg
    # model); without this placement the mask desyncs from the (full-frame)
    # image under the shared size-matcher/crop and the GT is silently
    # corrupted. Clamp to frame bounds (decoded extent can be +/-1 px).
    mask_np = decode_mask_to_image_res(meta["mask_obj"])
    if mask_np.shape[:2] != (orig_img_height, orig_img_width):
        full = np.zeros((orig_img_height, orig_img_width), dtype=bool)
        h0 = min(mask_np.shape[0], orig_img_height)
        w0 = min(mask_np.shape[1], orig_img_width)
        full[:h0, :w0] = mask_np[:h0, :w0]
        mask_np = full
    mask_t = torch.from_numpy(np.ascontiguousarray(mask_np, dtype=np.float32))[
        None, None
    ]

    if self.ensure_rgb:
        image = convert_to_rgb(image)
    elif self.ensure_grayscale:
        image = convert_to_grayscale(image)

    # Size matcher (apply the SAME geometry to image and mask).
    image, eff_scale = apply_sizematcher(
        image, max_height=self.max_hw[0], max_width=self.max_hw[1]
    )
    mask_t, _ = apply_sizematcher(
        mask_t, max_height=self.max_hw[0], max_width=self.max_hw[1]
    )
    instances = instances * eff_scale

    centroids = generate_centroids(
        instances,
        anchor_ind=self.anchor_ind,
        method=self.centroid_method,
        fallback=self.centroid_fallback,
    )
    instance, centroid = instances[0], centroids[0]

    # Oversized (sqrt(2)) crop for rotation headroom — kept for parity with
    # CenteredInstanceDataset even though geometric aug is skipped here.
    crop_size_aug = (
        (np.array([self.crop_size, self.crop_size]) * np.sqrt(2))
        .astype(np.int32)
        .tolist()
    )
    sample = generate_crops(image, instance, centroid, crop_size_aug)
    # Carry the mask through the IDENTICAL crop bbox.
    mask_t = crop_and_resize(
        mask_t, boxes=sample["instance_bbox"], size=crop_size_aug
    )

    sample["frame_idx"] = torch.tensor(lf_frame_idx, dtype=torch.int32)
    sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
    sample["num_instances"] = num_instances
    sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
        0
    )
    sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

    # Intensity augmentation on the image only (mask is intensity-invariant).
    if self.apply_aug and self.intensity_aug is not None:
        (
            sample["instance_image"],
            sample["instance"],
        ) = apply_intensity_augmentation(
            sample["instance_image"], sample["instance"], **self.intensity_aug
        )
    # Geometric augmentation: co-transform the centered-instance mask with the
    # SAME flip/affine matrix as the image + keypoints (nearest-neighbor, then
    # re-binarized). The oversized sqrt(2) crop above provides the rotation
    # headroom; the re-crop below trims back to crop_size so no out-of-frame
    # corners reach the model.
    if self.apply_aug and self.geometric_aug is not None:
        (
            sample["instance_image"],
            sample["instance"],
            mask_t,
        ) = apply_geometric_augmentation(
            sample["instance_image"],
            sample["instance"],
            symmetric_inds=self.symmetric_inds,
            masks=mask_t,
            **self.geometric_aug,
        )

    # Re-crop to the exact crop size (same bbox for image and mask).
    sample["instance_bbox"] = torch.unsqueeze(
        make_centered_bboxes(sample["centroid"][0], self.crop_size, self.crop_size),
        0,
    )
    sample["instance_image"] = crop_and_resize(
        sample["instance_image"],
        boxes=sample["instance_bbox"],
        size=(self.crop_size, self.crop_size),
    )
    mask_t = crop_and_resize(
        mask_t, boxes=sample["instance_bbox"], size=(self.crop_size, self.crop_size)
    )
    point = sample["instance_bbox"][0][0]
    sample["instance"] = sample["instance"] - point
    sample["centroid"] = sample["centroid"] - point

    # Resize image + keypoints by scale; match the mask to the image size.
    sample["instance_image"], sample["instance"] = apply_resizer(
        sample["instance_image"], sample["instance"], scale=self.scale
    )
    tgt_hw = sample["instance_image"].shape[-2:]
    if tuple(mask_t.shape[-2:]) != tuple(tgt_hw):
        # Antialiased bilinear, matching the image's ``apply_resizer`` /
        # ``tvf.resize`` above so the mask stays registered to the image
        # (standardized across every seg mask-geometry resize).
        mask_t = F.interpolate(
            mask_t,
            size=(int(tgt_hw[0]), int(tgt_hw[1])),
            mode="bilinear",
            align_corners=False,
            antialias=True,
        )

    # Pad both to the model's max stride (bottom-right).
    sample["instance_image"] = apply_pad_to_stride(
        sample["instance_image"], max_stride=self.max_stride
    )
    mask_t = apply_pad_to_stride(mask_t, max_stride=self.max_stride)

    img_hw = sample["instance_image"].shape[-2:]
    mask_bool = mask_t[0, 0].numpy() > 0.5
    # Single-element list -> the centered instance's mask only (no union),
    # downsampled + thresholded to the seg head's output stride.
    sample["foreground_mask"] = generate_foreground_mask(
        [mask_bool],
        img_hw=img_hw,
        output_stride=self.seg_head_config.output_stride,
    )
    sample["labels_idx"] = labels_idx

    return sample

__init__(labels, crop_size, seg_head_config, max_stride, anchor_ind=None, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    crop_size: int,
    seg_head_config: DictConfig,
    max_stride: int,
    anchor_ind: Optional[int] = None,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
) -> None:
    """Initialize class attributes."""
    self.seg_head_config = seg_head_config
    super().__init__(
        labels=labels,
        crop_size=crop_size,
        # The base class only uses confmap_head_config in its (overridden)
        # __getitem__; reuse the seg head config (it carries `output_stride`).
        confmap_head_config=seg_head_config,
        max_stride=max_stride,
        anchor_ind=anchor_ind,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
    )

CentroidDataset

Bases: BaseDataset

Dataset class for centroid models.

Attributes:

Name	Type	Description
max_stride		Scalar integer specifying the maximum stride that the image must be divisible by.
anchor_ind		Index of the node to use as the anchor point, based on its index in the ordered list of skeleton nodes.
use_user_centroids		Selects the single dataset-wide centroid source. True trains on user-annotated centroids (UserCentroid); False computes every centroid from instance keypoints (anchor_ind node, else mean of visible nodes). Frames that cannot supply a target in the chosen mode are dropped so the head never sees a per-frame mix of the two sources. If None, the mode is inferred from the labels (user centroids present -> True) and a warning is emitted; callers that care about train/val consistency should pass it explicitly (get_train_val_datasets does, via centroid_source).
user_instances_only		True if only user labeled instances should be used for training. If False, both user labeled and predicted instances would be used.
ensure_rgb		(bool) True if the input image should have 3 channels (RGB image). If input has only one
is	replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default	False.
ensure_grayscale		(bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
image.	If the source image has only one channel and this is set to False, then we retain the single channel input. Default	False.
intensity_aug		Intensity augmentation configuration. Can be: - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness'] - List of strings: Multiple intensity augmentations from the allowed values - Dictionary: Custom intensity configuration - None: No intensity augmentation applied
geometric_aug		Geometric augmentation configuration. Can be: - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup'] - List of strings: Multiple geometric augmentations from the allowed values - Dictionary: Custom geometric configuration - None: No geometric augmentation applied
scale		Factor to resize the image dimensions by, specified as a float. Default: 1.0.
apply_aug		True if augmentations should be applied to the data pipeline, else False. Default: False.
max_hw		Maximum height and width of images across the labels file. If max_height and max_width in the config is None, then max_hw is used (computed with sleap_nn.data.providers.get_max_height_width). Else the values in the config are used.
cache_img		String to indicate which caching to use: memory or disk. If None, the images aren't cached and loaded from the .slp file on each access.
cache_img_path		Path to save the .jpg files. If None, current working dir is used.
use_existing_imgs		Use existing imgs/ chunks in the cache_img_path.
confmap_head_config		DictConfig object with all the keys in the head_config section.
(required	keys	sigma, output_stride and anchor_part depending on the model type ).
rank		Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.
labels_list		List of sio.Labels objects. Used to store the labels in the cache. (only used if cache_img is None)

Methods:

Name	Description
__getitem__	Return dict with image and confmaps for centroids for given index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class CentroidDataset(BaseDataset):
    """Dataset class for centroid models.

    Attributes:
        max_stride: Scalar integer specifying the maximum stride that the image must be
            divisible by.
        anchor_ind: Index of the node to use as the anchor point, based on its index in the
            ordered list of skeleton nodes.
        use_user_centroids: Selects the single dataset-wide centroid source.
            `True` trains on user-annotated centroids (``UserCentroid``); `False`
            computes every centroid from instance keypoints (``anchor_ind`` node,
            else mean of visible nodes). Frames that cannot supply a target in
            the chosen mode are dropped so the head never sees a per-frame mix of
            the two sources. If `None`, the mode is inferred from the labels
            (user centroids present -> `True`) and a warning is emitted; callers
            that care about train/val consistency should pass it explicitly
            (``get_train_val_datasets`` does, via ``centroid_source``).
        user_instances_only: `True` if only user labeled instances should be used for training. If `False`,
            both user labeled and predicted instances would be used.
        ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one
        channel when this is set to `True`, then the images from single-channel
        is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default: `False`.
        ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
        is set to True, then we convert the image to grayscale (single-channel)
        image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
        intensity_aug: Intensity augmentation configuration. Can be:
            - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness']
            - List of strings: Multiple intensity augmentations from the allowed values
            - Dictionary: Custom intensity configuration
            - None: No intensity augmentation applied
        geometric_aug: Geometric augmentation configuration. Can be:
            - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup']
            - List of strings: Multiple geometric augmentations from the allowed values
            - Dictionary: Custom geometric configuration
            - None: No geometric augmentation applied
        scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
        apply_aug: `True` if augmentations should be applied to the data pipeline,
            else `False`. Default: `False`.
        max_hw: Maximum height and width of images across the labels file. If `max_height` and
           `max_width` in the config is None, then `max_hw` is used (computed with
            `sleap_nn.data.providers.get_max_height_width`). Else the values in the config
            are used.
        cache_img: String to indicate which caching to use: `memory` or `disk`. If `None`,
            the images aren't cached and loaded from the `.slp` file on each access.
        cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
        use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
        (required keys: `sigma`, `output_stride` and `anchor_part` depending on the model type ).
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        labels_list: List of `sio.Labels` objects. Used to store the labels in the cache. (only used if `cache_img` is `None`)
    """

    # Keep frames that have user centroids but no pose instances (the centroid
    # model can train on a bare centroid; other models can't). See BaseDataset.
    _include_centroid_only_frames: bool = True

    def __init__(
        self,
        labels: List[sio.Labels],
        confmap_head_config: DictConfig,
        max_stride: int,
        anchor_ind: Optional[int] = None,
        use_user_centroids: Optional[bool] = None,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
    ) -> None:
        """Initialize class attributes."""
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        self.anchor_ind = anchor_ind
        self.confmap_head_config = confmap_head_config
        # (method, fallback) for the centroid TARGET, resolved once from the head
        # config so the trained target and inference agree (#586).
        self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
            *centroid_method_from_config(confmap_head_config), anchor_ind
        )

        # Resolve ONE centroid source for the whole dataset so the head is never
        # trained against a per-frame mix of user-annotated and computed
        # centroids. ``get_train_val_datasets`` resolves this (respecting the
        # ``centroid_source`` config and sharing the decision across splits) and
        # passes an explicit bool. When constructed directly (e.g. in tests)
        # without it, infer from this dataset's frames and warn loudly — a
        # silently-chosen target is a subtle training footgun.
        if use_user_centroids is None:
            self.use_user_centroids = any(
                entry.get("user_centroids") for entry in self.lf_idx_list
            )
            src = "user-annotated" if self.use_user_centroids else "computed"
            logger.warning(
                "CentroidDataset: centroid source not specified; inferred "
                "'%s' centroids from the labels for ALL frames. Pass "
                "use_user_centroids explicitly (or set centroid_source in the "
                "config) to silence this.",
                src,
            )
        else:
            self.use_user_centroids = bool(use_user_centroids)

        # Enforce the resolved mode by dropping frames that cannot supply a
        # target in that mode (keeping them would force the per-frame fallback,
        # i.e. the mix we are eliminating). The two modes are mirror images:
        #   - user mode:     keep frames with a user centroid (+ negatives);
        #                    drop pose-only frames that have no user centroid.
        #   - computed mode: keep frames with a pose instance (+ negatives);
        #                    drop centroid-only frames that have no pose.
        # Negative frames (empty target) are valid in both modes.
        def _keeps(entry: Dict) -> bool:
            if entry.get("is_negative"):
                return True
            if self.use_user_centroids:
                return bool(entry.get("user_centroids"))
            return bool(entry.get("has_pose_instances"))

        n_before = len(self.lf_idx_list)
        self.lf_idx_list = [e for e in self.lf_idx_list if _keeps(e)]
        n_dropped = n_before - len(self.lf_idx_list)
        if n_dropped:
            if self.use_user_centroids:
                reason = (
                    "have pose instances but no UserCentroid annotation; annotate "
                    "centroids on them or set centroid_source='computed'"
                )
            else:
                reason = (
                    "have UserCentroid annotations but no pose instance; add pose "
                    "labels or set centroid_source='user'"
                )
            logger.warning(
                "CentroidDataset: dropped %d/%d frame(s) that %s.",
                n_dropped,
                n_before,
                reason,
            )

        # First-class centroid annotations may outnumber the pose instances in a
        # frame. The per-sample centroid target must have a fixed slot count so
        # the default collate can stack a batch, so grow ``max_instances`` to
        # cover the largest user-centroid count (cheap: reads the plain-float
        # lists already stored on ``lf_idx_list``, no h5py access).
        max_user_centroids = 0
        for entry in self.lf_idx_list:
            uc = entry.get("user_centroids")
            if uc is not None and len(uc) > max_user_centroids:
                max_user_centroids = len(uc)
        if max_user_centroids > self.max_instances:
            self.max_instances = max_user_centroids

        # Node count of the pose skeleton, used to shape the all-NaN placeholder
        # instances tensor for centroid-only frames (frames with user centroids
        # but no pose instance) so a batch mixing those with normal frames still
        # collates to a uniform instances shape.
        self._n_nodes = 1
        if labels and labels[0].skeletons and labels[0].skeletons[0].nodes:
            self._n_nodes = len(labels[0].skeletons[0].nodes)

    def _build_imageonly_sample(
        self, img: np.ndarray, frame_idx: int, video_idx: int
    ) -> Dict:
        """Build a sample dict for a centroid-only frame (no pose instances).

        Mirrors ``process_lf``'s output but with an all-NaN placeholder
        ``instances`` tensor of shape ``(1, max_instances, n_nodes, 2)`` — there
        are no pose instances, yet the shape must match normal frames so a batch
        collates uniformly. ``process_lf`` returns ``None`` for zero instances,
        so this path handles the pure-centroid case; the centroid target itself
        comes from ``user_centroids`` in ``__getitem__``.
        """
        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
        img_height, img_width = image.shape[-2:]
        image = np.expand_dims(image, axis=0)  # (n_samples=1, C, H, W)
        instances = torch.full(
            (1, self.max_instances, self._n_nodes, 2), torch.nan, dtype=torch.float32
        )
        return {
            "image": torch.from_numpy(image.copy()),
            "instances": instances,
            "video_idx": torch.tensor(video_idx, dtype=torch.int32),
            "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
            "orig_size": torch.Tensor([img_height, img_width]).unsqueeze(0),
            "num_instances": 0,
        }

    def _build_user_centroid_target(
        self, user_centroids: List[List[float]], scale: float, n_slots: int
    ) -> torch.Tensor:
        """Build the centroid confmap target from user-annotated centroids.

        Produces a tensor with the same shape/semantics ``generate_centroids``
        returns and that ``generate_multiconfmaps(..., is_centroids=True)``
        consumes: ``(n_samples=1, n_slots, 2)``, NaN-padded to ``n_slots`` so
        batches collate uniformly. ``scale`` maps the raw annotation
        coordinates into the preprocessed (size-matched + resized) image frame
        that the instance-derived path already operates in.
        """
        uc = torch.tensor(user_centroids, dtype=torch.float32) * scale  # (n_user, 2)
        n_user = uc.shape[0]
        if n_user < n_slots:
            pad = torch.full((n_slots - n_user, 2), torch.nan, dtype=torch.float32)
            uc = torch.cat([uc, pad], dim=0)
        elif n_user > n_slots:  # defensive; __init__ grows max_instances to avoid this
            uc = uc[:n_slots]
        return uc.unsqueeze(0)  # (1, n_slots, 2)

    def __getitem__(self, index) -> Dict:
        """Return dict with image and confmaps for centroids for given index."""
        sample = self.lf_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        video_idx = sample["video_idx"]
        lf_frame_idx = sample["frame_idx"]
        # Captured before ``sample`` is reassigned to the processed-frame dict.
        user_centroids = sample.get("user_centroids")

        if sample.get("is_negative", False):
            sample = self._load_negative_sample(sample)
        else:
            if self.cache_img is not None:
                instances = sample["instances"]
                if self.cache_img == "disk":
                    img = np.array(
                        Image.open(
                            f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                        )
                    )
                elif self.cache_img == "memory":
                    img = self.cache[(labels_idx, lf_idx)].copy()
            else:
                lf = self.labels_list[labels_idx][lf_idx]
                instances = lf.instances
                img = lf.image
            if img.ndim == 2:
                img = np.expand_dims(img, axis=2)

            # get dict
            sample = process_lf(
                instances_list=instances,
                img=img,
                frame_idx=lf_frame_idx,
                video_idx=video_idx,
                max_instances=self.max_instances,
                user_instances_only=self.user_instances_only,
            )
            # `process_lf` returns None when the frame has no (user) instances.
            # For the centroid model that is a valid pure-centroid frame (kept by
            # `_get_lf_idx_list` because it has user centroids); build an
            # image-only sample and let the centroid target come from
            # `user_centroids` below.
            if sample is None:
                sample = self._build_imageonly_sample(img, lf_frame_idx, video_idx)

        if self.ensure_rgb:
            sample["image"] = convert_to_rgb(sample["image"])
        elif self.ensure_grayscale:
            sample["image"] = convert_to_grayscale(sample["image"])

        # size matcher
        sample["image"], eff_scale = apply_sizematcher(
            sample["image"],
            max_height=self.max_hw[0],
            max_width=self.max_hw[1],
        )
        sample["instances"] = sample["instances"] * eff_scale
        sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

        # resize image
        sample["image"], sample["instances"] = apply_resizer(
            sample["image"],
            sample["instances"],
            scale=self.scale,
        )

        # Build the centroid target from the single dataset-wide source resolved
        # in ``__init__`` (``self.use_user_centroids``) — never a per-frame mix.
        # In user mode every kept non-negative frame carries a user centroid
        # (frame filtering guarantees it); negative frames have none and fall to
        # the computed path, which yields an all-NaN (empty) target as intended.
        if self.use_user_centroids and user_centroids:
            # Raw annotation coords -> preprocessed frame: instances were scaled
            # by ``eff_scale`` then by ``self.scale`` (apply_resizer) above.
            centroids = self._build_user_centroid_target(
                user_centroids,
                scale=eff_scale * self.scale,
                n_slots=sample["instances"].shape[1],
            )
            # The confmap generator and validation GT slice ``[:num_instances]``;
            # align the count with the number of user centroids.
            sample["num_instances"] = min(len(user_centroids), centroids.shape[1])
        else:
            # get the centroids based on the anchor idx
            centroids = generate_centroids(
                sample["instances"],
                anchor_ind=self.anchor_ind,
                method=self.centroid_method,
                fallback=self.centroid_fallback,
            )

        sample["centroids"] = centroids

        # Pad the image (if needed) according max stride
        sample["image"] = apply_pad_to_stride(
            sample["image"], max_stride=self.max_stride
        )

        # apply augmentation
        if self.apply_aug:
            if self.intensity_aug is not None:
                sample["image"], sample["centroids"] = apply_intensity_augmentation(
                    sample["image"],
                    sample["centroids"],
                    **self.intensity_aug,
                )

            if self.geometric_aug is not None:
                sample["image"], sample["centroids"] = apply_geometric_augmentation(
                    sample["image"],
                    sample["centroids"],
                    # Centroids are one point per instance (no node axis), so no
                    # symmetric swap applies — just mirror the coordinates.
                    symmetric_inds=[],
                    **self.geometric_aug,
                )

        img_hw = sample["image"].shape[-2:]

        # Generate confidence maps
        confidence_maps = generate_multiconfmaps(
            sample["centroids"],
            img_hw=img_hw,
            num_instances=sample["num_instances"],
            sigma=self.confmap_head_config.sigma,
            output_stride=self.confmap_head_config.output_stride,
            is_centroids=True,
        )

        sample["centroids_confidence_maps"] = confidence_maps
        sample["labels_idx"] = labels_idx
        if self.use_negative_frames:
            sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

        return sample

__getitem__(index)

Return dict with image and confmaps for centroids for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image and confmaps for centroids for given index."""
    sample = self.lf_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    video_idx = sample["video_idx"]
    lf_frame_idx = sample["frame_idx"]
    # Captured before ``sample`` is reassigned to the processed-frame dict.
    user_centroids = sample.get("user_centroids")

    if sample.get("is_negative", False):
        sample = self._load_negative_sample(sample)
    else:
        if self.cache_img is not None:
            instances = sample["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances = lf.instances
            img = lf.image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        # get dict
        sample = process_lf(
            instances_list=instances,
            img=img,
            frame_idx=lf_frame_idx,
            video_idx=video_idx,
            max_instances=self.max_instances,
            user_instances_only=self.user_instances_only,
        )
        # `process_lf` returns None when the frame has no (user) instances.
        # For the centroid model that is a valid pure-centroid frame (kept by
        # `_get_lf_idx_list` because it has user centroids); build an
        # image-only sample and let the centroid target come from
        # `user_centroids` below.
        if sample is None:
            sample = self._build_imageonly_sample(img, lf_frame_idx, video_idx)

    if self.ensure_rgb:
        sample["image"] = convert_to_rgb(sample["image"])
    elif self.ensure_grayscale:
        sample["image"] = convert_to_grayscale(sample["image"])

    # size matcher
    sample["image"], eff_scale = apply_sizematcher(
        sample["image"],
        max_height=self.max_hw[0],
        max_width=self.max_hw[1],
    )
    sample["instances"] = sample["instances"] * eff_scale
    sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

    # resize image
    sample["image"], sample["instances"] = apply_resizer(
        sample["image"],
        sample["instances"],
        scale=self.scale,
    )

    # Build the centroid target from the single dataset-wide source resolved
    # in ``__init__`` (``self.use_user_centroids``) — never a per-frame mix.
    # In user mode every kept non-negative frame carries a user centroid
    # (frame filtering guarantees it); negative frames have none and fall to
    # the computed path, which yields an all-NaN (empty) target as intended.
    if self.use_user_centroids and user_centroids:
        # Raw annotation coords -> preprocessed frame: instances were scaled
        # by ``eff_scale`` then by ``self.scale`` (apply_resizer) above.
        centroids = self._build_user_centroid_target(
            user_centroids,
            scale=eff_scale * self.scale,
            n_slots=sample["instances"].shape[1],
        )
        # The confmap generator and validation GT slice ``[:num_instances]``;
        # align the count with the number of user centroids.
        sample["num_instances"] = min(len(user_centroids), centroids.shape[1])
    else:
        # get the centroids based on the anchor idx
        centroids = generate_centroids(
            sample["instances"],
            anchor_ind=self.anchor_ind,
            method=self.centroid_method,
            fallback=self.centroid_fallback,
        )

    sample["centroids"] = centroids

    # Pad the image (if needed) according max stride
    sample["image"] = apply_pad_to_stride(
        sample["image"], max_stride=self.max_stride
    )

    # apply augmentation
    if self.apply_aug:
        if self.intensity_aug is not None:
            sample["image"], sample["centroids"] = apply_intensity_augmentation(
                sample["image"],
                sample["centroids"],
                **self.intensity_aug,
            )

        if self.geometric_aug is not None:
            sample["image"], sample["centroids"] = apply_geometric_augmentation(
                sample["image"],
                sample["centroids"],
                # Centroids are one point per instance (no node axis), so no
                # symmetric swap applies — just mirror the coordinates.
                symmetric_inds=[],
                **self.geometric_aug,
            )

    img_hw = sample["image"].shape[-2:]

    # Generate confidence maps
    confidence_maps = generate_multiconfmaps(
        sample["centroids"],
        img_hw=img_hw,
        num_instances=sample["num_instances"],
        sigma=self.confmap_head_config.sigma,
        output_stride=self.confmap_head_config.output_stride,
        is_centroids=True,
    )

    sample["centroids_confidence_maps"] = confidence_maps
    sample["labels_idx"] = labels_idx
    if self.use_negative_frames:
        sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

    return sample

__init__(labels, confmap_head_config, max_stride, anchor_ind=None, use_user_centroids=None, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    confmap_head_config: DictConfig,
    max_stride: int,
    anchor_ind: Optional[int] = None,
    use_user_centroids: Optional[bool] = None,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
) -> None:
    """Initialize class attributes."""
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=use_negative_frames,
    )
    self.anchor_ind = anchor_ind
    self.confmap_head_config = confmap_head_config
    # (method, fallback) for the centroid TARGET, resolved once from the head
    # config so the trained target and inference agree (#586).
    self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
        *centroid_method_from_config(confmap_head_config), anchor_ind
    )

    # Resolve ONE centroid source for the whole dataset so the head is never
    # trained against a per-frame mix of user-annotated and computed
    # centroids. ``get_train_val_datasets`` resolves this (respecting the
    # ``centroid_source`` config and sharing the decision across splits) and
    # passes an explicit bool. When constructed directly (e.g. in tests)
    # without it, infer from this dataset's frames and warn loudly — a
    # silently-chosen target is a subtle training footgun.
    if use_user_centroids is None:
        self.use_user_centroids = any(
            entry.get("user_centroids") for entry in self.lf_idx_list
        )
        src = "user-annotated" if self.use_user_centroids else "computed"
        logger.warning(
            "CentroidDataset: centroid source not specified; inferred "
            "'%s' centroids from the labels for ALL frames. Pass "
            "use_user_centroids explicitly (or set centroid_source in the "
            "config) to silence this.",
            src,
        )
    else:
        self.use_user_centroids = bool(use_user_centroids)

    # Enforce the resolved mode by dropping frames that cannot supply a
    # target in that mode (keeping them would force the per-frame fallback,
    # i.e. the mix we are eliminating). The two modes are mirror images:
    #   - user mode:     keep frames with a user centroid (+ negatives);
    #                    drop pose-only frames that have no user centroid.
    #   - computed mode: keep frames with a pose instance (+ negatives);
    #                    drop centroid-only frames that have no pose.
    # Negative frames (empty target) are valid in both modes.
    def _keeps(entry: Dict) -> bool:
        if entry.get("is_negative"):
            return True
        if self.use_user_centroids:
            return bool(entry.get("user_centroids"))
        return bool(entry.get("has_pose_instances"))

    n_before = len(self.lf_idx_list)
    self.lf_idx_list = [e for e in self.lf_idx_list if _keeps(e)]
    n_dropped = n_before - len(self.lf_idx_list)
    if n_dropped:
        if self.use_user_centroids:
            reason = (
                "have pose instances but no UserCentroid annotation; annotate "
                "centroids on them or set centroid_source='computed'"
            )
        else:
            reason = (
                "have UserCentroid annotations but no pose instance; add pose "
                "labels or set centroid_source='user'"
            )
        logger.warning(
            "CentroidDataset: dropped %d/%d frame(s) that %s.",
            n_dropped,
            n_before,
            reason,
        )

    # First-class centroid annotations may outnumber the pose instances in a
    # frame. The per-sample centroid target must have a fixed slot count so
    # the default collate can stack a batch, so grow ``max_instances`` to
    # cover the largest user-centroid count (cheap: reads the plain-float
    # lists already stored on ``lf_idx_list``, no h5py access).
    max_user_centroids = 0
    for entry in self.lf_idx_list:
        uc = entry.get("user_centroids")
        if uc is not None and len(uc) > max_user_centroids:
            max_user_centroids = len(uc)
    if max_user_centroids > self.max_instances:
        self.max_instances = max_user_centroids

    # Node count of the pose skeleton, used to shape the all-NaN placeholder
    # instances tensor for centroid-only frames (frames with user centroids
    # but no pose instance) so a batch mixing those with normal frames still
    # collates to a uniform instances shape.
    self._n_nodes = 1
    if labels and labels[0].skeletons and labels[0].skeletons[0].nodes:
        self._n_nodes = len(labels[0].skeletons[0].nodes)

EmbeddingDataset

Bases: BaseDataset

Dataset for the embedding (crop -> vector, re-ID) model type.

One sample per member detection of ONE carrier (see :class:EmbeddingMembership for which detections are members under each scope: tracked ones for tracklet, identified ones for global_id, every one for aug_view and for inference's include_untracked). The carrier is the detection_mode:

• mask: an lf.masks entry; the fixed-square crop is centered on the mask center-of-mass (or bounding-box midpoint) and carries the binary mask crop.
• pose: an lf.instances keypoint detection; the crop is centered on the pose centroid (anchor node with a per-instance mean-of-visible-nodes fallback) and carries an all-ones mask.

The detection_mode argument names the carrier the model trains / was trained on; it is used whenever it holds a member, otherwise the carrier holding more members is (see :func:resolve_embedding_detection_mode).

Returns the grayscale crop, a mask crop, and per-crop metadata (video_idx, frame_idx, group_id, global_group_id, item_id). When apply_aug is set, two independently-augmented views (instance_image / _view2) are produced for the two-view contrastive loss using the standard config-driven augmentation; otherwise a single un-augmented view is returned (val / inference).

The group_id keys the training groups: the global-identity index for global_id scope, a per-(labels, video, track) tracklet id for tracklet scope, or a per-detection id for aug_view. global_group_id is the grouping used for evaluation: the global-identity index, or, for a detection with no global identity, its own tracklet (untracked under aug_view: itself) in a separate id range (see :meth:EmbeddingMembership.group). The per-epoch validation and the post-training eval both group on it. The global identity of a detection is its real sio.Identity name when present, else its sio.Track name under track_names_are_global (see :func:_global_identity_label).

Attributes:

Name	Type	Description
crop_size		Side length of the square crop (should be divisible by max_stride).
class_names		Ordered global-identity vocabulary (the global_id group space; sio.Identity names, or track names under track_names_are_global).
embedding_head_config		The head leaf config (carries output_stride and the optional pose anchor_part).
id_scope		Training-group key: global_id | tracklet | aug_view.
track_names_are_global		Treat a sio.Track name as a global animal identity for detections lacking a sio.Identity (the pre-Identity convention).
membership		The :class:EmbeddingMembership deciding which detections are samples and their groups.
detection_mode		The carrier cropped: pose or mask.

Methods:

Name	Description
__getitem__	Return one grayscale crop + a mask crop + metadata.
__init__	Initialize class attributes.
__len__	Return the number of mask crops.

Source code in sleap_nn/data/custom_datasets.py

class EmbeddingDataset(BaseDataset):
    """Dataset for the ``embedding`` (crop -> vector, re-ID) model type.

    One sample per member detection of ONE carrier (see :class:`EmbeddingMembership`
    for which detections are members under each scope: tracked ones for ``tracklet``,
    identified ones for ``global_id``, every one for ``aug_view`` and for inference's
    ``include_untracked``). The carrier is the ``detection_mode``:

    - ``mask``: an ``lf.masks`` entry; the fixed-square crop is centered on the mask
      center-of-mass (or bounding-box midpoint) and carries the binary mask crop.
    - ``pose``: an ``lf.instances`` keypoint detection; the crop is centered on
      the pose centroid (anchor node with a per-instance mean-of-visible-nodes
      fallback) and carries an all-ones mask.

    The ``detection_mode`` argument names the carrier the model trains / was trained
    on; it is used whenever it holds a member, otherwise the carrier holding more
    members is (see :func:`resolve_embedding_detection_mode`).

    Returns the grayscale crop, a mask crop, and per-crop metadata (``video_idx``,
    ``frame_idx``, ``group_id``, ``global_group_id``, ``item_id``). When ``apply_aug``
    is set, two independently-augmented views (``instance_image`` / ``_view2``) are
    produced for the two-view contrastive loss using the standard config-driven
    augmentation; otherwise a single un-augmented view is returned (val / inference).

    The ``group_id`` keys the training groups: the global-identity index for
    ``global_id`` scope, a per-``(labels, video, track)`` tracklet id for ``tracklet``
    scope, or a per-detection id for ``aug_view``. ``global_group_id`` is the grouping
    used for evaluation: the global-identity index, or, for a detection with no global
    identity, its own tracklet (untracked under ``aug_view``: itself) in a separate id
    range (see :meth:`EmbeddingMembership.group`). The per-epoch validation and the
    post-training eval both group on it. The global identity of a detection is its real
    ``sio.Identity`` name when present, else its ``sio.Track`` name under
    ``track_names_are_global`` (see :func:`_global_identity_label`).

    Attributes:
        crop_size: Side length of the square crop (should be divisible by max_stride).
        class_names: Ordered global-identity vocabulary (the ``global_id`` group space;
            ``sio.Identity`` names, or track names under ``track_names_are_global``).
        embedding_head_config: The head leaf config (carries ``output_stride`` and the
            optional pose ``anchor_part``).
        id_scope: Training-group key: ``global_id`` | ``tracklet`` | ``aug_view``.
        track_names_are_global: Treat a ``sio.Track`` name as a global animal identity
            for detections lacking a ``sio.Identity`` (the pre-Identity convention).
        membership: The :class:`EmbeddingMembership` deciding which detections are
            samples and their groups.
        detection_mode: The carrier cropped: ``pose`` or ``mask``.
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        crop_size: int,
        class_names: List[str],
        embedding_head_config: DictConfig,
        max_stride: int,
        id_scope: str = "global_id",
        track_names_are_global: bool = True,
        crop_centering: str = "auto",
        include_untracked: bool = False,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = True,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        apply_aug: bool = False,
        scale: float = 1.0,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        frame_cache_size: int = 0,
        detection_mode: Optional[str] = None,
        min_mask_area: float = EMBEDDING_MIN_MASK_AREA,
    ) -> None:
        """Initialize class attributes.

        ``frame_cache_size`` keeps the last N decoded frames when ``cache_img`` is
        ``None``. Inference enumerates samples frame by frame, so ``1`` decodes each
        frame once instead of once per detection; ``0`` (default) keeps the
        decode-per-sample behavior for training, whose sampler has no frame
        locality.

        ``user_instances_only`` leaves out the predictions (``PredictedInstance`` and
        ``PredictedSegmentationMask``) that a user label on the same frame supersedes;
        it is ignored under
        ``include_untracked``, and the tracked-only inference path passes ``False``,
        so inference embeds predictions. ``detection_mode`` is the preferred carrier
        (``pose`` / ``mask``, ``None`` to count). ``min_mask_area`` is the smallest
        mask (in image pixels) that is a sample; inference passes ``1`` (skip only
        empty masks).
        """
        self.crop_size = crop_size
        self.class_names = list(class_names)
        self.embedding_head_config = embedding_head_config
        self._frame_cache = (
            _FrameLRU(int(frame_cache_size)) if frame_cache_size > 0 else None
        )
        if detection_mode is not None and detection_mode not in (
            EMBEDDING_DETECTION_MODES
        ):
            raise ValueError(
                f"Unknown embedding detection_mode {detection_mode!r}; choose one of "
                f"{'|'.join(EMBEDDING_DETECTION_MODES)}."
            )
        self._preferred_mode = detection_mode
        # `group_id` keying: global identity (global_id), per-(video, track) tracklet,
        # or per detection (aug_view); see `EmbeddingMembership.group`.
        self.id_scope = id_scope
        # Whether a `sio.Track` name doubles as a global animal identity for detections
        # without a real `sio.Identity` (the pre-Identity convention).
        self.track_names_are_global = bool(track_names_are_global)
        # Mask-mode crop center: `auto`/`mask_com` -> mask center-of-mass; `bbox` ->
        # mask bounding-box midpoint (robust to concave masks). Pose-mode centering is
        # driven by `anchor_part` regardless of this knob.
        if crop_centering not in ("auto", "mask_com", "bbox"):
            message = (
                f"Unknown crop_centering '{crop_centering}'; choose one of "
                f"auto|mask_com|bbox."
            )
            logger.error(message)
            raise ValueError(message)
        self.crop_centering = crop_centering
        # Inference re-tracking (WF2): when True, enumerate EVERY detection — tracked
        # or not — and assign a placeholder ``group_id=0`` (unused at inference). The
        # default (training / offline-retrieval) keeps the member-only enumeration,
        # where ``group_id`` is the real training-group key. Set before super().__init__
        # because the base ctor calls the overridden ``_get_lf_idx_list``.
        self.include_untracked = include_untracked
        self.membership = EmbeddingMembership(
            self.class_names,
            id_scope=id_scope,
            track_names_are_global=track_names_are_global,
            # Inference embeds predictions: `include_untracked` means every detection.
            user_instances_only=user_instances_only and not include_untracked,
            include_untracked=include_untracked,
            min_mask_area=min_mask_area,
            labels=labels,
        )
        # `scale` is NOT applied to embedding crops, at training or at inference: the
        # crop path sizes the frame via max_hw, crops crop_size around the detection,
        # and never resizes by `scale` (inference reuses this dataset, and
        # `EmbeddingLayer.predict` does no resizing of its own). Say so, since a
        # non-1.0 value otherwise looks like it shrinks the animal in the crop.
        if float(scale) != 1.0:
            logger.warning(
                f"data_config.preprocessing.scale={scale} is ignored for the embedding "
                "model: crops are sized by max_height/max_width and crop_size only, "
                "at training and inference alike. Use crop_size (or max_height/"
                "max_width) to change how large the animal appears in the crop."
            )
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            # Two-view contrastive aug now reuses the standard skia augmentation
            # (config-driven) per-crop in __getitem__ (CPU-side, like every other
            # dataset) instead of a bespoke GPU reimplementation in the LM.
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
        )
        # Anchor node for pose-centroid centering (Stage 1). Resolved against the
        # skeleton; falls back per-instance to the mean of visible nodes (the topdown
        # convention via `generate_centroids`).
        self.anchor_part = OmegaConf.select(
            embedding_head_config, "anchor_part", default=None
        )
        self.anchor_ind = None
        if self.anchor_part is not None:
            for label in labels:
                if label.skeletons:
                    names = label.skeletons[0].node_names
                    if self.anchor_part in names:
                        self.anchor_ind = names.index(self.anchor_part)
                    break
        # (method, fallback) for the crop center. Resolved after `anchor_ind` so an
        # `anchor_part` absent from the skeleton degrades to the fallback rather
        # than raising, matching this dataset's lenient anchor resolution (#586).
        self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
            *centroid_method_from_config(embedding_head_config), self.anchor_ind
        )

        # Detection mode: masks (crop on the mask center-of-mass) vs keypoint instances
        # (crop on the pose centroid, no mask) -- the preferred (trained) carrier when
        # it holds a member, else the one holding more.
        self.detection_mode = self._detect_mode(labels)
        if self.detection_mode == "pose":
            self.mask_idx_list = self._get_instance_idx_list(labels)
        else:
            self.mask_idx_list = self._get_mask_idx_list(labels)
        # Per-crop arrays for the group-aware batch sampler.
        self.group_ids = np.array([m["group_id"] for m in self.mask_idx_list], np.int64)
        # Video identity must be unique ACROSS labels files. `video_idx` indexes ONE
        # file's `labels.videos`, so video 0 of file A and video 0 of file B shared an
        # id: a cross-file pair at the same frame then looked like a same-video,
        # same-frame pair -- a "known negative" under `restrict_same_video=True` -- and
        # `GroupAwareBatchSampler`'s same-video guarantee silently spanned two files.
        # `video_idx` is kept as-is for image loading, which is per file.
        self._video_id_vocab: dict = {}
        for meta in self.mask_idx_list:
            key = (meta["labels_idx"], meta["video_idx"])
            meta["video_id"] = self._video_id_vocab.setdefault(
                key, len(self._video_id_vocab)
            )
        self.video_ids = np.array([m["video_id"] for m in self.mask_idx_list], np.int64)
        self.frame_ids = np.array(
            [m["frame_idx"] for m in self.mask_idx_list], np.int64
        )

    def _detect_mode(self, labels: List[sio.Labels]) -> str:
        """The carrier to crop: ``"mask"`` or ``"pose"`` (keypoints).

        Counts each carrier's MEMBERS (:meth:`EmbeddingMembership.is_member`, so
        identity-only masks count under ``global_id`` and every detection counts under
        ``include_untracked``) and applies :func:`resolve_embedding_detection_mode`:
        the preferred ``detection_mode`` when it holds a member, else the carrier with
        more members (ties to masks).

        Counting only TRACKED detections, as this once did, sent identity-only masks
        (the sleap-io #535 style) to pose mode and an empty dataset. Deciding on "any
        mask anywhere", as it did before that, let a single user-GT mask flip a whole
        pose-only ``.slp`` to mask mode.
        """
        return resolve_embedding_detection_mode(
            labels, self.membership.is_member, preferred=self._preferred_mode
        )

    def _is_member(self, det) -> bool:
        """Whether a detection is a sample (see :class:`EmbeddingMembership`)."""
        return self.membership.is_member(det)

    def _resolve_group(self, det, labels_idx, video_idx):
        """Return ``(group_id, global_group_id)`` for a detection, or ``None`` to skip."""
        return self.membership.group(det, labels_idx, video_idx)

    def _warn_skipped(self, skipped: Dict[str, int], carrier: str) -> None:
        """Say how many detections were left out, and why."""
        if not skipped:
            return
        reasons = ", ".join(f"{n} {why}" for why, n in sorted(skipped.items()))
        logger.warning(
            f"EmbeddingDataset: skipped {sum(skipped.values())} {carrier} "
            f"detection(s) under scope='{self.id_scope}': {reasons}."
        )

    def _get_instance_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Index per member keypoint instance (pose mode).

        Centroid via the topdown :func:`generate_centroids` (anchor node with a
        per-instance fallback to the mean of visible nodes).
        """
        idx_list = []
        skipped: Dict[str, int] = {}
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                for inst_idx, inst in enumerate(lf.instances):
                    video_idx = labels[labels_idx].videos.index(lf.video)
                    reason = self.membership.skip_reason(inst)
                    if reason is not None:
                        skipped[reason] = skipped.get(reason, 0) + 1
                        continue
                    pts = torch.from_numpy(inst.numpy()).to(
                        torch.float32
                    )  # (n_nodes,2)
                    centroid = generate_centroids(
                        pts.unsqueeze(0),
                        anchor_ind=self.anchor_ind,
                        method=self.centroid_method,
                        fallback=self.centroid_fallback,
                    )[
                        0
                    ]  # (x, y) in original image coords
                    # NaN or inf (a node at inf, or past float32's range): no crop
                    # window exists, so the detection is skipped like a NaN one.
                    if not torch.isfinite(centroid).all():
                        skipped["no centroid"] = skipped.get("no centroid", 0) + 1
                        continue
                    group_id, global_group_id = self._resolve_group(
                        inst, labels_idx, video_idx
                    )
                    centroid = centroid.numpy().astype(np.float32)
                    idx_list.append(
                        {
                            "labels_idx": labels_idx,
                            "lf_idx": lf_idx,
                            "instance_idx": inst_idx,
                            "video_idx": video_idx,
                            "frame_idx": lf.frame_idx,
                            "centroid": centroid,
                            # Object-exact source detection (same ``sio.Instance`` held
                            # by ``labels``), parallel to the mask path's ``mask_obj``,
                            # so the embedding writer can attach the vector to it.
                            "mask_obj": inst,
                            "group_id": group_id,
                            "global_group_id": global_group_id,
                        }
                    )
        self._warn_skipped(skipped, "pose")
        return idx_list

    def _get_lf_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Index frames carrying >=1 member mask OR instance (so the image cache covers them).

        Mask-only (gerbil) data carries no user *instances* and pose (fly) data carries
        no *masks*, so the base ``_get_lf_idx_list`` (which filters on user instances)
        can leave the image cache empty. Index on either a member ``lf.masks`` or a
        member ``lf.instances``. Runs before ``detection_mode`` is set, so it is
        mode-agnostic.
        """
        lf_idx_list = []
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                lf_masks = getattr(lf, "masks", None) or []
                has_member = any(self._is_member(m) for m in lf_masks) or any(
                    self._is_member(inst) for inst in lf.instances
                )
                if has_member:
                    video_idx = labels[labels_idx].videos.index(lf.video)
                    lf_idx_list.append(
                        {
                            "labels_idx": labels_idx,
                            "lf_idx": lf_idx,
                            "video_idx": video_idx,
                            "frame_idx": lf.frame_idx,
                            "is_negative": False,
                            "instances": None,
                        }
                    )
        return lf_idx_list

    def _get_mask_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Index per member mask, capturing the (picklable) mask object + its group id."""
        mask_idx_list = []
        skipped: Dict[str, int] = {}
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                lf_masks = getattr(lf, "masks", None) or []
                for mask_idx, mask_obj in enumerate(lf_masks):
                    video_idx = labels[labels_idx].videos.index(lf.video)
                    reason = self.membership.skip_reason(mask_obj)
                    if reason is not None:
                        skipped[reason] = skipped.get(reason, 0) + 1
                        continue
                    group_id, global_group_id = self._resolve_group(
                        mask_obj, labels_idx, video_idx
                    )
                    # The crop center, once, from the run-length encoding (see
                    # `_mask_crop_center`); `mask_image_hw` is how far the mask
                    # reaches, to tell whether a frame clips it.
                    rects = mask_foreground_rects(mask_obj)
                    has_rects = rects.shape[0] > 0
                    mask_idx_list.append(
                        {
                            "labels_idx": labels_idx,
                            "lf_idx": lf_idx,
                            "mask_idx": mask_idx,
                            "video_idx": video_idx,
                            "frame_idx": lf.frame_idx,
                            "mask_obj": mask_obj,
                            "group_id": group_id,
                            "global_group_id": global_group_id,
                            "mask_center": (
                                mask_rects_center(rects, self.crop_centering)
                                if has_rects
                                else None
                            ),
                            "mask_image_hw": (
                                (int(rects[:, 1].max()), int(rects[:, 3].max()))
                                if has_rects
                                else None
                            ),
                        }
                    )
        self._warn_skipped(skipped, "mask")
        return mask_idx_list

    def __len__(self) -> int:
        """Return the number of mask crops."""
        return len(self.mask_idx_list)

    def __getitem__(self, index) -> Dict:
        """Return one grayscale crop + a mask crop + metadata.

        Mask mode: crop centered on the mask COM, carrying the binary mask crop.
        Pose mode: crop centered on the pose centroid (mean of visible keypoints),
        carrying an all-ones mask (no segmentation; train maskless via
        ``preprocessing.burn_in=false``).
        """
        meta = self.mask_idx_list[index]
        labels_idx, lf_idx = meta["labels_idx"], meta["lf_idx"]

        # The frame is shared with the image caches: nothing below writes to it.
        if self.cache_img is not None and self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)]
        elif self.cache_img is not None and self.cache_img == "disk":
            img = np.array(
                Image.open(f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg")
            )
        else:
            img = self._read_frame(labels_idx, lf_idx)
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        if self._is_size_matched(img.shape[0], img.shape[1]):
            # The size matcher leaves this frame as it is: crop from the window the
            # crop reads (the same bytes, see `_crop_window`).
            if "centroid" in meta:  # pose mode
                instance_image, instance_mask = self._crop_pose_window(img, meta)
            else:  # mask mode
                instance_image, instance_mask = self._crop_mask_window(img, meta)
            return self._pack_sample(instance_image, instance_mask, meta, index)

        image = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)  # (1, C, H, W)
        image = torch.from_numpy(image.copy())

        if "centroid" in meta:  # pose mode
            instance_image, instance_mask = self._crop_pose(image, meta)
        else:  # mask mode
            instance_image, instance_mask = self._crop_mask(image, meta)
        return self._pack_sample(instance_image, instance_mask, meta, index)

    def _read_frame(self, labels_idx: int, lf_idx: int) -> np.ndarray:
        """Decode one frame, reusing the last ``frame_cache_size`` decodes.

        The returned array is shared with the cache; callers must not modify it in
        place (``__getitem__`` copies before any transform).
        """
        if self._frame_cache is None:
            return self.labels_list[labels_idx][lf_idx].image
        key = (labels_idx, lf_idx)
        img = self._frame_cache.get(key)
        if img is None:
            img = self.labels_list[labels_idx][lf_idx].image
            self._frame_cache.put(key, img)
        return img

    def _is_size_matched(self, height: int, width: int) -> bool:
        """Whether the size matcher (``max_hw``) leaves a ``height x width`` frame as is."""
        max_h, max_w = self.max_hw
        return (max_h is None or int(max_h) == int(height)) and (
            max_w is None or int(max_w) == int(width)
        )

    def _crop_window(self, frame: np.ndarray, cx: float, cy: float, mask_rects=None):
        """Crop ``crop_size`` around ``(cx, cy)``, reading only the window it samples.

        Returns what the size-matched case of :meth:`_crop_pose` / :meth:`_crop_mask`
        returned, byte for byte: the same box and the same Skia transform, fed a
        window of the frame (:func:`crop_and_resize_skia_window`), and the mask
        rasterized over that window only. A crop no longer costs a full-frame
        tensor copy, mask decode, colour conversion and RGBA copy.

        The exception is an RGB frame converted to grey. torchvision's conversion
        accumulates with ``add_(..., alpha=...)``, whose vectorized loop may use a
        fused multiply-add and its scalar tail loop not, so a pixel's grey level can
        depend on where it falls in the tensor. The conversion still runs on the
        whole frame, as before, and the window is taken after it.

        Args:
            frame: The ``(H, W, C)`` frame; read, never written.
            cx: Crop center x, in frame pixels.
            cy: Crop center y, in frame pixels.
            mask_rects: The mask's :func:`mask_foreground_rects` (mask mode), or
                ``None`` for pose mode's all-ones mask.
        """
        frame_hw = (int(frame.shape[0]), int(frame.shape[1]))
        size = (self.crop_size, self.crop_size)
        bbox = make_centered_bboxes(
            torch.tensor([cx, cy], dtype=torch.float32),
            self.crop_size,
            self.crop_size,
        ).unsqueeze(0)
        y0, y1, x0, x1 = crop_window_bounds(bbox, frame_hw)
        if frame.shape[2] == 3 and self.ensure_grayscale and not self.ensure_rgb:
            full = np.expand_dims(np.transpose(frame, (2, 0, 1)), axis=0)
            window = convert_to_grayscale(torch.from_numpy(full.copy()))[
                ..., y0:y1, x0:x1
            ]
        else:
            # A copy, never a view: the frame can be a cache entry, and a decoded
            # RGB frame has a negative channel stride (BGR -> RGB), which
            # `ascontiguousarray` keeps on an EMPTY window (a crop off the frame).
            window = np.transpose(frame[y0:y1, x0:x1], (2, 0, 1)).copy(order="C")
            window = torch.from_numpy(window).unsqueeze(0)
            if self.ensure_rgb:
                window = convert_to_rgb(window)
        instance_image = crop_and_resize_skia_window(
            window, (x0, y0), frame_hw, bbox, size
        )
        if mask_rects is None:
            # No segmentation in pose mode: an all-ones mask.
            instance_mask = torch.ones((1, 1) + size, dtype=torch.float32)
        else:
            mask_window = rasterize_rects(mask_rects, y0, y1, x0, x1)
            mask_window = torch.from_numpy(mask_window.astype(np.float32))[None, None]
            instance_mask = crop_and_resize_skia_window(
                mask_window, (x0, y0), frame_hw, bbox, size
            )
        return instance_image, instance_mask

    def _mask_crop_center(self, meta, rects, frame_hw) -> Tuple[float, float]:
        """The mask crop's center on a size-matched frame.

        The one computed at index time, unless the frame clips the mask (a top-down
        crop mask spilling off the frame, a mask grid larger than the frame): then
        the center of the part inside the frame, as the full-frame path computes it.
        """
        reach = meta.get("mask_image_hw")
        if (
            meta.get("mask_center") is not None
            and reach[0] <= frame_hw[0]
            and reach[1] <= frame_hw[1]
        ):
            return meta["mask_center"]
        return mask_rects_center(
            clip_rects(rects, frame_hw), self.crop_centering, frame_hw
        )

    def _crop_mask_window(self, frame: np.ndarray, meta):
        """:meth:`_crop_mask` of a size-matched ``(H, W, C)`` frame, from its window."""
        rects = mask_foreground_rects(meta["mask_obj"])
        cx, cy = self._mask_crop_center(meta, rects, frame.shape[:2])
        return self._crop_window(frame, cx, cy, mask_rects=rects)

    def _crop_pose_window(self, frame: np.ndarray, meta):
        """:meth:`_crop_pose` of a size-matched ``(H, W, C)`` frame, from its window."""
        cx, cy = float(meta["centroid"][0]), float(meta["centroid"][1])
        return self._crop_window(frame, cx, cy)

    def _crop_mask(self, image, meta):
        """Crop centered on the mask COM; return ``(image_crop, mask_crop)``.

        ``image`` is the ``(1, C, H, W)`` frame. A frame the size matcher leaves as
        is is cropped from its window (:meth:`_crop_mask_window`); one it resizes
        is resized and cropped whole.
        """
        from sleap_nn.inference.segmentation_convert import decode_mask_to_image_res

        orig_h, orig_w = image.shape[-2:]
        if self._is_size_matched(orig_h, orig_w):
            return self._crop_mask_window(image[0].permute(1, 2, 0).numpy(), meta)
        mask_np = decode_mask_to_image_res(meta["mask_obj"])
        if mask_np.shape[:2] != (orig_h, orig_w):
            full = np.zeros((orig_h, orig_w), dtype=bool)
            h0 = min(mask_np.shape[0], orig_h)
            w0 = min(mask_np.shape[1], orig_w)
            full[:h0, :w0] = mask_np[:h0, :w0]
            mask_np = full
        mask_t = torch.from_numpy(np.ascontiguousarray(mask_np, dtype=np.float32))[
            None, None
        ]

        if self.ensure_rgb:
            image = convert_to_rgb(image)
        elif self.ensure_grayscale:
            image = convert_to_grayscale(image)

        image, _ = apply_sizematcher(
            image, max_height=self.max_hw[0], max_width=self.max_hw[1]
        )
        mask_t, _ = apply_sizematcher(
            mask_t, max_height=self.max_hw[0], max_width=self.max_hw[1]
        )

        mask_bool = mask_t[0, 0].numpy() > 0.5
        if self.crop_centering == "bbox":
            cx, cy = _mask_bbox_midpoint(mask_bool)
        else:  # "auto" / "mask_com"
            cx, cy = _compute_mask_centroids([mask_bool])[0]
        bbox = make_centered_bboxes(
            torch.tensor([cx, cy], dtype=torch.float32),
            self.crop_size,
            self.crop_size,
        ).unsqueeze(0)
        instance_image = crop_and_resize(
            image, boxes=bbox, size=(self.crop_size, self.crop_size)
        )
        instance_mask = crop_and_resize(
            mask_t, boxes=bbox, size=(self.crop_size, self.crop_size)
        )
        return instance_image, instance_mask

    def _crop_pose(self, image, meta):
        """Crop centered on the pose centroid; return ``(image_crop, ones_mask)``.

        ``image`` is the ``(1, C, H, W)`` frame; a size-matched one is cropped from
        its window (:meth:`_crop_pose_window`).
        """
        if self._is_size_matched(*image.shape[-2:]):
            return self._crop_pose_window(image[0].permute(1, 2, 0).numpy(), meta)
        if self.ensure_rgb:
            image = convert_to_rgb(image)
        elif self.ensure_grayscale:
            image = convert_to_grayscale(image)

        image, ratio = apply_sizematcher(
            image, max_height=self.max_hw[0], max_width=self.max_hw[1]
        )
        cx = float(meta["centroid"][0]) * ratio
        cy = float(meta["centroid"][1]) * ratio
        bbox = make_centered_bboxes(
            torch.tensor([cx, cy], dtype=torch.float32),
            self.crop_size,
            self.crop_size,
        ).unsqueeze(0)
        instance_image = crop_and_resize(
            image, boxes=bbox, size=(self.crop_size, self.crop_size)
        )
        # No segmentation in pose mode: an all-ones mask (burn-in is a no-op; train
        # maskless via preprocessing.burn_in=false).
        instance_mask = torch.ones(
            (1, 1, self.crop_size, self.crop_size), dtype=torch.float32
        )
        return instance_image, instance_mask

    def _apply_crop_aug(self, image, mask):
        """Apply the standard config-driven skia aug to a crop + its mask.

        Reuses sleap-nn's :func:`apply_intensity_augmentation` /
        :func:`apply_geometric_augmentation` (the same functions every other dataset
        uses) instead of a bespoke GPU reimplementation: the mask co-transforms under
        the SAME affine via the ``masks=`` arg, and a placeholder keypoint rides along
        (the crop has no keypoints) and is discarded. Returns ``(image, mask)``.
        """
        # (n_samples=1, n_inst=1, n_nodes=1, 2) center placeholder for the keypoint
        # co-transform the skia aug expects; its output is ignored.
        dummy = torch.full((1, 1, 1, 2), float(self.crop_size) / 2.0)
        if self.intensity_aug is not None:
            image, _ = apply_intensity_augmentation(image, dummy, **self.intensity_aug)
        if self.geometric_aug is not None:
            image, _, mask = apply_geometric_augmentation(
                image, dummy, masks=mask, symmetric_inds=None, **self.geometric_aug
            )
        return image, mask

    def _pack_sample(self, instance_image, instance_mask, meta, index):
        """Pack one crop into a sample dict.

        Training (``apply_aug=True``): emit TWO independently-augmented views
        (``instance_image``/``instance_image_view2`` + masks) for the two-view
        contrastive loss. Val / inference (``apply_aug=False``): emit one un-augmented
        view (``instance_image``) only.
        """
        sample = {
            "group_id": torch.tensor(meta["group_id"], dtype=torch.int64),
            "global_group_id": torch.tensor(
                meta.get("global_group_id", meta["group_id"]), dtype=torch.int64
            ),
            "video_idx": torch.tensor(meta["video_idx"], dtype=torch.int64),
            # `video_id` (global, see __init__) is what the contrastive masks key
            # same-frame / same-video on; `video_idx` stays per-file for image loading.
            "video_id": torch.tensor(
                meta.get("video_id", meta["video_idx"]), dtype=torch.int64
            ),
            "frame_idx": torch.tensor(meta["frame_idx"], dtype=torch.int64),
            "item_id": torch.tensor(index, dtype=torch.int64),
            "labels_idx": meta["labels_idx"],
        }
        if self.apply_aug:
            img1, m1 = self._apply_crop_aug(instance_image, instance_mask)
            img2, m2 = self._apply_crop_aug(instance_image, instance_mask)
            sample["instance_image"] = img1.to(torch.float32)
            sample["instance_mask"] = (m1 > 0.5).to(torch.float32)
            sample["instance_image_view2"] = img2.to(torch.float32)
            sample["instance_mask_view2"] = (m2 > 0.5).to(torch.float32)
        else:
            sample["instance_image"] = instance_image.to(torch.float32)
            sample["instance_mask"] = (instance_mask > 0.5).to(torch.float32)
        return sample

__getitem__(index)

Return one grayscale crop + a mask crop + metadata.

Mask mode: crop centered on the mask COM, carrying the binary mask crop. Pose mode: crop centered on the pose centroid (mean of visible keypoints), carrying an all-ones mask (no segmentation; train maskless via preprocessing.burn_in=false).

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return one grayscale crop + a mask crop + metadata.

    Mask mode: crop centered on the mask COM, carrying the binary mask crop.
    Pose mode: crop centered on the pose centroid (mean of visible keypoints),
    carrying an all-ones mask (no segmentation; train maskless via
    ``preprocessing.burn_in=false``).
    """
    meta = self.mask_idx_list[index]
    labels_idx, lf_idx = meta["labels_idx"], meta["lf_idx"]

    # The frame is shared with the image caches: nothing below writes to it.
    if self.cache_img is not None and self.cache_img == "memory":
        img = self.cache[(labels_idx, lf_idx)]
    elif self.cache_img is not None and self.cache_img == "disk":
        img = np.array(
            Image.open(f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg")
        )
    else:
        img = self._read_frame(labels_idx, lf_idx)
    if img.ndim == 2:
        img = np.expand_dims(img, axis=2)

    if self._is_size_matched(img.shape[0], img.shape[1]):
        # The size matcher leaves this frame as it is: crop from the window the
        # crop reads (the same bytes, see `_crop_window`).
        if "centroid" in meta:  # pose mode
            instance_image, instance_mask = self._crop_pose_window(img, meta)
        else:  # mask mode
            instance_image, instance_mask = self._crop_mask_window(img, meta)
        return self._pack_sample(instance_image, instance_mask, meta, index)

    image = np.expand_dims(np.transpose(img, (2, 0, 1)), axis=0)  # (1, C, H, W)
    image = torch.from_numpy(image.copy())

    if "centroid" in meta:  # pose mode
        instance_image, instance_mask = self._crop_pose(image, meta)
    else:  # mask mode
        instance_image, instance_mask = self._crop_mask(image, meta)
    return self._pack_sample(instance_image, instance_mask, meta, index)

__init__(labels, crop_size, class_names, embedding_head_config, max_stride, id_scope='global_id', track_names_are_global=True, crop_centering='auto', include_untracked=False, user_instances_only=True, ensure_rgb=False, ensure_grayscale=True, intensity_aug=None, geometric_aug=None, apply_aug=False, scale=1.0, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, frame_cache_size=0, detection_mode=None, min_mask_area=EMBEDDING_MIN_MASK_AREA)

Initialize class attributes.

frame_cache_size keeps the last N decoded frames when cache_img is None. Inference enumerates samples frame by frame, so 1 decodes each frame once instead of once per detection; 0 (default) keeps the decode-per-sample behavior for training, whose sampler has no frame locality.

user_instances_only leaves out the predictions (PredictedInstance and PredictedSegmentationMask) that a user label on the same frame supersedes; it is ignored under include_untracked, and the tracked-only inference path passes False, so inference embeds predictions. detection_mode is the preferred carrier (pose / mask, None to count). min_mask_area is the smallest mask (in image pixels) that is a sample; inference passes 1 (skip only empty masks).

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    crop_size: int,
    class_names: List[str],
    embedding_head_config: DictConfig,
    max_stride: int,
    id_scope: str = "global_id",
    track_names_are_global: bool = True,
    crop_centering: str = "auto",
    include_untracked: bool = False,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = True,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    apply_aug: bool = False,
    scale: float = 1.0,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    frame_cache_size: int = 0,
    detection_mode: Optional[str] = None,
    min_mask_area: float = EMBEDDING_MIN_MASK_AREA,
) -> None:
    """Initialize class attributes.

    ``frame_cache_size`` keeps the last N decoded frames when ``cache_img`` is
    ``None``. Inference enumerates samples frame by frame, so ``1`` decodes each
    frame once instead of once per detection; ``0`` (default) keeps the
    decode-per-sample behavior for training, whose sampler has no frame
    locality.

    ``user_instances_only`` leaves out the predictions (``PredictedInstance`` and
    ``PredictedSegmentationMask``) that a user label on the same frame supersedes;
    it is ignored under
    ``include_untracked``, and the tracked-only inference path passes ``False``,
    so inference embeds predictions. ``detection_mode`` is the preferred carrier
    (``pose`` / ``mask``, ``None`` to count). ``min_mask_area`` is the smallest
    mask (in image pixels) that is a sample; inference passes ``1`` (skip only
    empty masks).
    """
    self.crop_size = crop_size
    self.class_names = list(class_names)
    self.embedding_head_config = embedding_head_config
    self._frame_cache = (
        _FrameLRU(int(frame_cache_size)) if frame_cache_size > 0 else None
    )
    if detection_mode is not None and detection_mode not in (
        EMBEDDING_DETECTION_MODES
    ):
        raise ValueError(
            f"Unknown embedding detection_mode {detection_mode!r}; choose one of "
            f"{'|'.join(EMBEDDING_DETECTION_MODES)}."
        )
    self._preferred_mode = detection_mode
    # `group_id` keying: global identity (global_id), per-(video, track) tracklet,
    # or per detection (aug_view); see `EmbeddingMembership.group`.
    self.id_scope = id_scope
    # Whether a `sio.Track` name doubles as a global animal identity for detections
    # without a real `sio.Identity` (the pre-Identity convention).
    self.track_names_are_global = bool(track_names_are_global)
    # Mask-mode crop center: `auto`/`mask_com` -> mask center-of-mass; `bbox` ->
    # mask bounding-box midpoint (robust to concave masks). Pose-mode centering is
    # driven by `anchor_part` regardless of this knob.
    if crop_centering not in ("auto", "mask_com", "bbox"):
        message = (
            f"Unknown crop_centering '{crop_centering}'; choose one of "
            f"auto|mask_com|bbox."
        )
        logger.error(message)
        raise ValueError(message)
    self.crop_centering = crop_centering
    # Inference re-tracking (WF2): when True, enumerate EVERY detection — tracked
    # or not — and assign a placeholder ``group_id=0`` (unused at inference). The
    # default (training / offline-retrieval) keeps the member-only enumeration,
    # where ``group_id`` is the real training-group key. Set before super().__init__
    # because the base ctor calls the overridden ``_get_lf_idx_list``.
    self.include_untracked = include_untracked
    self.membership = EmbeddingMembership(
        self.class_names,
        id_scope=id_scope,
        track_names_are_global=track_names_are_global,
        # Inference embeds predictions: `include_untracked` means every detection.
        user_instances_only=user_instances_only and not include_untracked,
        include_untracked=include_untracked,
        min_mask_area=min_mask_area,
        labels=labels,
    )
    # `scale` is NOT applied to embedding crops, at training or at inference: the
    # crop path sizes the frame via max_hw, crops crop_size around the detection,
    # and never resizes by `scale` (inference reuses this dataset, and
    # `EmbeddingLayer.predict` does no resizing of its own). Say so, since a
    # non-1.0 value otherwise looks like it shrinks the animal in the crop.
    if float(scale) != 1.0:
        logger.warning(
            f"data_config.preprocessing.scale={scale} is ignored for the embedding "
            "model: crops are sized by max_height/max_width and crop_size only, "
            "at training and inference alike. Use crop_size (or max_height/"
            "max_width) to change how large the animal appears in the crop."
        )
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        # Two-view contrastive aug now reuses the standard skia augmentation
        # (config-driven) per-crop in __getitem__ (CPU-side, like every other
        # dataset) instead of a bespoke GPU reimplementation in the LM.
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
    )
    # Anchor node for pose-centroid centering (Stage 1). Resolved against the
    # skeleton; falls back per-instance to the mean of visible nodes (the topdown
    # convention via `generate_centroids`).
    self.anchor_part = OmegaConf.select(
        embedding_head_config, "anchor_part", default=None
    )
    self.anchor_ind = None
    if self.anchor_part is not None:
        for label in labels:
            if label.skeletons:
                names = label.skeletons[0].node_names
                if self.anchor_part in names:
                    self.anchor_ind = names.index(self.anchor_part)
                break
    # (method, fallback) for the crop center. Resolved after `anchor_ind` so an
    # `anchor_part` absent from the skeleton degrades to the fallback rather
    # than raising, matching this dataset's lenient anchor resolution (#586).
    self.centroid_method, self.centroid_fallback = degrade_anchor_if_unresolved(
        *centroid_method_from_config(embedding_head_config), self.anchor_ind
    )

    # Detection mode: masks (crop on the mask center-of-mass) vs keypoint instances
    # (crop on the pose centroid, no mask) -- the preferred (trained) carrier when
    # it holds a member, else the one holding more.
    self.detection_mode = self._detect_mode(labels)
    if self.detection_mode == "pose":
        self.mask_idx_list = self._get_instance_idx_list(labels)
    else:
        self.mask_idx_list = self._get_mask_idx_list(labels)
    # Per-crop arrays for the group-aware batch sampler.
    self.group_ids = np.array([m["group_id"] for m in self.mask_idx_list], np.int64)
    # Video identity must be unique ACROSS labels files. `video_idx` indexes ONE
    # file's `labels.videos`, so video 0 of file A and video 0 of file B shared an
    # id: a cross-file pair at the same frame then looked like a same-video,
    # same-frame pair -- a "known negative" under `restrict_same_video=True` -- and
    # `GroupAwareBatchSampler`'s same-video guarantee silently spanned two files.
    # `video_idx` is kept as-is for image loading, which is per file.
    self._video_id_vocab: dict = {}
    for meta in self.mask_idx_list:
        key = (meta["labels_idx"], meta["video_idx"])
        meta["video_id"] = self._video_id_vocab.setdefault(
            key, len(self._video_id_vocab)
        )
    self.video_ids = np.array([m["video_id"] for m in self.mask_idx_list], np.int64)
    self.frame_ids = np.array(
        [m["frame_idx"] for m in self.mask_idx_list], np.int64
    )

__len__()

Return the number of mask crops.

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Return the number of mask crops."""
    return len(self.mask_idx_list)

EmbeddingMembership

Which detections an embedding dataset uses, and the groups they fall in.

The one statement of the rules, shared by :class:EmbeddingDataset (its index and its carrier choice) and by the trainer's setup-time checks (crop sizing, "can any batch hold a negative?"), which run before any dataset is built.

A detection is usable unless it is a prediction that a user label on the same frame supersedes, under user_instances_only (:func:~sleap_nn.data.utils.superseded_predictions: a PredictedInstance whose track a user Instance carries, a PredictedSegmentationMask a user mask adopted), an empty pose (every point NaN), or a mask with fewer than min_mask_area image pixels. A usable detection is a member when the scope can group it:

• tracklet: it carries a sio.Track.
• global_id: it has a global identity in class_names (a sio.Identity name, or its track name under track_names_are_global).
• aug_view: always. The scope is self-supervised: each detection is its own group, so unlabeled data trains.
• include_untracked (inference): always, with the placeholder group 0.

Parameters:

Name	Type	Description	Default
class_names	List[str]	The global-identity vocabulary (:func:resolve_embedding_class_names).	required
id_scope	str	global_id | tracklet | aug_view.	'global_id'
track_names_are_global	bool	A track name counts as a global identity.	True
user_instances_only	bool	Leave out predictions a user label on the same frame supersedes.	False
include_untracked	bool	Inference: every usable detection is a member.	False
min_mask_area	float	Masks with fewer image pixels are not usable.	EMBEDDING_MIN_MASK_AREA
labels	Optional[List[Labels]]	The labels the detections come from. Required with user_instances_only: whether a prediction is superseded is a property of its frame, so it is resolved here, once per frame.	None

Methods:

Name	Description
__init__	Store the rules; group ids are assigned as detections are grouped.
group	Return (group_id, global_group_id) for a member, or None to skip.
is_member	Whether a detection is a sample (side-effect free).
skip_reason	Why a detection is not a member, or None if it is one.

Source code in sleap_nn/data/custom_datasets.py

class EmbeddingMembership:
    """Which detections an embedding dataset uses, and the groups they fall in.

    The one statement of the rules, shared by :class:`EmbeddingDataset` (its index and
    its carrier choice) and by the trainer's setup-time checks (crop sizing, "can any
    batch hold a negative?"), which run before any dataset is built.

    A detection is **usable** unless it is a prediction that a user label on the same
    frame supersedes, under ``user_instances_only``
    (:func:`~sleap_nn.data.utils.superseded_predictions`: a ``PredictedInstance`` whose
    track a user ``Instance`` carries, a ``PredictedSegmentationMask`` a user mask
    adopted), an empty pose (every point NaN), or a mask with fewer than
    ``min_mask_area`` image pixels. A usable detection is a **member** when the scope
    can group it:

    - ``tracklet``: it carries a ``sio.Track``.
    - ``global_id``: it has a global identity in ``class_names`` (a ``sio.Identity``
      name, or its track name under ``track_names_are_global``).
    - ``aug_view``: always. The scope is self-supervised: each detection is its own
      group, so unlabeled data trains.
    - ``include_untracked`` (inference): always, with the placeholder group ``0``.

    Args:
        class_names: The global-identity vocabulary
            (:func:`resolve_embedding_class_names`).
        id_scope: ``global_id`` | ``tracklet`` | ``aug_view``.
        track_names_are_global: A track name counts as a global identity.
        user_instances_only: Leave out predictions a user label on the same frame
            supersedes.
        include_untracked: Inference: every usable detection is a member.
        min_mask_area: Masks with fewer image pixels are not usable.
        labels: The labels the detections come from. Required with
            ``user_instances_only``: whether a prediction is superseded is a property
            of its frame, so it is resolved here, once per frame.
    """

    def __init__(
        self,
        class_names: List[str],
        id_scope: str = "global_id",
        track_names_are_global: bool = True,
        user_instances_only: bool = False,
        include_untracked: bool = False,
        min_mask_area: float = EMBEDDING_MIN_MASK_AREA,
        labels: Optional[List[sio.Labels]] = None,
    ) -> None:
        """Store the rules; group ids are assigned as detections are grouped."""
        self.class_names = list(class_names)
        self._class_index = {name: i for i, name in enumerate(self.class_names)}
        self.id_scope = id_scope
        self.track_names_are_global = bool(track_names_are_global)
        self.user_instances_only = bool(user_instances_only)
        self.include_untracked = bool(include_untracked)
        self.min_mask_area = float(min_mask_area)
        # `id`s of the predictions a user label on their frame supersedes.
        self._superseded: set = set()
        if self.user_instances_only:
            if labels is None:
                raise ValueError(
                    "EmbeddingMembership(user_instances_only=True) needs the labels, "
                    "to find the predictions a user label supersedes."
                )
            for label in labels:
                for lf in label:
                    self._superseded |= superseded_predictions(lf)
        # Dense ids for tracklets (`tracklet` training groups; the eval fallback of a
        # detection with no global identity) and per-detection `aug_view` groups.
        self._tracklet_vocab: Dict[tuple, int] = {}
        self._n_detection_groups = 0

    def skip_reason(self, det) -> Optional[str]:
        """Why a detection is not a member, or ``None`` if it is one."""
        if id(det) in self._superseded:
            return "superseded by a user label (data_config.user_instances_only)"
        if isinstance(det, sio.SegmentationMask):
            if _mask_image_area(det) < max(self.min_mask_area, 1.0):
                return "empty mask" if det.area == 0 else "tiny mask"
        elif getattr(det, "is_empty", False):
            return "empty pose"
        if self.include_untracked or self.id_scope == "aug_view":
            return None
        if self.id_scope == "tracklet":
            return None if getattr(det, "track", None) is not None else "no track"
        label = _global_identity_label(det, self.track_names_are_global)
        if label is None or label not in self._class_index:
            return "no global identity"
        return None

    def is_member(self, det) -> bool:
        """Whether a detection is a sample (side-effect free)."""
        return self.skip_reason(det) is None

    def group(self, det, labels_idx: int, video_idx: int) -> Optional[Tuple[int, int]]:
        """Return ``(group_id, global_group_id)`` for a member, or ``None`` to skip.

        ``group_id`` is the TRAINING positive key: the global-identity index for
        ``global_id``, a dense per-``(labels, video, track)`` tracklet id for
        ``tracklet``, or a per-detection id for ``aug_view``.

        ``global_group_id`` is the eval grouping: the global-identity index in
        ``[0, len(class_names))``; for a detection with no global identity, its own
        tracklet (or, untracked under ``aug_view``, itself) keyed
        ``len(class_names) + id`` so it cannot share an id with an identity.

        Assigns ids, so call it once per indexed detection, in index order.
        """
        if not self.is_member(det):
            return None
        if self.include_untracked:
            # Inference: grouping is a training-only concern.
            return 0, 0
        global_label = _global_identity_label(det, self.track_names_are_global)
        gid = self._class_index.get(global_label) if global_label is not None else None
        track = getattr(det, "track", None)
        track_name = track.name if track is not None else None
        if self.id_scope == "tracklet":
            tid = self._tracklet_vocab.setdefault(
                (labels_idx, video_idx, track_name), len(self._tracklet_vocab)
            )
            # Identity indices and tracklet ids both count up from 0. Shift the
            # tracklet fallback past the identities: returning a bare `tid` made
            # tracklet k and identity k one group in the retrieval metrics.
            return tid, (gid if gid is not None else len(self.class_names) + tid)
        if self.id_scope == "aug_view":
            group_id = self._n_detection_groups
            self._n_detection_groups += 1
            if gid is not None:
                return group_id, gid
            key = (
                (labels_idx, video_idx, track_name)
                if track is not None
                else ("detection", group_id)
            )
            fallback = self._tracklet_vocab.setdefault(key, len(self._tracklet_vocab))
            return group_id, len(self.class_names) + fallback
        return gid, gid

__init__(class_names, id_scope='global_id', track_names_are_global=True, user_instances_only=False, include_untracked=False, min_mask_area=EMBEDDING_MIN_MASK_AREA, labels=None)

Store the rules; group ids are assigned as detections are grouped.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    class_names: List[str],
    id_scope: str = "global_id",
    track_names_are_global: bool = True,
    user_instances_only: bool = False,
    include_untracked: bool = False,
    min_mask_area: float = EMBEDDING_MIN_MASK_AREA,
    labels: Optional[List[sio.Labels]] = None,
) -> None:
    """Store the rules; group ids are assigned as detections are grouped."""
    self.class_names = list(class_names)
    self._class_index = {name: i for i, name in enumerate(self.class_names)}
    self.id_scope = id_scope
    self.track_names_are_global = bool(track_names_are_global)
    self.user_instances_only = bool(user_instances_only)
    self.include_untracked = bool(include_untracked)
    self.min_mask_area = float(min_mask_area)
    # `id`s of the predictions a user label on their frame supersedes.
    self._superseded: set = set()
    if self.user_instances_only:
        if labels is None:
            raise ValueError(
                "EmbeddingMembership(user_instances_only=True) needs the labels, "
                "to find the predictions a user label supersedes."
            )
        for label in labels:
            for lf in label:
                self._superseded |= superseded_predictions(lf)
    # Dense ids for tracklets (`tracklet` training groups; the eval fallback of a
    # detection with no global identity) and per-detection `aug_view` groups.
    self._tracklet_vocab: Dict[tuple, int] = {}
    self._n_detection_groups = 0

group(det, labels_idx, video_idx)

Return (group_id, global_group_id) for a member, or None to skip.

group_id is the TRAINING positive key: the global-identity index for global_id, a dense per-(labels, video, track) tracklet id for tracklet, or a per-detection id for aug_view.

global_group_id is the eval grouping: the global-identity index in [0, len(class_names)); for a detection with no global identity, its own tracklet (or, untracked under aug_view, itself) keyed len(class_names) + id so it cannot share an id with an identity.

Assigns ids, so call it once per indexed detection, in index order.

Source code in sleap_nn/data/custom_datasets.py

def group(self, det, labels_idx: int, video_idx: int) -> Optional[Tuple[int, int]]:
    """Return ``(group_id, global_group_id)`` for a member, or ``None`` to skip.

    ``group_id`` is the TRAINING positive key: the global-identity index for
    ``global_id``, a dense per-``(labels, video, track)`` tracklet id for
    ``tracklet``, or a per-detection id for ``aug_view``.

    ``global_group_id`` is the eval grouping: the global-identity index in
    ``[0, len(class_names))``; for a detection with no global identity, its own
    tracklet (or, untracked under ``aug_view``, itself) keyed
    ``len(class_names) + id`` so it cannot share an id with an identity.

    Assigns ids, so call it once per indexed detection, in index order.
    """
    if not self.is_member(det):
        return None
    if self.include_untracked:
        # Inference: grouping is a training-only concern.
        return 0, 0
    global_label = _global_identity_label(det, self.track_names_are_global)
    gid = self._class_index.get(global_label) if global_label is not None else None
    track = getattr(det, "track", None)
    track_name = track.name if track is not None else None
    if self.id_scope == "tracklet":
        tid = self._tracklet_vocab.setdefault(
            (labels_idx, video_idx, track_name), len(self._tracklet_vocab)
        )
        # Identity indices and tracklet ids both count up from 0. Shift the
        # tracklet fallback past the identities: returning a bare `tid` made
        # tracklet k and identity k one group in the retrieval metrics.
        return tid, (gid if gid is not None else len(self.class_names) + tid)
    if self.id_scope == "aug_view":
        group_id = self._n_detection_groups
        self._n_detection_groups += 1
        if gid is not None:
            return group_id, gid
        key = (
            (labels_idx, video_idx, track_name)
            if track is not None
            else ("detection", group_id)
        )
        fallback = self._tracklet_vocab.setdefault(key, len(self._tracklet_vocab))
        return group_id, len(self.class_names) + fallback
    return gid, gid

is_member(det)

Whether a detection is a sample (side-effect free).

Source code in sleap_nn/data/custom_datasets.py

def is_member(self, det) -> bool:
    """Whether a detection is a sample (side-effect free)."""
    return self.skip_reason(det) is None

skip_reason(det)

Why a detection is not a member, or None if it is one.

Source code in sleap_nn/data/custom_datasets.py

def skip_reason(self, det) -> Optional[str]:
    """Why a detection is not a member, or ``None`` if it is one."""
    if id(det) in self._superseded:
        return "superseded by a user label (data_config.user_instances_only)"
    if isinstance(det, sio.SegmentationMask):
        if _mask_image_area(det) < max(self.min_mask_area, 1.0):
            return "empty mask" if det.area == 0 else "tiny mask"
    elif getattr(det, "is_empty", False):
        return "empty pose"
    if self.include_untracked or self.id_scope == "aug_view":
        return None
    if self.id_scope == "tracklet":
        return None if getattr(det, "track", None) is not None else "no track"
    label = _global_identity_label(det, self.track_names_are_global)
    if label is None or label not in self._class_index:
        return "no global identity"
    return None

GroupAwareBatchSampler

Bases: Sampler

Group-aware batch sampler for contrastive embedding training.

Its only job is to make the wanted positives/negatives co-occur in a batch. Modes: - pk: P groups x K crops (the contrastive-standard sampler; guarantees K positives per group). - within_video: pick ONE video (among those holding >= 2 groups, in proportion to their crops), then min(P, groups in it) tracks x K crops from it — every in-batch pair is same-video so its relationship is KNOWN (the correct video-local sampler; cross-video pairs never co-occur). A batch is therefore smaller than P x K when its video holds fewer than P groups. - random: plain random batch (aug-view-only / self-supervised objectives).

Yields lists of dataset indices. __len__ = batches_per_epoch.

Methods:

Name	Description
__init__	Initialize the sampler from the dataset's per-crop arrays.
__iter__	Yield batches_per_epoch lists of dataset indices.
__len__	Number of batches per epoch.

Attributes:

Name	Type	Description
samples_per_batch	int	Crops one batch draws, on average.

Source code in sleap_nn/data/custom_datasets.py

class GroupAwareBatchSampler(torch.utils.data.Sampler):
    """Group-aware batch sampler for contrastive embedding training.

    Its only job is to make the wanted positives/negatives co-occur in a batch. Modes:
      - ``pk``: P groups x K crops (the contrastive-standard sampler; guarantees K
        positives per group).
      - ``within_video``: pick ONE video (among those holding >= 2 groups, in
        proportion to their crops), then ``min(P, groups in it)`` tracks x K crops
        from it — every in-batch pair is same-video so its relationship is KNOWN (the
        correct video-local sampler; cross-video pairs never co-occur). A batch is
        therefore smaller than ``P x K`` when its video holds fewer than P groups.
      - ``random``: plain random batch (aug-view-only / self-supervised objectives).

    Yields lists of dataset indices. ``__len__`` = ``batches_per_epoch``.
    """

    def __init__(
        self,
        group_ids: np.ndarray,
        video_ids: np.ndarray,
        frame_ids: np.ndarray,
        kind: str = "pk",
        P: int = 8,
        K: int = 16,
        batches_per_epoch: Optional[int] = None,
        seed: int = 0,
        rank: int = 0,
        world_size: int = 1,
    ):
        """Initialize the sampler from the dataset's per-crop arrays.

        Under multi-GPU (DDP) training, ``rank`` / ``world_size`` make each replica
        draw a DIFFERENT batch stream (the RNG is seeded per ``(seed, rank)``). With the
        per-replica ``batches_per_epoch``, the replicas process distinct batches whose
        gradients are all-reduced — a genuinely larger effective batch instead of every
        rank recomputing the identical batch. ``rank=0`` / ``world_size=1`` reproduces
        the single-GPU stream exactly.
        """
        if kind not in EMBEDDING_SAMPLER_KINDS:
            raise ValueError(
                f"Unknown sampler kind: {kind!r}; choose one of "
                f"{'|'.join(EMBEDDING_SAMPLER_KINDS)}."
            )
        self.group_ids = np.asarray(group_ids)
        self.video_ids = np.asarray(video_ids)
        self.frame_ids = np.asarray(frame_ids)
        self.kind = kind
        self.P = P
        self.K = K
        self.rank = int(rank)
        self.world_size = int(world_size)
        # Per-rank stream: offsetting the seed by the rank gives each replica an
        # independent batch sequence (SeedSequence decorrelates adjacent seeds), while
        # rank 0 reproduces the single-GPU stream (seed + 0 == seed) byte-for-byte.
        self.rng = np.random.default_rng(seed + self.rank)
        n = len(self.group_ids)
        self.all_idx = np.arange(n)

        self.uniq_groups = np.unique(self.group_ids)
        self.by_group = {g: self.all_idx[self.group_ids == g] for g in self.uniq_groups}

        self.uniq_videos = np.unique(self.video_ids)
        self.groups_in_video = {}
        self.by_video_group = {}
        for v in self.uniq_videos:
            vmask = self.video_ids == v
            vg = np.unique(self.group_ids[vmask])
            self.groups_in_video[v] = vg
            for g in vg:
                self.by_video_group[(v, g)] = self.all_idx[
                    vmask & (self.group_ids == g)
                ]
        elig = [v for v in self.uniq_videos if len(self.groups_in_video[v]) >= 2]
        self.elig_videos = np.array(elig) if elig else self.uniq_videos
        w = np.array([(self.video_ids == v).sum() for v in self.elig_videos], float)
        self.video_w = w / w.sum()

        self.batches_per_epoch = batches_per_epoch or max(
            1, int(np.ceil(n / self.samples_per_batch))
        )

    @property
    def samples_per_batch(self) -> int:
        """Crops one batch draws, on average.

        ``pk``: ``min(P, groups) x K``. ``random``: ``min(P x K, n)``.
        ``within_video``: each batch holds ``min(P, groups in its video) x K`` crops,
        so this is that size averaged over the videos the way they are drawn (in
        proportion to their crops). A 2-animal video yields ``2 x K`` crops, not
        ``P x K``: counting it like ``pk`` made an "epoch" cover only a fraction of
        the data (a quarter with ``P=8``).

        This is the unit that converts a dataset size into a number of training
        steps (see :func:`get_train_samples_per_step`).
        """
        if self.kind == "random":
            return max(1, min(self.P * self.K, len(self.all_idx)))
        if self.kind == "within_video" and len(self.elig_videos):
            groups = np.array(
                [min(self.P, len(self.groups_in_video[v])) for v in self.elig_videos],
                float,
            )
            return max(1, int(round(float(self.video_w @ groups) * self.K)))
        return max(1, min(self.P, len(self.uniq_groups)) * self.K)

    def __len__(self) -> int:
        """Number of batches per epoch."""
        return self.batches_per_epoch

    def _pk_batch(self):
        P = min(self.P, len(self.uniq_groups))
        groups = self.rng.choice(self.uniq_groups, size=P, replace=False)
        batch = []
        for g in groups:
            pool = self.by_group[g]
            replace = len(pool) < self.K
            batch.extend(self.rng.choice(pool, size=self.K, replace=replace).tolist())
        return batch

    def _within_video_batch(self):
        v = self.elig_videos[self.rng.choice(len(self.elig_videos), p=self.video_w)]
        gs = self.groups_in_video[v]
        P = min(self.P, len(gs))
        groups = self.rng.choice(gs, size=P, replace=False)
        batch = []
        for g in groups:
            pool = self.by_video_group[(v, g)]
            replace = len(pool) < self.K
            batch.extend(self.rng.choice(pool, size=self.K, replace=replace).tolist())
        return batch

    def _random_batch(self):
        size = min(self.P * self.K, len(self.all_idx))
        return self.rng.choice(self.all_idx, size=size, replace=False).tolist()

    def __iter__(self) -> Iterator:
        """Yield ``batches_per_epoch`` lists of dataset indices."""
        for _ in range(self.batches_per_epoch):
            if self.kind == "pk":
                yield self._pk_batch()
            elif self.kind == "within_video":
                yield self._within_video_batch()
            elif self.kind == "random":
                yield self._random_batch()
            else:
                raise ValueError(f"Unknown sampler kind: {self.kind}")

samples_per_batch property

Crops one batch draws, on average.

pk: min(P, groups) x K. random: min(P x K, n). within_video: each batch holds min(P, groups in its video) x K crops, so this is that size averaged over the videos the way they are drawn (in proportion to their crops). A 2-animal video yields 2 x K crops, not P x K: counting it like pk made an "epoch" cover only a fraction of the data (a quarter with P=8).

This is the unit that converts a dataset size into a number of training steps (see :func:get_train_samples_per_step).

__init__(group_ids, video_ids, frame_ids, kind='pk', P=8, K=16, batches_per_epoch=None, seed=0, rank=0, world_size=1)

Initialize the sampler from the dataset's per-crop arrays.

Under multi-GPU (DDP) training, rank / world_size make each replica draw a DIFFERENT batch stream (the RNG is seeded per (seed, rank)). With the per-replica batches_per_epoch, the replicas process distinct batches whose gradients are all-reduced — a genuinely larger effective batch instead of every rank recomputing the identical batch. rank=0 / world_size=1 reproduces the single-GPU stream exactly.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    group_ids: np.ndarray,
    video_ids: np.ndarray,
    frame_ids: np.ndarray,
    kind: str = "pk",
    P: int = 8,
    K: int = 16,
    batches_per_epoch: Optional[int] = None,
    seed: int = 0,
    rank: int = 0,
    world_size: int = 1,
):
    """Initialize the sampler from the dataset's per-crop arrays.

    Under multi-GPU (DDP) training, ``rank`` / ``world_size`` make each replica
    draw a DIFFERENT batch stream (the RNG is seeded per ``(seed, rank)``). With the
    per-replica ``batches_per_epoch``, the replicas process distinct batches whose
    gradients are all-reduced — a genuinely larger effective batch instead of every
    rank recomputing the identical batch. ``rank=0`` / ``world_size=1`` reproduces
    the single-GPU stream exactly.
    """
    if kind not in EMBEDDING_SAMPLER_KINDS:
        raise ValueError(
            f"Unknown sampler kind: {kind!r}; choose one of "
            f"{'|'.join(EMBEDDING_SAMPLER_KINDS)}."
        )
    self.group_ids = np.asarray(group_ids)
    self.video_ids = np.asarray(video_ids)
    self.frame_ids = np.asarray(frame_ids)
    self.kind = kind
    self.P = P
    self.K = K
    self.rank = int(rank)
    self.world_size = int(world_size)
    # Per-rank stream: offsetting the seed by the rank gives each replica an
    # independent batch sequence (SeedSequence decorrelates adjacent seeds), while
    # rank 0 reproduces the single-GPU stream (seed + 0 == seed) byte-for-byte.
    self.rng = np.random.default_rng(seed + self.rank)
    n = len(self.group_ids)
    self.all_idx = np.arange(n)

    self.uniq_groups = np.unique(self.group_ids)
    self.by_group = {g: self.all_idx[self.group_ids == g] for g in self.uniq_groups}

    self.uniq_videos = np.unique(self.video_ids)
    self.groups_in_video = {}
    self.by_video_group = {}
    for v in self.uniq_videos:
        vmask = self.video_ids == v
        vg = np.unique(self.group_ids[vmask])
        self.groups_in_video[v] = vg
        for g in vg:
            self.by_video_group[(v, g)] = self.all_idx[
                vmask & (self.group_ids == g)
            ]
    elig = [v for v in self.uniq_videos if len(self.groups_in_video[v]) >= 2]
    self.elig_videos = np.array(elig) if elig else self.uniq_videos
    w = np.array([(self.video_ids == v).sum() for v in self.elig_videos], float)
    self.video_w = w / w.sum()

    self.batches_per_epoch = batches_per_epoch or max(
        1, int(np.ceil(n / self.samples_per_batch))
    )

__iter__()

Yield batches_per_epoch lists of dataset indices.

Source code in sleap_nn/data/custom_datasets.py

def __iter__(self) -> Iterator:
    """Yield ``batches_per_epoch`` lists of dataset indices."""
    for _ in range(self.batches_per_epoch):
        if self.kind == "pk":
            yield self._pk_batch()
        elif self.kind == "within_video":
            yield self._within_video_batch()
        elif self.kind == "random":
            yield self._random_batch()
        else:
            raise ValueError(f"Unknown sampler kind: {self.kind}")

__len__()

Number of batches per epoch.

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Number of batches per epoch."""
    return self.batches_per_epoch

InfiniteDataLoader

Bases: DataLoader

Dataloader that reuses workers for infinite iteration.

This dataloader extends the PyTorch DataLoader to provide infinite recycling of workers, which improves efficiency for training loops that need to iterate through the dataset multiple times without recreating workers.

Attributes:

Name	Type	Description
batch_sampler	_RepeatSampler	A sampler that repeats indefinitely.
iterator	Iterator	The iterator from the parent DataLoader.
len_dataloader	Optional[int]	Number of minibatches to be generated. If None, this is set to len(dataset)/batch_size.

Methods:

Name	Description
__len__	Return the length of the batch sampler's sampler.
__iter__	Create a sampler that repeats indefinitely.
__del__	Ensure workers are properly terminated.
reset	Reset the iterator, useful when modifying dataset settings during training.

Examples:

Create an infinite dataloader for training

>>> dataset = CenteredInstanceDataset(...)
>>> dataloader = InfiniteDataLoader(dataset, batch_size=16, shuffle=True)
>>> for batch in dataloader:  # Infinite iteration
>>>     train_step(batch)

Source: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/build.py

Source code in sleap_nn/data/custom_datasets.py

class InfiniteDataLoader(DataLoader):
    """Dataloader that reuses workers for infinite iteration.

    This dataloader extends the PyTorch DataLoader to provide infinite recycling of workers, which improves efficiency
    for training loops that need to iterate through the dataset multiple times without recreating workers.

    Attributes:
        batch_sampler (_RepeatSampler): A sampler that repeats indefinitely.
        iterator (Iterator): The iterator from the parent DataLoader.
        len_dataloader (Optional[int]): Number of minibatches to be generated. If `None`, this is set to len(dataset)/batch_size.

    Methods:
        __len__: Return the length of the batch sampler's sampler.
        __iter__: Create a sampler that repeats indefinitely.
        __del__: Ensure workers are properly terminated.
        reset: Reset the iterator, useful when modifying dataset settings during training.

    Examples:
        Create an infinite dataloader for training
        >>> dataset = CenteredInstanceDataset(...)
        >>> dataloader = InfiniteDataLoader(dataset, batch_size=16, shuffle=True)
        >>> for batch in dataloader:  # Infinite iteration
        >>>     train_step(batch)

    Source: https://github.com/ultralytics/ultralytics/blob/main/ultralytics/data/build.py
    """

    def __init__(self, len_dataloader: Optional[int] = None, *args: Any, **kwargs: Any):
        """Initialize the InfiniteDataLoader with the same arguments as DataLoader."""
        super().__init__(*args, **kwargs)
        object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler))
        self.iterator = super().__iter__()
        self.len_dataloader = len_dataloader

    def __len__(self) -> int:
        """Return the length of the batch sampler's sampler."""
        # set the len to required number of steps per epoch as Lightning Trainer
        # doesn't use the `__iter__` directly but instead uses the length to set
        # the number of steps per epoch. If this is just set to len(sampler), then
        # it only iterates through the samples in the dataset (and doesn't cycle through)
        # if the required steps per epoch is more than batches in dataset.
        return (
            self.len_dataloader
            if self.len_dataloader is not None
            else len(self.batch_sampler.sampler)
        )

    def __iter__(self) -> Iterator:
        """Create an iterator that yields indefinitely from the underlying iterator."""
        while True:
            yield next(self.iterator)

    def __del__(self):
        """Ensure that workers are properly terminated when the dataloader is deleted."""
        try:
            if not hasattr(self.iterator, "_workers"):
                return
            for w in self.iterator._workers:  # force terminate
                if w.is_alive():
                    w.terminate()
            self.iterator._shutdown_workers()  # cleanup
        except Exception:
            pass

    def reset(self):
        """Reset the iterator to allow modifications to the dataset during training."""
        self.iterator = self._get_iterator()

__del__()

Ensure that workers are properly terminated when the dataloader is deleted.

Source code in sleap_nn/data/custom_datasets.py

def __del__(self):
    """Ensure that workers are properly terminated when the dataloader is deleted."""
    try:
        if not hasattr(self.iterator, "_workers"):
            return
        for w in self.iterator._workers:  # force terminate
            if w.is_alive():
                w.terminate()
        self.iterator._shutdown_workers()  # cleanup
    except Exception:
        pass

__init__(len_dataloader=None, *args, **kwargs)

Initialize the InfiniteDataLoader with the same arguments as DataLoader.

Source code in sleap_nn/data/custom_datasets.py

def __init__(self, len_dataloader: Optional[int] = None, *args: Any, **kwargs: Any):
    """Initialize the InfiniteDataLoader with the same arguments as DataLoader."""
    super().__init__(*args, **kwargs)
    object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler))
    self.iterator = super().__iter__()
    self.len_dataloader = len_dataloader

__iter__()

Create an iterator that yields indefinitely from the underlying iterator.

Source code in sleap_nn/data/custom_datasets.py

def __iter__(self) -> Iterator:
    """Create an iterator that yields indefinitely from the underlying iterator."""
    while True:
        yield next(self.iterator)

__len__()

Return the length of the batch sampler's sampler.

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Return the length of the batch sampler's sampler."""
    # set the len to required number of steps per epoch as Lightning Trainer
    # doesn't use the `__iter__` directly but instead uses the length to set
    # the number of steps per epoch. If this is just set to len(sampler), then
    # it only iterates through the samples in the dataset (and doesn't cycle through)
    # if the required steps per epoch is more than batches in dataset.
    return (
        self.len_dataloader
        if self.len_dataloader is not None
        else len(self.batch_sampler.sampler)
    )

reset()

Reset the iterator to allow modifications to the dataset during training.

Source code in sleap_nn/data/custom_datasets.py

def reset(self):
    """Reset the iterator to allow modifications to the dataset during training."""
    self.iterator = self._get_iterator()

ParallelCacheFiller

Parallel implementation of image caching using thread-local video copies.

This class uses ThreadPoolExecutor to parallelize I/O-bound operations when caching images to disk or memory. Each worker thread gets its own copy of video objects to ensure thread safety.

Attributes:

Name	Type	Description
labels		List of sio.Labels objects containing the data.
lf_idx_list		List of dictionaries with labeled frame indices.
cache_type		Either "disk" or "memory".
cache_path		Path to save cached images (for disk caching).
num_workers		Number of worker threads.

Methods:

Name	Description
__init__	Initialize the parallel cache filler.
fill_cache	Fill the cache in parallel.

Source code in sleap_nn/data/custom_datasets.py

class ParallelCacheFiller:
    """Parallel implementation of image caching using thread-local video copies.

    This class uses ThreadPoolExecutor to parallelize I/O-bound operations when
    caching images to disk or memory. Each worker thread gets its own copy of
    video objects to ensure thread safety.

    Attributes:
        labels: List of sio.Labels objects containing the data.
        lf_idx_list: List of dictionaries with labeled frame indices.
        cache_type: Either "disk" or "memory".
        cache_path: Path to save cached images (for disk caching).
        num_workers: Number of worker threads.
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        lf_idx_list: List[Dict],
        cache_type: str,
        cache_path: Optional[Path] = None,
        num_workers: int = 4,
    ):
        """Initialize the parallel cache filler.

        Args:
            labels: List of sio.Labels objects.
            lf_idx_list: List of sample dictionaries with frame indices.
            cache_type: Either "disk" or "memory".
            cache_path: Path for disk caching.
            num_workers: Number of worker threads.
        """
        self.labels = labels
        self.lf_idx_list = lf_idx_list
        self.cache_type = cache_type
        self.cache_path = cache_path
        self.num_workers = num_workers

        self.cache: Dict = {}
        self._cache_lock = threading.Lock()
        self._local = threading.local()
        self._video_info: Dict = {}

        # Prepare video copies for thread-local access
        self._prepare_video_copies()

    def _prepare_video_copies(self):
        """Close original videos and prepare for thread-local copies."""
        for label in self.labels:
            for video in label.videos:
                vid_id = id(video)
                if vid_id not in self._video_info:
                    # Store original state
                    original_open_backend = video.open_backend

                    # Close the video backend
                    video.close()
                    video.open_backend = False

                    self._video_info[vid_id] = {
                        "video": video,
                        "original_open_backend": original_open_backend,
                    }

    def _get_thread_local_video(self, video: sio.Video) -> sio.Video:
        """Get or create a thread-local video copy.

        Args:
            video: The original video object.

        Returns:
            A thread-local copy of the video that is safe to use.
        """
        vid_id = id(video)

        if not hasattr(self._local, "videos"):
            self._local.videos = {}

        if vid_id not in self._local.videos:
            # Create a thread-local copy
            video_copy = deepcopy(video)
            video_copy.open_backend = True
            self._local.videos[vid_id] = video_copy

        return self._local.videos[vid_id]

    def _process_sample(
        self, sample: Dict
    ) -> Tuple[int, int, Optional[np.ndarray], Optional[str]]:
        """Process a single sample (read image, optionally save/cache).

        Args:
            sample: Dictionary with labels_idx, lf_idx, etc.

        Returns:
            Tuple of (labels_idx, lf_idx, image_or_none, error_or_none).
        """
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]

        try:
            if sample.get("is_negative", False):
                # Negative frames: read directly from video by index
                video_idx = sample["video_idx"]
                frame_idx = sample["frame_idx"]
                video = self._get_thread_local_video(
                    self.labels[labels_idx].videos[video_idx]
                )
                img = video[frame_idx]
            else:
                # Positive frames: read from labeled frame
                lf = self.labels[labels_idx][lf_idx]
                video = self._get_thread_local_video(lf.video)
                img = video[lf.frame_idx]

            if img.shape[-1] == 1:
                img = np.squeeze(img)

            if self.cache_type == "disk":
                f_name = self.cache_path / f"sample_{labels_idx}_{lf_idx}.jpg"
                Image.fromarray(img).save(str(f_name), format="JPEG")
                return labels_idx, lf_idx, None, None
            elif self.cache_type == "memory":
                return labels_idx, lf_idx, img, None

        except Exception as e:
            return labels_idx, lf_idx, None, f"{type(e).__name__}: {str(e)}"

    def fill_cache(
        self, progress_callback=None
    ) -> Tuple[Dict, List[Tuple[int, int, str]]]:
        """Fill the cache in parallel.

        Args:
            progress_callback: Optional callback(completed_count) for progress updates.

        Returns:
            Tuple of (cache_dict, list_of_errors).
        """
        errors = []
        completed = 0

        with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
            futures = {
                executor.submit(self._process_sample, sample): sample
                for sample in self.lf_idx_list
            }

            for future in as_completed(futures):
                labels_idx, lf_idx, img, error = future.result()

                if error:
                    errors.append((labels_idx, lf_idx, error))
                elif self.cache_type == "memory" and img is not None:
                    with self._cache_lock:
                        self.cache[(labels_idx, lf_idx)] = img

                completed += 1
                if progress_callback:
                    progress_callback(completed)

        # Restore original video states
        self._restore_videos()

        return self.cache, errors

    def _restore_videos(self):
        """Restore original video states after caching is complete."""
        for vid_info in self._video_info.values():
            video = vid_info["video"]
            video.open_backend = vid_info["original_open_backend"]
            if video.open_backend:
                try:
                    video.open()
                except Exception:
                    pass

__init__(labels, lf_idx_list, cache_type, cache_path=None, num_workers=4)

Initialize the parallel cache filler.

Parameters:

Name	Type	Description	Default
labels	List[Labels]	List of sio.Labels objects.	required
lf_idx_list	List[Dict]	List of sample dictionaries with frame indices.	required
cache_type	str	Either "disk" or "memory".	required
cache_path	Optional[Path]	Path for disk caching.	None
num_workers	int	Number of worker threads.	4

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    lf_idx_list: List[Dict],
    cache_type: str,
    cache_path: Optional[Path] = None,
    num_workers: int = 4,
):
    """Initialize the parallel cache filler.

    Args:
        labels: List of sio.Labels objects.
        lf_idx_list: List of sample dictionaries with frame indices.
        cache_type: Either "disk" or "memory".
        cache_path: Path for disk caching.
        num_workers: Number of worker threads.
    """
    self.labels = labels
    self.lf_idx_list = lf_idx_list
    self.cache_type = cache_type
    self.cache_path = cache_path
    self.num_workers = num_workers

    self.cache: Dict = {}
    self._cache_lock = threading.Lock()
    self._local = threading.local()
    self._video_info: Dict = {}

    # Prepare video copies for thread-local access
    self._prepare_video_copies()

fill_cache(progress_callback=None)

Fill the cache in parallel.

Parameters:

Name	Type	Description	Default
progress_callback		Optional callback(completed_count) for progress updates.	None

Returns:

Type	Description
Tuple[Dict, List[Tuple[int, int, str]]]	Tuple of (cache_dict, list_of_errors).

Source code in sleap_nn/data/custom_datasets.py

def fill_cache(
    self, progress_callback=None
) -> Tuple[Dict, List[Tuple[int, int, str]]]:
    """Fill the cache in parallel.

    Args:
        progress_callback: Optional callback(completed_count) for progress updates.

    Returns:
        Tuple of (cache_dict, list_of_errors).
    """
    errors = []
    completed = 0

    with ThreadPoolExecutor(max_workers=self.num_workers) as executor:
        futures = {
            executor.submit(self._process_sample, sample): sample
            for sample in self.lf_idx_list
        }

        for future in as_completed(futures):
            labels_idx, lf_idx, img, error = future.result()

            if error:
                errors.append((labels_idx, lf_idx, error))
            elif self.cache_type == "memory" and img is not None:
                with self._cache_lock:
                    self.cache[(labels_idx, lf_idx)] = img

            completed += 1
            if progress_callback:
                progress_callback(completed)

    # Restore original video states
    self._restore_videos()

    return self.cache, errors

SemanticSegmentationDataset

Bases: BaseDataset

Dataset class for whole-frame semantic (foreground/background) segmentation.

Loads per-instance segmentation masks from LabeledFrame.masks and reduces them to a SINGLE binary foreground mask per frame (the union of every instance mask, area-downsampled to the segmentation head's output stride). There is NO instance grouping: unlike :class:BottomUpSegmentationDataset, no center heatmap, per-pixel offset field, or offset weight mask is generated — the model predicts one whole-frame foreground channel and decoding is a plain threshold.

Masks are captured into the sample index at construction time (mirroring how keypoint instances are captured for caching), so __getitem__ never needs a live Labels handle — this keeps it correct under the memory/disk image caching paths (where self.labels_list is None).

Note

Augmentation: intensity aug is applied to the image; geometric aug (rotation/scale/translate/flip) co-transforms every per-instance mask with the SAME affine matrix as the image (nearest-neighbor, re-binarized) at the preprocessed resolution, before the foreground mask is derived. Erase/mixup stay image-only. Mask resizing to the preprocessed image size handles scaling but is not pad-aware; for v1 train with scale=1.0 and input dims divisible by max_stride (and prefer small rotation ranges, since a full-frame rotation can clip instances at the frame edge, as it does for bottom-up pose).

Attributes:

Name	Type	Description
seg_head_config		Configuration for the segmentation head (output_stride).

Methods:

Name	Description
__getitem__	Return dict with image and whole-frame foreground mask for given index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class SemanticSegmentationDataset(BaseDataset):
    """Dataset class for whole-frame semantic (foreground/background) segmentation.

    Loads per-instance segmentation masks from ``LabeledFrame.masks`` and reduces
    them to a SINGLE binary foreground mask per frame (the union of every instance
    mask, area-downsampled to the segmentation head's output stride). There is NO
    instance grouping: unlike :class:`BottomUpSegmentationDataset`, no center
    heatmap, per-pixel offset field, or offset weight mask is generated — the model
    predicts one whole-frame foreground channel and decoding is a plain threshold.

    Masks are captured into the sample index at construction time (mirroring how
    keypoint instances are captured for caching), so ``__getitem__`` never needs
    a live ``Labels`` handle — this keeps it correct under the memory/disk image
    caching paths (where ``self.labels_list`` is ``None``).

    Note:
        Augmentation: intensity aug is applied to the image; geometric aug
        (rotation/scale/translate/flip) co-transforms every per-instance mask with the
        SAME affine matrix as the image (nearest-neighbor, re-binarized) at the
        preprocessed resolution, before the foreground mask is derived. Erase/mixup
        stay image-only. Mask resizing to the preprocessed image size handles scaling
        but is not pad-aware; for v1 train with ``scale=1.0`` and input dims divisible
        by ``max_stride`` (and prefer small rotation ranges, since a full-frame rotation
        can clip instances at the frame edge, as it does for bottom-up pose).

    Attributes:
        seg_head_config: Configuration for the segmentation head (``output_stride``).
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        seg_head_config: DictConfig,
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
    ) -> None:
        """Initialize class attributes."""
        self.seg_head_config = seg_head_config
        # Segmentation never uses negative frames (degenerate num_nodes/instances).
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=False,
        )

    def _apply_common_preprocessing(self, sample: Dict) -> Dict:
        """Apply common preprocessing with geometric augmentation deferred.

        Geometric augmentation must co-transform the segmentation masks, but those
        masks are not present in ``sample`` here (they are loaded separately), so the
        base method would warp only the image. We disable geometric aug for the base
        call (intensity aug still applies) and re-apply it in ``__getitem__`` once the
        masks have been resized to the preprocessed resolution, co-transforming image
        and masks with the same matrix.

        Args:
            sample: Sample dict with at least ``image`` and ``instances`` keys.

        Returns:
            The sample dict with preprocessing applied in-place.
        """
        saved_geometric_aug = self.geometric_aug
        self.geometric_aug = None
        try:
            sample = super()._apply_common_preprocessing(sample)
        finally:
            self.geometric_aug = saved_geometric_aug
        return sample

    def _get_lf_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return samples for frames that have segmentation masks.

        Overrides the base class to index frames by their masks rather than
        keypoint instances. The decoded mask arrays are captured into each
        sample so ``__getitem__`` does not depend on a live ``Labels`` handle.
        """
        lf_idx_list = []
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                lf_masks = getattr(lf, "masks", None)
                if not lf_masks:
                    continue
                # Scale-aware decode: masks written by the segmentation inference
                # layer are encoded at output-stride (non-identity scale); decode
                # them up to the IMAGE-pixel grid so self-training / pseudo-label
                # ``.slp`` files yield correctly-scaled targets (a stride-res
                # ``m.data`` would silently mis-scale the GT, since __getitem__'s
                # resize branch only fires on a preprocessing size change). Scale-1
                # GT masks take the zero-copy fast path.
                from sleap_nn.inference.segmentation_convert import (
                    decode_mask_to_image_res,
                )

                mask_arrays = [decode_mask_to_image_res(m) for m in lf_masks]
                if len(mask_arrays) == 0:
                    continue
                video_idx = label.videos.index(lf.video)
                sample = {
                    "labels_idx": labels_idx,
                    "lf_idx": lf_idx,
                    "video_idx": video_idx,
                    "frame_idx": lf.frame_idx,
                    "is_negative": False,
                    "instances": None,
                    "masks": mask_arrays,
                }
                lf_idx_list.append(sample)

        return lf_idx_list

    def __getitem__(self, index) -> Dict:
        """Return dict with image and whole-frame foreground mask for given index."""
        sample = self.lf_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        video_idx = sample["video_idx"]
        frame_idx = sample["frame_idx"]

        # Load image
        if self.cache_img is not None:
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            img = lf.image

        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
        image = np.expand_dims(image, axis=0)  # (1, C, H, W)
        image = torch.from_numpy(image.copy())

        # Dummy instances tensor (needed for preprocessing compatibility)
        instances = torch.zeros((1, 1, 1, 2), dtype=torch.float32)

        sample_dict = {
            "image": image,
            "instances": instances,
            "video_idx": torch.tensor(video_idx, dtype=torch.int32),
            "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
            "orig_size": torch.Tensor([image.shape[-2], image.shape[-1]]).unsqueeze(0),
            "num_instances": 0,
        }

        # Masks captured at index-build time (decoded bool arrays at orig res).
        # Place them on the full-frame image grid as one float (1, K, H, W) tensor
        # so they ride the IDENTICAL size-match / scale / stride-pad chain as the
        # image inside ``_apply_common_preprocessing`` (kept as a float tensor
        # end-to-end and binarized ONCE just before target generation). This is
        # what keeps the whole-frame target registered to the padded image and
        # makes ragged/offset-carrying decoded masks well-defined (see
        # ``_masks_to_frame_canvas``).
        mask_arrays = [np.asarray(m, dtype=bool) for m in sample["masks"]]
        orig_img_hw = (image.shape[-2], image.shape[-1])
        sample_dict["masks"] = _masks_to_frame_canvas(mask_arrays, orig_img_hw)

        # Apply common preprocessing (RGB/grayscale, size matching, scaling, padding).
        # Co-transforms ``sample_dict["masks"]`` with the same geometry as the image.
        sample_dict = self._apply_common_preprocessing(sample_dict)

        img_hw = sample_dict["image"].shape[-2:]

        # Geometric augmentation: co-transform the per-instance masks with the SAME
        # flip/affine matrix as the image (nearest-neighbor). Applied here (post
        # size-match / resize / pad) so image and masks share a resolution, and
        # BEFORE the foreground mask is generated so it is derived from the augmented
        # masks. Semantic segmentation has no keypoints, so a dummy instances tensor
        # rides along; erase/mixup stay image-only.
        if (
            self.apply_aug
            and self.geometric_aug is not None
            and sample_dict["masks"].shape[1] > 0
        ):
            (
                sample_dict["image"],
                _,
                sample_dict["masks"],
            ) = apply_geometric_augmentation(
                sample_dict["image"],
                torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                masks=sample_dict["masks"],
                **self.geometric_aug,
            )

        # Single re-binarization to bool arrays right before target generation.
        masks_t = sample_dict.pop("masks")
        mask_arrays = [masks_t[0, k].numpy() > 0.5 for k in range(masks_t.shape[1])]

        # Generate the single whole-frame foreground mask (union of all instance
        # masks, area-downsampled + re-binarized to the seg head's output stride).
        # No center/offset/weight targets: decoding is a plain threshold with no
        # instance grouping.
        foreground_mask = generate_foreground_mask(
            mask_arrays,
            img_hw=img_hw,
            output_stride=self.seg_head_config.output_stride,
            maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
        )

        sample_dict["foreground_mask"] = foreground_mask
        sample_dict["labels_idx"] = labels_idx

        return sample_dict

__getitem__(index)

Return dict with image and whole-frame foreground mask for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image and whole-frame foreground mask for given index."""
    sample = self.lf_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    video_idx = sample["video_idx"]
    frame_idx = sample["frame_idx"]

    # Load image
    if self.cache_img is not None:
        if self.cache_img == "disk":
            img = np.array(
                Image.open(
                    f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                )
            )
        elif self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)].copy()
    else:
        lf = self.labels_list[labels_idx][lf_idx]
        img = lf.image

    if img.ndim == 2:
        img = np.expand_dims(img, axis=2)

    image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
    image = np.expand_dims(image, axis=0)  # (1, C, H, W)
    image = torch.from_numpy(image.copy())

    # Dummy instances tensor (needed for preprocessing compatibility)
    instances = torch.zeros((1, 1, 1, 2), dtype=torch.float32)

    sample_dict = {
        "image": image,
        "instances": instances,
        "video_idx": torch.tensor(video_idx, dtype=torch.int32),
        "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
        "orig_size": torch.Tensor([image.shape[-2], image.shape[-1]]).unsqueeze(0),
        "num_instances": 0,
    }

    # Masks captured at index-build time (decoded bool arrays at orig res).
    # Place them on the full-frame image grid as one float (1, K, H, W) tensor
    # so they ride the IDENTICAL size-match / scale / stride-pad chain as the
    # image inside ``_apply_common_preprocessing`` (kept as a float tensor
    # end-to-end and binarized ONCE just before target generation). This is
    # what keeps the whole-frame target registered to the padded image and
    # makes ragged/offset-carrying decoded masks well-defined (see
    # ``_masks_to_frame_canvas``).
    mask_arrays = [np.asarray(m, dtype=bool) for m in sample["masks"]]
    orig_img_hw = (image.shape[-2], image.shape[-1])
    sample_dict["masks"] = _masks_to_frame_canvas(mask_arrays, orig_img_hw)

    # Apply common preprocessing (RGB/grayscale, size matching, scaling, padding).
    # Co-transforms ``sample_dict["masks"]`` with the same geometry as the image.
    sample_dict = self._apply_common_preprocessing(sample_dict)

    img_hw = sample_dict["image"].shape[-2:]

    # Geometric augmentation: co-transform the per-instance masks with the SAME
    # flip/affine matrix as the image (nearest-neighbor). Applied here (post
    # size-match / resize / pad) so image and masks share a resolution, and
    # BEFORE the foreground mask is generated so it is derived from the augmented
    # masks. Semantic segmentation has no keypoints, so a dummy instances tensor
    # rides along; erase/mixup stay image-only.
    if (
        self.apply_aug
        and self.geometric_aug is not None
        and sample_dict["masks"].shape[1] > 0
    ):
        (
            sample_dict["image"],
            _,
            sample_dict["masks"],
        ) = apply_geometric_augmentation(
            sample_dict["image"],
            torch.zeros((1, 1, 1, 2), dtype=torch.float32),
            masks=sample_dict["masks"],
            **self.geometric_aug,
        )

    # Single re-binarization to bool arrays right before target generation.
    masks_t = sample_dict.pop("masks")
    mask_arrays = [masks_t[0, k].numpy() > 0.5 for k in range(masks_t.shape[1])]

    # Generate the single whole-frame foreground mask (union of all instance
    # masks, area-downsampled + re-binarized to the seg head's output stride).
    # No center/offset/weight targets: decoding is a plain threshold with no
    # instance grouping.
    foreground_mask = generate_foreground_mask(
        mask_arrays,
        img_hw=img_hw,
        output_stride=self.seg_head_config.output_stride,
        maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
    )

    sample_dict["foreground_mask"] = foreground_mask
    sample_dict["labels_idx"] = labels_idx

    return sample_dict

__init__(labels, seg_head_config, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    seg_head_config: DictConfig,
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
) -> None:
    """Initialize class attributes."""
    self.seg_head_config = seg_head_config
    # Segmentation never uses negative frames (degenerate num_nodes/instances).
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=False,
    )

SemanticSegmentationTiledDataset

Bases: BaseDataset

Whole-frame semantic (fg/bg) segmentation dataset emitting fixed-size tiles.

Tiling analogue of :class:SemanticSegmentationDataset and a fg-only sibling of :class:BottomUpSegmentationTiledDataset: each frame is decomposed into overlapping square tiles (foreground-aware random draws for training, a deterministic grid for validation). A frame is decoded / channel-coerced / scaled once (cached in a per-worker LRU together with its decoded per-instance masks) and reused across all of its tiles. Each tile is then cut out with a constant-zero pad; when geometric augmentation is enabled the tile is taken via a sqrt(2) halo so a rotation has valid context, co-transforming every per-instance mask with the SAME affine matrix as the image (nearest-neighbor, re-binarized).

Unlike the bottom-up center-offset pipeline there is NO instance grouping: the single foreground_mask GT is the union of ALL masks that touch the tile (no centroid-ownership filter, no center heatmap, no offsets, no foreground weight).

Emits one sample per (frame, tile-slot); __len__ is the total number of tile slots. Returned samples match the SemanticSegmentationDataset key contract (plus an int32 tile_origin of shape (2,)), so the default collate and the SemanticSegmentationLightningModule apply with no changes.

Attributes:

Name	Type	Description
seg_head_config		Configuration for the segmentation (foreground) head.

Methods:

Name	Description
__getitem__	Return dict with image + foreground GT for one tile of one frame.
__init__	Initialize class attributes.
__len__	Return the number of tile samples (frames x tiles-per-frame).

Source code in sleap_nn/data/custom_datasets.py

class SemanticSegmentationTiledDataset(BaseDataset):
    """Whole-frame semantic (fg/bg) segmentation dataset emitting fixed-size tiles.

    Tiling analogue of :class:`SemanticSegmentationDataset` and a fg-only sibling of
    :class:`BottomUpSegmentationTiledDataset`: each frame is decomposed into
    overlapping square tiles (foreground-aware random draws for training, a
    deterministic grid for validation). A frame is decoded / channel-coerced / scaled
    once (cached in a per-worker LRU together with its decoded per-instance masks) and
    reused across all of its tiles. Each tile is then cut out with a constant-zero pad;
    when geometric augmentation is enabled the tile is taken via a ``sqrt(2)`` halo so a
    rotation has valid context, co-transforming every per-instance mask with the SAME
    affine matrix as the image (nearest-neighbor, re-binarized).

    Unlike the bottom-up center-offset pipeline there is NO instance grouping: the
    single ``foreground_mask`` GT is the union of ALL masks that touch the tile (no
    centroid-ownership filter, no center heatmap, no offsets, no foreground weight).

    Emits one sample per ``(frame, tile-slot)``; ``__len__`` is the total number of
    tile slots. Returned samples match the ``SemanticSegmentationDataset`` key contract
    (plus an ``int32`` ``tile_origin`` of shape ``(2,)``), so the default collate and
    the ``SemanticSegmentationLightningModule`` apply with no changes.

    Attributes:
        seg_head_config: Configuration for the segmentation (foreground) head.
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        seg_head_config: DictConfig,
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
        tiling: Optional[Union[DictConfig, Any]] = None,
        base_seed: int = 0,
    ) -> None:
        """Initialize class attributes."""
        self.seg_head_config = seg_head_config
        # Segmentation never uses negative frames (degenerate num_nodes/instances).
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=False,
            tiling=tiling,
            output_stride=seg_head_config.output_stride,
            base_seed=base_seed,
        )

        # Per-(frame, tile-slot) descriptors + contiguous per-frame index blocks.
        self.tile_idx_list = self._get_tile_idx_list(labels)
        self.frame_blocks = self._build_frame_blocks(self.tile_idx_list)

    def _get_lf_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return per-frame samples for frames that have segmentation masks.

        Mirrors ``SemanticSegmentationDataset._get_lf_idx_list``: indexes frames by
        their masks (not keypoint instances) and captures the decoded mask arrays so
        ``__getitem__`` never needs a live ``Labels`` handle. This is what the base
        ``__init__`` stores as ``self.lf_idx_list`` (used for image caching);
        ``_get_tile_idx_list`` explodes it into per-tile descriptors.
        """
        from sleap_nn.inference.segmentation_convert import decode_mask_to_image_res

        lf_idx_list = []
        for labels_idx, label in enumerate(labels):
            for lf_idx, lf in enumerate(label):
                lf_masks = getattr(lf, "masks", None)
                if not lf_masks:
                    continue
                # Scale-aware decode up to the IMAGE-pixel grid (scale-1 GT masks
                # take the zero-copy fast path); see SemanticSegmentationDataset.
                mask_arrays = [decode_mask_to_image_res(m) for m in lf_masks]
                if len(mask_arrays) == 0:
                    continue
                video_idx = label.videos.index(lf.video)
                lf_idx_list.append(
                    {
                        "labels_idx": labels_idx,
                        "lf_idx": lf_idx,
                        "video_idx": video_idx,
                        "frame_idx": lf.frame_idx,
                        "is_negative": False,
                        "instances": None,
                        "masks": mask_arrays,
                    }
                )
        return lf_idx_list

    def _get_tile_idx_list(self, labels: List[sio.Labels]) -> List[Dict]:
        """Return per-(frame, tile-slot) descriptors for the tiled seg dataset.

        Mirrors :meth:`BaseDataset._get_tile_idx_list` but keys off the mask-indexed
        per-frame list (``self.lf_idx_list``, built by :meth:`_get_lf_idx_list`) so the
        decoded ``mask_arrays`` ride along on every descriptor. Grid (val) pins one
        descriptor per :func:`generate_tile_grid` origin (sized-frame ``H, W`` taken
        from the decoded mask shape); foreground (train) emits ``samples_per_frame``
        slots with ``tile_origin=None`` (drawn at runtime). A frame's descriptors form
        a contiguous run (the block the sampler groups on).
        """
        tile_idx_list: List[Dict] = []
        for f in self.lf_idx_list:
            mask_arrays = f["masks"]
            # Masks are at image resolution; the sized (post-scale) frame H, W match
            # `apply_resizer`'s int(dim * scale) truncation used by `_frame_sized_hw`.
            mh, mw = mask_arrays[0].shape[:2]
            if self.scale != 1.0:
                sized_hw = (int(mh * self.scale), int(mw * self.scale))
            else:
                sized_hw = (int(mh), int(mw))

            if self.tile_sampling == "grid":
                origins = generate_tile_grid(
                    sized_hw,
                    tile_size=self.tile_size,
                    overlap=self.overlap,
                    output_stride=self.output_stride,
                    max_stride=self.max_stride,
                    min_overlap_fraction=self.min_overlap_fraction,
                )
            else:
                origins = [None] * self.samples_per_frame

            for sample_k, origin in enumerate(origins):
                tile_idx_list.append(
                    {
                        "labels_idx": f["labels_idx"],
                        "lf_idx": f["lf_idx"],
                        "video_idx": f["video_idx"],
                        "frame_idx": f["frame_idx"],
                        "masks": mask_arrays,
                        "sample_k": sample_k,
                        "tile_origin": origin,
                        "is_grid": self.tile_sampling == "grid",
                        "is_negative": False,
                    }
                )
        return tile_idx_list

    def __len__(self) -> int:
        """Return the number of tile samples (frames x tiles-per-frame)."""
        return len(self.tile_idx_list)

    def _read_image(self, d: Dict) -> np.ndarray:
        """Read a frame's raw HWC image (cache/disk/labels), restoring 2D -> 3D."""
        labels_idx = d["labels_idx"]
        lf_idx = d["lf_idx"]
        if self.cache_img is not None:
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            img = lf.image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)
        return img

    def _load_sized_frame(
        self, d: Dict
    ) -> Tuple[torch.Tensor, List[np.ndarray], tuple]:
        """Decode a frame once: channel-coerce + scale the image, size-match masks.

        Returns ``(image, mask_arrays, orig_hw)`` where ``image`` is a sized
        ``(1, C, H, W)`` tensor, ``mask_arrays`` are bool arrays at the sized image
        resolution, and ``orig_hw`` is the raw full-frame ``(H, W)`` before scaling.
        """
        img = self._read_image(d)  # HWC
        orig_hw = (int(img.shape[0]), int(img.shape[1]))

        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW
        image = np.expand_dims(image, axis=0)  # (1, C, H, W)
        image = torch.from_numpy(image.copy())

        if self.ensure_rgb:
            image = convert_to_rgb(image)
        elif self.ensure_grayscale:
            image = convert_to_grayscale(image)

        dummy = torch.zeros((1, 1, 1, 2), dtype=torch.float32)
        image, _ = apply_resizer(image, dummy, scale=self.scale)

        mask_arrays = [np.asarray(m, dtype=bool) for m in d["masks"]]
        sized_hw = (image.shape[-2], image.shape[-1])
        if (sized_hw != orig_hw) and len(mask_arrays) > 0:
            # Resize masks to the sized image resolution with antialiased bilinear,
            # matching the image's ``apply_resizer`` / ``tvf.resize`` above
            # (standardized across every seg mask-geometry resize).
            target_h, target_w = sized_hw
            resized = []
            for m in mask_arrays:
                m_tensor = (
                    torch.from_numpy(m.astype(np.float32)).unsqueeze(0).unsqueeze(0)
                )
                m_resized = F.interpolate(
                    m_tensor,
                    size=(target_h, target_w),
                    mode="bilinear",
                    align_corners=False,
                    antialias=True,
                )
                resized.append(m_resized.squeeze().numpy() > 0.5)
            mask_arrays = resized

        return image, mask_arrays, orig_hw

    def _slice_halo(
        self,
        image: torch.Tensor,
        mask_arrays: List[np.ndarray],
        hy0: int,
        hx0: int,
        side: int,
    ) -> Tuple[torch.Tensor, List[np.ndarray]]:
        """Slice a ``side x side`` window (top-left ``hy0, hx0``) with constant-zero pad.

        Slices both the image and every per-instance mask at the SAME offsets so they
        stay pixel-aligned. Out-of-bounds regions are zero.
        """
        _, C, H, W = image.shape
        ys, xs = max(0, hy0), max(0, hx0)
        ye, xe = min(H, hy0 + side), min(W, hx0 + side)
        win_img = image.new_zeros((1, C, side, side))
        win_masks = [np.zeros((side, side), dtype=bool) for _ in mask_arrays]
        if ye > ys and xe > xs:
            win_img[:, :, ys - hy0 : ye - hy0, xs - hx0 : xe - hx0] = image[
                :, :, ys:ye, xs:xe
            ]
            for k, m in enumerate(mask_arrays):
                win_masks[k][ys - hy0 : ye - hy0, xs - hx0 : xe - hx0] = m[ys:ye, xs:xe]
        return win_img, win_masks

    def __getitem__(self, index) -> Dict:
        """Return dict with image + foreground GT for one tile of one frame."""
        d = self.tile_idx_list[index]
        labels_idx = d["labels_idx"]
        video_idx = d["video_idx"]
        frame_idx = d["frame_idx"]
        epoch = int(self._epoch)
        ts = self.tile_size

        # 1. Decode the full frame once per (labels_idx, lf_idx), via per-worker LRU.
        cached = self._frame_lru().get((labels_idx, d["lf_idx"]))
        if cached is None:
            cached = self._load_sized_frame(d)
            self._frame_lru().put((labels_idx, d["lf_idx"]), cached)
        image, mask_arrays, orig_hw = cached
        sized_hw = (image.shape[-2], image.shape[-1])

        # 2. Resolve the tile origin: pinned for grid/val, drawn for train. The
        #    foreground-aware draw is seeded from instance centroids (placement only;
        #    semantic GT does not use them), so tiles still land on foreground.
        if d["is_grid"]:
            tile_origin = tuple(int(v) for v in d["tile_origin"])
            aug_seed = None
        else:
            rng = np.random.default_rng(
                tile_sample_seed(
                    self.base_seed, epoch, video_idx, frame_idx, d["sample_k"]
                )
            )
            cents = _compute_mask_centroids(mask_arrays)  # list of (x, y)
            centers = (
                torch.tensor(cents, dtype=torch.float32).reshape(-1, 2)
                if len(cents) > 0
                else torch.zeros((0, 2), dtype=torch.float32)
            )
            tile_origin = draw_tile_origin(
                centers,
                sized_hw,
                ts,
                d["sample_k"],
                self.samples_per_frame,
                self.tile_fg_fraction,
                self.center_jitter,
                rng,
            )
            aug_seed = tile_sample_seed(
                self.base_seed, epoch, video_idx, frame_idx, d["sample_k"], salt=1
            )

        y0, x0 = tile_origin
        apply_geo = self.apply_aug and self.geometric_aug is not None

        # 3. Cut the tile out of the frame (image + co-transformed masks). Under
        #    geometric aug, take a sqrt(2) halo centered on the tile center so the
        #    rotation has valid context, augment image + masks with the SAME matrix
        #    (nearest-neighbor, re-binarized), then trim the center tile back out.
        if apply_geo:
            halo = int(math.ceil(ts * math.sqrt(2)))
            hy0 = y0 - (halo - ts) // 2
            hx0 = x0 - (halo - ts) // 2
            halo_img, halo_masks = self._slice_halo(image, mask_arrays, hy0, hx0, halo)

            if len(halo_masks) > 0:
                halo_masks_t = torch.from_numpy(
                    np.stack([hm.astype(np.float32) for hm in halo_masks])
                ).unsqueeze(
                    0
                )  # (1, K, halo, halo)
                # The skia geometric backend samples its transform from the GLOBAL
                # numpy RNG (and torch); seed both so the halo path is reproducible.
                np.random.seed(aug_seed & 0xFFFFFFFF)
                torch.manual_seed(aug_seed)
                halo_img, _, halo_masks_t = apply_geometric_augmentation(
                    halo_img,
                    torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                    masks=halo_masks_t,
                    **dict(self.geometric_aug),
                )
                halo_masks = [
                    halo_masks_t[0, k].numpy() > 0.5
                    for k in range(halo_masks_t.shape[1])
                ]

            # Trim the augmented halo back to `ts`, centered on the halo center:
            # crop_and_resize for the image (codebase convention), an equivalent
            # integer center-slice for the (axis-aligned) mask arrays.
            c = halo / 2.0
            bbox = make_centered_bboxes(
                torch.tensor([[c, c]], dtype=torch.float32), ts, ts
            )
            tile_image = crop_and_resize(halo_img, boxes=bbox, size=(ts, ts))
            off = (halo - ts) // 2
            tile_masks = [hm[off : off + ts, off : off + ts] for hm in halo_masks]
        else:
            # Fast path (no aug): direct slice + constant-zero pad, byte-identical.
            tile_image, tile_masks = self._slice_halo(image, mask_arrays, y0, x0, ts)

        # 4. Intensity aug (image only) + pad to stride (no-op when ts % max_stride==0).
        if self.apply_aug and self.intensity_aug is not None:
            tile_image, _ = apply_intensity_augmentation(
                tile_image,
                torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                **self.intensity_aug,
            )
        tile_image = apply_pad_to_stride(tile_image, max_stride=self.max_stride)

        # 5. Foreground GT: union of ALL masks touching the tile (NO ownership filter,
        #    NO center/offset heads). Empty (fully-zero) tile masks contribute nothing.
        img_hw = tile_image.shape[-2:]
        foreground_mask = generate_foreground_mask(
            tile_masks,
            img_hw=img_hw,
            output_stride=self.seg_head_config.output_stride,
            maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
        )
        # Count of masks with any foreground in this tile (logging/eval convenience).
        num_masks_in_tile = sum(1 for m in tile_masks if int(m.sum()) > 0)

        return {
            "image": tile_image,
            "instances": torch.zeros((1, 1, 1, 2), dtype=torch.float32),
            "video_idx": torch.tensor(video_idx, dtype=torch.int32),
            "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
            "orig_size": torch.Tensor([orig_hw[0], orig_hw[1]]).unsqueeze(0),
            "num_instances": num_masks_in_tile,
            # Tiles are extracted in the model's input space (sizematcher bypassed),
            # so the effective scale is 1.0 (matches SingleInstanceTiledDataset).
            "eff_scale": torch.tensor(1.0, dtype=torch.float32),
            "foreground_mask": foreground_mask,
            "labels_idx": labels_idx,
            "tile_origin": torch.tensor(tile_origin, dtype=torch.int32),
        }

__getitem__(index)

Return dict with image + foreground GT for one tile of one frame.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image + foreground GT for one tile of one frame."""
    d = self.tile_idx_list[index]
    labels_idx = d["labels_idx"]
    video_idx = d["video_idx"]
    frame_idx = d["frame_idx"]
    epoch = int(self._epoch)
    ts = self.tile_size

    # 1. Decode the full frame once per (labels_idx, lf_idx), via per-worker LRU.
    cached = self._frame_lru().get((labels_idx, d["lf_idx"]))
    if cached is None:
        cached = self._load_sized_frame(d)
        self._frame_lru().put((labels_idx, d["lf_idx"]), cached)
    image, mask_arrays, orig_hw = cached
    sized_hw = (image.shape[-2], image.shape[-1])

    # 2. Resolve the tile origin: pinned for grid/val, drawn for train. The
    #    foreground-aware draw is seeded from instance centroids (placement only;
    #    semantic GT does not use them), so tiles still land on foreground.
    if d["is_grid"]:
        tile_origin = tuple(int(v) for v in d["tile_origin"])
        aug_seed = None
    else:
        rng = np.random.default_rng(
            tile_sample_seed(
                self.base_seed, epoch, video_idx, frame_idx, d["sample_k"]
            )
        )
        cents = _compute_mask_centroids(mask_arrays)  # list of (x, y)
        centers = (
            torch.tensor(cents, dtype=torch.float32).reshape(-1, 2)
            if len(cents) > 0
            else torch.zeros((0, 2), dtype=torch.float32)
        )
        tile_origin = draw_tile_origin(
            centers,
            sized_hw,
            ts,
            d["sample_k"],
            self.samples_per_frame,
            self.tile_fg_fraction,
            self.center_jitter,
            rng,
        )
        aug_seed = tile_sample_seed(
            self.base_seed, epoch, video_idx, frame_idx, d["sample_k"], salt=1
        )

    y0, x0 = tile_origin
    apply_geo = self.apply_aug and self.geometric_aug is not None

    # 3. Cut the tile out of the frame (image + co-transformed masks). Under
    #    geometric aug, take a sqrt(2) halo centered on the tile center so the
    #    rotation has valid context, augment image + masks with the SAME matrix
    #    (nearest-neighbor, re-binarized), then trim the center tile back out.
    if apply_geo:
        halo = int(math.ceil(ts * math.sqrt(2)))
        hy0 = y0 - (halo - ts) // 2
        hx0 = x0 - (halo - ts) // 2
        halo_img, halo_masks = self._slice_halo(image, mask_arrays, hy0, hx0, halo)

        if len(halo_masks) > 0:
            halo_masks_t = torch.from_numpy(
                np.stack([hm.astype(np.float32) for hm in halo_masks])
            ).unsqueeze(
                0
            )  # (1, K, halo, halo)
            # The skia geometric backend samples its transform from the GLOBAL
            # numpy RNG (and torch); seed both so the halo path is reproducible.
            np.random.seed(aug_seed & 0xFFFFFFFF)
            torch.manual_seed(aug_seed)
            halo_img, _, halo_masks_t = apply_geometric_augmentation(
                halo_img,
                torch.zeros((1, 1, 1, 2), dtype=torch.float32),
                masks=halo_masks_t,
                **dict(self.geometric_aug),
            )
            halo_masks = [
                halo_masks_t[0, k].numpy() > 0.5
                for k in range(halo_masks_t.shape[1])
            ]

        # Trim the augmented halo back to `ts`, centered on the halo center:
        # crop_and_resize for the image (codebase convention), an equivalent
        # integer center-slice for the (axis-aligned) mask arrays.
        c = halo / 2.0
        bbox = make_centered_bboxes(
            torch.tensor([[c, c]], dtype=torch.float32), ts, ts
        )
        tile_image = crop_and_resize(halo_img, boxes=bbox, size=(ts, ts))
        off = (halo - ts) // 2
        tile_masks = [hm[off : off + ts, off : off + ts] for hm in halo_masks]
    else:
        # Fast path (no aug): direct slice + constant-zero pad, byte-identical.
        tile_image, tile_masks = self._slice_halo(image, mask_arrays, y0, x0, ts)

    # 4. Intensity aug (image only) + pad to stride (no-op when ts % max_stride==0).
    if self.apply_aug and self.intensity_aug is not None:
        tile_image, _ = apply_intensity_augmentation(
            tile_image,
            torch.zeros((1, 1, 1, 2), dtype=torch.float32),
            **self.intensity_aug,
        )
    tile_image = apply_pad_to_stride(tile_image, max_stride=self.max_stride)

    # 5. Foreground GT: union of ALL masks touching the tile (NO ownership filter,
    #    NO center/offset heads). Empty (fully-zero) tile masks contribute nothing.
    img_hw = tile_image.shape[-2:]
    foreground_mask = generate_foreground_mask(
        tile_masks,
        img_hw=img_hw,
        output_stride=self.seg_head_config.output_stride,
        maxpool=bool(getattr(self.seg_head_config, "target_maxpool", False)),
    )
    # Count of masks with any foreground in this tile (logging/eval convenience).
    num_masks_in_tile = sum(1 for m in tile_masks if int(m.sum()) > 0)

    return {
        "image": tile_image,
        "instances": torch.zeros((1, 1, 1, 2), dtype=torch.float32),
        "video_idx": torch.tensor(video_idx, dtype=torch.int32),
        "frame_idx": torch.tensor(frame_idx, dtype=torch.int32),
        "orig_size": torch.Tensor([orig_hw[0], orig_hw[1]]).unsqueeze(0),
        "num_instances": num_masks_in_tile,
        # Tiles are extracted in the model's input space (sizematcher bypassed),
        # so the effective scale is 1.0 (matches SingleInstanceTiledDataset).
        "eff_scale": torch.tensor(1.0, dtype=torch.float32),
        "foreground_mask": foreground_mask,
        "labels_idx": labels_idx,
        "tile_origin": torch.tensor(tile_origin, dtype=torch.int32),
    }

__init__(labels, seg_head_config, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False, tiling=None, base_seed=0)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    seg_head_config: DictConfig,
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
    tiling: Optional[Union[DictConfig, Any]] = None,
    base_seed: int = 0,
) -> None:
    """Initialize class attributes."""
    self.seg_head_config = seg_head_config
    # Segmentation never uses negative frames (degenerate num_nodes/instances).
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=False,
        tiling=tiling,
        output_stride=seg_head_config.output_stride,
        base_seed=base_seed,
    )

    # Per-(frame, tile-slot) descriptors + contiguous per-frame index blocks.
    self.tile_idx_list = self._get_tile_idx_list(labels)
    self.frame_blocks = self._build_frame_blocks(self.tile_idx_list)

__len__()

Return the number of tile samples (frames x tiles-per-frame).

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Return the number of tile samples (frames x tiles-per-frame)."""
    return len(self.tile_idx_list)

SingleInstanceDataset

Bases: BaseDataset

Dataset class for single-instance models.

Attributes:

Name	Type	Description
max_stride		Scalar integer specifying the maximum stride that the image must be divisible by.
user_instances_only		True if only user labeled instances should be used for training. If False, both user labeled and predicted instances would be used.
ensure_rgb		(bool) True if the input image should have 3 channels (RGB image). If input has only one
is	replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default	False.
ensure_grayscale		(bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
image.	If the source image has only one channel and this is set to False, then we retain the single channel input. Default	False.
intensity_aug		Intensity augmentation configuration. Can be: - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness'] - List of strings: Multiple intensity augmentations from the allowed values - Dictionary: Custom intensity configuration - None: No intensity augmentation applied
geometric_aug		Geometric augmentation configuration. Can be: - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup'] - List of strings: Multiple geometric augmentations from the allowed values - Dictionary: Custom geometric configuration - None: No geometric augmentation applied
scale		Factor to resize the image dimensions by, specified as a float. Default: 1.0.
apply_aug		True if augmentations should be applied to the data pipeline, else False. Default: False.
max_hw		Maximum height and width of images across the labels file. If max_height and max_width in the config is None, then max_hw is used (computed with sleap_nn.data.providers.get_max_height_width). Else the values in the config are used.
cache_img		String to indicate which caching to use: memory or disk. If None, the images aren't cached and loaded from the .slp file on each access.
cache_img_path		Path to save the .jpg files. If None, current working dir is used.
use_existing_imgs		Use existing imgs/ chunks in the cache_img_path.
confmap_head_config		DictConfig object with all the keys in the head_config section.
(required	keys	sigma, output_stride and part_names depending on the model type ).
rank		Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.
labels_list		List of sio.Labels objects. Used to store the labels in the cache. (only used if cache_img is None)

Methods:

Name	Description
__getitem__	Return dict with image and confmaps for instance for given index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class SingleInstanceDataset(BaseDataset):
    """Dataset class for single-instance models.

    Attributes:
        max_stride: Scalar integer specifying the maximum stride that the image must be
            divisible by.
        user_instances_only: `True` if only user labeled instances should be used for training. If `False`,
            both user labeled and predicted instances would be used.
        ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one
        channel when this is set to `True`, then the images from single-channel
        is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default: `False`.
        ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
        is set to True, then we convert the image to grayscale (single-channel)
        image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
        intensity_aug: Intensity augmentation configuration. Can be:
            - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness']
            - List of strings: Multiple intensity augmentations from the allowed values
            - Dictionary: Custom intensity configuration
            - None: No intensity augmentation applied
        geometric_aug: Geometric augmentation configuration. Can be:
            - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup']
            - List of strings: Multiple geometric augmentations from the allowed values
            - Dictionary: Custom geometric configuration
            - None: No geometric augmentation applied
        scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
        apply_aug: `True` if augmentations should be applied to the data pipeline,
            else `False`. Default: `False`.
        max_hw: Maximum height and width of images across the labels file. If `max_height` and
           `max_width` in the config is None, then `max_hw` is used (computed with
            `sleap_nn.data.providers.get_max_height_width`). Else the values in the config
            are used.
        cache_img: String to indicate which caching to use: `memory` or `disk`. If `None`,
            the images aren't cached and loaded from the `.slp` file on each access.
        cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
        use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
        (required keys: `sigma`, `output_stride` and `part_names` depending on the model type ).
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        labels_list: List of `sio.Labels` objects. Used to store the labels in the cache. (only used if `cache_img` is `None`)
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        confmap_head_config: DictConfig,
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
    ) -> None:
        """Initialize class attributes."""
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        self.confmap_head_config = confmap_head_config

    def __getitem__(self, index) -> Dict:
        """Return dict with image and confmaps for instance for given index."""
        sample = self.lf_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        video_idx = sample["video_idx"]
        lf_frame_idx = sample["frame_idx"]

        if sample.get("is_negative", False):
            sample = self._load_negative_sample(sample)
        else:
            if self.cache_img is not None:
                instances = sample["instances"]
                if self.cache_img == "disk":
                    img = np.array(
                        Image.open(
                            f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                        )
                    )
                elif self.cache_img == "memory":
                    img = self.cache[(labels_idx, lf_idx)].copy()
            else:
                lf = self.labels_list[labels_idx][lf_idx]
                instances = lf.instances
                img = lf.image
            if img.ndim == 2:
                img = np.expand_dims(img, axis=2)

            # get dict
            sample = process_lf(
                instances_list=instances,
                img=img,
                frame_idx=lf_frame_idx,
                video_idx=video_idx,
                max_instances=self.max_instances,
                user_instances_only=self.user_instances_only,
            )

        sample = self._apply_common_preprocessing(sample)

        img_hw = sample["image"].shape[-2:]

        # Drop keypoints pushed outside the image by augmentation so their target
        # confidence map is empty rather than a partial blob at the image edge.
        sample["instances"] = filter_oob_points(
            sample["instances"], img_hw[0], img_hw[1]
        )

        # Generate confidence maps
        confidence_maps = generate_confmaps(
            sample["instances"],
            img_hw=img_hw,
            sigma=self.confmap_head_config.sigma,
            output_stride=self.confmap_head_config.output_stride,
        )

        sample["confidence_maps"] = confidence_maps
        sample["labels_idx"] = labels_idx
        if self.use_negative_frames:
            sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

        return sample

__getitem__(index)

Return dict with image and confmaps for instance for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image and confmaps for instance for given index."""
    sample = self.lf_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    video_idx = sample["video_idx"]
    lf_frame_idx = sample["frame_idx"]

    if sample.get("is_negative", False):
        sample = self._load_negative_sample(sample)
    else:
        if self.cache_img is not None:
            instances = sample["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances = lf.instances
            img = lf.image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        # get dict
        sample = process_lf(
            instances_list=instances,
            img=img,
            frame_idx=lf_frame_idx,
            video_idx=video_idx,
            max_instances=self.max_instances,
            user_instances_only=self.user_instances_only,
        )

    sample = self._apply_common_preprocessing(sample)

    img_hw = sample["image"].shape[-2:]

    # Drop keypoints pushed outside the image by augmentation so their target
    # confidence map is empty rather than a partial blob at the image edge.
    sample["instances"] = filter_oob_points(
        sample["instances"], img_hw[0], img_hw[1]
    )

    # Generate confidence maps
    confidence_maps = generate_confmaps(
        sample["instances"],
        img_hw=img_hw,
        sigma=self.confmap_head_config.sigma,
        output_stride=self.confmap_head_config.output_stride,
    )

    sample["confidence_maps"] = confidence_maps
    sample["labels_idx"] = labels_idx
    if self.use_negative_frames:
        sample["is_negative"] = self.lf_idx_list[index].get("is_negative", False)

    return sample

__init__(labels, confmap_head_config, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    confmap_head_config: DictConfig,
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
) -> None:
    """Initialize class attributes."""
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=use_negative_frames,
    )
    self.confmap_head_config = confmap_head_config

SingleInstanceTiledDataset

Bases: BaseDataset

Single-instance dataset that emits fixed-size tiles instead of whole frames.

Phase-A tiled training dataset. Each frame is decomposed into overlapping square tiles (foreground-aware random draws for training, a deterministic grid for validation). A frame is decoded/process_lf/channel-coerced/scaled once (cached in a per-worker LRU) and reused across all of its tiles; each tile is then sliced out (extract_tile, sizematcher bypassed), optionally geometrically augmented via the halo path, and its per-tile confidence maps are generated on tile-local coordinates.

Emits one sample per (frame, tile-slot); __len__ is the total number of tile slots. Returned samples match the SingleInstanceDataset key contract (plus an int32 tile_origin of shape (2,)), so the default collate applies with no custom collate_fn.

Single-instance keeps one pose per frame: on multi-instance labels a one-time warning is emitted (foreground sampling uses all keypoints and inference decodes a single global peak per node).

Methods:

Name	Description
__getitem__	Return dict with image + confmaps for one tile of one frame.
__init__	Initialize class attributes.
__len__	Return the number of tile samples (frames x tiles-per-frame).

Source code in sleap_nn/data/custom_datasets.py

class SingleInstanceTiledDataset(BaseDataset):
    """Single-instance dataset that emits fixed-size tiles instead of whole frames.

    Phase-A tiled training dataset. Each frame is decomposed into overlapping
    square tiles (foreground-aware random draws for training, a deterministic
    grid for validation). A frame is decoded/`process_lf`/channel-coerced/scaled
    once (cached in a per-worker LRU) and reused across all of its tiles; each
    tile is then sliced out (`extract_tile`, sizematcher bypassed), optionally
    geometrically augmented via the halo path, and its per-tile confidence maps
    are generated on tile-local coordinates.

    Emits one sample per ``(frame, tile-slot)``; ``__len__`` is the total number
    of tile slots. Returned samples match the ``SingleInstanceDataset`` key
    contract (plus an ``int32`` ``tile_origin`` of shape ``(2,)``), so the
    default collate applies with no custom ``collate_fn``.

    Single-instance keeps one pose per frame: on multi-instance labels a one-time
    warning is emitted (foreground sampling uses all keypoints and inference
    decodes a single global peak per node).
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        confmap_head_config: DictConfig,
        max_stride: int,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        use_negative_frames: bool = False,
        tiling: Optional[Union[DictConfig, Any]] = None,
        base_seed: int = 0,
    ) -> None:
        """Initialize class attributes."""
        super().__init__(
            labels=labels,
            max_stride=max_stride,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
            tiling=tiling,
            output_stride=confmap_head_config.output_stride,
            base_seed=base_seed,
        )
        self.confmap_head_config = confmap_head_config

        # Per-(frame, tile-slot) descriptors + contiguous per-frame index blocks.
        self.tile_idx_list = self._get_tile_idx_list(labels)
        self.frame_blocks = self._build_frame_blocks(self.tile_idx_list)

        if self.max_instances > 1:
            logger.warning(
                "SingleInstanceTiledDataset received labels with more than one "
                "instance per frame. Single-instance models keep one pose per "
                "frame: all keypoints seed foreground tile sampling and inference "
                "decodes a single global peak per node."
            )

    def __len__(self) -> int:
        """Return the number of tile samples (frames x tiles-per-frame)."""
        return len(self.tile_idx_list)

    def _read_frame(self, d: Dict) -> Tuple[np.ndarray, List[sio.Instance]]:
        """Read a frame's raw image + instances (cache/disk/labels), restoring 2D->3D.

        Mirrors the cache/disk/labels-list read in ``SingleInstanceDataset``.
        """
        labels_idx = d["labels_idx"]
        lf_idx = d["lf_idx"]
        if self.cache_img is not None:
            instances = d["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances = lf.instances
            img = lf.image
        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)
        return img, instances

    def _to_sized_frame(self, frame: Dict) -> Dict:
        """Channel-coerce + `apply_resizer(scale)` so the LRU stores sized frames.

        Applies the ensure_rgb/grayscale + ``scale`` prefix once, up front, so the
        cached frame (and every tile sliced from it) is already in the model's
        input space. ``_apply_common_preprocessing`` therefore does not re-scale.
        """
        if self.ensure_rgb:
            frame["image"] = convert_to_rgb(frame["image"])
        elif self.ensure_grayscale:
            frame["image"] = convert_to_grayscale(frame["image"])
        frame["image"], frame["instances"] = apply_resizer(
            frame["image"], frame["instances"], scale=self.scale
        )
        return frame

    def __getitem__(self, index) -> Dict:
        """Return dict with image + confmaps for one tile of one frame."""
        d = self.tile_idx_list[index]
        labels_idx = d["labels_idx"]
        epoch = int(self._epoch)

        # Decode the full frame once per (labels_idx, lf_idx), via per-worker LRU.
        # The cached value is the fully sized frame dict (image, instances +
        # per-frame metadata); tiles clone from it so the cache stays pristine.
        frame = (
            self._frame_lru().get((labels_idx, d["lf_idx"]))
            if not d["is_negative"]
            else None
        )
        if frame is None:
            if d["is_negative"]:
                frame = self._load_negative_sample(d)
            else:
                img, instances = self._read_frame(d)
                frame = process_lf(
                    instances_list=instances,
                    img=img,
                    frame_idx=d["frame_idx"],
                    video_idx=d["video_idx"],
                    max_instances=self.max_instances,
                    user_instances_only=self.user_instances_only,
                )
            frame = self._to_sized_frame(frame)
            if not d["is_negative"]:
                self._frame_lru().put((labels_idx, d["lf_idx"]), frame)

        sample = {
            "image": frame["image"].clone(),
            "instances": frame["instances"].clone(),
            "video_idx": frame["video_idx"],
            "frame_idx": frame["frame_idx"],
            "orig_size": frame["orig_size"],
            "num_instances": frame["num_instances"],
        }

        # Resolve the tile origin: pinned for grid/val, drawn for train.
        if d["is_grid"]:
            sample["tile_origin"] = d["tile_origin"]
        else:
            rng = np.random.default_rng(
                tile_sample_seed(
                    self.base_seed,
                    epoch,
                    d["video_idx"],
                    d["frame_idx"],
                    d["sample_k"],
                )
            )
            centers = frame_foreground_centers(sample["instances"])
            sample["tile_origin"] = draw_tile_origin(
                centers,
                sample["image"].shape[-2:],
                self.tile_size,
                d["sample_k"],
                self.samples_per_frame,
                self.tile_fg_fraction,
                self.center_jitter,
                rng,
                pos_ratio=0.0 if d["is_negative"] else 1.0,
            )
            sample["aug_seed"] = tile_sample_seed(
                self.base_seed,
                epoch,
                d["video_idx"],
                d["frame_idx"],
                d["sample_k"],
                salt=1,
            )

        sample = self._apply_common_preprocessing(sample)

        img_hw = sample["image"].shape[-2:]

        # min_visible_keypoints: drop instances with too few in-tile keypoints
        # BEFORE OOB/confmap generation so a barely-clipped instance at a seam
        # does not seed a partial blob.
        inst = sample["instances"]
        inside = (
            (inst[..., 0] >= 0)
            & (inst[..., 0] < img_hw[1])
            & (inst[..., 1] >= 0)
            & (inst[..., 1] < img_hw[0])
        )
        keep = inside.sum(dim=-1) >= self.min_visible_keypoints
        inst[~keep] = torch.nan
        sample["instances"] = inst

        # NaN out remaining OOB keypoints at final tile resolution, then confmaps.
        sample["instances"] = filter_oob_points(
            sample["instances"], img_hw[0], img_hw[1]
        )
        sample["confidence_maps"] = generate_confmaps(
            sample["instances"],
            img_hw=img_hw,
            sigma=self.confmap_head_config.sigma,
            output_stride=self.confmap_head_config.output_stride,
        )

        sample["labels_idx"] = labels_idx
        if self.use_negative_frames:
            sample["is_negative"] = d["is_negative"]

        # Drop the transient aug seed so the batch collates uniformly.
        sample.pop("aug_seed", None)

        return sample

__getitem__(index)

Return dict with image + confmaps for one tile of one frame.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with image + confmaps for one tile of one frame."""
    d = self.tile_idx_list[index]
    labels_idx = d["labels_idx"]
    epoch = int(self._epoch)

    # Decode the full frame once per (labels_idx, lf_idx), via per-worker LRU.
    # The cached value is the fully sized frame dict (image, instances +
    # per-frame metadata); tiles clone from it so the cache stays pristine.
    frame = (
        self._frame_lru().get((labels_idx, d["lf_idx"]))
        if not d["is_negative"]
        else None
    )
    if frame is None:
        if d["is_negative"]:
            frame = self._load_negative_sample(d)
        else:
            img, instances = self._read_frame(d)
            frame = process_lf(
                instances_list=instances,
                img=img,
                frame_idx=d["frame_idx"],
                video_idx=d["video_idx"],
                max_instances=self.max_instances,
                user_instances_only=self.user_instances_only,
            )
        frame = self._to_sized_frame(frame)
        if not d["is_negative"]:
            self._frame_lru().put((labels_idx, d["lf_idx"]), frame)

    sample = {
        "image": frame["image"].clone(),
        "instances": frame["instances"].clone(),
        "video_idx": frame["video_idx"],
        "frame_idx": frame["frame_idx"],
        "orig_size": frame["orig_size"],
        "num_instances": frame["num_instances"],
    }

    # Resolve the tile origin: pinned for grid/val, drawn for train.
    if d["is_grid"]:
        sample["tile_origin"] = d["tile_origin"]
    else:
        rng = np.random.default_rng(
            tile_sample_seed(
                self.base_seed,
                epoch,
                d["video_idx"],
                d["frame_idx"],
                d["sample_k"],
            )
        )
        centers = frame_foreground_centers(sample["instances"])
        sample["tile_origin"] = draw_tile_origin(
            centers,
            sample["image"].shape[-2:],
            self.tile_size,
            d["sample_k"],
            self.samples_per_frame,
            self.tile_fg_fraction,
            self.center_jitter,
            rng,
            pos_ratio=0.0 if d["is_negative"] else 1.0,
        )
        sample["aug_seed"] = tile_sample_seed(
            self.base_seed,
            epoch,
            d["video_idx"],
            d["frame_idx"],
            d["sample_k"],
            salt=1,
        )

    sample = self._apply_common_preprocessing(sample)

    img_hw = sample["image"].shape[-2:]

    # min_visible_keypoints: drop instances with too few in-tile keypoints
    # BEFORE OOB/confmap generation so a barely-clipped instance at a seam
    # does not seed a partial blob.
    inst = sample["instances"]
    inside = (
        (inst[..., 0] >= 0)
        & (inst[..., 0] < img_hw[1])
        & (inst[..., 1] >= 0)
        & (inst[..., 1] < img_hw[0])
    )
    keep = inside.sum(dim=-1) >= self.min_visible_keypoints
    inst[~keep] = torch.nan
    sample["instances"] = inst

    # NaN out remaining OOB keypoints at final tile resolution, then confmaps.
    sample["instances"] = filter_oob_points(
        sample["instances"], img_hw[0], img_hw[1]
    )
    sample["confidence_maps"] = generate_confmaps(
        sample["instances"],
        img_hw=img_hw,
        sigma=self.confmap_head_config.sigma,
        output_stride=self.confmap_head_config.output_stride,
    )

    sample["labels_idx"] = labels_idx
    if self.use_negative_frames:
        sample["is_negative"] = d["is_negative"]

    # Drop the transient aug seed so the batch collates uniformly.
    sample.pop("aug_seed", None)

    return sample

__init__(labels, confmap_head_config, max_stride, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, use_negative_frames=False, tiling=None, base_seed=0)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    confmap_head_config: DictConfig,
    max_stride: int,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    use_negative_frames: bool = False,
    tiling: Optional[Union[DictConfig, Any]] = None,
    base_seed: int = 0,
) -> None:
    """Initialize class attributes."""
    super().__init__(
        labels=labels,
        max_stride=max_stride,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        use_negative_frames=use_negative_frames,
        tiling=tiling,
        output_stride=confmap_head_config.output_stride,
        base_seed=base_seed,
    )
    self.confmap_head_config = confmap_head_config

    # Per-(frame, tile-slot) descriptors + contiguous per-frame index blocks.
    self.tile_idx_list = self._get_tile_idx_list(labels)
    self.frame_blocks = self._build_frame_blocks(self.tile_idx_list)

    if self.max_instances > 1:
        logger.warning(
            "SingleInstanceTiledDataset received labels with more than one "
            "instance per frame. Single-instance models keep one pose per "
            "frame: all keypoints seed foreground tile sampling and inference "
            "decodes a single global peak per node."
        )

__len__()

Return the number of tile samples (frames x tiles-per-frame).

Source code in sleap_nn/data/custom_datasets.py

def __len__(self) -> int:
    """Return the number of tile samples (frames x tiles-per-frame)."""
    return len(self.tile_idx_list)

TopDownCenteredInstanceMultiClassDataset

Bases: CenteredInstanceDataset

Dataset class for instance-centered confidence map ID models.

Attributes:

Name	Type	Description
max_stride		Scalar integer specifying the maximum stride that the image must be divisible by.
anchor_ind		Index of the node to use as the anchor point, based on its index in the ordered list of skeleton nodes.
user_instances_only		True if only user labeled instances should be used for training. If False, both user labeled and predicted instances would be used.
ensure_rgb		(bool) True if the input image should have 3 channels (RGB image). If input has only one
is	replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default	False.
ensure_grayscale		(bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
image.	If the source image has only one channel and this is set to False, then we retain the single channel input. Default	False.
intensity_aug		Intensity augmentation configuration. Can be: - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness'] - List of strings: Multiple intensity augmentations from the allowed values - Dictionary: Custom intensity configuration - None: No intensity augmentation applied
geometric_aug		Geometric augmentation configuration. Can be: - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup'] - List of strings: Multiple geometric augmentations from the allowed values - Dictionary: Custom geometric configuration - None: No geometric augmentation applied
scale		Factor to resize the image dimensions by, specified as a float. Default: 1.0.
apply_aug		True if augmentations should be applied to the data pipeline, else False. Default: False.
max_hw		Maximum height and width of images across the labels file. If max_height and max_width in the config is None, then max_hw is used (computed with sleap_nn.data.providers.get_max_height_width). Else the values in the config are used.
cache_img		String to indicate which caching to use: memory or disk. If None, the images aren't cached and loaded from the .slp file on each access.
cache_img_path		Path to save the .jpg files. If None, current working dir is used.
use_existing_imgs		Use existing imgs/ chunks in the cache_img_path.
crop_size		Crop size of each instance for centered-instance model. If scale is provided, then the cropped image will be resized according to scale.
rank		Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.
confmap_head_config		DictConfig object with all the keys in the head_config section. (required keys: sigma, output_stride, part_names and anchor_part depending on the model type ).
class_vectors_head_config		DictConfig object with all the keys in the head_config section. (required keys: classes, num_fc_layers, num_fc_units, output_stride, loss_weight).
labels_list		List of sio.Labels objects. Used to store the labels in the cache. (only used if cache_img is None)

Methods:

Name	Description
__getitem__	Return dict with cropped image and confmaps of instance for given index.
__init__	Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

class TopDownCenteredInstanceMultiClassDataset(CenteredInstanceDataset):
    """Dataset class for instance-centered confidence map ID models.

    Attributes:
        max_stride: Scalar integer specifying the maximum stride that the image must be
            divisible by.
        anchor_ind: Index of the node to use as the anchor point, based on its index in the
            ordered list of skeleton nodes.
        user_instances_only: `True` if only user labeled instances should be used for training. If `False`,
            both user labeled and predicted instances would be used.
        ensure_rgb: (bool) True if the input image should have 3 channels (RGB image). If input has only one
        channel when this is set to `True`, then the images from single-channel
        is replicated along the channel axis. If the image has three channels and this is set to False, then we retain the three channels. Default: `False`.
        ensure_grayscale: (bool) True if the input image should only have a single channel. If input has three channels (RGB) and this
        is set to True, then we convert the image to grayscale (single-channel)
        image. If the source image has only one channel and this is set to False, then we retain the single channel input. Default: `False`.
        intensity_aug: Intensity augmentation configuration. Can be:
            - String: One of ['uniform_noise', 'gaussian_noise', 'contrast', 'brightness']
            - List of strings: Multiple intensity augmentations from the allowed values
            - Dictionary: Custom intensity configuration
            - None: No intensity augmentation applied
        geometric_aug: Geometric augmentation configuration. Can be:
            - String: One of ['rotation', 'scale', 'translate', 'erase_scale', 'mixup']
            - List of strings: Multiple geometric augmentations from the allowed values
            - Dictionary: Custom geometric configuration
            - None: No geometric augmentation applied
        scale: Factor to resize the image dimensions by, specified as a float. Default: 1.0.
        apply_aug: `True` if augmentations should be applied to the data pipeline,
            else `False`. Default: `False`.
        max_hw: Maximum height and width of images across the labels file. If `max_height` and
           `max_width` in the config is None, then `max_hw` is used (computed with
            `sleap_nn.data.providers.get_max_height_width`). Else the values in the config
            are used.
        cache_img: String to indicate which caching to use: `memory` or `disk`. If `None`,
            the images aren't cached and loaded from the `.slp` file on each access.
        cache_img_path: Path to save the `.jpg` files. If `None`, current working dir is used.
        use_existing_imgs: Use existing imgs/ chunks in the `cache_img_path`.
        crop_size: Crop size of each instance for centered-instance model. If `scale` is provided, then the cropped image will be resized according to `scale`.
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        confmap_head_config: DictConfig object with all the keys in the `head_config` section.
            (required keys: `sigma`, `output_stride`, `part_names` and `anchor_part` depending on the model type ).
        class_vectors_head_config: DictConfig object with all the keys in the `head_config` section.
            (required keys: `classes`, `num_fc_layers`, `num_fc_units`, `output_stride`, `loss_weight`).
        labels_list: List of `sio.Labels` objects. Used to store the labels in the cache. (only used if `cache_img` is `None`)
    """

    def __init__(
        self,
        labels: List[sio.Labels],
        crop_size: int,
        confmap_head_config: DictConfig,
        class_vectors_head_config: DictConfig,
        max_stride: int,
        anchor_ind: Optional[int] = None,
        user_instances_only: bool = True,
        ensure_rgb: bool = False,
        ensure_grayscale: bool = False,
        intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
        scale: float = 1.0,
        apply_aug: bool = False,
        max_hw: Tuple[Optional[int]] = (None, None),
        cache_img: Optional[str] = None,
        cache_img_path: Optional[str] = None,
        use_existing_imgs: bool = False,
        rank: Optional[int] = None,
        parallel_caching: bool = True,
        cache_workers: int = 0,
        bbox_crop: Optional[BBoxCropSpec] = None,
    ) -> None:
        """Initialize class attributes."""
        super().__init__(
            labels=labels,
            crop_size=crop_size,
            confmap_head_config=confmap_head_config,
            max_stride=max_stride,
            anchor_ind=anchor_ind,
            user_instances_only=user_instances_only,
            ensure_rgb=ensure_rgb,
            ensure_grayscale=ensure_grayscale,
            intensity_aug=intensity_aug,
            geometric_aug=geometric_aug,
            scale=scale,
            apply_aug=apply_aug,
            max_hw=max_hw,
            cache_img=cache_img,
            cache_img_path=cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            bbox_crop=bbox_crop,
        )
        self.class_vectors_head_config = class_vectors_head_config
        self.class_names = self.class_vectors_head_config.classes

    def __getitem__(self, index) -> Dict:
        """Return dict with cropped image and confmaps of instance for given index."""
        sample = self.instance_idx_list[index]
        labels_idx = sample["labels_idx"]
        lf_idx = sample["lf_idx"]
        inst_idx = sample["inst_idx"]
        video_idx = sample["video_idx"]
        lf_frame_idx = sample["frame_idx"]
        crop_box = sample.get("bbox")

        if self.cache_img is not None:
            instances_list = sample["instances"]
            if self.cache_img == "disk":
                img = np.array(
                    Image.open(
                        f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                    )
                )
            elif self.cache_img == "memory":
                img = self.cache[(labels_idx, lf_idx)].copy()
        else:
            lf = self.labels_list[labels_idx][lf_idx]
            instances_list = lf.instances
            img = lf.image

        if img.ndim == 2:
            img = np.expand_dims(img, axis=2)

        image = np.transpose(img, (2, 0, 1))  # HWC -> CHW

        instances = []
        for inst in instances_list:
            instances.append(
                inst.numpy()
            )  # no need to filter empty instance (handled while creating instance_idx)
        instances = np.stack(instances, axis=0)

        # Add singleton time dimension for single frames.
        image = np.expand_dims(image, axis=0)  # (n_samples=1, C, H, W)
        instances = np.expand_dims(
            instances, axis=0
        )  # (n_samples=1, num_instances, num_nodes, 2)

        instances = torch.from_numpy(instances.astype("float32"))
        image = torch.from_numpy(image.copy())

        num_instances, _ = instances.shape[1:3]
        orig_img_height, orig_img_width = image.shape[-2:]

        instances = instances[:, inst_idx]

        if self.ensure_rgb:
            image = convert_to_rgb(image)
        elif self.ensure_grayscale:
            image = convert_to_grayscale(image)

        # size matcher
        image, eff_scale = apply_sizematcher(
            image,
            max_height=self.max_hw[0],
            max_width=self.max_hw[1],
        )
        instances = instances * eff_scale

        # get class vectors
        track_ids = torch.Tensor(
            [
                (
                    self.class_names.index(instances_list[idx].track.name)
                    if instances_list[idx].track is not None
                    else -1
                )
                for idx in range(num_instances)
            ]
        ).to(torch.int32)
        class_vectors = make_class_vectors(
            class_inds=track_ids,
            n_classes=torch.tensor(len(self.class_names), dtype=torch.int32),
        )

        # get the centroids based on the anchor idx
        centroids = generate_centroids(
            instances,
            anchor_ind=self.anchor_ind,
            method=self.centroid_method,
            fallback=self.centroid_fallback,
        )

        instance, centroid = instances[0], centroids[0]  # (n_samples=1)

        sample = self._crop_for_augmentation(
            image, instance, centroid, crop_box, eff_scale
        )

        sample["frame_idx"] = torch.tensor(lf_frame_idx, dtype=torch.int32)
        sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
        sample["num_instances"] = num_instances
        sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
            0
        )
        sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

        # apply augmentation
        if self.apply_aug:
            if self.intensity_aug is not None:
                (
                    sample["instance_image"],
                    sample["instance"],
                ) = apply_intensity_augmentation(
                    sample["instance_image"],
                    sample["instance"],
                    **self.intensity_aug,
                )

            if self.geometric_aug is not None:
                (
                    sample["instance_image"],
                    sample["instance"],
                ) = apply_geometric_augmentation(
                    sample["instance_image"],
                    sample["instance"],
                    symmetric_inds=self.symmetric_inds,
                    **self.geometric_aug,
                )

        # re-crop to original crop size
        sample = self._final_crop(sample)

        # resize image
        sample["instance_image"], sample["instance"] = apply_resizer(
            sample["instance_image"],
            sample["instance"],
            scale=self.scale,
        )

        # Pad the image (if needed) according max stride
        sample["instance_image"] = apply_pad_to_stride(
            sample["instance_image"], max_stride=self.max_stride
        )

        img_hw = sample["instance_image"].shape[-2:]

        # Drop keypoints pushed outside the crop by augmentation so their target
        # confidence map is empty rather than a partial blob at the crop edge.
        sample["instance"] = filter_oob_points(sample["instance"], img_hw[0], img_hw[1])

        # Generate confidence maps
        confidence_maps = generate_confmaps(
            sample["instance"],
            img_hw=img_hw,
            sigma=self.confmap_head_config.sigma,
            output_stride=self.confmap_head_config.output_stride,
        )

        sample["class_vectors"] = class_vectors[inst_idx].to(torch.float32)

        sample["confidence_maps"] = confidence_maps
        sample["labels_idx"] = labels_idx

        return sample

__getitem__(index)

Return dict with cropped image and confmaps of instance for given index.

Source code in sleap_nn/data/custom_datasets.py

def __getitem__(self, index) -> Dict:
    """Return dict with cropped image and confmaps of instance for given index."""
    sample = self.instance_idx_list[index]
    labels_idx = sample["labels_idx"]
    lf_idx = sample["lf_idx"]
    inst_idx = sample["inst_idx"]
    video_idx = sample["video_idx"]
    lf_frame_idx = sample["frame_idx"]
    crop_box = sample.get("bbox")

    if self.cache_img is not None:
        instances_list = sample["instances"]
        if self.cache_img == "disk":
            img = np.array(
                Image.open(
                    f"{self.cache_img_path}/sample_{labels_idx}_{lf_idx}.jpg"
                )
            )
        elif self.cache_img == "memory":
            img = self.cache[(labels_idx, lf_idx)].copy()
    else:
        lf = self.labels_list[labels_idx][lf_idx]
        instances_list = lf.instances
        img = lf.image

    if img.ndim == 2:
        img = np.expand_dims(img, axis=2)

    image = np.transpose(img, (2, 0, 1))  # HWC -> CHW

    instances = []
    for inst in instances_list:
        instances.append(
            inst.numpy()
        )  # no need to filter empty instance (handled while creating instance_idx)
    instances = np.stack(instances, axis=0)

    # Add singleton time dimension for single frames.
    image = np.expand_dims(image, axis=0)  # (n_samples=1, C, H, W)
    instances = np.expand_dims(
        instances, axis=0
    )  # (n_samples=1, num_instances, num_nodes, 2)

    instances = torch.from_numpy(instances.astype("float32"))
    image = torch.from_numpy(image.copy())

    num_instances, _ = instances.shape[1:3]
    orig_img_height, orig_img_width = image.shape[-2:]

    instances = instances[:, inst_idx]

    if self.ensure_rgb:
        image = convert_to_rgb(image)
    elif self.ensure_grayscale:
        image = convert_to_grayscale(image)

    # size matcher
    image, eff_scale = apply_sizematcher(
        image,
        max_height=self.max_hw[0],
        max_width=self.max_hw[1],
    )
    instances = instances * eff_scale

    # get class vectors
    track_ids = torch.Tensor(
        [
            (
                self.class_names.index(instances_list[idx].track.name)
                if instances_list[idx].track is not None
                else -1
            )
            for idx in range(num_instances)
        ]
    ).to(torch.int32)
    class_vectors = make_class_vectors(
        class_inds=track_ids,
        n_classes=torch.tensor(len(self.class_names), dtype=torch.int32),
    )

    # get the centroids based on the anchor idx
    centroids = generate_centroids(
        instances,
        anchor_ind=self.anchor_ind,
        method=self.centroid_method,
        fallback=self.centroid_fallback,
    )

    instance, centroid = instances[0], centroids[0]  # (n_samples=1)

    sample = self._crop_for_augmentation(
        image, instance, centroid, crop_box, eff_scale
    )

    sample["frame_idx"] = torch.tensor(lf_frame_idx, dtype=torch.int32)
    sample["video_idx"] = torch.tensor(video_idx, dtype=torch.int32)
    sample["num_instances"] = num_instances
    sample["orig_size"] = torch.Tensor([orig_img_height, orig_img_width]).unsqueeze(
        0
    )
    sample["eff_scale"] = torch.tensor(eff_scale, dtype=torch.float32)

    # apply augmentation
    if self.apply_aug:
        if self.intensity_aug is not None:
            (
                sample["instance_image"],
                sample["instance"],
            ) = apply_intensity_augmentation(
                sample["instance_image"],
                sample["instance"],
                **self.intensity_aug,
            )

        if self.geometric_aug is not None:
            (
                sample["instance_image"],
                sample["instance"],
            ) = apply_geometric_augmentation(
                sample["instance_image"],
                sample["instance"],
                symmetric_inds=self.symmetric_inds,
                **self.geometric_aug,
            )

    # re-crop to original crop size
    sample = self._final_crop(sample)

    # resize image
    sample["instance_image"], sample["instance"] = apply_resizer(
        sample["instance_image"],
        sample["instance"],
        scale=self.scale,
    )

    # Pad the image (if needed) according max stride
    sample["instance_image"] = apply_pad_to_stride(
        sample["instance_image"], max_stride=self.max_stride
    )

    img_hw = sample["instance_image"].shape[-2:]

    # Drop keypoints pushed outside the crop by augmentation so their target
    # confidence map is empty rather than a partial blob at the crop edge.
    sample["instance"] = filter_oob_points(sample["instance"], img_hw[0], img_hw[1])

    # Generate confidence maps
    confidence_maps = generate_confmaps(
        sample["instance"],
        img_hw=img_hw,
        sigma=self.confmap_head_config.sigma,
        output_stride=self.confmap_head_config.output_stride,
    )

    sample["class_vectors"] = class_vectors[inst_idx].to(torch.float32)

    sample["confidence_maps"] = confidence_maps
    sample["labels_idx"] = labels_idx

    return sample

__init__(labels, crop_size, confmap_head_config, class_vectors_head_config, max_stride, anchor_ind=None, user_instances_only=True, ensure_rgb=False, ensure_grayscale=False, intensity_aug=None, geometric_aug=None, scale=1.0, apply_aug=False, max_hw=(None, None), cache_img=None, cache_img_path=None, use_existing_imgs=False, rank=None, parallel_caching=True, cache_workers=0, bbox_crop=None)

Initialize class attributes.

Source code in sleap_nn/data/custom_datasets.py

def __init__(
    self,
    labels: List[sio.Labels],
    crop_size: int,
    confmap_head_config: DictConfig,
    class_vectors_head_config: DictConfig,
    max_stride: int,
    anchor_ind: Optional[int] = None,
    user_instances_only: bool = True,
    ensure_rgb: bool = False,
    ensure_grayscale: bool = False,
    intensity_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    geometric_aug: Optional[Union[str, List[str], Dict[str, Any]]] = None,
    scale: float = 1.0,
    apply_aug: bool = False,
    max_hw: Tuple[Optional[int]] = (None, None),
    cache_img: Optional[str] = None,
    cache_img_path: Optional[str] = None,
    use_existing_imgs: bool = False,
    rank: Optional[int] = None,
    parallel_caching: bool = True,
    cache_workers: int = 0,
    bbox_crop: Optional[BBoxCropSpec] = None,
) -> None:
    """Initialize class attributes."""
    super().__init__(
        labels=labels,
        crop_size=crop_size,
        confmap_head_config=confmap_head_config,
        max_stride=max_stride,
        anchor_ind=anchor_ind,
        user_instances_only=user_instances_only,
        ensure_rgb=ensure_rgb,
        ensure_grayscale=ensure_grayscale,
        intensity_aug=intensity_aug,
        geometric_aug=geometric_aug,
        scale=scale,
        apply_aug=apply_aug,
        max_hw=max_hw,
        cache_img=cache_img,
        cache_img_path=cache_img_path,
        use_existing_imgs=use_existing_imgs,
        rank=rank,
        parallel_caching=parallel_caching,
        cache_workers=cache_workers,
        bbox_crop=bbox_crop,
    )
    self.class_vectors_head_config = class_vectors_head_config
    self.class_names = self.class_vectors_head_config.classes

embedding_membership_from_config(config, labels, *, user_instances_only=None)

The training :class:EmbeddingMembership of an embedding config.

Reads the same settings :func:get_train_val_datasets builds the datasets from (identity grouping, user_instances_only) and the vocabulary of labels (the train + val labels), for checks that must run before the datasets exist.

Source code in sleap_nn/data/custom_datasets.py

def embedding_membership_from_config(
    config: DictConfig,
    labels: List[sio.Labels],
    *,
    user_instances_only: Optional[bool] = None,
) -> EmbeddingMembership:
    """The training :class:`EmbeddingMembership` of an embedding config.

    Reads the same settings :func:`get_train_val_datasets` builds the datasets from
    (identity grouping, ``user_instances_only``) and the vocabulary of ``labels`` (the
    train + val labels), for checks that must run before the datasets exist.
    """
    track_names_are_global, id_scope = resolve_embedding_grouping(config)
    if user_instances_only is None:
        user_instances_only = bool(
            OmegaConf.select(config, "data_config.user_instances_only", default=True)
        )
    class_names = resolve_embedding_class_names(
        labels,
        track_names_are_global=track_names_are_global,
        user_instances_only=user_instances_only,
    )
    return EmbeddingMembership(
        class_names,
        id_scope=id_scope,
        track_names_are_global=track_names_are_global,
        user_instances_only=user_instances_only,
        labels=labels,
    )

embedding_preferred_carrier(config)

The carrier an embedding model was trained on (pose / mask), if known.

A file can carry both carriers (top-down segmentation and SAM output are poses with linked masks), and the two crop differently (pose centroid vs mask center, all-ones vs real mask). A model should embed the carrier it was trained on, so training records it in model_config.head_configs.embedding.embedding.detection_mode. For a config saved before that field existed, data_config.preprocessing.burn_in=True means the model was trained on masks (burn-in needs one); otherwise the carrier is unknown and None is returned (the dataset then counts, see :func:resolve_embedding_detection_mode).

Source code in sleap_nn/data/custom_datasets.py

def embedding_preferred_carrier(config: DictConfig) -> Optional[str]:
    """The carrier an embedding model was trained on (``pose`` / ``mask``), if known.

    A file can carry both carriers (top-down segmentation and SAM output are poses
    with linked masks), and the two crop differently (pose centroid vs mask center,
    all-ones vs real mask). A model should embed the carrier it was trained on, so
    training records it in
    ``model_config.head_configs.embedding.embedding.detection_mode``. For a config
    saved before that field existed, ``data_config.preprocessing.burn_in=True`` means
    the model was trained on masks (burn-in needs one); otherwise the carrier is
    unknown and ``None`` is returned (the dataset then counts, see
    :func:`resolve_embedding_detection_mode`).
    """
    explicit = OmegaConf.select(
        config,
        "model_config.head_configs.embedding.embedding.detection_mode",
        default=None,
    )
    if explicit is not None:
        return str(explicit)
    if bool(
        OmegaConf.select(config, "data_config.preprocessing.burn_in", default=False)
    ):
        return "mask"
    return None

embedding_training_groups(labels, membership, detection_mode)

(group_ids, video_ids, frame_ids) of the crops an embedding dataset would index.

The per-crop arrays :class:GroupAwareBatchSampler draws batches from, computed from the labels alone (no image decode) so the objective can be checked against them before the dataset is built and cached. video_ids are unique across labels files, like :attr:EmbeddingDataset.video_ids.

Source code in sleap_nn/data/custom_datasets.py

def embedding_training_groups(
    labels: List[sio.Labels], membership: EmbeddingMembership, detection_mode: str
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
    """``(group_ids, video_ids, frame_ids)`` of the crops an embedding dataset would index.

    The per-crop arrays :class:`GroupAwareBatchSampler` draws batches from, computed
    from the labels alone (no image decode) so the objective can be checked against
    them before the dataset is built and cached. ``video_ids`` are unique across
    labels files, like :attr:`EmbeddingDataset.video_ids`.
    """
    groups, videos, frames = [], [], []
    video_vocab: Dict[tuple, int] = {}
    for labels_idx, label in enumerate(labels):
        video_index = {id(v): i for i, v in enumerate(label.videos)}
        for lf in label:
            dets = (
                (getattr(lf, "masks", None) or [])
                if detection_mode == "mask"
                else lf.instances
            )
            video_idx = video_index[id(lf.video)]
            for det in dets:
                group = membership.group(det, labels_idx, video_idx)
                if group is None:
                    continue
                groups.append(group[0])
                videos.append(
                    video_vocab.setdefault((labels_idx, video_idx), len(video_vocab))
                )
                frames.append(lf.frame_idx)
    return (
        np.asarray(groups, np.int64),
        np.asarray(videos, np.int64),
        np.asarray(frames, np.int64),
    )

get_embedding_batch_sampler(train_dataset, config, batches_per_epoch=None, rank=0, world_size=1)

Build the group-aware batch sampler an embedding training run uses.

The one place the objective.sampler config is read, so the training loader and the steps-per-epoch count (:func:get_train_samples_per_step) always agree on the batch composition.

Source code in sleap_nn/data/custom_datasets.py

def get_embedding_batch_sampler(
    train_dataset: "EmbeddingDataset",
    config: DictConfig,
    batches_per_epoch: Optional[int] = None,
    rank: int = 0,
    world_size: int = 1,
) -> GroupAwareBatchSampler:
    """Build the group-aware batch sampler an ``embedding`` training run uses.

    The one place the ``objective.sampler`` config is read, so the training loader
    and the steps-per-epoch count (:func:`get_train_samples_per_step`) always agree
    on the batch composition.
    """
    sampler = resolve_embedding_objective(
        OmegaConf.select(
            config, "model_config.head_configs.embedding.embedding.objective"
        )
    ).sampler
    return GroupAwareBatchSampler(
        group_ids=train_dataset.group_ids,
        video_ids=train_dataset.video_ids,
        frame_ids=train_dataset.frame_ids,
        kind=sampler.kind,
        P=sampler.groups_per_batch,
        K=sampler.samples_per_group,
        batches_per_epoch=batches_per_epoch,
        seed=OmegaConf.select(config, "trainer_config.seed", default=0) or 0,
        rank=rank,
        world_size=world_size,
    )

get_steps_per_epoch(dataset, batch_size)

Compute the number of steps (iterations) per epoch for the given dataset.

Source code in sleap_nn/data/custom_datasets.py

def get_steps_per_epoch(dataset: BaseDataset, batch_size: int):
    """Compute the number of steps (iterations) per epoch for the given dataset."""
    return (len(dataset) // batch_size) + (1 if (len(dataset) % batch_size) else 0)

get_train_samples_per_step(train_dataset, config)

Number of training samples one optimizer step consumes (per device).

train_data_loader.batch_size for every model type except embedding, whose loader ignores it: :class:GroupAwareBatchSampler draws P groups x K crops per step instead. Dividing the dataset by batch_size there overstated the steps one pass needs by P x K / batch_size (with the defaults P=8, K=16 and batch_size=4, an "epoch" walked the dataset 32 times).

Source code in sleap_nn/data/custom_datasets.py

def get_train_samples_per_step(train_dataset: BaseDataset, config: DictConfig) -> int:
    """Number of training samples one optimizer step consumes (per device).

    ``train_data_loader.batch_size`` for every model type except ``embedding``, whose
    loader ignores it: :class:`GroupAwareBatchSampler` draws P groups x K crops per
    step instead. Dividing the dataset by ``batch_size`` there overstated the steps
    one pass needs by ``P x K / batch_size`` (with the defaults P=8, K=16 and
    ``batch_size=4``, an "epoch" walked the dataset 32 times).
    """
    if isinstance(train_dataset, EmbeddingDataset):
        return get_embedding_batch_sampler(train_dataset, config).samples_per_batch
    return config.trainer_config.train_data_loader.batch_size

get_train_val_dataloaders(train_dataset, val_dataset, config, train_steps_per_epoch=None, val_steps_per_epoch=None, rank=None, trainer_devices=1)

Return the train and val dataloaders.

Parameters:

Name	Type	Description	Default
train_dataset	BaseDataset	Train dataset-instance of one of the dataset classes [SingleInstanceDataset, CentroidDataset, CenteredInstanceDataset, BottomUpDataset, BottomUpMultiClassDataset, TopDownCenteredInstanceMultiClassDataset].	required
val_dataset	BaseDataset	Val dataset-instance of one of the dataset classes [SingleInstanceDataset, CentroidDataset, CenteredInstanceDataset, BottomUpDataset, BottomUpMultiClassDataset, TopDownCenteredInstanceMultiClassDataset].	required
config	DictConfig	Sleap-nn config.	required
train_steps_per_epoch	Optional[int]	Number of minibatches (steps) to train for in an epoch. If set to None, this is set to the number of batches in the training data. Note: In a multi-gpu training setup, the effective steps during training would be the trainer_steps_per_epoch / trainer_devices.	None
val_steps_per_epoch	Optional[int]	Number of minibatches (steps) to run validation for in an epoch. If set to None, this is set to the number of batches in the val data.	None
rank	Optional[int]	Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.	None
trainer_devices	int	Number of devices to use for training.	1

Returns:

Type	Description
	A tuple (train_dataloader, val_dataloader).

Source code in sleap_nn/data/custom_datasets.py

def get_train_val_dataloaders(
    train_dataset: BaseDataset,
    val_dataset: BaseDataset,
    config: DictConfig,
    train_steps_per_epoch: Optional[int] = None,
    val_steps_per_epoch: Optional[int] = None,
    rank: Optional[int] = None,
    trainer_devices: int = 1,
):
    """Return the train and val dataloaders.

    Args:
        train_dataset: Train dataset-instance of one of the dataset classes [SingleInstanceDataset, CentroidDataset, CenteredInstanceDataset, BottomUpDataset, BottomUpMultiClassDataset, TopDownCenteredInstanceMultiClassDataset].
        val_dataset: Val dataset-instance of one of the dataset classes [SingleInstanceDataset, CentroidDataset, CenteredInstanceDataset, BottomUpDataset, BottomUpMultiClassDataset, TopDownCenteredInstanceMultiClassDataset].
        config: Sleap-nn config.
        train_steps_per_epoch: Number of minibatches (steps) to train for in an epoch. If set to `None`, this is set to the number of batches in the training data. **Note**: In a multi-gpu training setup, the effective steps during training would be the `trainer_steps_per_epoch` / `trainer_devices`.
        val_steps_per_epoch: Number of minibatches (steps) to run validation for in an epoch. If set to `None`, this is set to the number of batches in the val data.
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.
        trainer_devices: Number of devices to use for training.

    Returns:
        A tuple (train_dataloader, val_dataloader).
    """
    pin_memory = (
        config.trainer_config.train_data_loader.pin_memory
        if "pin_memory" in config.trainer_config.train_data_loader
        and config.trainer_config.train_data_loader.pin_memory is not None
        else True
    )

    if train_steps_per_epoch is None:
        train_steps_per_epoch = config.trainer_config.train_steps_per_epoch
        if train_steps_per_epoch is None:
            train_steps_per_epoch = get_steps_per_epoch(
                dataset=train_dataset,
                batch_size=get_train_samples_per_step(train_dataset, config),
            )

    if val_steps_per_epoch is None:
        val_steps_per_epoch = get_steps_per_epoch(
            dataset=val_dataset,
            batch_size=config.trainer_config.val_data_loader.batch_size,
        )

    # Embedding training uses a group-aware BATCH sampler (PK / within_video) so the
    # wanted positives/negatives co-occur in each batch. The batch_sampler is mutually
    # exclusive with batch_size/shuffle/sampler, so build a distinct loader here.
    if isinstance(train_dataset, EmbeddingDataset):
        batch_sampler = get_embedding_batch_sampler(
            train_dataset,
            config,
            batches_per_epoch=max(1, round(train_steps_per_epoch / trainer_devices)),
            # DDP: each rank draws an INDEPENDENT (seed + rank) batch stream of the same
            # length over the full dataset — not a partition — so the all-reduced gradient
            # aggregates roughly world_size x P x K decorrelated crops per step.
            rank=rank if rank is not None else 0,
            world_size=trainer_devices,
        )
        # Use a plain DataLoader (NOT InfiniteDataLoader): the GroupAwareBatchSampler
        # is already epoch-bounded (yields `batches_per_epoch` batches per __iter__,
        # re-randomized each epoch via its rng), so the infinite-recycling wrapper is
        # unnecessary. Critically, InfiniteDataLoader eagerly creates its worker
        # iterator + wraps the sampler in an infinite _RepeatSampler, which deadlocks
        # under Lightning with num_workers>0; the plain loader iterates with workers
        # correctly (so heavy CPU/skia aug can be parallelized).
        train_nw = config.trainer_config.train_data_loader.num_workers
        val_nw = config.trainer_config.val_data_loader.num_workers
        train_data_loader = DataLoader(
            dataset=train_dataset,
            batch_sampler=batch_sampler,
            num_workers=train_nw,
            pin_memory=pin_memory,
            persistent_workers=train_nw > 0,
        )
        val_data_loader = DataLoader(
            dataset=val_dataset,
            shuffle=False,
            batch_size=config.trainer_config.val_data_loader.batch_size,
            num_workers=val_nw,
            pin_memory=pin_memory,
            persistent_workers=val_nw > 0,
        )
        return train_data_loader, val_data_loader

    tiling = OmegaConf.select(config, "data_config.preprocessing.tiling", default=None)
    tiling_enabled = tiling is not None and tiling.enabled

    # Under tiling, a frame-grouped block sampler REPLACES DistributedSampler (it
    # shards whole frame blocks so a frame's tiles stay together and DDP-disjoint),
    # and a per-worker RNG init de-correlates the halo augmentation streams.
    worker_init_fn = tiling_worker_init_fn if tiling_enabled else None

    if tiling_enabled:
        train_sampler = FrameGroupedTileSampler(
            train_dataset.frame_blocks,
            batch_size=config.trainer_config.train_data_loader.batch_size,
            shuffle=config.trainer_config.train_data_loader.shuffle,
            seed=config.trainer_config.seed or 0,
            num_replicas=trainer_devices,
            rank=(rank if rank is not None else 0),
        )
    else:
        train_sampler = (
            DistributedSampler(
                dataset=train_dataset,
                shuffle=config.trainer_config.train_data_loader.shuffle,
                rank=rank if rank is not None else 0,
                num_replicas=trainer_devices,
            )
            if trainer_devices > 1
            else None
        )

    train_data_loader = InfiniteDataLoader(
        dataset=train_dataset,
        sampler=train_sampler,
        len_dataloader=max(1, round(train_steps_per_epoch / trainer_devices)),
        shuffle=(
            config.trainer_config.train_data_loader.shuffle
            if train_sampler is None
            else None
        ),
        batch_size=config.trainer_config.train_data_loader.batch_size,
        num_workers=config.trainer_config.train_data_loader.num_workers,
        pin_memory=pin_memory,
        worker_init_fn=worker_init_fn,
        persistent_workers=(
            True if config.trainer_config.train_data_loader.num_workers > 0 else None
        ),
        prefetch_factor=(
            config.trainer_config.train_data_loader.batch_size
            if config.trainer_config.train_data_loader.num_workers > 0
            else None
        ),
    )

    if tiling_enabled:
        val_sampler = FrameGroupedTileSampler(
            val_dataset.frame_blocks,
            batch_size=config.trainer_config.val_data_loader.batch_size,
            shuffle=False,
            seed=config.trainer_config.seed or 0,
            num_replicas=trainer_devices,
            rank=(rank if rank is not None else 0),
        )
    else:
        val_sampler = (
            DistributedSampler(
                dataset=val_dataset,
                shuffle=False,
                rank=rank if rank is not None else 0,
                num_replicas=trainer_devices,
            )
            if trainer_devices > 1
            else None
        )
    val_data_loader = InfiniteDataLoader(
        dataset=val_dataset,
        shuffle=False if val_sampler is None else None,
        sampler=val_sampler,
        len_dataloader=(
            max(1, round(val_steps_per_epoch / trainer_devices))
            if trainer_devices > 1
            else None
        ),
        batch_size=config.trainer_config.val_data_loader.batch_size,
        num_workers=config.trainer_config.val_data_loader.num_workers,
        pin_memory=pin_memory,
        worker_init_fn=worker_init_fn,
        persistent_workers=(
            True if config.trainer_config.val_data_loader.num_workers > 0 else None
        ),
        prefetch_factor=(
            config.trainer_config.val_data_loader.batch_size
            if config.trainer_config.val_data_loader.num_workers > 0
            else None
        ),
    )

    return train_data_loader, val_data_loader

get_train_val_datasets(train_labels, val_labels, config, rank=None)

Return the train and val datasets.

Parameters:

Name	Type	Description	Default
train_labels	List[Labels]	List of train labels.	required
val_labels	List[Labels]	List of val labels.	required
config	DictConfig	Sleap-nn config.	required
rank	Optional[int]	Indicates the rank of the process. Used during distributed training to ensure that image storage to disk occurs only once across all workers.	None

Returns:

Type	Description
	A tuple (train_dataset, val_dataset).

Source code in sleap_nn/data/custom_datasets.py

def get_train_val_datasets(
    train_labels: List[sio.Labels],
    val_labels: List[sio.Labels],
    config: DictConfig,
    rank: Optional[int] = None,
):
    """Return the train and val datasets.

    Args:
        train_labels: List of train labels.
        val_labels: List of val labels.
        config: Sleap-nn config.
        rank: Indicates the rank of the process. Used during distributed training to ensure that image storage to
            disk occurs only once across all workers.

    Returns:
        A tuple (train_dataset, val_dataset).
    """
    cache_imgs = (
        config.data_config.data_pipeline_fw.split("_")[-1]
        if "cache_img" in config.data_config.data_pipeline_fw
        else None
    )
    base_cache_img_path = config.data_config.cache_img_path
    train_cache_img_path, val_cache_img_path = None, None

    if cache_imgs == "disk":
        train_cache_img_path = Path(base_cache_img_path) / "train_imgs"
        val_cache_img_path = Path(base_cache_img_path) / "val_imgs"
    use_existing_imgs = config.data_config.use_existing_imgs

    # Parallel caching configuration
    parallel_caching = getattr(config.data_config, "parallel_caching", True)
    cache_workers = getattr(config.data_config, "cache_workers", 0)

    use_negative_frames = getattr(config.data_config, "use_negative_frames", False)

    model_type = get_model_type_from_cfg(config=config)
    backbone_type = get_backbone_type_from_cfg(config=config)

    # Gate the lossy disk cache for the embedding (appearance / re-ID) model (SPEC §5.1):
    # `torch_dataset_cache_img_disk` writes source frames as JPEG, which silently
    # degrades an appearance model. The disk cache has no lossless format, so refuse it
    # for embedding and require the in-memory cache (or an uncached fw) instead.
    if model_type == "embedding" and cache_imgs == "disk":
        raise ValueError(
            "data_pipeline_fw='torch_dataset_cache_img_disk' is not supported for the "
            "`embedding` model type: the disk cache stores frames as JPEG, which "
            "silently degrades an appearance / re-ID model. Use "
            "data_pipeline_fw='torch_dataset_cache_img_memory' (lossless, recommended) "
            "or 'torch_dataset' (uncached)."
        )

    if use_negative_frames and model_type in (
        "centered_instance",
        "multi_class_topdown",
        "centered_instance_segmentation",
        "embedding",
    ):
        logger.warning(
            f"use_negative_frames is enabled but model_type='{model_type}' "
            f"operates at instance-crop level and does not support frame-level "
            f"negatives. Negative frames will be disabled."
        )
        use_negative_frames = False

    if cache_imgs == "disk" and use_existing_imgs:
        if not (
            train_cache_img_path.exists()
            and train_cache_img_path.is_dir()
            and any(train_cache_img_path.glob("*.jpg"))
        ):
            message = f"There are no images in the path: {train_cache_img_path}"
            logger.error(message)
            raise Exception(message)

        if not (
            val_cache_img_path.exists()
            and val_cache_img_path.is_dir()
            and any(val_cache_img_path.glob("*.jpg"))
        ):
            message = f"There are no images in the path: {val_cache_img_path}"
            logger.error(message)
            raise Exception(message)

    if model_type == "bottomup":
        train_dataset = BottomUpDataset(
            labels=train_labels,
            confmap_head_config=config.model_config.head_configs.bottomup.confmaps,
            pafs_head_config=config.model_config.head_configs.bottomup.pafs,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            scale=config.data_config.preprocessing.scale,
            apply_aug=config.data_config.use_augmentations_train,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=train_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        val_dataset = BottomUpDataset(
            labels=val_labels,
            confmap_head_config=config.model_config.head_configs.bottomup.confmaps,
            pafs_head_config=config.model_config.head_configs.bottomup.pafs,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=None,
            geometric_aug=None,
            scale=config.data_config.preprocessing.scale,
            apply_aug=False,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=val_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )

    elif model_type == "multi_class_bottomup":
        train_dataset = BottomUpMultiClassDataset(
            labels=train_labels,
            confmap_head_config=config.model_config.head_configs.multi_class_bottomup.confmaps,
            class_maps_head_config=config.model_config.head_configs.multi_class_bottomup.class_maps,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            scale=config.data_config.preprocessing.scale,
            apply_aug=config.data_config.use_augmentations_train,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=train_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        val_dataset = BottomUpMultiClassDataset(
            labels=val_labels,
            confmap_head_config=config.model_config.head_configs.multi_class_bottomup.confmaps,
            class_maps_head_config=config.model_config.head_configs.multi_class_bottomup.class_maps,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=None,
            geometric_aug=None,
            scale=config.data_config.preprocessing.scale,
            apply_aug=False,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=val_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )

    elif model_type == "centered_instance":
        bbox_crop = BBoxCropSpec.from_config(config)
        nodes = config.model_config.head_configs.centered_instance.confmaps.part_names
        anchor_part = (
            config.model_config.head_configs.centered_instance.confmaps.anchor_part
        )
        anchor_ind = nodes.index(anchor_part) if anchor_part is not None else None
        train_dataset = CenteredInstanceDataset(
            labels=train_labels,
            confmap_head_config=config.model_config.head_configs.centered_instance.confmaps,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            scale=config.data_config.preprocessing.scale,
            apply_aug=config.data_config.use_augmentations_train,
            crop_size=config.data_config.preprocessing.crop_size,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=train_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            bbox_crop=bbox_crop,
        )
        val_dataset = CenteredInstanceDataset(
            labels=val_labels,
            confmap_head_config=config.model_config.head_configs.centered_instance.confmaps,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=None,
            geometric_aug=None,
            scale=config.data_config.preprocessing.scale,
            apply_aug=False,
            crop_size=config.data_config.preprocessing.crop_size,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=val_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            bbox_crop=bbox_crop,
        )

    elif model_type == "multi_class_topdown":
        bbox_crop = BBoxCropSpec.from_config(config)
        nodes = config.model_config.head_configs.multi_class_topdown.confmaps.part_names
        anchor_part = (
            config.model_config.head_configs.multi_class_topdown.confmaps.anchor_part
        )
        anchor_ind = nodes.index(anchor_part) if anchor_part is not None else None
        train_dataset = TopDownCenteredInstanceMultiClassDataset(
            labels=train_labels,
            confmap_head_config=config.model_config.head_configs.multi_class_topdown.confmaps,
            class_vectors_head_config=config.model_config.head_configs.multi_class_topdown.class_vectors,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            scale=config.data_config.preprocessing.scale,
            apply_aug=config.data_config.use_augmentations_train,
            crop_size=config.data_config.preprocessing.crop_size,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=train_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            bbox_crop=bbox_crop,
        )
        val_dataset = TopDownCenteredInstanceMultiClassDataset(
            labels=val_labels,
            confmap_head_config=config.model_config.head_configs.multi_class_topdown.confmaps,
            class_vectors_head_config=config.model_config.head_configs.multi_class_topdown.class_vectors,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=None,
            geometric_aug=None,
            scale=config.data_config.preprocessing.scale,
            apply_aug=False,
            crop_size=config.data_config.preprocessing.crop_size,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=val_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            bbox_crop=bbox_crop,
        )

    elif model_type == "centroid":
        # Mask-only labels carry no skeleton at all (the centroid target comes from
        # `UserCentroid` annotations, possibly derived from masks via
        # `data_config.centroids_from_masks`). Indexing `skeletons[0]` crashed with
        # a bare `ConfigIndexError: list index out of range`; an absent skeleton
        # just means there is no anchor node to resolve.
        skeletons_cfg = OmegaConf.select(config, "data_config.skeletons", default=None)
        nodes = [x["name"] for x in skeletons_cfg[0]["nodes"]] if skeletons_cfg else []
        anchor_part = config.model_config.head_configs.centroid.confmaps.anchor_part
        # The anchor/instance-keypoint path is now only a FALLBACK for frames
        # without user centroids, so an anchor_part that is None OR absent from
        # the pose skeleton must NOT crash (it did: ``nodes.index`` raised
        # ValueError). Resolve when possible, else leave ``anchor_ind=None`` and
        # the fallback derives the centroid from the mean of visible nodes.
        anchor_ind = (
            nodes.index(anchor_part)
            if anchor_part is not None and anchor_part in nodes
            else None
        )
        # Resolve ONE centroid source for the whole run (no per-frame mix of
        # user-annotated and computed centroids). Inference (when unset) reads
        # the TRAIN labels and the same decision is applied to both splits so
        # train and val can never disagree on the centroid definition.
        centroid_source = OmegaConf.select(
            config,
            "model_config.head_configs.centroid.confmaps.centroid_source",
            default=None,
        )
        use_user_centroids = resolve_centroid_source(centroid_source, train_labels)
        train_dataset = CentroidDataset(
            labels=train_labels,
            confmap_head_config=config.model_config.head_configs.centroid.confmaps,
            use_user_centroids=use_user_centroids,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            scale=config.data_config.preprocessing.scale,
            apply_aug=config.data_config.use_augmentations_train,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=train_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )
        val_dataset = CentroidDataset(
            labels=val_labels,
            confmap_head_config=config.model_config.head_configs.centroid.confmaps,
            use_user_centroids=use_user_centroids,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=None,
            geometric_aug=None,
            scale=config.data_config.preprocessing.scale,
            apply_aug=False,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=val_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            use_negative_frames=use_negative_frames,
        )

    elif model_type == "bottomup_segmentation":
        seg_cfg = config.model_config.head_configs.bottomup_segmentation
        tiling = OmegaConf.select(
            config, "data_config.preprocessing.tiling", default=None
        )
        if tiling is not None and tiling.enabled:
            # Tiled bottom-up segmentation training. Train draws foreground-aware
            # tiles (aug on); val always uses a deterministic full-coverage grid
            # (no aug). Masks are co-transformed with the image under the halo path.
            base_seed = config.trainer_config.seed or 0
            train_dataset = BottomUpSegmentationTiledDataset(
                labels=train_labels,
                seg_head_config=seg_cfg.segmentation,
                center_head_config=seg_cfg.center,
                offset_head_config=seg_cfg.offsets,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=(
                    config.data_config.augmentation_config.intensity
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                geometric_aug=(
                    config.data_config.augmentation_config.geometric
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                scale=config.data_config.preprocessing.scale,
                apply_aug=config.data_config.use_augmentations_train,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=train_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
                tiling=OmegaConf.merge(tiling, {"sampling": tiling.sampling}),
                base_seed=base_seed,
            )
            val_dataset = BottomUpSegmentationTiledDataset(
                labels=val_labels,
                seg_head_config=seg_cfg.segmentation,
                center_head_config=seg_cfg.center,
                offset_head_config=seg_cfg.offsets,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=None,
                geometric_aug=None,
                scale=config.data_config.preprocessing.scale,
                apply_aug=False,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=val_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
                tiling=OmegaConf.merge(tiling, {"sampling": "grid"}),
                base_seed=base_seed,
            )
        else:
            train_dataset = BottomUpSegmentationDataset(
                labels=train_labels,
                seg_head_config=seg_cfg.segmentation,
                center_head_config=seg_cfg.center,
                offset_head_config=seg_cfg.offsets,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=(
                    config.data_config.augmentation_config.intensity
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                geometric_aug=(
                    config.data_config.augmentation_config.geometric
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                scale=config.data_config.preprocessing.scale,
                apply_aug=config.data_config.use_augmentations_train,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=train_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
            )
            val_dataset = BottomUpSegmentationDataset(
                labels=val_labels,
                seg_head_config=seg_cfg.segmentation,
                center_head_config=seg_cfg.center,
                offset_head_config=seg_cfg.offsets,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=None,
                geometric_aug=None,
                scale=config.data_config.preprocessing.scale,
                apply_aug=False,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=val_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
            )

    elif model_type == "semantic_segmentation":
        seg_cfg = config.model_config.head_configs.semantic_segmentation
        tiling = OmegaConf.select(
            config, "data_config.preprocessing.tiling", default=None
        )
        if tiling is not None and tiling.enabled:
            # Tiled whole-frame semantic segmentation. Train draws foreground-aware
            # tiles (aug on); val always uses a deterministic full-coverage grid (no
            # aug). Masks are co-transformed with the image under the halo path.
            base_seed = config.trainer_config.seed or 0
            train_dataset = SemanticSegmentationTiledDataset(
                labels=train_labels,
                seg_head_config=seg_cfg.segmentation,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=(
                    config.data_config.augmentation_config.intensity
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                geometric_aug=(
                    config.data_config.augmentation_config.geometric
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                scale=config.data_config.preprocessing.scale,
                apply_aug=config.data_config.use_augmentations_train,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=train_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
                tiling=OmegaConf.merge(tiling, {"sampling": tiling.sampling}),
                base_seed=base_seed,
            )
            val_dataset = SemanticSegmentationTiledDataset(
                labels=val_labels,
                seg_head_config=seg_cfg.segmentation,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=None,
                geometric_aug=None,
                scale=config.data_config.preprocessing.scale,
                apply_aug=False,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=val_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
                tiling=OmegaConf.merge(tiling, {"sampling": "grid"}),
                base_seed=base_seed,
            )
        else:
            train_dataset = SemanticSegmentationDataset(
                labels=train_labels,
                seg_head_config=seg_cfg.segmentation,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=(
                    config.data_config.augmentation_config.intensity
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                geometric_aug=(
                    config.data_config.augmentation_config.geometric
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                scale=config.data_config.preprocessing.scale,
                apply_aug=config.data_config.use_augmentations_train,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=train_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
            )
            val_dataset = SemanticSegmentationDataset(
                labels=val_labels,
                seg_head_config=seg_cfg.segmentation,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=None,
                geometric_aug=None,
                scale=config.data_config.preprocessing.scale,
                apply_aug=False,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=val_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=False,
            )

    elif model_type in ("bbox_centernet", "bbox_ssd"):
        head_configs = config.model_config.head_configs
        target_encoder = get_bbox_target_encoder(model_type, head_configs)
        categories = bbox_head_categories(model_type, head_configs)
        common = dict(
            target_encoder=target_encoder,
            categories=categories,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            scale=config.data_config.preprocessing.scale,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
        )
        train_dataset = BBoxDataset(
            labels=train_labels,
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            apply_aug=config.data_config.use_augmentations_train,
            cache_img_path=train_cache_img_path,
            use_negative_frames=use_negative_frames,
            **common,
        )
        val_dataset = BBoxDataset(
            labels=val_labels,
            intensity_aug=None,
            geometric_aug=None,
            apply_aug=False,
            cache_img_path=val_cache_img_path,
            use_negative_frames=False,
            **common,
        )

    elif model_type == "centered_instance_segmentation":
        seg_cfg = config.model_config.head_configs.centered_instance_segmentation
        anchor_part = seg_cfg.segmentation.anchor_part
        if anchor_part is not None:
            nodes = [x["name"] for x in config.data_config.skeletons[0]["nodes"]]
            anchor_ind = nodes.index(anchor_part)
        else:
            anchor_ind = None
        train_dataset = CenteredInstanceSegmentationDataset(
            labels=train_labels,
            crop_size=config.data_config.preprocessing.crop_size,
            seg_head_config=seg_cfg.segmentation,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            scale=config.data_config.preprocessing.scale,
            apply_aug=config.data_config.use_augmentations_train,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=train_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
        )
        val_dataset = CenteredInstanceSegmentationDataset(
            labels=val_labels,
            crop_size=config.data_config.preprocessing.crop_size,
            seg_head_config=seg_cfg.segmentation,
            max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                "max_stride"
            ],
            anchor_ind=anchor_ind,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=config.data_config.preprocessing.ensure_rgb,
            ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
            intensity_aug=None,
            geometric_aug=None,
            scale=config.data_config.preprocessing.scale,
            apply_aug=False,
            max_hw=(
                config.data_config.preprocessing.max_height,
                config.data_config.preprocessing.max_width,
            ),
            cache_img=cache_imgs,
            cache_img_path=val_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
        )

    elif model_type == "embedding":
        emb_cfg = config.model_config.head_configs.embedding.embedding
        # Whether a per-video track name may stand in as a global animal identity for
        # detections without a real `sio.Identity` (the pre-Identity convention), and
        # the training-group key (global identity vs per-video tracklet). The
        # post-training eval reads the same pair, so both group alike.
        track_names_are_global, id_scope = resolve_embedding_grouping(config)
        user_instances_only = bool(config.data_config.user_instances_only)
        # One global-identity vocabulary shared by train + val (the group_id space):
        # `sio.Identity` names, else track names under `track_names_are_global`.
        class_names = resolve_embedding_class_names(
            train_labels + val_labels,
            track_names_are_global=track_names_are_global,
            user_instances_only=user_instances_only,
        )
        # The carrier to crop. `ModelTrainer` records the one it resolved from the
        # train labels in `detection_mode` before this runs, so train and val crop
        # the same carrier and the saved config says which one inference must use.
        detection_mode = embedding_preferred_carrier(config)
        max_stride = config.model_config.backbone_config[f"{backbone_type}"][
            "max_stride"
        ]
        crop_size = config.data_config.preprocessing.crop_size
        max_hw = (
            config.data_config.preprocessing.max_height,
            config.data_config.preprocessing.max_width,
        )
        # Grayscale is the embedding default (helps the cross-video gap), but the user
        # can opt into RGB via `preprocessing.ensure_rgb`. A 3ch ImageNet backbone
        # repeats the gray channel in Model.forward, so grayscale data still works with
        # convnext/swint. When neither flag is set, default to grayscale.
        emb_ensure_rgb = bool(config.data_config.preprocessing.ensure_rgb)
        emb_ensure_grayscale = (
            bool(config.data_config.preprocessing.ensure_grayscale)
            or not emb_ensure_rgb
        )
        crop_centering = OmegaConf.select(
            config, "data_config.preprocessing.crop_centering", default="auto"
        )
        train_dataset = EmbeddingDataset(
            labels=train_labels,
            crop_size=crop_size,
            class_names=class_names,
            embedding_head_config=emb_cfg,
            max_stride=max_stride,
            id_scope=id_scope,
            track_names_are_global=track_names_are_global,
            crop_centering=crop_centering,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=emb_ensure_rgb,
            ensure_grayscale=emb_ensure_grayscale,
            # Two contrastive views via the standard config-driven skia aug (per crop).
            intensity_aug=(
                config.data_config.augmentation_config.intensity
                if config.data_config.augmentation_config is not None
                else None
            ),
            geometric_aug=(
                config.data_config.augmentation_config.geometric
                if config.data_config.augmentation_config is not None
                else None
            ),
            apply_aug=config.data_config.use_augmentations_train,
            scale=config.data_config.preprocessing.scale,
            max_hw=max_hw,
            cache_img=cache_imgs,
            cache_img_path=train_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            detection_mode=detection_mode,
        )
        val_dataset = EmbeddingDataset(
            labels=val_labels,
            crop_size=crop_size,
            class_names=class_names,
            embedding_head_config=emb_cfg,
            max_stride=max_stride,
            id_scope=id_scope,
            track_names_are_global=track_names_are_global,
            crop_centering=crop_centering,
            user_instances_only=config.data_config.user_instances_only,
            ensure_rgb=emb_ensure_rgb,
            ensure_grayscale=emb_ensure_grayscale,
            scale=config.data_config.preprocessing.scale,
            max_hw=max_hw,
            cache_img=cache_imgs,
            cache_img_path=val_cache_img_path,
            use_existing_imgs=use_existing_imgs,
            rank=rank,
            parallel_caching=parallel_caching,
            cache_workers=cache_workers,
            detection_mode=detection_mode,
        )

        if len(train_dataset) == 0:
            needs = {
                "tracklet": "a sio.Track",
                "global_id": (
                    "a sio.Identity (or a track name, under "
                    "data_config.identity.track_names_are_global)"
                ),
                "aug_view": "nothing but a non-empty pose or mask",
            }[id_scope]
            message = (
                "The embedding train dataset is empty: no detection is a sample under "
                f"positives.scope='{id_scope}', which needs {needs} "
                f"({len(class_names)} global identit(ies) found)"
                + (
                    ", and data_config.user_instances_only=True leaves out predictions "
                    "that a user label on the same frame supersedes"
                    if user_instances_only
                    else ""
                )
                + "."
            )
            logger.error(message)
            raise ValueError(message)

    else:
        tiling = OmegaConf.select(
            config, "data_config.preprocessing.tiling", default=None
        )
        if tiling is not None and tiling.enabled:
            # Tiled single-instance training. Train draws foreground-aware tiles;
            # val always uses a deterministic full-coverage grid (no aug).
            base_seed = config.trainer_config.seed or 0
            train_dataset = SingleInstanceTiledDataset(
                labels=train_labels,
                confmap_head_config=config.model_config.head_configs.single_instance.confmaps,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=(
                    config.data_config.augmentation_config.intensity
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                geometric_aug=(
                    config.data_config.augmentation_config.geometric
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                scale=config.data_config.preprocessing.scale,
                apply_aug=config.data_config.use_augmentations_train,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=train_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=use_negative_frames,
                tiling=OmegaConf.merge(tiling, {"sampling": tiling.sampling}),
                base_seed=base_seed,
            )
            val_dataset = SingleInstanceTiledDataset(
                labels=val_labels,
                confmap_head_config=config.model_config.head_configs.single_instance.confmaps,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=None,
                geometric_aug=None,
                scale=config.data_config.preprocessing.scale,
                apply_aug=False,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=val_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=use_negative_frames,
                tiling=OmegaConf.merge(tiling, {"sampling": "grid"}),
                base_seed=base_seed,
            )
        else:
            train_dataset = SingleInstanceDataset(
                labels=train_labels,
                confmap_head_config=config.model_config.head_configs.single_instance.confmaps,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=(
                    config.data_config.augmentation_config.intensity
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                geometric_aug=(
                    config.data_config.augmentation_config.geometric
                    if config.data_config.augmentation_config is not None
                    else None
                ),
                scale=config.data_config.preprocessing.scale,
                apply_aug=config.data_config.use_augmentations_train,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=train_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=use_negative_frames,
            )
            val_dataset = SingleInstanceDataset(
                labels=val_labels,
                confmap_head_config=config.model_config.head_configs.single_instance.confmaps,
                max_stride=config.model_config.backbone_config[f"{backbone_type}"][
                    "max_stride"
                ],
                user_instances_only=config.data_config.user_instances_only,
                ensure_rgb=config.data_config.preprocessing.ensure_rgb,
                ensure_grayscale=config.data_config.preprocessing.ensure_grayscale,
                intensity_aug=None,
                geometric_aug=None,
                scale=config.data_config.preprocessing.scale,
                apply_aug=False,
                max_hw=(
                    config.data_config.preprocessing.max_height,
                    config.data_config.preprocessing.max_width,
                ),
                cache_img=cache_imgs,
                cache_img_path=val_cache_img_path,
                use_existing_imgs=use_existing_imgs,
                rank=rank,
                parallel_caching=parallel_caching,
                cache_workers=cache_workers,
                use_negative_frames=use_negative_frames,
            )

    if cache_imgs == "disk" and use_existing_imgs:
        _validate_existing_disk_cache_complete(
            train_dataset, train_cache_img_path, "train"
        )
        _validate_existing_disk_cache_complete(val_dataset, val_cache_img_path, "val")

    # If using caching, close the videos to prevent `h5py objects can't be pickled error` when num_workers > 0.
    if "cache_img" in config.data_config.data_pipeline_fw:
        for train, val in zip(train_labels, val_labels):
            for video in train.videos:
                if video.is_open:
                    video.close()
            for video in val.videos:
                if video.is_open:
                    video.close()

    return train_dataset, val_dataset

labels_have_user_centroids(labels)

Return True if any labeled frame carries a usable UserCentroid.

Mirrors the per-frame extraction used to build the centroid sample list (BaseDataset._extract_user_centroid_xy): predicted centroids and NaN coordinates do not count, and an installed sleap-io without first-class centroids yields False.

Source code in sleap_nn/data/custom_datasets.py

def labels_have_user_centroids(labels: List[sio.Labels]) -> bool:
    """Return True if any labeled frame carries a usable ``UserCentroid``.

    Mirrors the per-frame extraction used to build the centroid sample list
    (``BaseDataset._extract_user_centroid_xy``): predicted centroids and NaN
    coordinates do not count, and an installed sleap-io without first-class
    centroids yields False.
    """
    for label in labels:
        for lf in label:
            if BaseDataset._extract_user_centroid_xy(lf):
                return True
    return False

predict_embedding_carrier(carriers, preferred=None)

The carrier :func:resolve_embedding_detection_mode will pick, before the labels exist.

Knowing only WHICH carriers will hold detections (e.g. from the model types of a detection stack), not how many: the preferred (trained) carrier if it is among them, else the only one there is.

Parameters:

Name	Type	Description	Default
carriers		The carriers (pose / mask) the labels will hold.	required
preferred	Optional[str]	:func:embedding_preferred_carrier of the model.	None

Returns:

Type	Description
	pose / mask, or None when the choice depends on the counts (both carriers and no preference) or there will be no carrier at all.

Source code in sleap_nn/data/custom_datasets.py

def predict_embedding_carrier(carriers, preferred: Optional[str] = None):
    """The carrier :func:`resolve_embedding_detection_mode` will pick, before the labels exist.

    Knowing only WHICH carriers will hold detections (e.g. from the model types of a
    detection stack), not how many: the preferred (trained) carrier if it is among
    them, else the only one there is.

    Args:
        carriers: The carriers (``pose`` / ``mask``) the labels will hold.
        preferred: :func:`embedding_preferred_carrier` of the model.

    Returns:
        ``pose`` / ``mask``, or ``None`` when the choice depends on the counts (both
        carriers and no preference) or there will be no carrier at all.
    """
    carriers = set(carriers)
    if preferred in carriers:
        return preferred
    if len(carriers) == 1:
        return next(iter(carriers))
    return None

resolve_centroid_source(centroid_source, train_labels)

Resolve the centroid target source to one dataset-wide mode.

Returns True to train on user-annotated centroids, False to compute centroids from instance keypoints (the anchor node, else the mean of visible nodes). The centroid model must use ONE source for the whole dataset; mixing the two trains the head against two different definitions of "centroid".

centroid_source (from CentroidConfMapsConfig) selects the mode explicitly — "user" or "computed" ("anchor" is accepted as an alias for "computed"). When it is None the mode is INFERRED from the training labels and a loud warning is emitted, because a silently-chosen target is a subtle training footgun.

Raises:

Type	Description
ValueError	if centroid_source is a non-empty unrecognized string.

Source code in sleap_nn/data/custom_datasets.py

def resolve_centroid_source(
    centroid_source: Optional[str], train_labels: List[sio.Labels]
) -> bool:
    """Resolve the centroid target source to one dataset-wide mode.

    Returns ``True`` to train on user-annotated centroids, ``False`` to compute
    centroids from instance keypoints (the anchor node, else the mean of
    visible nodes). The centroid model must use ONE source for the whole
    dataset; mixing the two trains the head against two different definitions of
    "centroid".

    ``centroid_source`` (from ``CentroidConfMapsConfig``) selects the mode
    explicitly — ``"user"`` or ``"computed"`` (``"anchor"`` is accepted as an
    alias for ``"computed"``). When it is ``None`` the mode is INFERRED from the
    training labels and a loud warning is emitted, because a silently-chosen
    target is a subtle training footgun.

    Raises:
        ValueError: if ``centroid_source`` is a non-empty unrecognized string.
    """
    if centroid_source is not None:
        source = str(centroid_source).strip().lower()
        if source == CENTROID_SOURCE_USER:
            logger.info(
                "Centroid target source: user-annotated centroids "
                "(centroid_source='user')."
            )
            return True
        if source in (CENTROID_SOURCE_COMPUTED, "anchor"):
            logger.info(
                "Centroid target source: computed from instance keypoints "
                "(centroid_source='computed')."
            )
            return False
        raise ValueError(
            f"Invalid centroid_source={centroid_source!r}. Expected 'user', "
            f"'computed', or None (infer from the training labels)."
        )

    has_user = labels_have_user_centroids(train_labels)
    chosen = (
        "user-annotated centroids (UserCentroid)"
        if has_user
        else "computed centroids (anchor node / mean of visible nodes)"
    )
    bar = "=" * 76
    logger.warning(
        "\n%s\n"
        "centroid_source is NOT set: INFERRING the centroid target from the "
        "training labels.\n"
        "  -> Training on %s for ALL frames.\n"
        "Set model_config.head_configs.centroid.confmaps.centroid_source to "
        "'user' or\n'computed' to make this explicit and silence this warning.\n"
        "%s",
        bar,
        chosen,
        bar,
    )
    return has_user

resolve_embedding_class_names(labels, track_names_are_global=True, user_instances_only=False)

Collect the sorted global-identity vocabulary (the global_id / eval grouping).

The vocabulary is the set of GLOBAL animal identities — a real sio.Identity name when a detection carries one, else its sio.Track name under the track_names_are_global promise (see :func:_global_identity_label). Resolving per detection (not per track) means an animal with both an Identity and a per-video Track contributes only its identity, so identity- and track-labelled data share one coherent vocabulary. Scanning BOTH train + val labels and sorting gives one consistent vocabulary shared by the train and val datasets.

With user_instances_only (data_config.user_instances_only), predictions a user label on the same frame supersedes are left out, as they are left out of the dataset itself (:func:~sleap_nn.data.utils.embedding_detections).

Source code in sleap_nn/data/custom_datasets.py

def resolve_embedding_class_names(
    labels: List[sio.Labels],
    track_names_are_global: bool = True,
    user_instances_only: bool = False,
) -> List[str]:
    """Collect the sorted global-identity vocabulary (the ``global_id`` / eval grouping).

    The vocabulary is the set of GLOBAL animal identities — a real ``sio.Identity``
    name when a detection carries one, else its ``sio.Track`` name under the
    ``track_names_are_global`` promise (see :func:`_global_identity_label`). Resolving
    per detection (not per track) means an animal with both an ``Identity`` and a
    per-video ``Track`` contributes only its identity, so identity- and track-labelled
    data share one coherent vocabulary. Scanning BOTH train + val labels and sorting
    gives one consistent vocabulary shared by the train and val datasets.

    With ``user_instances_only`` (``data_config.user_instances_only``), predictions a
    user label on the same frame supersedes are left out, as they are left out of the
    dataset itself (:func:`~sleap_nn.data.utils.embedding_detections`).
    """
    names = set()
    for label in labels:
        for lf in label:
            instances, masks = embedding_detections(lf, user_instances_only)
            for det in instances + masks:
                lab = _global_identity_label(det, track_names_are_global)
                if lab is not None:
                    names.add(lab)
    return sorted(names)

resolve_embedding_detection_mode(labels, is_member, preferred=None)

Pick the carrier (pose / mask) an embedding dataset crops.

Counts the detections on each carrier that is_member accepts (the dataset's own membership rule, so identity-only masks and untracked detections count the way the dataset will index them) with :func:sleap_nn.inference.tracking.count_carriers.

• preferred (the carrier the model trains / was trained on, see :func:embedding_preferred_carrier) wins whenever it holds any member; if it holds none the other carrier is used.
• With no preference, the carrier holding MORE members wins. Ties go to masks: the historical rule, and so the carrier every model saved before detection_mode was recorded was trained on when its file held both. (count_carriers' own tie rule, pose, answers a different question: which carrier to TRACK when both carry vectors.)

Source code in sleap_nn/data/custom_datasets.py

def resolve_embedding_detection_mode(
    labels: List[sio.Labels], is_member, preferred: Optional[str] = None
) -> str:
    """Pick the carrier (``pose`` / ``mask``) an embedding dataset crops.

    Counts the detections on each carrier that ``is_member`` accepts (the dataset's
    own membership rule, so identity-only masks and untracked detections count the
    way the dataset will index them) with
    :func:`sleap_nn.inference.tracking.count_carriers`.

    - ``preferred`` (the carrier the model trains / was trained on, see
      :func:`embedding_preferred_carrier`) wins whenever it holds any member; if it
      holds none the other carrier is used.
    - With no preference, the carrier holding MORE members wins. Ties go to masks:
      the historical rule, and so the carrier every model saved before
      ``detection_mode`` was recorded was trained on when its file held both.
      (``count_carriers``' own tie rule, pose, answers a different question: which
      carrier to TRACK when both carry vectors.)
    """
    from sleap_nn.inference.tracking import count_carriers  # lazy: no import cycle

    n_pose = n_mask = 0
    for label in labels:
        counts = count_carriers(label, is_member)
        n_pose += counts.n_pose_with
        n_mask += counts.n_mask_with
    n_with = {"pose": n_pose, "mask": n_mask}
    if preferred is not None:
        if preferred not in EMBEDDING_DETECTION_MODES:
            raise ValueError(
                f"Unknown embedding detection_mode {preferred!r}; choose one of "
                f"{'|'.join(EMBEDDING_DETECTION_MODES)}."
            )
        if n_with[preferred] > 0:
            return preferred
        other = "mask" if preferred == "pose" else "pose"
        if n_with[other] > 0:
            logger.warning(
                f"Embedding carrier: the model uses the {preferred} carrier "
                "(head_configs.embedding.embedding.detection_mode, or masks under "
                f"burn_in), but these labels hold no {preferred} sample; using their "
                f"{n_with[other]} {other} detection(s) instead."
            )
            return other
        return preferred
    if n_mask == 0:
        return "pose"
    return "mask" if n_mask >= n_pose else "pose"

resolve_embedding_grouping(config)

Return (track_names_are_global, id_scope) for an embedding training config.

These two settings decide which detections an :class:EmbeddingDataset keeps and how it groups them into identities. The per-epoch validation set (built by :func:get_train_val_datasets) and the post-training retrieval eval (:func:sleap_nn.inference.embedding.embed_labels_for_eval) both read them here, so the metric that selects the checkpoint and the reported one group detections the same way.

Parameters:

Name	Type	Description	Default
config	DictConfig	The training config (data_config + model_config), e.g. the training_config.yaml saved next to a trained model.	required

Returns:

Type	Description
Tuple[bool, str]	track_names_are_global (data_config.identity, default False) and the objective's positives.scope as resolved by :func:~sleap_nn.config.model_config.resolve_embedding_objective.

Source code in sleap_nn/data/custom_datasets.py

def resolve_embedding_grouping(config: DictConfig) -> Tuple[bool, str]:
    """Return ``(track_names_are_global, id_scope)`` for an embedding training config.

    These two settings decide which detections an :class:`EmbeddingDataset` keeps and
    how it groups them into identities. The per-epoch validation set (built by
    :func:`get_train_val_datasets`) and the post-training retrieval eval
    (:func:`sleap_nn.inference.embedding.embed_labels_for_eval`) both read them here,
    so the metric that selects the checkpoint and the reported one group detections
    the same way.

    Args:
        config: The training config (``data_config`` + ``model_config``), e.g. the
            ``training_config.yaml`` saved next to a trained model.

    Returns:
        ``track_names_are_global`` (``data_config.identity``, default ``False``) and
        the objective's ``positives.scope`` as resolved by
        :func:`~sleap_nn.config.model_config.resolve_embedding_objective`.
    """
    track_names_are_global = bool(
        OmegaConf.select(
            config, "data_config.identity.track_names_are_global", default=False
        )
    )
    objective = OmegaConf.select(
        config,
        "model_config.head_configs.embedding.embedding.objective",
        default=None,
    )
    id_scope = resolve_embedding_objective(objective).positives.scope
    return track_names_are_global, id_scope