lightning_modules

sleap_nn.training.lightning_modules

This module has the LightningModule classes for all model types.

Classes:

Name	Description
BBoxCenterNetLightningModule	Lightning Module for the bbox_centernet (CenterNet-style) detector.
BBoxSSDLightningModule	Lightning Module for the bbox_ssd (SSD-style, anchor-based) detector.
BottomUpLightningModule	Lightning Module for BottomUp Model.
BottomUpMultiClassLightningModule	Lightning Module for BottomUp ID Model.
BottomUpSegmentationLightningModule	Lightning Module for Bottom-Up Instance Segmentation.
CentroidLightningModule	Lightning Module for Centroid Model.
EmbeddingLightningModule	Lightning Module for the embedding (crop -> vector, re-ID) model type.
LightningModel	Base PyTorch Lightning Module for all sleap-nn models.
SemanticSegmentationLightningModule	Lightning Module for whole-frame semantic (foreground/background) segmentation.
SingleInstanceLightningModule	Lightning Module for SingleInstance Model.
TopDownCenteredInstanceLightningModule	Lightning Module for TopDownCenteredInstance Model.
TopDownCenteredInstanceMultiClassLightningModule	Lightning Module for TopDownCenteredInstance ID Model.
TopDownCenteredInstanceSegmentationLightningModule	Lightning Module for top-down (crop-centered) instance segmentation (#622).

Functions:

Name	Description
set_embedding_burn_in_from_config	Honor data_config.preprocessing.burn_in on an EmbeddingLightningModule.
validate_embedding_identity	Enforce the identity-equality gates for the embedding objective.
validate_embedding_views	Check that the two contrastive views of a crop can differ.

BBoxCenterNetLightningModule

Bases: _BBoxLightningModel

Lightning Module for the bbox_centernet (CenterNet-style) detector.

Three heads at one stride -- a per-category center heatmap (logits), box size (grid cells) and sub-cell center offset -- trained with CenterNet's objective: the penalty-reduced focal loss on the heatmap (normalized by the number of objects) plus L1 on size and offset at the center cells, weighted by the heads' loss_weight (CenterNet's 1 / 0.1 / 1 by default). forward returns the heatmap as probabilities (sigmoid applied) for the decoder, mirroring :class:SemanticSegmentationLightningModule.

Methods:

Name	Description
__init__	Initialise the configs, the model and the loss knobs.

Source code in sleap_nn/training/lightning_modules.py

class BBoxCenterNetLightningModule(_BBoxLightningModel):
    """Lightning Module for the ``bbox_centernet`` (CenterNet-style) detector.

    Three heads at one stride -- a per-category center heatmap (logits), box size
    (grid cells) and sub-cell center offset -- trained with CenterNet's objective:
    the penalty-reduced focal loss on the heatmap (normalized by the number of
    objects) plus L1 on size and offset at the center cells, weighted by the heads'
    ``loss_weight`` (CenterNet's 1 / 0.1 / 1 by default). ``forward`` returns the
    heatmap as probabilities (sigmoid applied) for the decoder, mirroring
    :class:`SemanticSegmentationLightningModule`.
    """

    def __init__(self, **kwargs):
        """Initialise the configs, the model and the loss knobs."""
        super().__init__(**kwargs)
        cfg = self.head_configs.bbox_centernet
        self.focal_alpha = float(cfg.center.focal_alpha)
        self.focal_beta = float(cfg.center.focal_beta)
        self.center_weight = float(cfg.center.loss_weight)
        self.size_weight = float(cfg.size.loss_weight)
        self.offset_weight = float(cfg.offset.loss_weight)
        # Focal-loss prior-probability bias init (RetinaNet / CenterNet): start the
        # heatmap near p=0.1 instead of 0.5 so >99%-background targets do not
        # swamp the first steps (see `CentroidLightningModule`).
        bias_value = -math.log((1.0 - 0.1) / 0.1)
        for head, head_layer in zip(self.model.heads, self.model.head_layers):
            if head.name == "BBoxCenterHead":
                nn.init.constant_(getattr(head_layer, head.name)[0].bias, bias_value)

    def _decoder_inputs(self, raw):
        """Sigmoid the center heatmap; size/offset pass through."""
        out = dict(raw)
        out["BBoxCenterHead"] = torch.sigmoid(raw["BBoxCenterHead"])
        return out

    def _viz_heatmap(self, dec_in):
        heatmap = dec_in["BBoxCenterHead"][0].max(dim=0, keepdim=True).values
        return heatmap.cpu().numpy().transpose(1, 2, 0)

    def _compute_losses(self, raw, batch):
        # `BBoxDataset` targets are (B, C, h, w) -- no n_samples axis to squeeze.
        # (Squeezing dim 1 of a 1-class target would broadcast it against the
        # (B, 1, h, w) logits into a (B, B, h, w) cross-sample loss.)
        heatmap_target = batch["center_heatmap"]
        reg_mask = batch["reg_mask"]
        center_loss = compute_centernet_focal_loss(
            raw["BBoxCenterHead"],
            heatmap_target,
            alpha=self.focal_alpha,
            beta=self.focal_beta,
        )
        size_loss = compute_masked_l1(
            raw["BBoxSizeHead"], batch["size_target"], reg_mask
        )
        offset_loss = compute_masked_l1(
            raw["BBoxOffsetHead"], batch["offset_target"], reg_mask
        )
        loss = (
            self.center_weight * center_loss
            + self.size_weight * size_loss
            + self.offset_weight * offset_loss
        )
        return loss, {
            "center_loss": center_loss,
            "size_loss": size_loss,
            "offset_loss": offset_loss,
        }

__init__(**kwargs)

Initialise the configs, the model and the loss knobs.

Source code in sleap_nn/training/lightning_modules.py

def __init__(self, **kwargs):
    """Initialise the configs, the model and the loss knobs."""
    super().__init__(**kwargs)
    cfg = self.head_configs.bbox_centernet
    self.focal_alpha = float(cfg.center.focal_alpha)
    self.focal_beta = float(cfg.center.focal_beta)
    self.center_weight = float(cfg.center.loss_weight)
    self.size_weight = float(cfg.size.loss_weight)
    self.offset_weight = float(cfg.offset.loss_weight)
    # Focal-loss prior-probability bias init (RetinaNet / CenterNet): start the
    # heatmap near p=0.1 instead of 0.5 so >99%-background targets do not
    # swamp the first steps (see `CentroidLightningModule`).
    bias_value = -math.log((1.0 - 0.1) / 0.1)
    for head, head_layer in zip(self.model.heads, self.model.head_layers):
        if head.name == "BBoxCenterHead":
            nn.init.constant_(getattr(head_layer, head.name)[0].bias, bias_value)

BBoxSSDLightningModule

Bases: _BBoxLightningModel

Lightning Module for the bbox_ssd (SSD-style, anchor-based) detector.

One multi-scale head predicts, per anchor of every pyramid level, class logits (background first) and SSD box deltas. Trained with SSD's multibox loss: softmax cross-entropy over positives plus the hardest neg_pos_ratio background anchors per positive, and smooth-L1 on the positives' deltas (:func:~sleap_nn.training.losses.compute_ssd_loss). forward returns the raw SSDHead tensor; the decoder applies the softmax.

Methods:

Name	Description
__init__	Initialise the configs, the model and the loss knobs.

Source code in sleap_nn/training/lightning_modules.py

class BBoxSSDLightningModule(_BBoxLightningModel):
    """Lightning Module for the ``bbox_ssd`` (SSD-style, anchor-based) detector.

    One multi-scale head predicts, per anchor of every pyramid level, class logits
    (background first) and SSD box deltas. Trained with SSD's multibox loss: softmax
    cross-entropy over positives plus the hardest ``neg_pos_ratio`` background
    anchors per positive, and smooth-L1 on the positives' deltas
    (:func:`~sleap_nn.training.losses.compute_ssd_loss`). ``forward`` returns the
    raw ``SSDHead`` tensor; the decoder applies the softmax.
    """

    def __init__(self, **kwargs):
        """Initialise the configs, the model and the loss knobs."""
        super().__init__(**kwargs)
        ssd = self.head_configs.bbox_ssd.ssd
        self.neg_pos_ratio = float(ssd.neg_pos_ratio)
        self.box_weight = float(ssd.loss_weight)

    def _compute_losses(self, raw, batch):
        loss, cls_loss, box_loss = compute_ssd_loss(
            raw["SSDHead"],
            batch["ssd_cls_target"],
            batch["ssd_box_target"],
            neg_pos_ratio=self.neg_pos_ratio,
            box_weight=self.box_weight,
        )
        return loss, {"cls_loss": cls_loss, "box_loss": box_loss}

__init__(**kwargs)

Initialise the configs, the model and the loss knobs.

Source code in sleap_nn/training/lightning_modules.py

def __init__(self, **kwargs):
    """Initialise the configs, the model and the loss knobs."""
    super().__init__(**kwargs)
    ssd = self.head_configs.bbox_ssd.ssd
    self.neg_pos_ratio = float(ssd.neg_pos_ratio)
    self.box_weight = float(ssd.loss_weight)

BottomUpLightningModule

Bases: LightningModel

Lightning Module for BottomUp Model.

This is a subclass of the LightningModel to configure the training/ validation steps and forward pass specific to BottomUp model. Bottom-Up models predict all keypoints simultaneously and use Part Affinity Fields (PAFs) to group keypoints into individual animals.

Parameters:

Name	Type	Description	Default
model_type	str	Type of the model. One of single_instance, centered_instance, centroid, bottomup, multi_class_bottomup, multi_class_topdown.	required
backbone_type	str	Backbone model. One of unet, convnext and swint.	required
backbone_config	Union[str, Dict[str, Any], DictConfig]	Backbone configuration. Can be: - String: One of the preset backbone types: - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"] - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"] - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"] - Dictionary: Custom configuration with structure: { "unet": {UNetConfig parameters}, "convnext": {ConvNextConfig parameters}, "swint": {SwinTConfig parameters} } Only one backbone type should be specified in the dictionary. - DictConfig: OmegaConf DictConfig object containing backbone configuration.	required
head_configs	DictConfig	Head configuration dictionary containing model-specific parameters. For Single Instance: confmaps with part_names, sigma, output_stride. For Centroid: confmaps with anchor_part, sigma, output_stride. For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride. For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight. For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight. For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.	required
pretrained_backbone_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.	None
pretrained_head_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.	None
init_weights	Optional[str]	Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.	'xavier'
lr_scheduler	Optional[Union[str, DictConfig]]	Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.	None
online_mining	Optional[bool]	If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).	False
hard_to_easy_ratio	Optional[float]	Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.	2.0
min_hard_keypoints	Optional[int]	Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.	2
max_hard_keypoints	Optional[int]	Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.	None
loss_scale	Optional[float]	Factor to scale hard keypoint losses by. Default: 5.0.	5.0
optimizer	Optional[str]	Optimizer name. One of ["Adam", "AdamW"].	'Adam'
learning_rate	Optional[float]	Learning rate for the optimizer. Default: 1e-3.	0.001
amsgrad	Optional[bool]	Enable AMSGrad with the optimizer. Default: False.	False

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model.
get_visualization_data	Extract visualization data from a sample.
training_step	Training step.
validation_step	Validation step.
visualize_example	Visualize predictions during training (used with callbacks).
visualize_pafs_example	Visualize PAF predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

class BottomUpLightningModule(LightningModel):
    """Lightning Module for BottomUp Model.

    This is a subclass of the `LightningModel` to configure the training/ validation steps
    and forward pass specific to BottomUp model. Bottom-Up models predict all keypoints
    simultaneously and use Part Affinity Fields (PAFs) to group keypoints into individual animals.

    Args:
        model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
        backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
        backbone_config: Backbone configuration. Can be:
            - String: One of the preset backbone types:
                - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"]
                - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"]
                - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"]
            - Dictionary: Custom configuration with structure:
                {
                    "unet": {UNetConfig parameters},
                    "convnext": {ConvNextConfig parameters},
                    "swint": {SwinTConfig parameters}
                }
                Only one backbone type should be specified in the dictionary.
            - DictConfig: OmegaConf DictConfig object containing backbone configuration.
        head_configs: Head configuration dictionary containing model-specific parameters.
            For Single Instance: confmaps with part_names, sigma, output_stride.
            For Centroid: confmaps with anchor_part, sigma, output_stride.
            For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride.
            For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight.
            For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight.
            For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.
        pretrained_backbone_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.
        pretrained_head_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.
        init_weights: Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.
        lr_scheduler: Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.
        online_mining: If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).
        hard_to_easy_ratio: Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.
        min_hard_keypoints: Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.
        max_hard_keypoints: Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.
        loss_scale: Factor to scale hard keypoint losses by. Default: 5.0.
        optimizer: Optimizer name. One of ["Adam", "AdamW"].
        learning_rate: Learning rate for the optimizer. Default: 1e-3.
        amsgrad: Enable AMSGrad with the optimizer. Default: False.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )

        paf_scorer = PAFScorer(
            part_names=self.head_configs.bottomup.confmaps.part_names,
            edges=self.head_configs.bottomup.pafs.edges,
            pafs_stride=self.head_configs.bottomup.pafs.output_stride,
        )
        self.bottomup_inf_layer = BottomUpInferenceModel(
            torch_model=self.forward,
            paf_scorer=paf_scorer,
            peak_threshold=0.1,  # Lower threshold for epoch-end eval during training
            input_scale=1.0,
            return_confmaps=True,
            return_pafs=True,
            cms_output_stride=self.head_configs.bottomup.confmaps.output_stride,
            pafs_output_stride=self.head_configs.bottomup.pafs.output_stride,
            max_peaks_per_node=100,  # Prevents combinatorial explosion in early training
        )
        self.node_names = list(self.head_configs.bottomup.confmaps.part_names)

    def get_visualization_data(
        self, sample, include_pafs: bool = False
    ) -> VisualizationData:
        """Extract visualization data from a sample."""
        ex = sample.copy()
        ex["eff_scale"] = torch.tensor([1.0])
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["image"] = ex["image"].unsqueeze(dim=0)
        output = self.bottomup_inf_layer(ex)[0]

        peaks = output["pred_instance_peaks"][0].cpu().numpy()
        peak_values = output["pred_peak_values"][0].cpu().numpy()
        img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        gt_instances = ex["instances"][0].cpu().numpy()
        confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

        pred_pafs = None
        if include_pafs:
            pafs = output["pred_part_affinity_fields"].cpu().numpy()[0]
            pred_pafs = pafs  # (h, w, 2*edges)

        return VisualizationData(
            image=img,
            pred_confmaps=confmaps,
            pred_peaks=peaks,
            pred_peak_values=peak_values,
            gt_instances=gt_instances,
            node_names=self.node_names,
            output_scale=confmaps.shape[0] / img.shape[0],
            is_paired=False,
            pred_pafs=pred_pafs,
        )

    def visualize_example(self, sample):
        """Visualize predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plt.xlim(plt.xlim())
        plt.ylim(plt.ylim())
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def visualize_pafs_example(self, sample):
        """Visualize PAF predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample, include_pafs=True)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)

        pafs = data.pred_pafs
        pafs = pafs.reshape((pafs.shape[0], pafs.shape[1], -1, 2))
        pafs_mag = np.sqrt(pafs[..., 0] ** 2 + pafs[..., 1] ** 2)
        plot_confmaps(pafs_mag, output_scale=pafs_mag.shape[0] / data.image.shape[0])
        return fig

    def forward(self, img):
        """Forward pass of the model."""
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        output = self.model(img)
        return {
            "MultiInstanceConfmapsHead": output["MultiInstanceConfmapsHead"],
            "PartAffinityFieldsHead": output["PartAffinityFieldsHead"],
        }

    def training_step(self, batch, batch_idx):
        """Training step."""
        X = torch.squeeze(batch["image"], dim=1)
        y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
        y_paf = batch["part_affinity_fields"]
        X = normalize_on_gpu(X)
        preds = self.model(X)
        pafs = preds["PartAffinityFieldsHead"]
        confmaps = preds["MultiInstanceConfmapsHead"]

        confmap_loss = self._compute_negative_weighted_loss(confmaps, y_confmap, batch)
        pafs_loss = self._compute_negative_weighted_loss(pafs, y_paf, batch)
        self._log_negative_split_metrics(
            [
                ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
                ("paf", pafs, y_paf, self.loss_weights[1]),
            ],
            batch,
            stage="train",
        )

        if self.online_mining is not None and self.online_mining:
            confmap_ohkm_loss = compute_ohkm_loss(
                y_gt=y_confmap,
                y_pr=confmaps,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            pafs_ohkm_loss = compute_ohkm_loss(
                y_gt=y_paf,
                y_pr=pafs,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            confmap_loss += confmap_ohkm_loss
            pafs_loss += pafs_ohkm_loss

        losses = {
            "MultiInstanceConfmapsHead": confmap_loss,
            "PartAffinityFieldsHead": pafs_loss,
        }
        loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
        # Log step-level loss (every batch, uses global_step x-axis)
        self.log(
            "loss",
            loss,
            prog_bar=True,
            on_step=True,
            on_epoch=False,
            sync_dist=True,
        )

        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
        self._accumulate_loss(loss)
        self.log(
            "train/confmaps_loss",
            confmap_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "train/paf_loss",
            pafs_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step."""
        X = torch.squeeze(batch["image"], dim=1)
        y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
        y_paf = batch["part_affinity_fields"]
        X = normalize_on_gpu(X)

        preds = self.model(X)
        pafs = preds["PartAffinityFieldsHead"]
        confmaps = preds["MultiInstanceConfmapsHead"]

        confmap_loss = self._compute_negative_weighted_loss(
            confmaps, y_confmap, batch, stage="val"
        )
        pafs_loss = self._compute_negative_weighted_loss(
            pafs, y_paf, batch, stage="val"
        )
        self._log_negative_split_metrics(
            [
                ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
                ("paf", pafs, y_paf, self.loss_weights[1]),
            ],
            batch,
            stage="val",
        )

        if self.online_mining is not None and self.online_mining:
            confmap_ohkm_loss = compute_ohkm_loss(
                y_gt=y_confmap,
                y_pr=confmaps,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            pafs_ohkm_loss = compute_ohkm_loss(
                y_gt=y_paf,
                y_pr=pafs,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            confmap_loss += confmap_ohkm_loss
            pafs_loss += pafs_ohkm_loss

        losses = {
            "MultiInstanceConfmapsHead": confmap_loss,
            "PartAffinityFieldsHead": pafs_loss,
        }

        val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "val/confmaps_loss",
            confmap_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "val/paf_loss",
            pafs_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Collect predictions for epoch-end evaluation if enabled
        if self._collect_val_predictions:
            with torch.no_grad():
                # Note: Do NOT squeeze the image here - the forward() method expects
                # (batch, n_samples, C, H, W) and handles the n_samples squeeze internally
                inference_output = self.bottomup_inf_layer(batch)
                if isinstance(inference_output, list):
                    inference_output = inference_output[0]

            batch_size = len(batch["frame_idx"])
            for i in range(batch_size):
                eff = batch["eff_scale"][i].cpu().numpy()

                # Predictions are already in original space (variable number of instances)
                pred_peaks = inference_output["pred_instance_peaks"][i]
                pred_scores = inference_output["pred_peak_values"][i]
                if torch.is_tensor(pred_peaks):
                    pred_peaks = pred_peaks.cpu().numpy()
                if torch.is_tensor(pred_scores):
                    pred_scores = pred_scores.cpu().numpy()

                # Transform GT to original space
                # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
                gt_prep = batch["instances"][i].cpu().numpy()
                if gt_prep.ndim == 4:
                    gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
                gt_orig = gt_prep / eff
                num_inst = batch["num_instances"][i].item()
                gt_orig = gt_orig[:num_inst]  # Only valid instances

                self.val_predictions.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "pred_peaks": pred_peaks,  # Original space, variable instances
                        "pred_scores": pred_scores,
                    }
                )
                self.val_ground_truth.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "gt_instances": gt_orig,  # Original space
                        "num_instances": num_inst,
                    }
                )

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )

    paf_scorer = PAFScorer(
        part_names=self.head_configs.bottomup.confmaps.part_names,
        edges=self.head_configs.bottomup.pafs.edges,
        pafs_stride=self.head_configs.bottomup.pafs.output_stride,
    )
    self.bottomup_inf_layer = BottomUpInferenceModel(
        torch_model=self.forward,
        paf_scorer=paf_scorer,
        peak_threshold=0.1,  # Lower threshold for epoch-end eval during training
        input_scale=1.0,
        return_confmaps=True,
        return_pafs=True,
        cms_output_stride=self.head_configs.bottomup.confmaps.output_stride,
        pafs_output_stride=self.head_configs.bottomup.pafs.output_stride,
        max_peaks_per_node=100,  # Prevents combinatorial explosion in early training
    )
    self.node_names = list(self.head_configs.bottomup.confmaps.part_names)

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    output = self.model(img)
    return {
        "MultiInstanceConfmapsHead": output["MultiInstanceConfmapsHead"],
        "PartAffinityFieldsHead": output["PartAffinityFieldsHead"],
    }

get_visualization_data(sample, include_pafs=False)

Extract visualization data from a sample.

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(
    self, sample, include_pafs: bool = False
) -> VisualizationData:
    """Extract visualization data from a sample."""
    ex = sample.copy()
    ex["eff_scale"] = torch.tensor([1.0])
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["image"] = ex["image"].unsqueeze(dim=0)
    output = self.bottomup_inf_layer(ex)[0]

    peaks = output["pred_instance_peaks"][0].cpu().numpy()
    peak_values = output["pred_peak_values"][0].cpu().numpy()
    img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    gt_instances = ex["instances"][0].cpu().numpy()
    confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

    pred_pafs = None
    if include_pafs:
        pafs = output["pred_part_affinity_fields"].cpu().numpy()[0]
        pred_pafs = pafs  # (h, w, 2*edges)

    return VisualizationData(
        image=img,
        pred_confmaps=confmaps,
        pred_peaks=peaks,
        pred_peak_values=peak_values,
        gt_instances=gt_instances,
        node_names=self.node_names,
        output_scale=confmaps.shape[0] / img.shape[0],
        is_paired=False,
        pred_pafs=pred_pafs,
    )

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    X = torch.squeeze(batch["image"], dim=1)
    y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
    y_paf = batch["part_affinity_fields"]
    X = normalize_on_gpu(X)
    preds = self.model(X)
    pafs = preds["PartAffinityFieldsHead"]
    confmaps = preds["MultiInstanceConfmapsHead"]

    confmap_loss = self._compute_negative_weighted_loss(confmaps, y_confmap, batch)
    pafs_loss = self._compute_negative_weighted_loss(pafs, y_paf, batch)
    self._log_negative_split_metrics(
        [
            ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
            ("paf", pafs, y_paf, self.loss_weights[1]),
        ],
        batch,
        stage="train",
    )

    if self.online_mining is not None and self.online_mining:
        confmap_ohkm_loss = compute_ohkm_loss(
            y_gt=y_confmap,
            y_pr=confmaps,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        pafs_ohkm_loss = compute_ohkm_loss(
            y_gt=y_paf,
            y_pr=pafs,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        confmap_loss += confmap_ohkm_loss
        pafs_loss += pafs_ohkm_loss

    losses = {
        "MultiInstanceConfmapsHead": confmap_loss,
        "PartAffinityFieldsHead": pafs_loss,
    }
    loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
    # Log step-level loss (every batch, uses global_step x-axis)
    self.log(
        "loss",
        loss,
        prog_bar=True,
        on_step=True,
        on_epoch=False,
        sync_dist=True,
    )

    # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
    self._accumulate_loss(loss)
    self.log(
        "train/confmaps_loss",
        confmap_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "train/paf_loss",
        pafs_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    return loss

validation_step(batch, batch_idx)

Validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step."""
    X = torch.squeeze(batch["image"], dim=1)
    y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
    y_paf = batch["part_affinity_fields"]
    X = normalize_on_gpu(X)

    preds = self.model(X)
    pafs = preds["PartAffinityFieldsHead"]
    confmaps = preds["MultiInstanceConfmapsHead"]

    confmap_loss = self._compute_negative_weighted_loss(
        confmaps, y_confmap, batch, stage="val"
    )
    pafs_loss = self._compute_negative_weighted_loss(
        pafs, y_paf, batch, stage="val"
    )
    self._log_negative_split_metrics(
        [
            ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
            ("paf", pafs, y_paf, self.loss_weights[1]),
        ],
        batch,
        stage="val",
    )

    if self.online_mining is not None and self.online_mining:
        confmap_ohkm_loss = compute_ohkm_loss(
            y_gt=y_confmap,
            y_pr=confmaps,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        pafs_ohkm_loss = compute_ohkm_loss(
            y_gt=y_paf,
            y_pr=pafs,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        confmap_loss += confmap_ohkm_loss
        pafs_loss += pafs_ohkm_loss

    losses = {
        "MultiInstanceConfmapsHead": confmap_loss,
        "PartAffinityFieldsHead": pafs_loss,
    }

    val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "val/confmaps_loss",
        confmap_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "val/paf_loss",
        pafs_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Collect predictions for epoch-end evaluation if enabled
    if self._collect_val_predictions:
        with torch.no_grad():
            # Note: Do NOT squeeze the image here - the forward() method expects
            # (batch, n_samples, C, H, W) and handles the n_samples squeeze internally
            inference_output = self.bottomup_inf_layer(batch)
            if isinstance(inference_output, list):
                inference_output = inference_output[0]

        batch_size = len(batch["frame_idx"])
        for i in range(batch_size):
            eff = batch["eff_scale"][i].cpu().numpy()

            # Predictions are already in original space (variable number of instances)
            pred_peaks = inference_output["pred_instance_peaks"][i]
            pred_scores = inference_output["pred_peak_values"][i]
            if torch.is_tensor(pred_peaks):
                pred_peaks = pred_peaks.cpu().numpy()
            if torch.is_tensor(pred_scores):
                pred_scores = pred_scores.cpu().numpy()

            # Transform GT to original space
            # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
            gt_prep = batch["instances"][i].cpu().numpy()
            if gt_prep.ndim == 4:
                gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
            gt_orig = gt_prep / eff
            num_inst = batch["num_instances"][i].item()
            gt_orig = gt_orig[:num_inst]  # Only valid instances

            self.val_predictions.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "pred_peaks": pred_peaks,  # Original space, variable instances
                    "pred_scores": pred_scores,
                }
            )
            self.val_ground_truth.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "gt_instances": gt_orig,  # Original space
                    "num_instances": num_inst,
                }
            )

visualize_example(sample)

Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plt.xlim(plt.xlim())
    plt.ylim(plt.ylim())
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

visualize_pafs_example(sample)

Visualize PAF predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_pafs_example(self, sample):
    """Visualize PAF predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample, include_pafs=True)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)

    pafs = data.pred_pafs
    pafs = pafs.reshape((pafs.shape[0], pafs.shape[1], -1, 2))
    pafs_mag = np.sqrt(pafs[..., 0] ** 2 + pafs[..., 1] ** 2)
    plot_confmaps(pafs_mag, output_scale=pafs_mag.shape[0] / data.image.shape[0])
    return fig

BottomUpMultiClassLightningModule

Bases: LightningModel

Lightning Module for BottomUp ID Model.

This is a subclass of the LightningModel to configure the training/ validation steps and forward pass specific to BottomUp ID model. Multi-Class Bottom-Up models predict all keypoints simultaneously and classify instances using class maps to identify individual animals across frames.

Parameters:

Name	Type	Description	Default
model_type	str	Type of the model. One of single_instance, centered_instance, centroid, bottomup, multi_class_bottomup, multi_class_topdown.	required
backbone_type	str	Backbone model. One of unet, convnext and swint.	required
backbone_config	Union[str, Dict[str, Any], DictConfig]	Backbone configuration. Can be: - String: One of the preset backbone types: - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"] - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"] - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"] - Dictionary: Custom configuration with structure: { "unet": {UNetConfig parameters}, "convnext": {ConvNextConfig parameters}, "swint": {SwinTConfig parameters} } Only one backbone type should be specified in the dictionary. - DictConfig: OmegaConf DictConfig object containing backbone configuration.	required
head_configs	DictConfig	Head configuration dictionary containing model-specific parameters. For Single Instance: confmaps with part_names, sigma, output_stride. For Centroid: confmaps with anchor_part, sigma, output_stride. For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride. For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight. For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight. For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.	required
pretrained_backbone_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.	None
pretrained_head_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.	None
init_weights	Optional[str]	Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.	'xavier'
lr_scheduler	Optional[Union[str, DictConfig]]	Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.	None
online_mining	Optional[bool]	If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).	False
hard_to_easy_ratio	Optional[float]	Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.	2.0
min_hard_keypoints	Optional[int]	Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.	2
max_hard_keypoints	Optional[int]	Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.	None
loss_scale	Optional[float]	Factor to scale hard keypoint losses by. Default: 5.0.	5.0
optimizer	Optional[str]	Optimizer name. One of ["Adam", "AdamW"].	'Adam'
learning_rate	Optional[float]	Learning rate for the optimizer. Default: 1e-3.	0.001
amsgrad	Optional[bool]	Enable AMSGrad with the optimizer. Default: False.	False

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model.
get_visualization_data	Extract visualization data from a sample.
training_step	Training step.
validation_step	Validation step.
visualize_class_maps_example	Visualize class map predictions during training (used with callbacks).
visualize_example	Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

class BottomUpMultiClassLightningModule(LightningModel):
    """Lightning Module for BottomUp ID Model.

    This is a subclass of the `LightningModel` to configure the training/ validation steps
    and forward pass specific to BottomUp ID model. Multi-Class Bottom-Up models predict
    all keypoints simultaneously and classify instances using class maps to identify
    individual animals across frames.

    Args:
        model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
        backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
        backbone_config: Backbone configuration. Can be:
            - String: One of the preset backbone types:
                - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"]
                - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"]
                - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"]
            - Dictionary: Custom configuration with structure:
                {
                    "unet": {UNetConfig parameters},
                    "convnext": {ConvNextConfig parameters},
                    "swint": {SwinTConfig parameters}
                }
                Only one backbone type should be specified in the dictionary.
            - DictConfig: OmegaConf DictConfig object containing backbone configuration.
        head_configs: Head configuration dictionary containing model-specific parameters.
            For Single Instance: confmaps with part_names, sigma, output_stride.
            For Centroid: confmaps with anchor_part, sigma, output_stride.
            For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride.
            For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight.
            For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight.
            For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.
        pretrained_backbone_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.
        pretrained_head_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.
        init_weights: Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.
        lr_scheduler: Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.
        online_mining: If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).
        hard_to_easy_ratio: Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.
        min_hard_keypoints: Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.
        max_hard_keypoints: Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.
        loss_scale: Factor to scale hard keypoint losses by. Default: 5.0.
        optimizer: Optimizer name. One of ["Adam", "AdamW"].
        learning_rate: Learning rate for the optimizer. Default: 1e-3.
        amsgrad: Enable AMSGrad with the optimizer. Default: False.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )
        self.bottomup_inf_layer = BottomUpMultiClassInferenceModel(
            torch_model=self.forward,
            peak_threshold=0.2,
            input_scale=1.0,
            return_confmaps=True,
            return_class_maps=True,
            cms_output_stride=self.head_configs.multi_class_bottomup.confmaps.output_stride,
            class_maps_output_stride=self.head_configs.multi_class_bottomup.class_maps.output_stride,
        )
        self.node_names = list(
            self.head_configs.multi_class_bottomup.confmaps.part_names
        )

    def get_visualization_data(
        self, sample, include_class_maps: bool = False
    ) -> VisualizationData:
        """Extract visualization data from a sample."""
        ex = sample.copy()
        ex["eff_scale"] = torch.tensor([1.0])
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["image"] = ex["image"].unsqueeze(dim=0)
        output = self.bottomup_inf_layer(ex)[0]

        peaks = output["pred_instance_peaks"][0].cpu().numpy()
        peak_values = output["pred_peak_values"][0].cpu().numpy()
        img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        gt_instances = ex["instances"][0].cpu().numpy()
        confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

        pred_class_maps = None
        if include_class_maps:
            pred_class_maps = (
                output["pred_class_maps"].cpu().numpy()[0].transpose(1, 2, 0)
            )

        return VisualizationData(
            image=img,
            pred_confmaps=confmaps,
            pred_peaks=peaks,
            pred_peak_values=peak_values,
            gt_instances=gt_instances,
            node_names=self.node_names,
            output_scale=confmaps.shape[0] / img.shape[0],
            is_paired=False,
            pred_class_maps=pred_class_maps,
        )

    def visualize_example(self, sample):
        """Visualize predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plt.xlim(plt.xlim())
        plt.ylim(plt.ylim())
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def visualize_class_maps_example(self, sample):
        """Visualize class map predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample, include_class_maps=True)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(
            data.pred_class_maps,
            output_scale=data.pred_class_maps.shape[0] / data.image.shape[0],
        )
        return fig

    def forward(self, img):
        """Forward pass of the model."""
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        output = self.model(img)
        return {
            "MultiInstanceConfmapsHead": output["MultiInstanceConfmapsHead"],
            "ClassMapsHead": output["ClassMapsHead"],
        }

    def training_step(self, batch, batch_idx):
        """Training step."""
        X = torch.squeeze(batch["image"], dim=1)
        y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
        y_classmap = torch.squeeze(batch["class_maps"], dim=1)
        X = normalize_on_gpu(X)
        preds = self.model(X)
        classmaps = preds["ClassMapsHead"]
        confmaps = preds["MultiInstanceConfmapsHead"]

        confmap_loss = self._compute_negative_weighted_loss(confmaps, y_confmap, batch)
        classmaps_loss = self._compute_negative_weighted_loss(
            classmaps, y_classmap, batch
        )
        self._log_negative_split_metrics(
            [
                ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
                ("classmap", classmaps, y_classmap, self.loss_weights[1]),
            ],
            batch,
            stage="train",
        )

        if self.online_mining is not None and self.online_mining:
            confmap_ohkm_loss = compute_ohkm_loss(
                y_gt=y_confmap,
                y_pr=confmaps,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            confmap_loss += confmap_ohkm_loss

        losses = {
            "MultiInstanceConfmapsHead": confmap_loss,
            "ClassMapsHead": classmaps_loss,
        }
        loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
        # Log step-level loss (every batch, uses global_step x-axis)
        self.log(
            "loss",
            loss,
            prog_bar=True,
            on_step=True,
            on_epoch=False,
            sync_dist=True,
        )

        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
        self._accumulate_loss(loss)
        self.log(
            "train/confmaps_loss",
            confmap_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "train/classmap_loss",
            classmaps_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Compute classification accuracy at GT keypoint locations
        with torch.no_grad():
            # Get output stride for class maps
            cms_stride = self.head_configs.multi_class_bottomup.class_maps.output_stride

            # Get GT instances and sample class maps at those locations
            instances = batch["instances"]  # (batch, n_samples, max_inst, n_nodes, 2)
            if instances.dim() == 5:
                instances = instances.squeeze(1)  # (batch, max_inst, n_nodes, 2)
            num_instances = batch["num_instances"]  # (batch,)

            correct = 0
            total = 0
            for b in range(instances.shape[0]):
                n_inst = num_instances[b].item()
                for inst_idx in range(n_inst):
                    for node_idx in range(instances.shape[2]):
                        # Get keypoint location (in input image space)
                        kp = instances[b, inst_idx, node_idx]  # (2,) = (x, y)
                        if torch.isnan(kp).any():
                            continue

                        # Convert to class map space
                        x_cm = (
                            (kp[0] / cms_stride)
                            .long()
                            .clamp(0, classmaps.shape[-1] - 1)
                        )
                        y_cm = (
                            (kp[1] / cms_stride)
                            .long()
                            .clamp(0, classmaps.shape[-2] - 1)
                        )

                        # Sample predicted and GT class at this location
                        pred_class = classmaps[b, :, y_cm, x_cm].argmax()
                        gt_class = y_classmap[b, :, y_cm, x_cm].argmax()

                        if pred_class == gt_class:
                            correct += 1
                        total += 1

            if total > 0:
                class_accuracy = torch.tensor(correct / total, device=X.device)
                self.log(
                    "train/class_accuracy",
                    class_accuracy,
                    on_step=False,
                    on_epoch=True,
                    sync_dist=True,
                )

        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step."""
        X = torch.squeeze(batch["image"], dim=1)
        y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
        y_classmap = torch.squeeze(batch["class_maps"], dim=1)
        X = normalize_on_gpu(X)

        preds = self.model(X)
        classmaps = preds["ClassMapsHead"]
        confmaps = preds["MultiInstanceConfmapsHead"]

        confmap_loss = self._compute_negative_weighted_loss(
            confmaps, y_confmap, batch, stage="val"
        )
        classmaps_loss = self._compute_negative_weighted_loss(
            classmaps, y_classmap, batch, stage="val"
        )
        self._log_negative_split_metrics(
            [
                ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
                ("classmap", classmaps, y_classmap, self.loss_weights[1]),
            ],
            batch,
            stage="val",
        )

        if self.online_mining is not None and self.online_mining:
            confmap_ohkm_loss = compute_ohkm_loss(
                y_gt=y_confmap,
                y_pr=confmaps,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            confmap_loss += confmap_ohkm_loss

        losses = {
            "MultiInstanceConfmapsHead": confmap_loss,
            "ClassMapsHead": classmaps_loss,
        }

        val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "val/confmaps_loss",
            confmap_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "val/classmap_loss",
            classmaps_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Compute classification accuracy at GT keypoint locations
        with torch.no_grad():
            # Get output stride for class maps
            cms_stride = self.head_configs.multi_class_bottomup.class_maps.output_stride

            # Get GT instances and sample class maps at those locations
            instances = batch["instances"]  # (batch, n_samples, max_inst, n_nodes, 2)
            if instances.dim() == 5:
                instances = instances.squeeze(1)  # (batch, max_inst, n_nodes, 2)
            num_instances = batch["num_instances"]  # (batch,)

            correct = 0
            total = 0
            for b in range(instances.shape[0]):
                n_inst = num_instances[b].item()
                for inst_idx in range(n_inst):
                    for node_idx in range(instances.shape[2]):
                        # Get keypoint location (in input image space)
                        kp = instances[b, inst_idx, node_idx]  # (2,) = (x, y)
                        if torch.isnan(kp).any():
                            continue

                        # Convert to class map space
                        x_cm = (
                            (kp[0] / cms_stride)
                            .long()
                            .clamp(0, classmaps.shape[-1] - 1)
                        )
                        y_cm = (
                            (kp[1] / cms_stride)
                            .long()
                            .clamp(0, classmaps.shape[-2] - 1)
                        )

                        # Sample predicted and GT class at this location
                        pred_class = classmaps[b, :, y_cm, x_cm].argmax()
                        gt_class = y_classmap[b, :, y_cm, x_cm].argmax()

                        if pred_class == gt_class:
                            correct += 1
                        total += 1

            if total > 0:
                class_accuracy = torch.tensor(correct / total, device=X.device)
                self.log(
                    "val/class_accuracy",
                    class_accuracy,
                    on_step=False,
                    on_epoch=True,
                    sync_dist=True,
                )

        # Collect predictions for epoch-end evaluation if enabled
        if self._collect_val_predictions:
            with torch.no_grad():
                # Note: Do NOT squeeze the image here - the forward() method expects
                # (batch, n_samples, C, H, W) and handles the n_samples squeeze internally
                inference_output = self.bottomup_inf_layer(batch)
                if isinstance(inference_output, list):
                    inference_output = inference_output[0]

            batch_size = len(batch["frame_idx"])
            for i in range(batch_size):
                eff = batch["eff_scale"][i].cpu().numpy()

                # Predictions are already in original space (variable number of instances)
                pred_peaks = inference_output["pred_instance_peaks"][i]
                pred_scores = inference_output["pred_peak_values"][i]
                if torch.is_tensor(pred_peaks):
                    pred_peaks = pred_peaks.cpu().numpy()
                if torch.is_tensor(pred_scores):
                    pred_scores = pred_scores.cpu().numpy()

                # Transform GT to original space
                # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
                gt_prep = batch["instances"][i].cpu().numpy()
                if gt_prep.ndim == 4:
                    gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
                gt_orig = gt_prep / eff
                num_inst = batch["num_instances"][i].item()
                gt_orig = gt_orig[:num_inst]  # Only valid instances

                self.val_predictions.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "pred_peaks": pred_peaks,  # Original space, variable instances
                        "pred_scores": pred_scores,
                    }
                )
                self.val_ground_truth.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "gt_instances": gt_orig,  # Original space
                        "num_instances": num_inst,
                    }
                )

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )
    self.bottomup_inf_layer = BottomUpMultiClassInferenceModel(
        torch_model=self.forward,
        peak_threshold=0.2,
        input_scale=1.0,
        return_confmaps=True,
        return_class_maps=True,
        cms_output_stride=self.head_configs.multi_class_bottomup.confmaps.output_stride,
        class_maps_output_stride=self.head_configs.multi_class_bottomup.class_maps.output_stride,
    )
    self.node_names = list(
        self.head_configs.multi_class_bottomup.confmaps.part_names
    )

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    output = self.model(img)
    return {
        "MultiInstanceConfmapsHead": output["MultiInstanceConfmapsHead"],
        "ClassMapsHead": output["ClassMapsHead"],
    }

get_visualization_data(sample, include_class_maps=False)

Extract visualization data from a sample.

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(
    self, sample, include_class_maps: bool = False
) -> VisualizationData:
    """Extract visualization data from a sample."""
    ex = sample.copy()
    ex["eff_scale"] = torch.tensor([1.0])
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["image"] = ex["image"].unsqueeze(dim=0)
    output = self.bottomup_inf_layer(ex)[0]

    peaks = output["pred_instance_peaks"][0].cpu().numpy()
    peak_values = output["pred_peak_values"][0].cpu().numpy()
    img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    gt_instances = ex["instances"][0].cpu().numpy()
    confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

    pred_class_maps = None
    if include_class_maps:
        pred_class_maps = (
            output["pred_class_maps"].cpu().numpy()[0].transpose(1, 2, 0)
        )

    return VisualizationData(
        image=img,
        pred_confmaps=confmaps,
        pred_peaks=peaks,
        pred_peak_values=peak_values,
        gt_instances=gt_instances,
        node_names=self.node_names,
        output_scale=confmaps.shape[0] / img.shape[0],
        is_paired=False,
        pred_class_maps=pred_class_maps,
    )

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    X = torch.squeeze(batch["image"], dim=1)
    y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
    y_classmap = torch.squeeze(batch["class_maps"], dim=1)
    X = normalize_on_gpu(X)
    preds = self.model(X)
    classmaps = preds["ClassMapsHead"]
    confmaps = preds["MultiInstanceConfmapsHead"]

    confmap_loss = self._compute_negative_weighted_loss(confmaps, y_confmap, batch)
    classmaps_loss = self._compute_negative_weighted_loss(
        classmaps, y_classmap, batch
    )
    self._log_negative_split_metrics(
        [
            ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
            ("classmap", classmaps, y_classmap, self.loss_weights[1]),
        ],
        batch,
        stage="train",
    )

    if self.online_mining is not None and self.online_mining:
        confmap_ohkm_loss = compute_ohkm_loss(
            y_gt=y_confmap,
            y_pr=confmaps,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        confmap_loss += confmap_ohkm_loss

    losses = {
        "MultiInstanceConfmapsHead": confmap_loss,
        "ClassMapsHead": classmaps_loss,
    }
    loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
    # Log step-level loss (every batch, uses global_step x-axis)
    self.log(
        "loss",
        loss,
        prog_bar=True,
        on_step=True,
        on_epoch=False,
        sync_dist=True,
    )

    # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
    self._accumulate_loss(loss)
    self.log(
        "train/confmaps_loss",
        confmap_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "train/classmap_loss",
        classmaps_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Compute classification accuracy at GT keypoint locations
    with torch.no_grad():
        # Get output stride for class maps
        cms_stride = self.head_configs.multi_class_bottomup.class_maps.output_stride

        # Get GT instances and sample class maps at those locations
        instances = batch["instances"]  # (batch, n_samples, max_inst, n_nodes, 2)
        if instances.dim() == 5:
            instances = instances.squeeze(1)  # (batch, max_inst, n_nodes, 2)
        num_instances = batch["num_instances"]  # (batch,)

        correct = 0
        total = 0
        for b in range(instances.shape[0]):
            n_inst = num_instances[b].item()
            for inst_idx in range(n_inst):
                for node_idx in range(instances.shape[2]):
                    # Get keypoint location (in input image space)
                    kp = instances[b, inst_idx, node_idx]  # (2,) = (x, y)
                    if torch.isnan(kp).any():
                        continue

                    # Convert to class map space
                    x_cm = (
                        (kp[0] / cms_stride)
                        .long()
                        .clamp(0, classmaps.shape[-1] - 1)
                    )
                    y_cm = (
                        (kp[1] / cms_stride)
                        .long()
                        .clamp(0, classmaps.shape[-2] - 1)
                    )

                    # Sample predicted and GT class at this location
                    pred_class = classmaps[b, :, y_cm, x_cm].argmax()
                    gt_class = y_classmap[b, :, y_cm, x_cm].argmax()

                    if pred_class == gt_class:
                        correct += 1
                    total += 1

        if total > 0:
            class_accuracy = torch.tensor(correct / total, device=X.device)
            self.log(
                "train/class_accuracy",
                class_accuracy,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )

    return loss

validation_step(batch, batch_idx)

Validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step."""
    X = torch.squeeze(batch["image"], dim=1)
    y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
    y_classmap = torch.squeeze(batch["class_maps"], dim=1)
    X = normalize_on_gpu(X)

    preds = self.model(X)
    classmaps = preds["ClassMapsHead"]
    confmaps = preds["MultiInstanceConfmapsHead"]

    confmap_loss = self._compute_negative_weighted_loss(
        confmaps, y_confmap, batch, stage="val"
    )
    classmaps_loss = self._compute_negative_weighted_loss(
        classmaps, y_classmap, batch, stage="val"
    )
    self._log_negative_split_metrics(
        [
            ("confmaps", confmaps, y_confmap, self.loss_weights[0]),
            ("classmap", classmaps, y_classmap, self.loss_weights[1]),
        ],
        batch,
        stage="val",
    )

    if self.online_mining is not None and self.online_mining:
        confmap_ohkm_loss = compute_ohkm_loss(
            y_gt=y_confmap,
            y_pr=confmaps,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        confmap_loss += confmap_ohkm_loss

    losses = {
        "MultiInstanceConfmapsHead": confmap_loss,
        "ClassMapsHead": classmaps_loss,
    }

    val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "val/confmaps_loss",
        confmap_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "val/classmap_loss",
        classmaps_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Compute classification accuracy at GT keypoint locations
    with torch.no_grad():
        # Get output stride for class maps
        cms_stride = self.head_configs.multi_class_bottomup.class_maps.output_stride

        # Get GT instances and sample class maps at those locations
        instances = batch["instances"]  # (batch, n_samples, max_inst, n_nodes, 2)
        if instances.dim() == 5:
            instances = instances.squeeze(1)  # (batch, max_inst, n_nodes, 2)
        num_instances = batch["num_instances"]  # (batch,)

        correct = 0
        total = 0
        for b in range(instances.shape[0]):
            n_inst = num_instances[b].item()
            for inst_idx in range(n_inst):
                for node_idx in range(instances.shape[2]):
                    # Get keypoint location (in input image space)
                    kp = instances[b, inst_idx, node_idx]  # (2,) = (x, y)
                    if torch.isnan(kp).any():
                        continue

                    # Convert to class map space
                    x_cm = (
                        (kp[0] / cms_stride)
                        .long()
                        .clamp(0, classmaps.shape[-1] - 1)
                    )
                    y_cm = (
                        (kp[1] / cms_stride)
                        .long()
                        .clamp(0, classmaps.shape[-2] - 1)
                    )

                    # Sample predicted and GT class at this location
                    pred_class = classmaps[b, :, y_cm, x_cm].argmax()
                    gt_class = y_classmap[b, :, y_cm, x_cm].argmax()

                    if pred_class == gt_class:
                        correct += 1
                    total += 1

        if total > 0:
            class_accuracy = torch.tensor(correct / total, device=X.device)
            self.log(
                "val/class_accuracy",
                class_accuracy,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )

    # Collect predictions for epoch-end evaluation if enabled
    if self._collect_val_predictions:
        with torch.no_grad():
            # Note: Do NOT squeeze the image here - the forward() method expects
            # (batch, n_samples, C, H, W) and handles the n_samples squeeze internally
            inference_output = self.bottomup_inf_layer(batch)
            if isinstance(inference_output, list):
                inference_output = inference_output[0]

        batch_size = len(batch["frame_idx"])
        for i in range(batch_size):
            eff = batch["eff_scale"][i].cpu().numpy()

            # Predictions are already in original space (variable number of instances)
            pred_peaks = inference_output["pred_instance_peaks"][i]
            pred_scores = inference_output["pred_peak_values"][i]
            if torch.is_tensor(pred_peaks):
                pred_peaks = pred_peaks.cpu().numpy()
            if torch.is_tensor(pred_scores):
                pred_scores = pred_scores.cpu().numpy()

            # Transform GT to original space
            # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
            gt_prep = batch["instances"][i].cpu().numpy()
            if gt_prep.ndim == 4:
                gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
            gt_orig = gt_prep / eff
            num_inst = batch["num_instances"][i].item()
            gt_orig = gt_orig[:num_inst]  # Only valid instances

            self.val_predictions.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "pred_peaks": pred_peaks,  # Original space, variable instances
                    "pred_scores": pred_scores,
                }
            )
            self.val_ground_truth.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "gt_instances": gt_orig,  # Original space
                    "num_instances": num_inst,
                }
            )

visualize_class_maps_example(sample)

Visualize class map predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_class_maps_example(self, sample):
    """Visualize class map predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample, include_class_maps=True)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(
        data.pred_class_maps,
        output_scale=data.pred_class_maps.shape[0] / data.image.shape[0],
    )
    return fig

visualize_example(sample)

Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plt.xlim(plt.xlim())
    plt.ylim(plt.ylim())
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

BottomUpSegmentationLightningModule

Bases: LightningModel

Lightning Module for Bottom-Up Instance Segmentation.

Predicts foreground masks, instance center heatmaps, and per-pixel offset vectors for grouping pixels into instances.

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model.
get_visualization_data	Extract visualization data from a sample.
training_step	Training step.
validation_step	Validation step.
visualize_example	Visualize segmentation predictions during training.

Source code in sleap_nn/training/lightning_modules.py

class BottomUpSegmentationLightningModule(LightningModel):
    """Lightning Module for Bottom-Up Instance Segmentation.

    Predicts foreground masks, instance center heatmaps, and per-pixel offset
    vectors for grouping pixels into instances.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )

        seg_cfg = self.head_configs[self.model_type]
        self.seg_inf_layer = BottomUpSegmentationInferenceModel(
            torch_model=self.forward,
            fg_threshold=0.5,
            peak_threshold=0.1,
            output_stride=seg_cfg.segmentation.output_stride,
        )
        # bce-dice foreground-loss knobs (defaults preserve the symmetric loss).
        self.fg_bce_weight = getattr(seg_cfg.segmentation, "bce_weight", 0.5)
        self.fg_dice_weight = getattr(seg_cfg.segmentation, "dice_weight", 0.5)
        self.fg_bce_pos_weight = getattr(seg_cfg.segmentation, "bce_pos_weight", None)

    def get_visualization_data(
        self,
        sample,
        include_center_heatmap: bool = False,
        include_offsets: bool = False,
        include_gt_mask: bool = False,
        include_instance_masks: bool = False,
    ) -> VisualizationData:
        """Extract visualization data from a sample.

        For segmentation models, the foreground probability map is used as
        the confidence map overlay. No keypoints/peaks are drawn on the
        predictions panel (segmentation has none); the optional
        ``instance_masks`` overlay shows the offset-grouped per-instance masks.

        Args:
            sample: A sample dictionary from the data pipeline.
            include_center_heatmap: If True, include the center heatmap in the
                returned data for separate visualization.
            include_offsets: If True, include the center-offset field for a
                separate offset-magnitude visualization.
            include_gt_mask: If True, include the ground-truth foreground mask
                for a GT-vs-prediction overlay.
            include_instance_masks: If True, keep the grouped per-instance masks
                (already computed by the offset-grouping below) for a colored
                instance-mask overlay. No extra forward/grouping pass is run.

        Returns:
            VisualizationData with foreground map and center locations.
        """
        ex = sample.copy()
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["image"] = ex["image"].unsqueeze(dim=0)

        # Run forward pass to get predictions
        with torch.no_grad():
            img = ex["image"].squeeze(1).to(self.device)
            img = normalize_on_gpu(img)
            preds = self.model(img)

        # Foreground probability as confmap overlay (H, W, 1)
        fg_prob = torch.sigmoid(preds["SegmentationHead"][0]).cpu().numpy()
        fg_prob = fg_prob.transpose(1, 2, 0)  # (H, W, 1)

        # Get image as (H, W, C)
        img_np = ex["image"][0, 0].cpu().numpy().transpose(1, 2, 0)

        # Colored per-instance mask overlay (separate `instance_masks` panel) is
        # produced by offset grouping. This is the ONLY place grouping is needed:
        # the `predictions` panel is a segmentation foreground map with NO
        # keypoints, so we do not peak-find just to draw center dots on it.
        instance_masks = None
        if include_instance_masks:
            from sleap_nn.inference.segmentation import group_instances_from_offsets

            seg_cfg = self.head_configs[self.model_type]
            output_stride = seg_cfg.segmentation.output_stride
            fg_sigmoid = torch.sigmoid(preds["SegmentationHead"])
            instances = group_instances_from_offsets(
                foreground=fg_sigmoid[0:1],
                center_heatmap=preds["InstanceCenterHead"][0:1],
                offsets=preds["CenterOffsetHead"][0:1],
                fg_threshold=0.5,
                peak_threshold=0.1,
                output_stride=output_stride,
            )
            # Each is a (H, W) bool at output-stride resolution (same grid as fg_prob).
            instance_masks = [inst["mask"] for inst in instances]

        # Segmentation has no keypoints -> draw no peak dots on the predictions
        # panel (matches TopDownCenteredInstanceSegmentationLightningModule).
        gt_pts = np.zeros((0, 1, 2))
        pred_pts = np.zeros((0, 1, 2))

        # Optionally include center heatmap for separate visualization
        center_hmap = None
        if include_center_heatmap:
            center_hmap = preds["InstanceCenterHead"][0].cpu().numpy()
            center_hmap = center_hmap.transpose(1, 2, 0)  # (H, W, 1)

        # Optionally include the center-offset field for an offset-magnitude viz
        offsets = None
        if include_offsets:
            offsets = preds["CenterOffsetHead"][0].cpu().numpy()
            offsets = offsets.transpose(1, 2, 0)  # (H, W, 2) -> (dx, dy)

        # Optionally include the GT foreground mask for a GT-vs-pred overlay
        gt_mask = None
        if include_gt_mask and "foreground_mask" in ex:
            gt_mask = ex["foreground_mask"].squeeze().cpu().numpy()  # (H, W)

        return VisualizationData(
            image=img_np,
            pred_confmaps=fg_prob,
            pred_peaks=pred_pts,
            pred_peak_values=np.zeros((0,)),
            gt_instances=gt_pts,
            node_names=["center"],
            output_scale=fg_prob.shape[0] / img_np.shape[0],
            is_paired=False,
            pred_center_heatmap=center_hmap,
            pred_offsets=offsets,
            gt_mask=gt_mask,
            instance_masks=instance_masks,
        )

    def visualize_example(self, sample):
        """Visualize segmentation predictions during training."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plt.xlim(plt.xlim())
        plt.ylim(plt.ylim())
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def forward(self, img):
        """Forward pass of the model."""
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        output = self.model(img)
        return {
            "SegmentationHead": torch.sigmoid(output["SegmentationHead"]),
            "InstanceCenterHead": output["InstanceCenterHead"],
            "CenterOffsetHead": output["CenterOffsetHead"],
        }

    def training_step(self, batch, batch_idx):
        """Training step."""
        X = torch.squeeze(batch["image"], dim=1)
        y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
        y_center = torch.squeeze(batch["center_heatmap"], dim=1)
        y_offsets = torch.squeeze(batch["center_offsets"], dim=1)
        y_weight = torch.squeeze(batch["foreground_weight"], dim=1)

        X = normalize_on_gpu(X)
        preds = self.model(X)

        pred_fg = preds["SegmentationHead"]
        pred_center = preds["InstanceCenterHead"]
        pred_offsets = preds["CenterOffsetHead"]

        fg_loss = compute_bce_dice_loss(
            pred_fg,
            y_fg,
            bce_weight=self.fg_bce_weight,
            dice_weight=self.fg_dice_weight,
            pos_weight=self.fg_bce_pos_weight,
        )
        center_loss = F.mse_loss(pred_center, y_center)
        offset_loss = compute_masked_smooth_l1(pred_offsets, y_offsets, y_weight)

        losses = {
            "SegmentationHead": fg_loss,
            "InstanceCenterHead": center_loss,
            "CenterOffsetHead": offset_loss,
        }
        seg_cfg = self.head_configs[self.model_type]
        loss = (
            seg_cfg.segmentation.loss_weight * losses["SegmentationHead"]
            + seg_cfg.center.loss_weight * losses["InstanceCenterHead"]
            + seg_cfg.offsets.loss_weight * losses["CenterOffsetHead"]
        )

        self.log(
            "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
        )
        self._accumulate_loss(loss)
        self.log("train/fg_loss", fg_loss, on_step=False, on_epoch=True, sync_dist=True)
        self.log(
            "train/center_loss",
            center_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "train/offset_loss",
            offset_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step."""
        X = torch.squeeze(batch["image"], dim=1)
        y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
        y_center = torch.squeeze(batch["center_heatmap"], dim=1)
        y_offsets = torch.squeeze(batch["center_offsets"], dim=1)
        y_weight = torch.squeeze(batch["foreground_weight"], dim=1)

        X = normalize_on_gpu(X)
        preds = self.model(X)

        pred_fg = preds["SegmentationHead"]
        pred_center = preds["InstanceCenterHead"]
        pred_offsets = preds["CenterOffsetHead"]

        fg_loss = compute_bce_dice_loss(
            pred_fg,
            y_fg,
            bce_weight=self.fg_bce_weight,
            dice_weight=self.fg_dice_weight,
            pos_weight=self.fg_bce_pos_weight,
        )
        center_loss = F.mse_loss(pred_center, y_center)
        offset_loss = compute_masked_smooth_l1(pred_offsets, y_offsets, y_weight)

        losses = {
            "SegmentationHead": fg_loss,
            "InstanceCenterHead": center_loss,
            "CenterOffsetHead": offset_loss,
        }
        seg_cfg = self.head_configs[self.model_type]
        val_loss = (
            seg_cfg.segmentation.loss_weight * losses["SegmentationHead"]
            + seg_cfg.center.loss_weight * losses["InstanceCenterHead"]
            + seg_cfg.offsets.loss_weight * losses["CenterOffsetHead"]
        )

        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log("val/fg_loss", fg_loss, on_step=False, on_epoch=True, sync_dist=True)
        self.log(
            "val/center_loss", center_loss, on_step=False, on_epoch=True, sync_dist=True
        )
        self.log(
            "val/offset_loss", offset_loss, on_step=False, on_epoch=True, sync_dist=True
        )

        # Foreground IoU averaged PER-SAMPLE (mean of per-image IoUs), not pooled
        # over the whole batch tensor (which over-weights large/foreground-heavy
        # images). Lightning's on_epoch aggregation then yields a true
        # mean-per-image IoU.
        pred_fg_binary = (pred_fg > 0.0).float()
        dims = (1, 2, 3)
        intersection = (pred_fg_binary * y_fg).sum(dim=dims)
        union = pred_fg_binary.sum(dim=dims) + y_fg.sum(dim=dims) - intersection
        iou = (intersection / (union + 1e-6)).mean()
        self.log("val/fg_iou", iou, on_step=False, on_epoch=True, sync_dist=True)

        # Optional instance-level mask eval: when enabled by the
        # SegmentationEvaluationCallback, recover per-instance masks by grouping the
        # predicted AND ground-truth heads on the SAME preprocessed stride grid (no
        # original-resolution remapping), so a mask-IoU mAP/precision/recall can be
        # computed against the instance-separated GT. Appended in lockstep with the
        # GT so the callback pairs them positionally.
        if self._collect_val_predictions:
            from sleap_nn.inference.segmentation import group_instances_from_offsets

            stride = seg_cfg.segmentation.output_stride
            fg_prob = torch.sigmoid(pred_fg)
            for i in range(X.shape[0]):
                pred_insts = group_instances_from_offsets(
                    foreground=fg_prob[i : i + 1],
                    center_heatmap=pred_center[i : i + 1],
                    offsets=pred_offsets[i : i + 1],
                    fg_threshold=0.5,
                    peak_threshold=0.1,
                    output_stride=stride,
                )
                gt_insts = group_instances_from_offsets(
                    foreground=y_fg[i : i + 1],
                    center_heatmap=y_center[i : i + 1],
                    offsets=y_offsets[i : i + 1],
                    fg_threshold=0.5,
                    peak_threshold=0.1,
                    output_stride=stride,
                )
                self.val_predictions.append({"masks": [d["mask"] for d in pred_insts]})
                self.val_ground_truth.append({"masks": [d["mask"] for d in gt_insts]})

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )

    seg_cfg = self.head_configs[self.model_type]
    self.seg_inf_layer = BottomUpSegmentationInferenceModel(
        torch_model=self.forward,
        fg_threshold=0.5,
        peak_threshold=0.1,
        output_stride=seg_cfg.segmentation.output_stride,
    )
    # bce-dice foreground-loss knobs (defaults preserve the symmetric loss).
    self.fg_bce_weight = getattr(seg_cfg.segmentation, "bce_weight", 0.5)
    self.fg_dice_weight = getattr(seg_cfg.segmentation, "dice_weight", 0.5)
    self.fg_bce_pos_weight = getattr(seg_cfg.segmentation, "bce_pos_weight", None)

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    output = self.model(img)
    return {
        "SegmentationHead": torch.sigmoid(output["SegmentationHead"]),
        "InstanceCenterHead": output["InstanceCenterHead"],
        "CenterOffsetHead": output["CenterOffsetHead"],
    }

get_visualization_data(sample, include_center_heatmap=False, include_offsets=False, include_gt_mask=False, include_instance_masks=False)

Extract visualization data from a sample.

For segmentation models, the foreground probability map is used as the confidence map overlay. No keypoints/peaks are drawn on the predictions panel (segmentation has none); the optional instance_masks overlay shows the offset-grouped per-instance masks.

Parameters:

Name	Type	Description	Default
sample		A sample dictionary from the data pipeline.	required
include_center_heatmap	bool	If True, include the center heatmap in the returned data for separate visualization.	False
include_offsets	bool	If True, include the center-offset field for a separate offset-magnitude visualization.	False
include_gt_mask	bool	If True, include the ground-truth foreground mask for a GT-vs-prediction overlay.	False
include_instance_masks	bool	If True, keep the grouped per-instance masks (already computed by the offset-grouping below) for a colored instance-mask overlay. No extra forward/grouping pass is run.	False

Returns:

Type	Description
VisualizationData	VisualizationData with foreground map and center locations.

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(
    self,
    sample,
    include_center_heatmap: bool = False,
    include_offsets: bool = False,
    include_gt_mask: bool = False,
    include_instance_masks: bool = False,
) -> VisualizationData:
    """Extract visualization data from a sample.

    For segmentation models, the foreground probability map is used as
    the confidence map overlay. No keypoints/peaks are drawn on the
    predictions panel (segmentation has none); the optional
    ``instance_masks`` overlay shows the offset-grouped per-instance masks.

    Args:
        sample: A sample dictionary from the data pipeline.
        include_center_heatmap: If True, include the center heatmap in the
            returned data for separate visualization.
        include_offsets: If True, include the center-offset field for a
            separate offset-magnitude visualization.
        include_gt_mask: If True, include the ground-truth foreground mask
            for a GT-vs-prediction overlay.
        include_instance_masks: If True, keep the grouped per-instance masks
            (already computed by the offset-grouping below) for a colored
            instance-mask overlay. No extra forward/grouping pass is run.

    Returns:
        VisualizationData with foreground map and center locations.
    """
    ex = sample.copy()
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["image"] = ex["image"].unsqueeze(dim=0)

    # Run forward pass to get predictions
    with torch.no_grad():
        img = ex["image"].squeeze(1).to(self.device)
        img = normalize_on_gpu(img)
        preds = self.model(img)

    # Foreground probability as confmap overlay (H, W, 1)
    fg_prob = torch.sigmoid(preds["SegmentationHead"][0]).cpu().numpy()
    fg_prob = fg_prob.transpose(1, 2, 0)  # (H, W, 1)

    # Get image as (H, W, C)
    img_np = ex["image"][0, 0].cpu().numpy().transpose(1, 2, 0)

    # Colored per-instance mask overlay (separate `instance_masks` panel) is
    # produced by offset grouping. This is the ONLY place grouping is needed:
    # the `predictions` panel is a segmentation foreground map with NO
    # keypoints, so we do not peak-find just to draw center dots on it.
    instance_masks = None
    if include_instance_masks:
        from sleap_nn.inference.segmentation import group_instances_from_offsets

        seg_cfg = self.head_configs[self.model_type]
        output_stride = seg_cfg.segmentation.output_stride
        fg_sigmoid = torch.sigmoid(preds["SegmentationHead"])
        instances = group_instances_from_offsets(
            foreground=fg_sigmoid[0:1],
            center_heatmap=preds["InstanceCenterHead"][0:1],
            offsets=preds["CenterOffsetHead"][0:1],
            fg_threshold=0.5,
            peak_threshold=0.1,
            output_stride=output_stride,
        )
        # Each is a (H, W) bool at output-stride resolution (same grid as fg_prob).
        instance_masks = [inst["mask"] for inst in instances]

    # Segmentation has no keypoints -> draw no peak dots on the predictions
    # panel (matches TopDownCenteredInstanceSegmentationLightningModule).
    gt_pts = np.zeros((0, 1, 2))
    pred_pts = np.zeros((0, 1, 2))

    # Optionally include center heatmap for separate visualization
    center_hmap = None
    if include_center_heatmap:
        center_hmap = preds["InstanceCenterHead"][0].cpu().numpy()
        center_hmap = center_hmap.transpose(1, 2, 0)  # (H, W, 1)

    # Optionally include the center-offset field for an offset-magnitude viz
    offsets = None
    if include_offsets:
        offsets = preds["CenterOffsetHead"][0].cpu().numpy()
        offsets = offsets.transpose(1, 2, 0)  # (H, W, 2) -> (dx, dy)

    # Optionally include the GT foreground mask for a GT-vs-pred overlay
    gt_mask = None
    if include_gt_mask and "foreground_mask" in ex:
        gt_mask = ex["foreground_mask"].squeeze().cpu().numpy()  # (H, W)

    return VisualizationData(
        image=img_np,
        pred_confmaps=fg_prob,
        pred_peaks=pred_pts,
        pred_peak_values=np.zeros((0,)),
        gt_instances=gt_pts,
        node_names=["center"],
        output_scale=fg_prob.shape[0] / img_np.shape[0],
        is_paired=False,
        pred_center_heatmap=center_hmap,
        pred_offsets=offsets,
        gt_mask=gt_mask,
        instance_masks=instance_masks,
    )

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    X = torch.squeeze(batch["image"], dim=1)
    y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
    y_center = torch.squeeze(batch["center_heatmap"], dim=1)
    y_offsets = torch.squeeze(batch["center_offsets"], dim=1)
    y_weight = torch.squeeze(batch["foreground_weight"], dim=1)

    X = normalize_on_gpu(X)
    preds = self.model(X)

    pred_fg = preds["SegmentationHead"]
    pred_center = preds["InstanceCenterHead"]
    pred_offsets = preds["CenterOffsetHead"]

    fg_loss = compute_bce_dice_loss(
        pred_fg,
        y_fg,
        bce_weight=self.fg_bce_weight,
        dice_weight=self.fg_dice_weight,
        pos_weight=self.fg_bce_pos_weight,
    )
    center_loss = F.mse_loss(pred_center, y_center)
    offset_loss = compute_masked_smooth_l1(pred_offsets, y_offsets, y_weight)

    losses = {
        "SegmentationHead": fg_loss,
        "InstanceCenterHead": center_loss,
        "CenterOffsetHead": offset_loss,
    }
    seg_cfg = self.head_configs[self.model_type]
    loss = (
        seg_cfg.segmentation.loss_weight * losses["SegmentationHead"]
        + seg_cfg.center.loss_weight * losses["InstanceCenterHead"]
        + seg_cfg.offsets.loss_weight * losses["CenterOffsetHead"]
    )

    self.log(
        "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
    )
    self._accumulate_loss(loss)
    self.log("train/fg_loss", fg_loss, on_step=False, on_epoch=True, sync_dist=True)
    self.log(
        "train/center_loss",
        center_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "train/offset_loss",
        offset_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    return loss

validation_step(batch, batch_idx)

Validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step."""
    X = torch.squeeze(batch["image"], dim=1)
    y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
    y_center = torch.squeeze(batch["center_heatmap"], dim=1)
    y_offsets = torch.squeeze(batch["center_offsets"], dim=1)
    y_weight = torch.squeeze(batch["foreground_weight"], dim=1)

    X = normalize_on_gpu(X)
    preds = self.model(X)

    pred_fg = preds["SegmentationHead"]
    pred_center = preds["InstanceCenterHead"]
    pred_offsets = preds["CenterOffsetHead"]

    fg_loss = compute_bce_dice_loss(
        pred_fg,
        y_fg,
        bce_weight=self.fg_bce_weight,
        dice_weight=self.fg_dice_weight,
        pos_weight=self.fg_bce_pos_weight,
    )
    center_loss = F.mse_loss(pred_center, y_center)
    offset_loss = compute_masked_smooth_l1(pred_offsets, y_offsets, y_weight)

    losses = {
        "SegmentationHead": fg_loss,
        "InstanceCenterHead": center_loss,
        "CenterOffsetHead": offset_loss,
    }
    seg_cfg = self.head_configs[self.model_type]
    val_loss = (
        seg_cfg.segmentation.loss_weight * losses["SegmentationHead"]
        + seg_cfg.center.loss_weight * losses["InstanceCenterHead"]
        + seg_cfg.offsets.loss_weight * losses["CenterOffsetHead"]
    )

    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log("val/fg_loss", fg_loss, on_step=False, on_epoch=True, sync_dist=True)
    self.log(
        "val/center_loss", center_loss, on_step=False, on_epoch=True, sync_dist=True
    )
    self.log(
        "val/offset_loss", offset_loss, on_step=False, on_epoch=True, sync_dist=True
    )

    # Foreground IoU averaged PER-SAMPLE (mean of per-image IoUs), not pooled
    # over the whole batch tensor (which over-weights large/foreground-heavy
    # images). Lightning's on_epoch aggregation then yields a true
    # mean-per-image IoU.
    pred_fg_binary = (pred_fg > 0.0).float()
    dims = (1, 2, 3)
    intersection = (pred_fg_binary * y_fg).sum(dim=dims)
    union = pred_fg_binary.sum(dim=dims) + y_fg.sum(dim=dims) - intersection
    iou = (intersection / (union + 1e-6)).mean()
    self.log("val/fg_iou", iou, on_step=False, on_epoch=True, sync_dist=True)

    # Optional instance-level mask eval: when enabled by the
    # SegmentationEvaluationCallback, recover per-instance masks by grouping the
    # predicted AND ground-truth heads on the SAME preprocessed stride grid (no
    # original-resolution remapping), so a mask-IoU mAP/precision/recall can be
    # computed against the instance-separated GT. Appended in lockstep with the
    # GT so the callback pairs them positionally.
    if self._collect_val_predictions:
        from sleap_nn.inference.segmentation import group_instances_from_offsets

        stride = seg_cfg.segmentation.output_stride
        fg_prob = torch.sigmoid(pred_fg)
        for i in range(X.shape[0]):
            pred_insts = group_instances_from_offsets(
                foreground=fg_prob[i : i + 1],
                center_heatmap=pred_center[i : i + 1],
                offsets=pred_offsets[i : i + 1],
                fg_threshold=0.5,
                peak_threshold=0.1,
                output_stride=stride,
            )
            gt_insts = group_instances_from_offsets(
                foreground=y_fg[i : i + 1],
                center_heatmap=y_center[i : i + 1],
                offsets=y_offsets[i : i + 1],
                fg_threshold=0.5,
                peak_threshold=0.1,
                output_stride=stride,
            )
            self.val_predictions.append({"masks": [d["mask"] for d in pred_insts]})
            self.val_ground_truth.append({"masks": [d["mask"] for d in gt_insts]})

visualize_example(sample)

Visualize segmentation predictions during training.

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize segmentation predictions during training."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plt.xlim(plt.xlim())
    plt.ylim(plt.ylim())
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

CentroidLightningModule

Bases: LightningModel

Lightning Module for Centroid Model.

This is a subclass of the LightningModel to configure the training/ validation steps and forward pass specific to centroid model. Centroid models detect the center points of animals in the image, which are then used by Top-Down models for keypoint prediction.

Parameters:

Name	Type	Description	Default
model_type	str	Type of the model. One of single_instance, centered_instance, centroid, bottomup, multi_class_bottomup, multi_class_topdown.	required
backbone_type	str	Backbone model. One of unet, convnext and swint.	required
backbone_config	Union[str, Dict[str, Any], DictConfig]	Backbone configuration. Can be: - String: One of the preset backbone types: - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"] - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"] - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"] - Dictionary: Custom configuration with structure: { "unet": {UNetConfig parameters}, "convnext": {ConvNextConfig parameters}, "swint": {SwinTConfig parameters} } Only one backbone type should be specified in the dictionary. - DictConfig: OmegaConf DictConfig object containing backbone configuration.	required
head_configs	DictConfig	Head configuration dictionary containing model-specific parameters. For Single Instance: confmaps with part_names, sigma, output_stride. For Centroid: confmaps with anchor_part, sigma, output_stride. For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride. For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight. For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight. For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.	required
pretrained_backbone_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.	None
pretrained_head_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.	None
init_weights	Optional[str]	Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.	'xavier'
lr_scheduler	Optional[Union[str, DictConfig]]	Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.	None
online_mining	Optional[bool]	If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).	False
hard_to_easy_ratio	Optional[float]	Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.	2.0
min_hard_keypoints	Optional[int]	Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.	2
max_hard_keypoints	Optional[int]	Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.	None
loss_scale	Optional[float]	Factor to scale hard keypoint losses by. Default: 5.0.	5.0
optimizer	Optional[str]	Optimizer name. One of ["Adam", "AdamW"].	'Adam'
learning_rate	Optional[float]	Learning rate for the optimizer. Default: 1e-3.	0.001
amsgrad	Optional[bool]	Enable AMSGrad with the optimizer. Default: False.	False

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model.
get_visualization_data	Extract visualization data from a sample.
training_step	Training step.
validation_step	Validation step.
visualize_example	Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

class CentroidLightningModule(LightningModel):
    """Lightning Module for Centroid Model.

    This is a subclass of the `LightningModel` to configure the training/ validation steps
    and forward pass specific to centroid model. Centroid models detect the center points
    of animals in the image, which are then used by Top-Down models for keypoint prediction.

    Args:
        model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
        backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
        backbone_config: Backbone configuration. Can be:
            - String: One of the preset backbone types:
                - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"]
                - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"]
                - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"]
            - Dictionary: Custom configuration with structure:
                {
                    "unet": {UNetConfig parameters},
                    "convnext": {ConvNextConfig parameters},
                    "swint": {SwinTConfig parameters}
                }
                Only one backbone type should be specified in the dictionary.
            - DictConfig: OmegaConf DictConfig object containing backbone configuration.
        head_configs: Head configuration dictionary containing model-specific parameters.
            For Single Instance: confmaps with part_names, sigma, output_stride.
            For Centroid: confmaps with anchor_part, sigma, output_stride.
            For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride.
            For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight.
            For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight.
            For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.
        pretrained_backbone_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.
        pretrained_head_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.
        init_weights: Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.
        lr_scheduler: Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.
        online_mining: If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).
        hard_to_easy_ratio: Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.
        min_hard_keypoints: Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.
        max_hard_keypoints: Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.
        loss_scale: Factor to scale hard keypoint losses by. Default: 5.0.
        optimizer: Optimizer name. One of ["Adam", "AdamW"].
        learning_rate: Learning rate for the optimizer. Default: 1e-3.
        amsgrad: Enable AMSGrad with the optimizer. Default: False.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
        centroid_focal_loss_alpha: Optional[float] = 0.0,
        centroid_focal_loss_beta: Optional[float] = 4.0,
        centroid_focal_loss_pos_threshold: Optional[float] = 0.5,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )
        # Centroid-only knobs, not threaded through the base `LightningModel`
        # since no other model type implements this loss -- see
        # `CentroidConfMapsConfig.focal_loss_alpha`.
        self.centroid_focal_loss_alpha = centroid_focal_loss_alpha
        self.centroid_focal_loss_beta = centroid_focal_loss_beta
        self.centroid_focal_loss_pos_threshold = centroid_focal_loss_pos_threshold

        # RetinaNet/CenterNet "prior probability" bias init (Lin et al. 2017,
        # Focal Loss for Dense Object Detection, sec 4.1). Without this, a
        # freshly-initialized sigmoid head starts at ~0.5 everywhere; with a
        # focal loss and a target that's >99% background pixels, that start
        # point gives weak, roughly-symmetric gradients that can leave
        # training stuck near its initial value for many epochs. Biasing the
        # pre-sigmoid logit so the head starts near a low constant
        # probability (matching the true class balance) gives the loss a much
        # stronger initial gradient toward learning the sparse foreground.
        if self.centroid_focal_loss_alpha != 0.0:
            prior_prob = 0.01
            bias_value = -math.log((1.0 - prior_prob) / prior_prob)
            for head, head_layer in zip(self.model.heads, self.model.head_layers):
                if head.name == "CentroidConfmapsHead":
                    nn.init.constant_(
                        getattr(head_layer, head.name)[0].bias, bias_value
                    )

        self.centroid_inf_layer = CentroidCrop(
            torch_model=self.forward,
            peak_threshold=0.2,
            return_confmaps=True,
            output_stride=self.head_configs.centroid.confmaps.output_stride,
            input_scale=1.0,
        )
        self.node_names = ["centroid"]

    def get_visualization_data(self, sample) -> VisualizationData:
        """Extract visualization data from a sample."""
        ex = sample.copy()
        ex["eff_scale"] = torch.tensor([1.0])
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["image"] = ex["image"].unsqueeze(dim=0)
        gt_centroids = ex["centroids"].cpu().numpy()
        output = self.centroid_inf_layer(ex)

        peaks = output["centroids"][0].cpu().numpy()
        centroid_vals = output["centroid_vals"][0].cpu().numpy()
        img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        confmaps = output["pred_centroid_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

        return VisualizationData(
            image=img,
            pred_confmaps=confmaps,
            pred_peaks=peaks,
            pred_peak_values=centroid_vals,
            gt_instances=gt_centroids,
            node_names=self.node_names,
            output_scale=confmaps.shape[0] / img.shape[0],
            is_paired=False,
        )

    def visualize_example(self, sample):
        """Visualize predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def forward(self, img):
        """Forward pass of the model."""
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        return self.model(img)["CentroidConfmapsHead"]

    def _compute_loss(
        self, y_preds: torch.Tensor, y: torch.Tensor, batch: Dict, stage: str = "train"
    ) -> torch.Tensor:
        """Negative-weighted MSE, or a focal loss in place of MSE.

        When ``centroid_focal_loss_alpha == 0`` this is exactly
        :meth:`_compute_negative_weighted_loss` (plain MSE, optionally
        negative-frame-weighted on train). When nonzero, replaces the base
        per-pixel MSE with
        :func:`sleap_nn.training.losses.compute_centroid_focal_loss` before
        applying the same negative-frame weighting on top -- see
        `CentroidConfMapsConfig.focal_loss_alpha`. Requires the head's output to
        be a calibrated ``(0, 1)`` probability (see
        ``CentroidConfmapsHead.use_sigmoid_activation``). ``val``/eval always
        uses plain unweighted MSE (matching
        ``_compute_negative_weighted_loss``'s own val-stage behavior), so
        ``ModelCheckpoint``/``EarlyStopping`` stay comparable across every
        experiment in this family.
        """
        focal_alpha = self.centroid_focal_loss_alpha
        if focal_alpha == 0.0:
            return self._compute_negative_weighted_loss(y_preds, y, batch, stage=stage)

        if stage != "train":
            return nn.MSELoss()(y_preds, y)

        per_sample = compute_centroid_focal_loss(
            y_preds,
            y,
            alpha=focal_alpha,
            beta=self.centroid_focal_loss_beta,
            pos_threshold=self.centroid_focal_loss_pos_threshold,
            reduction="none",
        ).mean(dim=list(range(1, y_preds.ndim)))

        is_negative = batch.get("is_negative", None)
        if is_negative is None or self.negative_loss_weight == 1.0:
            return per_sample.mean()

        is_neg = is_negative.to(y_preds.device)
        weights = torch.where(
            is_neg,
            torch.tensor(self.negative_loss_weight, device=y_preds.device),
            torch.tensor(1.0, device=y_preds.device),
        )
        return (per_sample * weights).mean()

    def training_step(self, batch, batch_idx):
        """Training step."""
        X, y = (
            torch.squeeze(batch["image"], dim=1),
            torch.squeeze(batch["centroids_confidence_maps"], dim=1),
        )
        X = normalize_on_gpu(X)

        y_preds = self.model(X)["CentroidConfmapsHead"]
        loss = self._compute_loss(y_preds, y, batch, stage="train")
        self._log_negative_split_metrics(
            [("confmaps", y_preds, y, 1.0)], batch, stage="train"
        )
        self._log_confmap_fg_bg_loss(y_preds, y, stage="train")
        # Log step-level loss (every batch, uses global_step x-axis)
        self.log(
            "loss",
            loss,
            prog_bar=True,
            on_step=True,
            on_epoch=False,
            sync_dist=True,
        )

        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
        self._accumulate_loss(loss)
        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step."""
        X, y = (
            torch.squeeze(batch["image"], dim=1),
            torch.squeeze(batch["centroids_confidence_maps"], dim=1),
        )
        X = normalize_on_gpu(X)

        y_preds = self.model(X)["CentroidConfmapsHead"]
        val_loss = self._compute_loss(y_preds, y, batch, stage="val")
        self._log_negative_split_metrics(
            [("confmaps", y_preds, y, 1.0)], batch, stage="val"
        )
        self._log_confmap_fg_bg_loss(y_preds, y, stage="val")
        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Collect predictions for epoch-end evaluation if enabled
        if self._collect_val_predictions:
            # Save GT centroids before inference (inference overwrites batch["centroids"])
            batch["gt_centroids"] = batch["centroids"].clone()

            with torch.no_grad():
                inference_output = self.centroid_inf_layer(batch)

            batch_size = len(batch["frame_idx"])
            for i in range(batch_size):
                eff = batch["eff_scale"][i].cpu().numpy()

                # Predictions are in original image space (inference divides by eff_scale)
                # centroids shape: (batch, 1, max_instances, 2) - squeeze to (max_instances, 2)
                pred_centroids = (
                    inference_output["centroids"][i].squeeze(0).cpu().numpy()
                )
                pred_vals = inference_output["centroid_vals"][i].cpu().numpy()

                # Transform GT centroids from preprocessed to original image space
                # Use "gt_centroids" since inference overwrites "centroids" with predictions
                gt_centroids_prep = (
                    batch["gt_centroids"][i].cpu().numpy()
                )  # (n_samples=1, max_inst, 2)
                gt_centroids_orig = gt_centroids_prep.squeeze(0) / eff  # (max_inst, 2)
                num_inst = batch["num_instances"][i].item()

                # Filter to valid instances (non-NaN)
                valid_pred_mask = ~np.isnan(pred_centroids).any(axis=1)
                pred_centroids = pred_centroids[valid_pred_mask]
                pred_vals = pred_vals[valid_pred_mask]

                gt_centroids_valid = gt_centroids_orig[:num_inst]

                self.val_predictions.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "pred_peaks": pred_centroids.reshape(
                            -1, 1, 2
                        ),  # (n_inst, 1, 2)
                        "pred_scores": pred_vals.reshape(-1, 1),  # (n_inst, 1)
                    }
                )
                self.val_ground_truth.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "gt_instances": gt_centroids_valid.reshape(
                            -1, 1, 2
                        ),  # (n_inst, 1, 2)
                        "num_instances": num_inst,
                    }
                )

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0, centroid_focal_loss_alpha=0.0, centroid_focal_loss_beta=4.0, centroid_focal_loss_pos_threshold=0.5)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
    centroid_focal_loss_alpha: Optional[float] = 0.0,
    centroid_focal_loss_beta: Optional[float] = 4.0,
    centroid_focal_loss_pos_threshold: Optional[float] = 0.5,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )
    # Centroid-only knobs, not threaded through the base `LightningModel`
    # since no other model type implements this loss -- see
    # `CentroidConfMapsConfig.focal_loss_alpha`.
    self.centroid_focal_loss_alpha = centroid_focal_loss_alpha
    self.centroid_focal_loss_beta = centroid_focal_loss_beta
    self.centroid_focal_loss_pos_threshold = centroid_focal_loss_pos_threshold

    # RetinaNet/CenterNet "prior probability" bias init (Lin et al. 2017,
    # Focal Loss for Dense Object Detection, sec 4.1). Without this, a
    # freshly-initialized sigmoid head starts at ~0.5 everywhere; with a
    # focal loss and a target that's >99% background pixels, that start
    # point gives weak, roughly-symmetric gradients that can leave
    # training stuck near its initial value for many epochs. Biasing the
    # pre-sigmoid logit so the head starts near a low constant
    # probability (matching the true class balance) gives the loss a much
    # stronger initial gradient toward learning the sparse foreground.
    if self.centroid_focal_loss_alpha != 0.0:
        prior_prob = 0.01
        bias_value = -math.log((1.0 - prior_prob) / prior_prob)
        for head, head_layer in zip(self.model.heads, self.model.head_layers):
            if head.name == "CentroidConfmapsHead":
                nn.init.constant_(
                    getattr(head_layer, head.name)[0].bias, bias_value
                )

    self.centroid_inf_layer = CentroidCrop(
        torch_model=self.forward,
        peak_threshold=0.2,
        return_confmaps=True,
        output_stride=self.head_configs.centroid.confmaps.output_stride,
        input_scale=1.0,
    )
    self.node_names = ["centroid"]

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    return self.model(img)["CentroidConfmapsHead"]

get_visualization_data(sample)

Extract visualization data from a sample.

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(self, sample) -> VisualizationData:
    """Extract visualization data from a sample."""
    ex = sample.copy()
    ex["eff_scale"] = torch.tensor([1.0])
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["image"] = ex["image"].unsqueeze(dim=0)
    gt_centroids = ex["centroids"].cpu().numpy()
    output = self.centroid_inf_layer(ex)

    peaks = output["centroids"][0].cpu().numpy()
    centroid_vals = output["centroid_vals"][0].cpu().numpy()
    img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    confmaps = output["pred_centroid_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

    return VisualizationData(
        image=img,
        pred_confmaps=confmaps,
        pred_peaks=peaks,
        pred_peak_values=centroid_vals,
        gt_instances=gt_centroids,
        node_names=self.node_names,
        output_scale=confmaps.shape[0] / img.shape[0],
        is_paired=False,
    )

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    X, y = (
        torch.squeeze(batch["image"], dim=1),
        torch.squeeze(batch["centroids_confidence_maps"], dim=1),
    )
    X = normalize_on_gpu(X)

    y_preds = self.model(X)["CentroidConfmapsHead"]
    loss = self._compute_loss(y_preds, y, batch, stage="train")
    self._log_negative_split_metrics(
        [("confmaps", y_preds, y, 1.0)], batch, stage="train"
    )
    self._log_confmap_fg_bg_loss(y_preds, y, stage="train")
    # Log step-level loss (every batch, uses global_step x-axis)
    self.log(
        "loss",
        loss,
        prog_bar=True,
        on_step=True,
        on_epoch=False,
        sync_dist=True,
    )

    # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
    self._accumulate_loss(loss)
    return loss

validation_step(batch, batch_idx)

Validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step."""
    X, y = (
        torch.squeeze(batch["image"], dim=1),
        torch.squeeze(batch["centroids_confidence_maps"], dim=1),
    )
    X = normalize_on_gpu(X)

    y_preds = self.model(X)["CentroidConfmapsHead"]
    val_loss = self._compute_loss(y_preds, y, batch, stage="val")
    self._log_negative_split_metrics(
        [("confmaps", y_preds, y, 1.0)], batch, stage="val"
    )
    self._log_confmap_fg_bg_loss(y_preds, y, stage="val")
    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Collect predictions for epoch-end evaluation if enabled
    if self._collect_val_predictions:
        # Save GT centroids before inference (inference overwrites batch["centroids"])
        batch["gt_centroids"] = batch["centroids"].clone()

        with torch.no_grad():
            inference_output = self.centroid_inf_layer(batch)

        batch_size = len(batch["frame_idx"])
        for i in range(batch_size):
            eff = batch["eff_scale"][i].cpu().numpy()

            # Predictions are in original image space (inference divides by eff_scale)
            # centroids shape: (batch, 1, max_instances, 2) - squeeze to (max_instances, 2)
            pred_centroids = (
                inference_output["centroids"][i].squeeze(0).cpu().numpy()
            )
            pred_vals = inference_output["centroid_vals"][i].cpu().numpy()

            # Transform GT centroids from preprocessed to original image space
            # Use "gt_centroids" since inference overwrites "centroids" with predictions
            gt_centroids_prep = (
                batch["gt_centroids"][i].cpu().numpy()
            )  # (n_samples=1, max_inst, 2)
            gt_centroids_orig = gt_centroids_prep.squeeze(0) / eff  # (max_inst, 2)
            num_inst = batch["num_instances"][i].item()

            # Filter to valid instances (non-NaN)
            valid_pred_mask = ~np.isnan(pred_centroids).any(axis=1)
            pred_centroids = pred_centroids[valid_pred_mask]
            pred_vals = pred_vals[valid_pred_mask]

            gt_centroids_valid = gt_centroids_orig[:num_inst]

            self.val_predictions.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "pred_peaks": pred_centroids.reshape(
                        -1, 1, 2
                    ),  # (n_inst, 1, 2)
                    "pred_scores": pred_vals.reshape(-1, 1),  # (n_inst, 1)
                }
            )
            self.val_ground_truth.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "gt_instances": gt_centroids_valid.reshape(
                        -1, 1, 2
                    ),  # (n_inst, 1, 2)
                    "num_instances": num_inst,
                }
            )

visualize_example(sample)

Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

EmbeddingLightningModule

Bases: LightningModel

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

The embedder (backbone + EmbeddingHead) maps a grayscale, mask-burned-in crop to an L2-normalized vector. Training is contrastive: the training_step makes two augmented views of each crop (GPU-side), builds a positive/negative mask from each item's (video, frame, group, item_id), and applies the configured contrastive loss on a train-only projection head. The objective (positives x negatives x loss) is read from the head config; swapping it requires no code change here.

Methods:

Name	Description
__init__	Initialise the configs, model, objective + projection head.
configure_optimizers	Optimizer over trainable params only (frozen backbone -> adapter only).
forward	Inference forward: image -> embedding (B, D).
training_step	Two-view contrastive training step (views augmented in the dataset).
validation_step	Embed the val batch (no aug); compute a val loss + collect for retrieval.

Source code in sleap_nn/training/lightning_modules.py

class EmbeddingLightningModule(LightningModel):
    """Lightning Module for the ``embedding`` (crop -> vector, re-ID) model type.

    The embedder (backbone + EmbeddingHead) maps a grayscale, mask-burned-in crop to an
    L2-normalized vector. Training is contrastive: the ``training_step`` makes two
    augmented views of each crop (GPU-side), builds a positive/negative mask from each
    item's ``(video, frame, group, item_id)``, and applies the configured contrastive
    loss on a train-only projection head. The objective (positives x negatives x loss)
    is read from the head config; swapping it requires no code change here.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs, model, objective + projection head."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )
        leaf = self.head_configs.embedding.embedding
        self.embedding_dim = leaf.embedding_dim
        self.freeze_backbone = bool(
            OmegaConf.select(leaf, "freeze_backbone", default=False)
        )
        # Grayscale input + mask burn-in + per-crop standardize. `burn_in` defaults OFF
        # (matching `PreprocessingConfig.burn_in=False` and the sample config); BOTH the
        # training factory and the inference loader override it from
        # `data_config.preprocessing.burn_in` (set_embedding_burn_in_from_config), so a
        # mask-based model opts in there and train/inference stay in lockstep. When
        # burn_in is off, `_standardize` falls back to whole-crop standardize.
        self.burn_in = False
        # Per-crop standardize is always on (nothing sets this False), so the ``/255``
        # branch of ``_build_input`` is currently unreachable and ``background_fill='mean'``
        # is identical to ``'black'`` (the foreground mean is 0 in standardized space) —
        # see PreprocessingConfig.background_fill docs.
        self.standardize = True
        # What the masked-out background is replaced with when burn_in is on. The
        # factory / inference loader override this from
        # ``data_config.preprocessing.background_fill`` (set_embedding_burn_in_from_config).
        self.background_fill = "black"

        # Resolve the objective with defaults (nested sub-configs may be None).
        obj = resolve_embedding_objective(
            OmegaConf.select(leaf, "objective", default=None)
        )
        self.loss_name = obj.loss.name
        self.loss_temperature = obj.loss.temperature
        self.loss_margin = obj.loss.margin
        if float(self.loss_temperature) <= 0:
            raise ValueError(
                "head_configs.embedding.embedding.objective.loss.temperature must be "
                f"> 0 (got {self.loss_temperature}); a non-positive temperature makes "
                "the contrastive logits inf/NaN or inverts the objective."
            )
        if float(self.loss_margin) < 0:
            raise ValueError(
                "head_configs.embedding.embedding.objective.loss.margin must be >= 0 "
                f"(got {self.loss_margin})."
            )
        self.pos_scope = obj.positives.scope
        aug_views = int(obj.positives.aug_views)
        if aug_views != 2:
            raise ValueError(
                "head_configs.embedding.embedding.objective.positives.aug_views must "
                f"be 2 (got {aug_views}); the contrastive training_step uses exactly "
                "two augmented views per crop. Other view counts are not supported "
                "in P1 — set aug_views=2 (the default)."
            )
        if self.pos_scope == "aug_view":
            # Validation is not doubled (single view per crop) and aug_view has no
            # identity positives, so the per-row positive set is empty -> val/loss is
            # structurally ~0 every epoch. Checkpoint selection uses the retrieval
            # metric (not val/loss), but a `reduce_lr_on_plateau` scheduler that
            # monitors val/loss would decay the LR spuriously.
            logger.warning(
                "Embedding positives.scope='aug_view': val/loss is uninformative "
                "(no in-batch validation positives). Checkpoint selection uses the "
                "retrieval metric; avoid lr_scheduler='reduce_lr_on_plateau' (it "
                "monitors val/loss) for this self-supervised regime."
            )
        self.neg_sources = list(obj.negatives.sources)
        self.neg_exclude_same_track = bool(obj.negatives.exclude_same_track)
        self.neg_restrict_same_video = bool(obj.negatives.restrict_same_video)
        self.use_projection = bool(obj.use_projection)
        projection_dim = int(obj.projection_dim)

        self.loss_fn = get_contrastive_loss(self.loss_name)
        # Train-only projection head (discarded at inference) for supcon/infonce.
        if self.use_projection and self.loss_name in ("supcon", "infonce"):
            self.projection = nn.Sequential(
                nn.Linear(self.embedding_dim, projection_dim),
                nn.ReLU(inplace=True),
                nn.Linear(projection_dim, projection_dim),
            )
        else:
            self.projection = None

        if self.freeze_backbone:
            # Freeze the pretrained ENCODER only, in eval mode for the whole run (the
            # mechanism `backbone_config.pretrained.freeze` uses). Setting
            # `requires_grad=False` on the whole backbone instead left BatchNorm in
            # train mode (running stats drifted, the train forward used batch stats)
            # and froze the randomly initialized middle blocks of a native
            # convnext/swint -- see `EmbeddingHeadConfig.freeze_backbone`.
            self.model.backbone.freeze_encoder()

    # ---- objective helpers ----
    def _project(self, e: torch.Tensor) -> torch.Tensor:
        if self.projection is None:
            return F.normalize(e, dim=1)
        return F.normalize(self.projection(e), dim=1)

    def _loss_kwargs(self):
        if self.loss_name in ("supcon", "infonce"):
            return {"temperature": self.loss_temperature}
        return {"margin": self.loss_margin}

    def _build_masks(self, item_id, video, frame, group):
        return build_contrastive_masks(
            item_id=item_id,
            video=video,
            frame=frame,
            group=group,
            positives_scope=self.pos_scope,
            negatives_sources=self.neg_sources,
            exclude_same_track=self.neg_exclude_same_track,
            restrict_same_video=self.neg_restrict_same_video,
        )

    # ---- crop intensity (Stage 5) + two-view aug (Stage 0, GPU) ----
    def _standardize(self, gray, mask, eps=1e-5, std_floor=EMBEDDING_STD_FLOOR):
        # Reduce over spatial dims ONLY (keep the channel axis) so each channel is
        # standardized independently. For grayscale (C=1) this is byte-identical to the
        # old whole-tensor reduction; for RGB (C=3) it yields a true per-channel
        # zero-mean/unit-std instead of a ~C x-scaled cross-channel "mean". `mask` is
        # single-channel and broadcasts across channels.
        m = mask if self.burn_in else torch.ones_like(mask)
        cnt = m.sum((2, 3), keepdim=True).clamp(min=1)
        mu = (gray * m).sum((2, 3), keepdim=True) / cnt
        std = (((gray - mu) ** 2 * m).sum((2, 3), keepdim=True) / cnt).sqrt() + eps
        # Floor the std at one grey level (inputs are on the 0-255 scale): `eps` alone
        # let a flat (or empty) foreground divide by ~1e-5, so a `grey` background fill
        # came out ~1e7. Real crops have a std of tens of grey levels, where the floor
        # is inactive and the result is bit-identical to before.
        std = std.clamp(min=std_floor)
        g = (gray - mu) / std
        if not self.burn_in:
            return g
        # Compose the standardized foreground with the configured background fill.
        return g * m + self._background_fill(g, m, mu, std) * (1 - m)

    def _background_fill(self, g, m, mu, std):
        """Standardized-space fill for the masked-out background (burn-in only).

        ``black`` (default) and ``mean`` are 0 (the foreground mean in standardized
        space — the original mask-multiply); ``grey`` is the standardized value of raw
        mid-grey for the crop; ``noise`` is per-pixel standard-normal noise. ``noise``
        is applied only during training (it is an augmentation): at eval / inference it
        falls back to the neutral 0 fill so the retrieval metric — and thus checkpoint
        selection — stays deterministic.
        """
        fill = getattr(self, "background_fill", "black")
        if fill == "grey":
            return (127.5 - mu) / std
        if fill == "noise" and self.training:
            return torch.randn_like(g)
        # "black" / "mean" / (noise at eval) -> foreground mean == 0 in standardized space.
        return torch.zeros_like(g)

    def _build_input(self, gray, mask):
        if self.standardize:
            return self._standardize(gray, mask)
        g = gray / 255.0
        if not self.burn_in:
            return g
        fill = getattr(self, "background_fill", "black")
        if fill == "grey":
            bg = torch.full_like(g, 0.5)
        elif fill == "mean":
            cnt = mask.sum((1, 2, 3), keepdim=True).clamp(min=1)
            bg = (g * mask).sum((1, 2, 3), keepdim=True) / cnt
        elif fill == "noise" and self.training:  # augmentation; deterministic at eval
            bg = torch.rand_like(g)
        else:  # "black" / (noise at eval)
            bg = torch.zeros_like(g)
        return g * mask + bg * (1 - mask)

    def _two_views(self, batch):
        """Return the two augmented (gray, mask) views for the contrastive step.

        The views are produced by the standard config-driven skia augmentation in
        ``EmbeddingDataset.__getitem__`` (``instance_image`` / ``instance_image_view2``).
        If a second view is absent (``apply_aug=False``), the first view is reused
        (degenerate — train with ``use_augmentations_train=True``).
        """
        gray1 = torch.squeeze(batch["instance_image"], dim=1).to(torch.float32)
        mask1 = torch.squeeze(batch["instance_mask"], dim=1).to(torch.float32)
        if "instance_image_view2" in batch:
            gray2 = torch.squeeze(batch["instance_image_view2"], dim=1).to(
                torch.float32
            )
            mask2 = torch.squeeze(batch["instance_mask_view2"], dim=1).to(torch.float32)
        else:
            gray2, mask2 = gray1, mask1
        return (gray1, mask1), (gray2, mask2)

    def forward(self, img, mask=None):
        """Inference forward: image -> embedding (B, D).

        Runs the SAME crop pipeline as training/validation (``_build_input``: mask
        burn-in + per-crop standardize), so pass ``mask`` (the instance mask) to
        reproduce a burn-in model's masked, foreground-only standardize and stay in
        lockstep with training. When ``mask is None`` — no mask available, e.g.
        centroid-driven raw-frame inference — a whole-crop (all-ones) standardize is
        used, which DIVERGES from a burn-in model's masked training standardize
        (callers on that path warn; see ``inference/embedding.py``). The mask-driven
        inference layer (``EmbeddingLayer``) calls ``_build_input`` with the real mask
        directly, so it is unaffected by this fallback.
        """
        img = torch.squeeze(img, dim=1).to(self.device).to(torch.float32)
        if mask is None:
            mask = torch.ones_like(img[:, :1])
        else:
            mask = mask.to(self.device).to(torch.float32)
            if mask.dim() == 5:
                mask = torch.squeeze(mask, dim=1)
        return self.model(self._build_input(img, mask))["EmbeddingHead"]

    def training_step(self, batch, batch_idx):
        """Two-view contrastive training step (views augmented in the dataset)."""
        (g1, m1), (g2, m2) = self._two_views(batch)
        x = torch.cat([self._build_input(g1, m1), self._build_input(g2, m2)], dim=0)
        e = self.model(x)["EmbeddingHead"]
        z = self._project(e)

        item_id = torch.cat([batch["item_id"], batch["item_id"]], dim=0)
        # `video_id` is unique across labels files; `video_idx` is per file (see
        # EmbeddingDataset.__init__). Fall back for batches built before that field
        # existed (e.g. a hand-built dict in a test).
        video_key = "video_id" if "video_id" in batch else "video_idx"
        video = torch.cat([batch[video_key], batch[video_key]], dim=0)
        frame = torch.cat([batch["frame_idx"], batch["frame_idx"]], dim=0)
        group = torch.cat([batch["group_id"], batch["group_id"]], dim=0)
        pos, neg = self._build_masks(item_id, video, frame, group)
        loss = self.loss_fn(z, pos, neg, **self._loss_kwargs())

        self.log(
            "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
        )
        self._accumulate_loss(loss)
        self.log("train/loss", loss, on_step=False, on_epoch=True, sync_dist=True)
        with torch.no_grad():
            self.log(
                "train/pos_per_anchor",
                (pos.sum(1).float().mean()),
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )
        return loss

    def validation_step(self, batch, batch_idx):
        """Embed the val batch (no aug); compute a val loss + collect for retrieval."""
        gray = torch.squeeze(batch["instance_image"], dim=1).to(torch.float32)
        mask = torch.squeeze(batch["instance_mask"], dim=1).to(torch.float32)
        x = self._build_input(gray, mask)
        e = self.model(x)["EmbeddingHead"]
        z = self._project(e)

        item_id = batch["item_id"]
        pos, neg = self._build_masks(
            item_id,
            batch["video_id" if "video_id" in batch else "video_idx"],
            batch["frame_idx"],
            batch["group_id"],
        )
        val_loss = self.loss_fn(z, pos, neg, **self._loss_kwargs())
        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        if self._collect_val_predictions:
            emb = e.detach().cpu()
            # Evaluate on the eval grouping (`global_group_id`: the global identity, or
            # a bare tracklet's own group), independent of the training group (e.g.
            # tracklet), so different objectives are comparable on one retrieval
            # metric. The post-training eval (train.py) groups on the same key.
            label_key = "global_group_id" if "global_group_id" in batch else "group_id"
            labels = batch[label_key].detach().cpu()
            for i in range(emb.shape[0]):
                self.val_predictions.append({"embedding": emb[i]})
                self.val_ground_truth.append({"label": int(labels[i])})

    def configure_optimizers(self):
        """Optimizer over trainable params only (frozen backbone -> adapter only).

        Builds the optimizer over ``requires_grad`` params, then delegates to the base
        scheduler construction with THAT optimizer so the scheduler is bound to the
        optimizer that is returned (val/loss monitor; checkpoint SELECTION still moves
        to the retrieval metric via ModelCheckpoint in the trainer).
        """
        optim = torch.optim.AdamW if self.optimizer == "AdamW" else torch.optim.Adam
        optimizer = optim(
            self._trainable_param_groups(), lr=self.lr, amsgrad=self.amsgrad
        )
        return LightningModel.configure_optimizers(self, optimizer=optimizer)

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs, model, objective + projection head.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs, model, objective + projection head."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )
    leaf = self.head_configs.embedding.embedding
    self.embedding_dim = leaf.embedding_dim
    self.freeze_backbone = bool(
        OmegaConf.select(leaf, "freeze_backbone", default=False)
    )
    # Grayscale input + mask burn-in + per-crop standardize. `burn_in` defaults OFF
    # (matching `PreprocessingConfig.burn_in=False` and the sample config); BOTH the
    # training factory and the inference loader override it from
    # `data_config.preprocessing.burn_in` (set_embedding_burn_in_from_config), so a
    # mask-based model opts in there and train/inference stay in lockstep. When
    # burn_in is off, `_standardize` falls back to whole-crop standardize.
    self.burn_in = False
    # Per-crop standardize is always on (nothing sets this False), so the ``/255``
    # branch of ``_build_input`` is currently unreachable and ``background_fill='mean'``
    # is identical to ``'black'`` (the foreground mean is 0 in standardized space) —
    # see PreprocessingConfig.background_fill docs.
    self.standardize = True
    # What the masked-out background is replaced with when burn_in is on. The
    # factory / inference loader override this from
    # ``data_config.preprocessing.background_fill`` (set_embedding_burn_in_from_config).
    self.background_fill = "black"

    # Resolve the objective with defaults (nested sub-configs may be None).
    obj = resolve_embedding_objective(
        OmegaConf.select(leaf, "objective", default=None)
    )
    self.loss_name = obj.loss.name
    self.loss_temperature = obj.loss.temperature
    self.loss_margin = obj.loss.margin
    if float(self.loss_temperature) <= 0:
        raise ValueError(
            "head_configs.embedding.embedding.objective.loss.temperature must be "
            f"> 0 (got {self.loss_temperature}); a non-positive temperature makes "
            "the contrastive logits inf/NaN or inverts the objective."
        )
    if float(self.loss_margin) < 0:
        raise ValueError(
            "head_configs.embedding.embedding.objective.loss.margin must be >= 0 "
            f"(got {self.loss_margin})."
        )
    self.pos_scope = obj.positives.scope
    aug_views = int(obj.positives.aug_views)
    if aug_views != 2:
        raise ValueError(
            "head_configs.embedding.embedding.objective.positives.aug_views must "
            f"be 2 (got {aug_views}); the contrastive training_step uses exactly "
            "two augmented views per crop. Other view counts are not supported "
            "in P1 — set aug_views=2 (the default)."
        )
    if self.pos_scope == "aug_view":
        # Validation is not doubled (single view per crop) and aug_view has no
        # identity positives, so the per-row positive set is empty -> val/loss is
        # structurally ~0 every epoch. Checkpoint selection uses the retrieval
        # metric (not val/loss), but a `reduce_lr_on_plateau` scheduler that
        # monitors val/loss would decay the LR spuriously.
        logger.warning(
            "Embedding positives.scope='aug_view': val/loss is uninformative "
            "(no in-batch validation positives). Checkpoint selection uses the "
            "retrieval metric; avoid lr_scheduler='reduce_lr_on_plateau' (it "
            "monitors val/loss) for this self-supervised regime."
        )
    self.neg_sources = list(obj.negatives.sources)
    self.neg_exclude_same_track = bool(obj.negatives.exclude_same_track)
    self.neg_restrict_same_video = bool(obj.negatives.restrict_same_video)
    self.use_projection = bool(obj.use_projection)
    projection_dim = int(obj.projection_dim)

    self.loss_fn = get_contrastive_loss(self.loss_name)
    # Train-only projection head (discarded at inference) for supcon/infonce.
    if self.use_projection and self.loss_name in ("supcon", "infonce"):
        self.projection = nn.Sequential(
            nn.Linear(self.embedding_dim, projection_dim),
            nn.ReLU(inplace=True),
            nn.Linear(projection_dim, projection_dim),
        )
    else:
        self.projection = None

    if self.freeze_backbone:
        # Freeze the pretrained ENCODER only, in eval mode for the whole run (the
        # mechanism `backbone_config.pretrained.freeze` uses). Setting
        # `requires_grad=False` on the whole backbone instead left BatchNorm in
        # train mode (running stats drifted, the train forward used batch stats)
        # and froze the randomly initialized middle blocks of a native
        # convnext/swint -- see `EmbeddingHeadConfig.freeze_backbone`.
        self.model.backbone.freeze_encoder()

configure_optimizers()

Optimizer over trainable params only (frozen backbone -> adapter only).

Builds the optimizer over requires_grad params, then delegates to the base scheduler construction with THAT optimizer so the scheduler is bound to the optimizer that is returned (val/loss monitor; checkpoint SELECTION still moves to the retrieval metric via ModelCheckpoint in the trainer).

Source code in sleap_nn/training/lightning_modules.py

def configure_optimizers(self):
    """Optimizer over trainable params only (frozen backbone -> adapter only).

    Builds the optimizer over ``requires_grad`` params, then delegates to the base
    scheduler construction with THAT optimizer so the scheduler is bound to the
    optimizer that is returned (val/loss monitor; checkpoint SELECTION still moves
    to the retrieval metric via ModelCheckpoint in the trainer).
    """
    optim = torch.optim.AdamW if self.optimizer == "AdamW" else torch.optim.Adam
    optimizer = optim(
        self._trainable_param_groups(), lr=self.lr, amsgrad=self.amsgrad
    )
    return LightningModel.configure_optimizers(self, optimizer=optimizer)

forward(img, mask=None)

Inference forward: image -> embedding (B, D).

Runs the SAME crop pipeline as training/validation (_build_input: mask burn-in + per-crop standardize), so pass mask (the instance mask) to reproduce a burn-in model's masked, foreground-only standardize and stay in lockstep with training. When mask is None — no mask available, e.g. centroid-driven raw-frame inference — a whole-crop (all-ones) standardize is used, which DIVERGES from a burn-in model's masked training standardize (callers on that path warn; see inference/embedding.py). The mask-driven inference layer (EmbeddingLayer) calls _build_input with the real mask directly, so it is unaffected by this fallback.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img, mask=None):
    """Inference forward: image -> embedding (B, D).

    Runs the SAME crop pipeline as training/validation (``_build_input``: mask
    burn-in + per-crop standardize), so pass ``mask`` (the instance mask) to
    reproduce a burn-in model's masked, foreground-only standardize and stay in
    lockstep with training. When ``mask is None`` — no mask available, e.g.
    centroid-driven raw-frame inference — a whole-crop (all-ones) standardize is
    used, which DIVERGES from a burn-in model's masked training standardize
    (callers on that path warn; see ``inference/embedding.py``). The mask-driven
    inference layer (``EmbeddingLayer``) calls ``_build_input`` with the real mask
    directly, so it is unaffected by this fallback.
    """
    img = torch.squeeze(img, dim=1).to(self.device).to(torch.float32)
    if mask is None:
        mask = torch.ones_like(img[:, :1])
    else:
        mask = mask.to(self.device).to(torch.float32)
        if mask.dim() == 5:
            mask = torch.squeeze(mask, dim=1)
    return self.model(self._build_input(img, mask))["EmbeddingHead"]

training_step(batch, batch_idx)

Two-view contrastive training step (views augmented in the dataset).

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Two-view contrastive training step (views augmented in the dataset)."""
    (g1, m1), (g2, m2) = self._two_views(batch)
    x = torch.cat([self._build_input(g1, m1), self._build_input(g2, m2)], dim=0)
    e = self.model(x)["EmbeddingHead"]
    z = self._project(e)

    item_id = torch.cat([batch["item_id"], batch["item_id"]], dim=0)
    # `video_id` is unique across labels files; `video_idx` is per file (see
    # EmbeddingDataset.__init__). Fall back for batches built before that field
    # existed (e.g. a hand-built dict in a test).
    video_key = "video_id" if "video_id" in batch else "video_idx"
    video = torch.cat([batch[video_key], batch[video_key]], dim=0)
    frame = torch.cat([batch["frame_idx"], batch["frame_idx"]], dim=0)
    group = torch.cat([batch["group_id"], batch["group_id"]], dim=0)
    pos, neg = self._build_masks(item_id, video, frame, group)
    loss = self.loss_fn(z, pos, neg, **self._loss_kwargs())

    self.log(
        "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
    )
    self._accumulate_loss(loss)
    self.log("train/loss", loss, on_step=False, on_epoch=True, sync_dist=True)
    with torch.no_grad():
        self.log(
            "train/pos_per_anchor",
            (pos.sum(1).float().mean()),
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
    return loss

validation_step(batch, batch_idx)

Embed the val batch (no aug); compute a val loss + collect for retrieval.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Embed the val batch (no aug); compute a val loss + collect for retrieval."""
    gray = torch.squeeze(batch["instance_image"], dim=1).to(torch.float32)
    mask = torch.squeeze(batch["instance_mask"], dim=1).to(torch.float32)
    x = self._build_input(gray, mask)
    e = self.model(x)["EmbeddingHead"]
    z = self._project(e)

    item_id = batch["item_id"]
    pos, neg = self._build_masks(
        item_id,
        batch["video_id" if "video_id" in batch else "video_idx"],
        batch["frame_idx"],
        batch["group_id"],
    )
    val_loss = self.loss_fn(z, pos, neg, **self._loss_kwargs())
    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    if self._collect_val_predictions:
        emb = e.detach().cpu()
        # Evaluate on the eval grouping (`global_group_id`: the global identity, or
        # a bare tracklet's own group), independent of the training group (e.g.
        # tracklet), so different objectives are comparable on one retrieval
        # metric. The post-training eval (train.py) groups on the same key.
        label_key = "global_group_id" if "global_group_id" in batch else "group_id"
        labels = batch[label_key].detach().cpu()
        for i in range(emb.shape[0]):
            self.val_predictions.append({"embedding": emb[i]})
            self.val_ground_truth.append({"label": int(labels[i])})

LightningModel

Bases: LightningModule

Base PyTorch Lightning Module for all sleap-nn models.

This class is a sub-class of Torch Lightning Module to configure the training and validation steps.

Parameters:

Name	Type	Description	Default
model_type	str	Type of the model. One of single_instance, centered_instance, centroid, bottomup, multi_class_bottomup, multi_class_topdown.	required
backbone_type	str	Backbone model. One of unet, convnext and swint.	required
backbone_config	Union[str, Dict[str, Any], DictConfig]	Backbone configuration. Can be: - String: One of the preset backbone types: - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"] - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"] - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"] - Dictionary: Custom configuration with structure: { "unet": {UNetConfig parameters}, "convnext": {ConvNextConfig parameters}, "swint": {SwinTConfig parameters} } Only one backbone type should be specified in the dictionary. - DictConfig: OmegaConf DictConfig object containing backbone configuration.	required
head_configs	DictConfig	Head configuration dictionary containing model-specific parameters. For Single Instance: confmaps with part_names, sigma, output_stride. For Centroid: confmaps with anchor_part, sigma, output_stride. For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride. For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight. For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight. For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.	required
pretrained_backbone_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.	None
pretrained_head_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.	None
init_weights	Optional[str]	Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.	'xavier'
lr_scheduler	Optional[Union[str, DictConfig]]	Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.	None
online_mining	Optional[bool]	If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).	False
hard_to_easy_ratio	Optional[float]	Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.	2.0
min_hard_keypoints	Optional[int]	Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.	2
max_hard_keypoints	Optional[int]	Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.	None
loss_scale	Optional[float]	Factor to scale hard keypoint losses by. Default: 5.0.	5.0
optimizer	Optional[str]	Optimizer name. One of ["Adam", "AdamW"].	'Adam'
learning_rate	Optional[float]	Learning rate for the optimizer. Default: 1e-3.	0.001
amsgrad	Optional[bool]	Enable AMSGrad with the optimizer. Default: False.	False

Methods:

Name	Description
__init__	Initialise the configs and the model.
configure_optimizers	Configure optimiser and learning rate scheduler.
forward	Forward pass of the model.
get_lightning_model_from_config	Get lightning model from config.
on_train_batch_end	Count per-device samples this epoch (for throughput logging).
on_train_epoch_end	Configure the train timer at the end of every epoch.
on_train_epoch_start	Configure the train timer at the beginning of each epoch.
on_validation_epoch_end	Configure the val timer at the end of every epoch.
on_validation_epoch_start	Configure the val timer at the beginning of each epoch.
training_step	Training step.
validation_step	Validation step.

Source code in sleap_nn/training/lightning_modules.py

class LightningModel(L.LightningModule):
    """Base PyTorch Lightning Module for all sleap-nn models.

    This class is a sub-class of Torch Lightning Module to configure the training and validation steps.

    Args:
        model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
        backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
        backbone_config: Backbone configuration. Can be:
            - String: One of the preset backbone types:
                - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"]
                - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"]
                - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"]
            - Dictionary: Custom configuration with structure:
                {
                    "unet": {UNetConfig parameters},
                    "convnext": {ConvNextConfig parameters},
                    "swint": {SwinTConfig parameters}
                }
                Only one backbone type should be specified in the dictionary.
            - DictConfig: OmegaConf DictConfig object containing backbone configuration.
        head_configs: Head configuration dictionary containing model-specific parameters.
            For Single Instance: confmaps with part_names, sigma, output_stride.
            For Centroid: confmaps with anchor_part, sigma, output_stride.
            For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride.
            For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight.
            For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight.
            For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.
        pretrained_backbone_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.
        pretrained_head_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.
        init_weights: Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.
        lr_scheduler: Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.
        online_mining: If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).
        hard_to_easy_ratio: Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.
        min_hard_keypoints: Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.
        max_hard_keypoints: Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.
        loss_scale: Factor to scale hard keypoint losses by. Default: 5.0.
        optimizer: Optimizer name. One of ["Adam", "AdamW"].
        learning_rate: Learning rate for the optimizer. Default: 1e-3.
        amsgrad: Enable AMSGrad with the optimizer. Default: False.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__()
        self.negative_loss_weight = negative_loss_weight
        self.model_type = model_type
        self.backbone_type = backbone_type
        if not isinstance(backbone_config, DictConfig):
            backbone_cfg = get_backbone_config(backbone_config)
            config = OmegaConf.structured(backbone_cfg)
            OmegaConf.to_container(config, resolve=True, throw_on_missing=True)
            config = DictConfig(config)
        else:
            config = backbone_config
        self.backbone_config = config
        self.head_configs = head_configs
        self.pretrained_backbone_weights = pretrained_backbone_weights
        self.pretrained_head_weights = pretrained_head_weights
        self.in_channels = self.backbone_config[f"{self.backbone_type}"]["in_channels"]
        self.input_expand_channels = self.in_channels
        self.init_weights = init_weights
        self.lr_scheduler = lr_scheduler
        self.online_mining = online_mining
        self.hard_to_easy_ratio = hard_to_easy_ratio
        self.min_hard_keypoints = min_hard_keypoints
        self.max_hard_keypoints = max_hard_keypoints
        self.loss_scale = loss_scale
        self.optimizer = optimizer
        self.lr = learning_rate
        self.amsgrad = amsgrad

        self.model = Model(
            backbone_type=self.backbone_type,
            backbone_config=self.backbone_config[f"{self.backbone_type}"],
            head_configs=self.head_configs[self.model_type],
            model_type=self.model_type,
        )

        if len(self.head_configs[self.model_type]) > 1:
            self.loss_weights = [
                (
                    self.head_configs[self.model_type][x].loss_weight
                    if self.head_configs[self.model_type][x].loss_weight is not None
                    else 1.0
                )
                for x in self.head_configs[self.model_type]
            ]

        self.training_loss = {}
        self.val_loss = {}
        self.learning_rate = {}

        # For epoch-averaged loss tracking
        self._epoch_loss_sum = 0.0
        self._epoch_loss_count = 0

        # For throughput logging (samples/frames per second).
        self._epoch_sample_count = 0
        self._samples_per_frame_cache = None

        # For epoch-end evaluation
        self.val_predictions: List[Dict] = []
        self.val_ground_truth: List[Dict] = []
        self._collect_val_predictions: bool = False

        # Initialization for encoder and decoder stacks.
        if self.init_weights == "xavier":
            self.model.apply(xavier_init_weights)

        # Pre-trained weights for the encoder stack - only for swint and convnext
        if self.backbone_type == "convnext" or self.backbone_type == "swint":
            if (
                self.backbone_config[f"{self.backbone_type}"]["pre_trained_weights"]
                is not None
            ):
                ckpt = MODEL_WEIGHTS[
                    self.backbone_config[f"{self.backbone_type}"]["pre_trained_weights"]
                ].DEFAULT.get_state_dict(progress=True, check_hash=True)
                self.model.backbone.enc.load_state_dict(ckpt, strict=False)

        # External pretrained (HuggingFace) backbone: the wrapper loaded its
        # weights in __init__, but the xavier init above clobbered them (it runs
        # on every Conv2d/Linear). Re-apply the snapshotted pretrained weights, then
        # freeze the encoder if requested. Mirrors the convnext/swint ordering.
        if self.backbone_type == "pretrained":
            self.model.backbone.reload_pretrained_weights()
            if getattr(self.model.backbone, "freeze", False):
                self.model.backbone.freeze_encoder()

        # Initializing backbone (encoder + decoder) with trained ckpts
        if self.pretrained_backbone_weights is not None:
            logger.info(
                f"Loading backbone weights from `{self.pretrained_backbone_weights}` ..."
            )
            if self.pretrained_backbone_weights.endswith(".ckpt"):
                ckpt = torch.load(
                    self.pretrained_backbone_weights,
                    map_location="cpu",
                    weights_only=False,
                )
                ckpt["state_dict"] = {
                    k: ckpt["state_dict"][k]
                    for k in ckpt["state_dict"].keys()
                    if ".backbone" in k
                }
                self.load_state_dict(ckpt["state_dict"], strict=False)

            elif self.pretrained_backbone_weights.endswith(".h5"):
                # load from sleap model weights
                load_legacy_model_weights(
                    self.model.backbone,
                    self.pretrained_backbone_weights,
                    component="backbone",
                )

            else:
                message = f"Unsupported file extension for pretrained backbone weights. Please provide a .ckpt or .h5 file."
                logger.error(message)
                raise ValueError(message)

        # Initializing head layers with trained ckpts.
        if self.pretrained_head_weights is not None:
            logger.info(
                f"Loading head weights from `{self.pretrained_head_weights}` ..."
            )
            if self.pretrained_head_weights.endswith(".ckpt"):
                ckpt = torch.load(
                    self.pretrained_head_weights,
                    map_location="cpu",
                    weights_only=False,
                )
                ckpt["state_dict"] = {
                    k: ckpt["state_dict"][k]
                    for k in ckpt["state_dict"].keys()
                    if ".head_layers" in k
                }
                self.load_state_dict(ckpt["state_dict"], strict=False)

            elif self.pretrained_head_weights.endswith(".h5"):
                # load from sleap model weights
                load_legacy_model_weights(
                    self.model.head_layers,
                    self.pretrained_head_weights,
                    component="head",
                )

            else:
                message = f"Unsupported file extension for pretrained head weights. Please provide a .ckpt or .h5 file."
                logger.error(message)
                raise ValueError(message)

    @classmethod
    def get_lightning_model_from_config(cls, config: DictConfig):
        """Get lightning model from config."""
        model_type = get_model_type_from_cfg(config)
        backbone_type = get_backbone_type_from_cfg(config)

        lightning_models = {
            "single_instance": SingleInstanceLightningModule,
            "centroid": CentroidLightningModule,
            "centered_instance": TopDownCenteredInstanceLightningModule,
            "bottomup": BottomUpLightningModule,
            "multi_class_bottomup": BottomUpMultiClassLightningModule,
            "multi_class_topdown": TopDownCenteredInstanceMultiClassLightningModule,
            "bottomup_segmentation": BottomUpSegmentationLightningModule,
            "centered_instance_segmentation": TopDownCenteredInstanceSegmentationLightningModule,
            "semantic_segmentation": SemanticSegmentationLightningModule,
            "embedding": EmbeddingLightningModule,
            "bbox_centernet": BBoxCenterNetLightningModule,
            "bbox_ssd": BBoxSSDLightningModule,
        }

        if model_type not in lightning_models:
            message = f"Incorrect model type. Please check if one of the following keys in the head configs is not None: [`single_instance`, `centroid`, `centered_instance`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`, `bottomup_segmentation`, `centered_instance_segmentation`, `semantic_segmentation`, `embedding`, `bbox_centernet`, `bbox_ssd`]"
            logger.error(message)
            raise ValueError(message)

        negative_loss_weight = getattr(config.data_config, "negative_loss_weight", 1.0)

        # See CentroidConfMapsConfig.focal_loss_alpha -- centroid-only.
        extra_kwargs = {}
        if model_type == "centroid":
            centroid_confmaps_config = (
                config.model_config.head_configs.centroid.confmaps
            )
            extra_kwargs["centroid_focal_loss_alpha"] = getattr(
                centroid_confmaps_config, "focal_loss_alpha", 0.0
            )
            extra_kwargs["centroid_focal_loss_beta"] = getattr(
                centroid_confmaps_config, "focal_loss_beta", 4.0
            )
            extra_kwargs["centroid_focal_loss_pos_threshold"] = getattr(
                centroid_confmaps_config, "focal_loss_pos_threshold", 0.5
            )

        lightning_model = lightning_models[model_type](
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=config.model_config.backbone_config,
            head_configs=config.model_config.head_configs,
            pretrained_backbone_weights=config.model_config.pretrained_backbone_weights,
            pretrained_head_weights=config.model_config.pretrained_head_weights,
            init_weights=config.model_config.init_weights,
            lr_scheduler=config.trainer_config.lr_scheduler,
            online_mining=config.trainer_config.online_hard_keypoint_mining.online_mining,
            hard_to_easy_ratio=config.trainer_config.online_hard_keypoint_mining.hard_to_easy_ratio,
            min_hard_keypoints=config.trainer_config.online_hard_keypoint_mining.min_hard_keypoints,
            max_hard_keypoints=config.trainer_config.online_hard_keypoint_mining.max_hard_keypoints,
            loss_scale=config.trainer_config.online_hard_keypoint_mining.loss_scale,
            optimizer=config.trainer_config.optimizer_name,
            learning_rate=config.trainer_config.optimizer.lr,
            amsgrad=config.trainer_config.optimizer.amsgrad,
            negative_loss_weight=negative_loss_weight,
            **extra_kwargs,
        )

        if model_type == "embedding":
            # Make `data_config.preprocessing.burn_in` live (mask-on vs centroid-crop).
            set_embedding_burn_in_from_config(lightning_model, config)

        return lightning_model

    def forward(self, img):
        """Forward pass of the model."""
        pass

    def on_train_epoch_start(self):
        """Configure the train timer at the beginning of each epoch."""
        self.train_start_time = time.time()
        # Reset epoch loss tracking
        self._epoch_loss_sum = 0.0
        self._epoch_loss_count = 0
        # Reset per-device sample count for throughput.
        self._epoch_sample_count = 0

    def _accumulate_loss(self, loss: torch.Tensor):
        """Accumulate loss for epoch-averaged logging. Call this in training_step."""
        self._epoch_loss_sum += loss.detach().item()
        self._epoch_loss_count += 1

    def on_train_batch_end(self, outputs, batch, batch_idx):
        """Count per-device samples this epoch (for throughput logging)."""
        try:
            if isinstance(batch, dict):
                if "image" in batch and hasattr(batch["image"], "shape"):
                    self._epoch_sample_count += int(batch["image"].shape[0])
                elif "frame_idx" in batch:
                    self._epoch_sample_count += len(batch["frame_idx"])
        except Exception:
            pass

    def _tiling_samples_per_frame(self) -> int:
        """Tiles sampled per source frame (1 when not tiling).

        Used to convert the sample (tile) throughput into a *source-frame*
        throughput so ``train/frames_per_sec`` reads as full frames, not crops.
        """
        if self._samples_per_frame_cache is not None:
            return self._samples_per_frame_cache
        spf = 1
        try:
            dl = getattr(self.trainer, "train_dataloader", None)
            ds = getattr(dl, "dataset", None)
            if ds is not None and getattr(ds, "tiling_enabled", False):
                spf = max(1, int(getattr(ds, "samples_per_frame", 1) or 1))
        except Exception:
            spf = 1
        self._samples_per_frame_cache = spf
        return spf

    def on_train_epoch_end(self):
        """Configure the train timer at the end of every epoch."""
        train_time = time.time() - self.train_start_time
        self.log(
            "train/time",
            train_time,
            prog_bar=False,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        # Log epoch explicitly for custom x-axis support in wandb
        self.log(
            "epoch",
            float(self.current_epoch),
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        # Log epoch-averaged training loss
        if self._epoch_loss_count > 0:
            avg_loss = self._epoch_loss_sum / self._epoch_loss_count
            self.log(
                "train/loss",
                avg_loss,
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )
        # Log current learning rate (useful for monitoring LR schedulers)
        if self.trainer.optimizers:
            lr = self.trainer.optimizers[0].param_groups[0]["lr"]
            self.log(
                "train/lr",
                lr,
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )

        # Throughput: optimizer steps/sec, sample (tile/crop) throughput, and
        # full source-frames/sec. Uses the GLOBAL batch (per-device samples x
        # world_size), so it is correct under multi-GPU. ``frames_per_sec`` divides
        # the sample rate by tiles-sampled-per-frame under tiling, so it reads as
        # source frames rather than crops/tiles.
        if train_time > 0:
            world = int(getattr(self.trainer, "world_size", 1) or 1)
            spf = self._tiling_samples_per_frame()
            steps_per_sec = self._epoch_loss_count / train_time
            samples_per_sec = (self._epoch_sample_count * world) / train_time
            frames_per_sec = samples_per_sec / spf
            for _name, _val in (
                ("train/steps_per_sec", steps_per_sec),
                ("train/samples_per_sec", samples_per_sec),
                ("train/frames_per_sec", frames_per_sec),
            ):
                self.log(
                    _name,
                    _val,
                    prog_bar=False,
                    on_step=False,
                    on_epoch=True,
                    sync_dist=False,
                )

    def on_validation_epoch_start(self):
        """Configure the val timer at the beginning of each epoch."""
        self.val_start_time = time.time()
        # Clear accumulated predictions for new epoch
        self.val_predictions = []
        self.val_ground_truth = []

    def on_validation_epoch_end(self):
        """Configure the val timer at the end of every epoch."""
        val_time = time.time() - self.val_start_time
        self.log(
            "val/time",
            val_time,
            prog_bar=False,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        # Log epoch explicitly so val/* metrics can use it as x-axis in wandb
        # (mirrors what on_train_epoch_end does for train/* metrics)
        self.log(
            "epoch",
            float(self.current_epoch),
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

    def _compute_negative_weighted_loss(
        self, y_preds: torch.Tensor, y: torch.Tensor, batch: Dict, stage: str = "train"
    ) -> torch.Tensor:
        """Compute MSE loss with optional negative sample weighting.

        This is a **pure loss function**: it computes the (optionally weighted)
        loss for backprop and logs **nothing**. Diagnostic split metrics
        (positive/negative loss and counts) are logged separately by
        :meth:`_log_negative_split_metrics`, which is called exactly once per
        step (so counts are not double-logged for multi-head models, which call
        this loss helper once per head).

        When ``is_negative`` is absent from the batch this returns plain
        ``nn.MSELoss()``. Otherwise per-sample MSE is computed and weighting
        depends on ``stage``:

        * ``stage == "train"``: negative samples are weighted by
          ``negative_loss_weight`` (only when it is not 1.0).
        * ``stage != "train"`` (e.g. ``"val"``): the loss is **always
          unweighted** (``per_sample.mean()``). This keeps ``val/loss``
          numerically identical to plain ``nn.MSELoss()`` so that
          ``ModelCheckpoint(monitor="val/loss")`` and ``EarlyStopping``
          behavior is unchanged.

        Args:
            y_preds: Predicted tensor.
            y: Ground truth tensor.
            batch: Batch dictionary, may contain ``is_negative`` key.
            stage: Weighting gate. ``"train"`` applies negative weighting; any
                other value returns the unweighted mean.

        Returns:
            The (optionally weighted) loss tensor.
        """
        is_negative = batch.get("is_negative", None)
        if is_negative is None:
            return nn.MSELoss()(y_preds, y)

        # Per-sample loss for weighting. Mean over all non-batch dims makes this
        # agnostic to head shape (confmaps vs PAF vs classmaps).
        per_sample = (
            (y_preds - y).pow(2).mean(dim=list(range(1, y_preds.ndim)))
        )  # (batch,)

        # Negative weighting is train-only; val/eval stages stay unweighted so
        # val/loss remains numerically identical to plain nn.MSELoss().
        if stage != "train" or self.negative_loss_weight == 1.0:
            return per_sample.mean()

        is_neg = is_negative.to(y_preds.device)
        weights = torch.where(
            is_neg,
            torch.tensor(self.negative_loss_weight, device=y_preds.device),
            torch.tensor(1.0, device=y_preds.device),
        )
        return (per_sample * weights).mean()

    def _log_negative_split_metrics(
        self,
        heads: List[Tuple[str, torch.Tensor, torch.Tensor, float]],
        batch: Dict,
        stage: str = "train",
    ) -> None:
        """Log positive/negative split diagnostics, exactly once per step.

        Decoupled from :meth:`_compute_negative_weighted_loss` so the counts
        are not double-logged for multi-head models (which call the loss helper
        once per head). **All logged loss values are unweighted with respect to
        ``negative_loss_weight``** (a train-only backprop trick) on both train
        and val. The word "weighted" below refers only to the per-head
        ``loss_weights``.

        Logged keys (only when ``is_negative`` is present in ``batch``):

        * ``{stage}/n_positive`` / ``{stage}/n_negative`` (``reduce_fx="sum"``)
          -- logged once per step.
        * ``{stage}/loss_positive`` / ``{stage}/loss_negative`` -- PRIMARY
          WEIGHTED cross-head aggregate ``sum(weight * head_split)``; mirrors
          the composition of the real optimized loss.
        * ``{stage}/loss_positive_unweighted`` /
          ``{stage}/loss_negative_unweighted`` -- SECONDARY plain MEAN across
          heads of the per-head splits.
        * ``{stage}/{name}_loss_positive`` / ``{stage}/{name}_loss_negative``
          -- per-head split, only emitted when there is more than one head.

        For single-head models the weighted and unweighted aggregates are equal
        (one head, weight 1.0) and no per-head keys are emitted.

        Args:
            heads: List of ``(name, y_pred, y, weight)`` tuples. ``weight`` is
                the per-head ``loss_weight`` (1.0 for single-head models).
            batch: Batch dict; logs nothing if ``is_negative`` is absent.
            stage: Key prefix, e.g. ``"train"`` or ``"val"``.
        """
        is_negative = batch.get("is_negative", None)
        if is_negative is None:
            return

        is_neg = is_negative.to(heads[0][1].device)
        is_pos = ~is_neg
        n_neg = int(is_neg.sum().item())
        n_pos = int(is_pos.sum().item())

        # Counts: logged once per step regardless of head count.
        self.log(
            f"{stage}/n_positive",
            float(n_pos),
            prog_bar=False,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
            reduce_fx="sum",
        )
        self.log(
            f"{stage}/n_negative",
            float(n_neg),
            prog_bar=False,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
            reduce_fx="sum",
        )

        multi_head = len(heads) > 1
        weighted_pos = 0.0
        weighted_neg = 0.0
        unweighted_pos = []
        unweighted_neg = []

        for name, y_pred, y, weight in heads:
            per_sample = (y_pred - y).pow(2).mean(dim=list(range(1, y_pred.ndim)))

            if n_pos > 0:
                head_pos = per_sample[is_pos].mean()
                weighted_pos = weighted_pos + weight * head_pos
                unweighted_pos.append(head_pos)
                if multi_head:
                    self.log(
                        f"{stage}/{name}_loss_positive",
                        head_pos,
                        prog_bar=False,
                        on_step=False,
                        on_epoch=True,
                        sync_dist=True,
                    )
            if n_neg > 0:
                head_neg = per_sample[is_neg].mean()
                weighted_neg = weighted_neg + weight * head_neg
                unweighted_neg.append(head_neg)
                if multi_head:
                    self.log(
                        f"{stage}/{name}_loss_negative",
                        head_neg,
                        prog_bar=False,
                        on_step=False,
                        on_epoch=True,
                        sync_dist=True,
                    )

        if n_pos > 0:
            self.log(
                f"{stage}/loss_positive",
                weighted_pos,
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )
            self.log(
                f"{stage}/loss_positive_unweighted",
                torch.stack(unweighted_pos).mean(),
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )
        if n_neg > 0:
            self.log(
                f"{stage}/loss_negative",
                weighted_neg,
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )
            self.log(
                f"{stage}/loss_negative_unweighted",
                torch.stack(unweighted_neg).mean(),
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )

    def _log_confmap_fg_bg_loss(
        self,
        y_pred: torch.Tensor,
        y: torch.Tensor,
        stage: str = "train",
        threshold: float = 0.5,
    ) -> None:
        """Log the confmap MSE split into foreground vs background pixels.

        DIAGNOSTIC / LOGGING ONLY -- these values are **not** added to the
        optimized loss. Gaussian confidence-map targets are dominated by
        near-zero background pixels (the foreground blob is typically ~1-2% of
        the map), so plain ``MSELoss`` is mostly the background term. Splitting
        the squared error by the GROUND-TRUTH confmap value lets us watch the
        foreground/background imbalance evolve over training and informs whether
        a weighted (or focal-style) confmap loss would help.

        Pixels are split by target value: foreground = ``y > threshold``,
        background = ``y < threshold`` (pixels exactly at ``threshold`` are
        ignored). Metrics are epoch-averaged and DDP-synced.

        Logged keys:

        * ``{stage}/confmap_loss_fg`` -- mean squared error over foreground pixels.
        * ``{stage}/confmap_loss_bg`` -- mean squared error over background pixels.
        * ``{stage}/confmap_fg_frac`` -- fraction of pixels that are foreground
          (a direct measure of the imbalance).

        Args:
            y_pred: Predicted confidence maps.
            y: Ground-truth confidence maps (Gaussian peaks in ``[0, 1]``), same
                shape as ``y_pred``.
            stage: Key prefix, e.g. ``"train"`` or ``"val"``.
            threshold: Foreground/background split on the target value.
                *Default*: ``0.5``.
        """
        with torch.no_grad():
            se = (y_pred - y).pow(2)
            fg = y > threshold
            bg = y < threshold
            zero = torch.zeros((), device=y_pred.device)
            fg_loss = se[fg].mean() if fg.any() else zero
            bg_loss = se[bg].mean() if bg.any() else zero
            fg_frac = fg.float().mean()
        for key, val in (
            (f"{stage}/confmap_loss_fg", fg_loss),
            (f"{stage}/confmap_loss_bg", bg_loss),
            (f"{stage}/confmap_fg_frac", fg_frac),
        ):
            self.log(
                key,
                val,
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )

    def training_step(self, batch, batch_idx):
        """Training step."""
        pass

    def validation_step(self, batch, batch_idx):
        """Validation step."""
        pass

    def _trainable_param_groups(self):
        """Parameters that require grad, split into LR groups when asked for.

        A ``pretrained`` backbone can train its encoder at a scaled (and, for a
        plain ViT, layer-wise decayed) learning rate; everything else -- neck,
        decoder, heads -- trains at ``self.lr``. Without those options this is the
        flat list of trainable parameters (a single optimizer group, as before).
        """
        trainable = [p for p in self.parameters() if p.requires_grad]
        backbone = getattr(self.model, "backbone", None)
        groups_fn = getattr(backbone, "encoder_param_groups", None)
        encoder_groups = groups_fn(self.lr) if groups_fn is not None else None
        if not encoder_groups:
            return trainable
        in_encoder = {id(p) for g in encoder_groups for p in g["params"]}
        rest = [p for p in trainable if id(p) not in in_encoder]
        return [{"params": rest, "lr": self.lr}] + encoder_groups

    def configure_optimizers(self, optimizer=None):
        """Configure optimiser and learning rate scheduler.

        Args:
            optimizer: Optional pre-built optimizer. Subclasses that need a custom
                parameter set (e.g. the embedding model's frozen-backbone filter) build
                their own optimizer and pass it here so the scheduler is bound to the
                SAME optimizer that is returned (Lightning rejects a scheduler attached
                to an optimizer that is not returned from ``configure_optimizers``).
        """
        if optimizer is None:
            if self.optimizer == "Adam":
                optim = torch.optim.Adam
            elif self.optimizer == "AdamW":
                optim = torch.optim.AdamW

            # Only optimize params that require gradients so a frozen pretrained
            # backbone (freeze=True) is excluded; a no-op for fully-trainable models.
            optimizer = optim(
                self._trainable_param_groups(),
                lr=self.lr,
                amsgrad=self.amsgrad,
            )

        lr_scheduler_cfg = LRSchedulerConfig()
        if self.lr_scheduler is None:
            return {
                "optimizer": optimizer,
            }

        # Encoder param groups train at a fraction of `self.lr` (encoder_lr_scale /
        # layer_decay). A schedule's absolute LRs (warmup start, floor, end) are set
        # against `self.lr`, so give each group the same fraction of them; applied
        # as-is, a floor above a deep-decayed group's LR would raise that group.
        group_ratios = [
            g["lr"] / self.lr if self.lr else 1.0 for g in optimizer.param_groups
        ]

        def _per_group_lr(value):
            if all(r == 1.0 for r in group_ratios):
                return value
            if isinstance(value, Sequence) and not isinstance(value, str):
                if len(value) == len(group_ratios):
                    return list(value)  # already one per param group
                if len(value) != 1:
                    raise ValueError(
                        f"An LR schedule value {list(value)} gives {len(value)} "
                        "entries, but encoder_lr_scale / layer_decay split the "
                        f"optimizer into {len(group_ratios)} param groups. Give one "
                        "value; it is scaled for each group."
                    )
                value = value[0]
            return [float(value) * r for r in group_ratios]

        scheduler = None
        if isinstance(self.lr_scheduler, str):
            if self.lr_scheduler == "step_lr":
                lr_scheduler_cfg.step_lr = StepLRConfig()
            elif self.lr_scheduler == "reduce_lr_on_plateau":
                lr_scheduler_cfg.reduce_lr_on_plateau = ReduceLROnPlateauConfig()
            elif self.lr_scheduler == "cosine_annealing_warmup":
                lr_scheduler_cfg.cosine_annealing_warmup = CosineAnnealingWarmupConfig()
            elif self.lr_scheduler == "linear_warmup_linear_decay":
                lr_scheduler_cfg.linear_warmup_linear_decay = (
                    LinearWarmupLinearDecayConfig()
                )
            else:
                # An unrecognized name left `lr_scheduler_cfg` at its default, whose
                # `reduce_lr_on_plateau` is populated (the other three default to
                # None) -- so a typo silently trained on ReduceLROnPlateau instead of
                # the schedule the user asked for. Name the valid choices instead.
                raise ValueError(
                    f"Unknown lr_scheduler {self.lr_scheduler!r}. Expected one of "
                    "'step_lr', 'reduce_lr_on_plateau', 'cosine_annealing_warmup', "
                    "'linear_warmup_linear_decay', a scheduler config, or None."
                )

        elif isinstance(self.lr_scheduler, dict) or OmegaConf.is_config(
            self.lr_scheduler
        ):
            # `isinstance(x, dict)` is False for an OmegaConf DictConfig, which is what
            # the trainer and every YAML config actually produce -- so a DictConfig fell
            # through to the default LRSchedulerConfig here. That went unnoticed only
            # because the checks below used to dereference `self.lr_scheduler` directly,
            # bypassing this branch entirely; routing them through `lr_scheduler_cfg`
            # (correct, and required for the string form) exposes it.
            lr_scheduler_cfg = self.lr_scheduler

        # Explicit priority order per LRSchedulerConfig's own docstring:
        # cosine_annealing_warmup > linear_warmup_linear_decay > step_lr >
        # reduce_lr_on_plateau. `reduce_lr_on_plateau` defaults to a populated
        # (non-None) config while the other three default to None, so a plain
        # `for k, v in self.lr_scheduler.items(): if v is not None: ... break`
        # (the previous implementation) picked whichever scheduler happened to
        # be first in the dataclass's FIELD DECLARATION order among the
        # non-None ones -- silently ignoring this documented priority and
        # defaulting to ReduceLROnPlateau for any user who set
        # cosine_annealing_warmup/linear_warmup_linear_decay without also
        # explicitly nulling reduce_lr_on_plateau. No error, no warning --
        # training just ran with the wrong LR schedule indefinitely.
        # Read the LOCAL `lr_scheduler_cfg`, not the raw ctor arg. The string branch
        # above populates `lr_scheduler_cfg` and the dict branch aliases it, but the
        # checks below previously dereferenced `self.lr_scheduler` directly -- so the
        # documented string shorthand (`lr_scheduler="step_lr"`) raised
        # `AttributeError: 'str' object has no attribute 'cosine_annealing_warmup'`,
        # and a PARTIAL dict (a user setting just one scheduler) raised
        # `ConfigAttributeError: Missing key cosine_annealing_warmup`. Only a dict with
        # all four keys present worked. `train()` normalizes the string upstream, which
        # is why the CLI path never hit this; direct LightningModule construction did.
        def _sched(name):
            """Scheduler sub-config by name, tolerating a partial dict."""
            if OmegaConf.is_config(lr_scheduler_cfg):
                return OmegaConf.select(lr_scheduler_cfg, name, default=None)
            # A plain Python dict has no attributes, so `getattr` alone returned
            # None for every name and produced NO scheduler -- silently, where
            # `main` raised. Sub-configs may themselves be dicts, so wrap them
            # for the attribute access the branches below do.
            if isinstance(lr_scheduler_cfg, Mapping):
                sub = lr_scheduler_cfg.get(name, None)
                return OmegaConf.create(sub) if isinstance(sub, Mapping) else sub
            return getattr(lr_scheduler_cfg, name, None)

        if _sched("cosine_annealing_warmup") is not None:
            cfg = _sched("cosine_annealing_warmup")
            # Use trainer's max_epochs if not specified in config
            max_epochs = (
                cfg.max_epochs
                if cfg.max_epochs is not None
                else self.trainer.max_epochs
            )
            scheduler = LinearWarmupCosineAnnealingLR(
                optimizer=optimizer,
                warmup_epochs=cfg.warmup_epochs,
                max_epochs=max_epochs,
                warmup_start_lr=_per_group_lr(cfg.warmup_start_lr),
                eta_min=_per_group_lr(cfg.eta_min),
                # Stepped once per epoch: without this the default
                # `warmup_start_lr=0` trains the whole first epoch at LR 0.
                skip_warmup_start=True,
            )
        elif _sched("linear_warmup_linear_decay") is not None:
            cfg = _sched("linear_warmup_linear_decay")
            # Use trainer's max_epochs if not specified in config
            max_epochs = (
                cfg.max_epochs
                if cfg.max_epochs is not None
                else self.trainer.max_epochs
            )
            scheduler = LinearWarmupLinearDecayLR(
                optimizer=optimizer,
                warmup_epochs=cfg.warmup_epochs,
                max_epochs=max_epochs,
                warmup_start_lr=_per_group_lr(cfg.warmup_start_lr),
                end_lr=_per_group_lr(cfg.end_lr),
                # Stepped once per epoch: without this the default
                # `warmup_start_lr=0` trains the whole first epoch at LR 0.
                skip_warmup_start=True,
            )
        elif _sched("step_lr") is not None:
            scheduler = torch.optim.lr_scheduler.StepLR(
                optimizer=optimizer,
                step_size=_sched("step_lr").step_size,
                gamma=_sched("step_lr").gamma,
            )
        elif _sched("reduce_lr_on_plateau") is not None:
            scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
                optimizer,
                mode="min",
                threshold=_sched("reduce_lr_on_plateau").threshold,
                threshold_mode=_sched("reduce_lr_on_plateau").threshold_mode,
                cooldown=_sched("reduce_lr_on_plateau").cooldown,
                patience=_sched("reduce_lr_on_plateau").patience,
                factor=_sched("reduce_lr_on_plateau").factor,
                min_lr=_per_group_lr(_sched("reduce_lr_on_plateau").min_lr),
            )
        if scheduler is None:
            return {
                "optimizer": optimizer,
            }

        return {
            "optimizer": optimizer,
            "lr_scheduler": {
                "scheduler": scheduler,
                "monitor": "val/loss",
            },
        }

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__()
    self.negative_loss_weight = negative_loss_weight
    self.model_type = model_type
    self.backbone_type = backbone_type
    if not isinstance(backbone_config, DictConfig):
        backbone_cfg = get_backbone_config(backbone_config)
        config = OmegaConf.structured(backbone_cfg)
        OmegaConf.to_container(config, resolve=True, throw_on_missing=True)
        config = DictConfig(config)
    else:
        config = backbone_config
    self.backbone_config = config
    self.head_configs = head_configs
    self.pretrained_backbone_weights = pretrained_backbone_weights
    self.pretrained_head_weights = pretrained_head_weights
    self.in_channels = self.backbone_config[f"{self.backbone_type}"]["in_channels"]
    self.input_expand_channels = self.in_channels
    self.init_weights = init_weights
    self.lr_scheduler = lr_scheduler
    self.online_mining = online_mining
    self.hard_to_easy_ratio = hard_to_easy_ratio
    self.min_hard_keypoints = min_hard_keypoints
    self.max_hard_keypoints = max_hard_keypoints
    self.loss_scale = loss_scale
    self.optimizer = optimizer
    self.lr = learning_rate
    self.amsgrad = amsgrad

    self.model = Model(
        backbone_type=self.backbone_type,
        backbone_config=self.backbone_config[f"{self.backbone_type}"],
        head_configs=self.head_configs[self.model_type],
        model_type=self.model_type,
    )

    if len(self.head_configs[self.model_type]) > 1:
        self.loss_weights = [
            (
                self.head_configs[self.model_type][x].loss_weight
                if self.head_configs[self.model_type][x].loss_weight is not None
                else 1.0
            )
            for x in self.head_configs[self.model_type]
        ]

    self.training_loss = {}
    self.val_loss = {}
    self.learning_rate = {}

    # For epoch-averaged loss tracking
    self._epoch_loss_sum = 0.0
    self._epoch_loss_count = 0

    # For throughput logging (samples/frames per second).
    self._epoch_sample_count = 0
    self._samples_per_frame_cache = None

    # For epoch-end evaluation
    self.val_predictions: List[Dict] = []
    self.val_ground_truth: List[Dict] = []
    self._collect_val_predictions: bool = False

    # Initialization for encoder and decoder stacks.
    if self.init_weights == "xavier":
        self.model.apply(xavier_init_weights)

    # Pre-trained weights for the encoder stack - only for swint and convnext
    if self.backbone_type == "convnext" or self.backbone_type == "swint":
        if (
            self.backbone_config[f"{self.backbone_type}"]["pre_trained_weights"]
            is not None
        ):
            ckpt = MODEL_WEIGHTS[
                self.backbone_config[f"{self.backbone_type}"]["pre_trained_weights"]
            ].DEFAULT.get_state_dict(progress=True, check_hash=True)
            self.model.backbone.enc.load_state_dict(ckpt, strict=False)

    # External pretrained (HuggingFace) backbone: the wrapper loaded its
    # weights in __init__, but the xavier init above clobbered them (it runs
    # on every Conv2d/Linear). Re-apply the snapshotted pretrained weights, then
    # freeze the encoder if requested. Mirrors the convnext/swint ordering.
    if self.backbone_type == "pretrained":
        self.model.backbone.reload_pretrained_weights()
        if getattr(self.model.backbone, "freeze", False):
            self.model.backbone.freeze_encoder()

    # Initializing backbone (encoder + decoder) with trained ckpts
    if self.pretrained_backbone_weights is not None:
        logger.info(
            f"Loading backbone weights from `{self.pretrained_backbone_weights}` ..."
        )
        if self.pretrained_backbone_weights.endswith(".ckpt"):
            ckpt = torch.load(
                self.pretrained_backbone_weights,
                map_location="cpu",
                weights_only=False,
            )
            ckpt["state_dict"] = {
                k: ckpt["state_dict"][k]
                for k in ckpt["state_dict"].keys()
                if ".backbone" in k
            }
            self.load_state_dict(ckpt["state_dict"], strict=False)

        elif self.pretrained_backbone_weights.endswith(".h5"):
            # load from sleap model weights
            load_legacy_model_weights(
                self.model.backbone,
                self.pretrained_backbone_weights,
                component="backbone",
            )

        else:
            message = f"Unsupported file extension for pretrained backbone weights. Please provide a .ckpt or .h5 file."
            logger.error(message)
            raise ValueError(message)

    # Initializing head layers with trained ckpts.
    if self.pretrained_head_weights is not None:
        logger.info(
            f"Loading head weights from `{self.pretrained_head_weights}` ..."
        )
        if self.pretrained_head_weights.endswith(".ckpt"):
            ckpt = torch.load(
                self.pretrained_head_weights,
                map_location="cpu",
                weights_only=False,
            )
            ckpt["state_dict"] = {
                k: ckpt["state_dict"][k]
                for k in ckpt["state_dict"].keys()
                if ".head_layers" in k
            }
            self.load_state_dict(ckpt["state_dict"], strict=False)

        elif self.pretrained_head_weights.endswith(".h5"):
            # load from sleap model weights
            load_legacy_model_weights(
                self.model.head_layers,
                self.pretrained_head_weights,
                component="head",
            )

        else:
            message = f"Unsupported file extension for pretrained head weights. Please provide a .ckpt or .h5 file."
            logger.error(message)
            raise ValueError(message)

configure_optimizers(optimizer=None)

Configure optimiser and learning rate scheduler.

Parameters:

Name	Type	Description	Default
optimizer		Optional pre-built optimizer. Subclasses that need a custom parameter set (e.g. the embedding model's frozen-backbone filter) build their own optimizer and pass it here so the scheduler is bound to the SAME optimizer that is returned (Lightning rejects a scheduler attached to an optimizer that is not returned from configure_optimizers).	None

Source code in sleap_nn/training/lightning_modules.py

def configure_optimizers(self, optimizer=None):
    """Configure optimiser and learning rate scheduler.

    Args:
        optimizer: Optional pre-built optimizer. Subclasses that need a custom
            parameter set (e.g. the embedding model's frozen-backbone filter) build
            their own optimizer and pass it here so the scheduler is bound to the
            SAME optimizer that is returned (Lightning rejects a scheduler attached
            to an optimizer that is not returned from ``configure_optimizers``).
    """
    if optimizer is None:
        if self.optimizer == "Adam":
            optim = torch.optim.Adam
        elif self.optimizer == "AdamW":
            optim = torch.optim.AdamW

        # Only optimize params that require gradients so a frozen pretrained
        # backbone (freeze=True) is excluded; a no-op for fully-trainable models.
        optimizer = optim(
            self._trainable_param_groups(),
            lr=self.lr,
            amsgrad=self.amsgrad,
        )

    lr_scheduler_cfg = LRSchedulerConfig()
    if self.lr_scheduler is None:
        return {
            "optimizer": optimizer,
        }

    # Encoder param groups train at a fraction of `self.lr` (encoder_lr_scale /
    # layer_decay). A schedule's absolute LRs (warmup start, floor, end) are set
    # against `self.lr`, so give each group the same fraction of them; applied
    # as-is, a floor above a deep-decayed group's LR would raise that group.
    group_ratios = [
        g["lr"] / self.lr if self.lr else 1.0 for g in optimizer.param_groups
    ]

    def _per_group_lr(value):
        if all(r == 1.0 for r in group_ratios):
            return value
        if isinstance(value, Sequence) and not isinstance(value, str):
            if len(value) == len(group_ratios):
                return list(value)  # already one per param group
            if len(value) != 1:
                raise ValueError(
                    f"An LR schedule value {list(value)} gives {len(value)} "
                    "entries, but encoder_lr_scale / layer_decay split the "
                    f"optimizer into {len(group_ratios)} param groups. Give one "
                    "value; it is scaled for each group."
                )
            value = value[0]
        return [float(value) * r for r in group_ratios]

    scheduler = None
    if isinstance(self.lr_scheduler, str):
        if self.lr_scheduler == "step_lr":
            lr_scheduler_cfg.step_lr = StepLRConfig()
        elif self.lr_scheduler == "reduce_lr_on_plateau":
            lr_scheduler_cfg.reduce_lr_on_plateau = ReduceLROnPlateauConfig()
        elif self.lr_scheduler == "cosine_annealing_warmup":
            lr_scheduler_cfg.cosine_annealing_warmup = CosineAnnealingWarmupConfig()
        elif self.lr_scheduler == "linear_warmup_linear_decay":
            lr_scheduler_cfg.linear_warmup_linear_decay = (
                LinearWarmupLinearDecayConfig()
            )
        else:
            # An unrecognized name left `lr_scheduler_cfg` at its default, whose
            # `reduce_lr_on_plateau` is populated (the other three default to
            # None) -- so a typo silently trained on ReduceLROnPlateau instead of
            # the schedule the user asked for. Name the valid choices instead.
            raise ValueError(
                f"Unknown lr_scheduler {self.lr_scheduler!r}. Expected one of "
                "'step_lr', 'reduce_lr_on_plateau', 'cosine_annealing_warmup', "
                "'linear_warmup_linear_decay', a scheduler config, or None."
            )

    elif isinstance(self.lr_scheduler, dict) or OmegaConf.is_config(
        self.lr_scheduler
    ):
        # `isinstance(x, dict)` is False for an OmegaConf DictConfig, which is what
        # the trainer and every YAML config actually produce -- so a DictConfig fell
        # through to the default LRSchedulerConfig here. That went unnoticed only
        # because the checks below used to dereference `self.lr_scheduler` directly,
        # bypassing this branch entirely; routing them through `lr_scheduler_cfg`
        # (correct, and required for the string form) exposes it.
        lr_scheduler_cfg = self.lr_scheduler

    # Explicit priority order per LRSchedulerConfig's own docstring:
    # cosine_annealing_warmup > linear_warmup_linear_decay > step_lr >
    # reduce_lr_on_plateau. `reduce_lr_on_plateau` defaults to a populated
    # (non-None) config while the other three default to None, so a plain
    # `for k, v in self.lr_scheduler.items(): if v is not None: ... break`
    # (the previous implementation) picked whichever scheduler happened to
    # be first in the dataclass's FIELD DECLARATION order among the
    # non-None ones -- silently ignoring this documented priority and
    # defaulting to ReduceLROnPlateau for any user who set
    # cosine_annealing_warmup/linear_warmup_linear_decay without also
    # explicitly nulling reduce_lr_on_plateau. No error, no warning --
    # training just ran with the wrong LR schedule indefinitely.
    # Read the LOCAL `lr_scheduler_cfg`, not the raw ctor arg. The string branch
    # above populates `lr_scheduler_cfg` and the dict branch aliases it, but the
    # checks below previously dereferenced `self.lr_scheduler` directly -- so the
    # documented string shorthand (`lr_scheduler="step_lr"`) raised
    # `AttributeError: 'str' object has no attribute 'cosine_annealing_warmup'`,
    # and a PARTIAL dict (a user setting just one scheduler) raised
    # `ConfigAttributeError: Missing key cosine_annealing_warmup`. Only a dict with
    # all four keys present worked. `train()` normalizes the string upstream, which
    # is why the CLI path never hit this; direct LightningModule construction did.
    def _sched(name):
        """Scheduler sub-config by name, tolerating a partial dict."""
        if OmegaConf.is_config(lr_scheduler_cfg):
            return OmegaConf.select(lr_scheduler_cfg, name, default=None)
        # A plain Python dict has no attributes, so `getattr` alone returned
        # None for every name and produced NO scheduler -- silently, where
        # `main` raised. Sub-configs may themselves be dicts, so wrap them
        # for the attribute access the branches below do.
        if isinstance(lr_scheduler_cfg, Mapping):
            sub = lr_scheduler_cfg.get(name, None)
            return OmegaConf.create(sub) if isinstance(sub, Mapping) else sub
        return getattr(lr_scheduler_cfg, name, None)

    if _sched("cosine_annealing_warmup") is not None:
        cfg = _sched("cosine_annealing_warmup")
        # Use trainer's max_epochs if not specified in config
        max_epochs = (
            cfg.max_epochs
            if cfg.max_epochs is not None
            else self.trainer.max_epochs
        )
        scheduler = LinearWarmupCosineAnnealingLR(
            optimizer=optimizer,
            warmup_epochs=cfg.warmup_epochs,
            max_epochs=max_epochs,
            warmup_start_lr=_per_group_lr(cfg.warmup_start_lr),
            eta_min=_per_group_lr(cfg.eta_min),
            # Stepped once per epoch: without this the default
            # `warmup_start_lr=0` trains the whole first epoch at LR 0.
            skip_warmup_start=True,
        )
    elif _sched("linear_warmup_linear_decay") is not None:
        cfg = _sched("linear_warmup_linear_decay")
        # Use trainer's max_epochs if not specified in config
        max_epochs = (
            cfg.max_epochs
            if cfg.max_epochs is not None
            else self.trainer.max_epochs
        )
        scheduler = LinearWarmupLinearDecayLR(
            optimizer=optimizer,
            warmup_epochs=cfg.warmup_epochs,
            max_epochs=max_epochs,
            warmup_start_lr=_per_group_lr(cfg.warmup_start_lr),
            end_lr=_per_group_lr(cfg.end_lr),
            # Stepped once per epoch: without this the default
            # `warmup_start_lr=0` trains the whole first epoch at LR 0.
            skip_warmup_start=True,
        )
    elif _sched("step_lr") is not None:
        scheduler = torch.optim.lr_scheduler.StepLR(
            optimizer=optimizer,
            step_size=_sched("step_lr").step_size,
            gamma=_sched("step_lr").gamma,
        )
    elif _sched("reduce_lr_on_plateau") is not None:
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            optimizer,
            mode="min",
            threshold=_sched("reduce_lr_on_plateau").threshold,
            threshold_mode=_sched("reduce_lr_on_plateau").threshold_mode,
            cooldown=_sched("reduce_lr_on_plateau").cooldown,
            patience=_sched("reduce_lr_on_plateau").patience,
            factor=_sched("reduce_lr_on_plateau").factor,
            min_lr=_per_group_lr(_sched("reduce_lr_on_plateau").min_lr),
        )
    if scheduler is None:
        return {
            "optimizer": optimizer,
        }

    return {
        "optimizer": optimizer,
        "lr_scheduler": {
            "scheduler": scheduler,
            "monitor": "val/loss",
        },
    }

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    pass

get_lightning_model_from_config(config) classmethod

Get lightning model from config.

Source code in sleap_nn/training/lightning_modules.py

@classmethod
def get_lightning_model_from_config(cls, config: DictConfig):
    """Get lightning model from config."""
    model_type = get_model_type_from_cfg(config)
    backbone_type = get_backbone_type_from_cfg(config)

    lightning_models = {
        "single_instance": SingleInstanceLightningModule,
        "centroid": CentroidLightningModule,
        "centered_instance": TopDownCenteredInstanceLightningModule,
        "bottomup": BottomUpLightningModule,
        "multi_class_bottomup": BottomUpMultiClassLightningModule,
        "multi_class_topdown": TopDownCenteredInstanceMultiClassLightningModule,
        "bottomup_segmentation": BottomUpSegmentationLightningModule,
        "centered_instance_segmentation": TopDownCenteredInstanceSegmentationLightningModule,
        "semantic_segmentation": SemanticSegmentationLightningModule,
        "embedding": EmbeddingLightningModule,
        "bbox_centernet": BBoxCenterNetLightningModule,
        "bbox_ssd": BBoxSSDLightningModule,
    }

    if model_type not in lightning_models:
        message = f"Incorrect model type. Please check if one of the following keys in the head configs is not None: [`single_instance`, `centroid`, `centered_instance`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`, `bottomup_segmentation`, `centered_instance_segmentation`, `semantic_segmentation`, `embedding`, `bbox_centernet`, `bbox_ssd`]"
        logger.error(message)
        raise ValueError(message)

    negative_loss_weight = getattr(config.data_config, "negative_loss_weight", 1.0)

    # See CentroidConfMapsConfig.focal_loss_alpha -- centroid-only.
    extra_kwargs = {}
    if model_type == "centroid":
        centroid_confmaps_config = (
            config.model_config.head_configs.centroid.confmaps
        )
        extra_kwargs["centroid_focal_loss_alpha"] = getattr(
            centroid_confmaps_config, "focal_loss_alpha", 0.0
        )
        extra_kwargs["centroid_focal_loss_beta"] = getattr(
            centroid_confmaps_config, "focal_loss_beta", 4.0
        )
        extra_kwargs["centroid_focal_loss_pos_threshold"] = getattr(
            centroid_confmaps_config, "focal_loss_pos_threshold", 0.5
        )

    lightning_model = lightning_models[model_type](
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=config.model_config.backbone_config,
        head_configs=config.model_config.head_configs,
        pretrained_backbone_weights=config.model_config.pretrained_backbone_weights,
        pretrained_head_weights=config.model_config.pretrained_head_weights,
        init_weights=config.model_config.init_weights,
        lr_scheduler=config.trainer_config.lr_scheduler,
        online_mining=config.trainer_config.online_hard_keypoint_mining.online_mining,
        hard_to_easy_ratio=config.trainer_config.online_hard_keypoint_mining.hard_to_easy_ratio,
        min_hard_keypoints=config.trainer_config.online_hard_keypoint_mining.min_hard_keypoints,
        max_hard_keypoints=config.trainer_config.online_hard_keypoint_mining.max_hard_keypoints,
        loss_scale=config.trainer_config.online_hard_keypoint_mining.loss_scale,
        optimizer=config.trainer_config.optimizer_name,
        learning_rate=config.trainer_config.optimizer.lr,
        amsgrad=config.trainer_config.optimizer.amsgrad,
        negative_loss_weight=negative_loss_weight,
        **extra_kwargs,
    )

    if model_type == "embedding":
        # Make `data_config.preprocessing.burn_in` live (mask-on vs centroid-crop).
        set_embedding_burn_in_from_config(lightning_model, config)

    return lightning_model

on_train_batch_end(outputs, batch, batch_idx)

Count per-device samples this epoch (for throughput logging).

Source code in sleap_nn/training/lightning_modules.py

def on_train_batch_end(self, outputs, batch, batch_idx):
    """Count per-device samples this epoch (for throughput logging)."""
    try:
        if isinstance(batch, dict):
            if "image" in batch and hasattr(batch["image"], "shape"):
                self._epoch_sample_count += int(batch["image"].shape[0])
            elif "frame_idx" in batch:
                self._epoch_sample_count += len(batch["frame_idx"])
    except Exception:
        pass

on_train_epoch_end()

Configure the train timer at the end of every epoch.

Source code in sleap_nn/training/lightning_modules.py

def on_train_epoch_end(self):
    """Configure the train timer at the end of every epoch."""
    train_time = time.time() - self.train_start_time
    self.log(
        "train/time",
        train_time,
        prog_bar=False,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    # Log epoch explicitly for custom x-axis support in wandb
    self.log(
        "epoch",
        float(self.current_epoch),
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    # Log epoch-averaged training loss
    if self._epoch_loss_count > 0:
        avg_loss = self._epoch_loss_sum / self._epoch_loss_count
        self.log(
            "train/loss",
            avg_loss,
            prog_bar=False,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
    # Log current learning rate (useful for monitoring LR schedulers)
    if self.trainer.optimizers:
        lr = self.trainer.optimizers[0].param_groups[0]["lr"]
        self.log(
            "train/lr",
            lr,
            prog_bar=False,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

    # Throughput: optimizer steps/sec, sample (tile/crop) throughput, and
    # full source-frames/sec. Uses the GLOBAL batch (per-device samples x
    # world_size), so it is correct under multi-GPU. ``frames_per_sec`` divides
    # the sample rate by tiles-sampled-per-frame under tiling, so it reads as
    # source frames rather than crops/tiles.
    if train_time > 0:
        world = int(getattr(self.trainer, "world_size", 1) or 1)
        spf = self._tiling_samples_per_frame()
        steps_per_sec = self._epoch_loss_count / train_time
        samples_per_sec = (self._epoch_sample_count * world) / train_time
        frames_per_sec = samples_per_sec / spf
        for _name, _val in (
            ("train/steps_per_sec", steps_per_sec),
            ("train/samples_per_sec", samples_per_sec),
            ("train/frames_per_sec", frames_per_sec),
        ):
            self.log(
                _name,
                _val,
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=False,
            )

on_train_epoch_start()

Configure the train timer at the beginning of each epoch.

Source code in sleap_nn/training/lightning_modules.py

def on_train_epoch_start(self):
    """Configure the train timer at the beginning of each epoch."""
    self.train_start_time = time.time()
    # Reset epoch loss tracking
    self._epoch_loss_sum = 0.0
    self._epoch_loss_count = 0
    # Reset per-device sample count for throughput.
    self._epoch_sample_count = 0

on_validation_epoch_end()

Configure the val timer at the end of every epoch.

Source code in sleap_nn/training/lightning_modules.py

def on_validation_epoch_end(self):
    """Configure the val timer at the end of every epoch."""
    val_time = time.time() - self.val_start_time
    self.log(
        "val/time",
        val_time,
        prog_bar=False,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    # Log epoch explicitly so val/* metrics can use it as x-axis in wandb
    # (mirrors what on_train_epoch_end does for train/* metrics)
    self.log(
        "epoch",
        float(self.current_epoch),
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

on_validation_epoch_start()

Configure the val timer at the beginning of each epoch.

Source code in sleap_nn/training/lightning_modules.py

def on_validation_epoch_start(self):
    """Configure the val timer at the beginning of each epoch."""
    self.val_start_time = time.time()
    # Clear accumulated predictions for new epoch
    self.val_predictions = []
    self.val_ground_truth = []

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    pass

validation_step(batch, batch_idx)

Validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step."""
    pass

SemanticSegmentationLightningModule

Bases: LightningModel

Lightning Module for whole-frame semantic (foreground/background) segmentation.

Predicts a single binary foreground mask over the ENTIRE frame with NO instance grouping. This is a hybrid of two existing segmentation modules:

• :class:BottomUpSegmentationLightningModule — whole-frame image input and tiling compatibility, but WITHOUT its instance-center / center-offset heads or the offset-based grouping (there is no per-instance separation here).
• :class:TopDownCenteredInstanceSegmentationLightningModule — the bce-dice foreground loss and foreground-IoU metric, but on the whole frame rather than a centroid crop.

A lone SegmentationHead (1-channel logits) is trained with compute_bce_dice_loss and decoded at inference by thresholding the foreground probability into ONE mask per frame (no grouping).

forward mirrors :class:BottomUpSegmentationLightningModule — it returns {"SegmentationHead": sigmoid(logits)} (a dict with sigmoid ALREADY applied). This is load-bearing: the whole-frame SemanticSegmentationLayer and its tiled wrapper read SegmentationHead as probabilities, and the tiled path Gaussian-averages the foreground map across tile overlaps (a blend that is only correct for probabilities, not logits). training_step / validation_step bypass forward and call self.model(X) directly to supervise the raw logits with compute_bce_dice_loss.

All-background frames are fully supported: a frame with no instances yields an all-zero foreground_mask (see generate_foreground_mask with an empty mask list), and bce-dice supervises the model to predict background everywhere. Enable such frames via data_config.use_negative_frames=True (consumed in the data pipeline, not here).

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass returning the foreground PROBABILITY map (sigmoid applied).
get_visualization_data	Whole-frame viz: image + predicted foreground vs GT mask overlay.
training_step	Training step (bce-dice on the whole-frame foreground mask).
validation_step	Validation step (val loss + whole-frame foreground IoU).
visualize_example	Visualize the predicted foreground mask over the frame during training.

Source code in sleap_nn/training/lightning_modules.py

class SemanticSegmentationLightningModule(LightningModel):
    """Lightning Module for whole-frame semantic (foreground/background) segmentation.

    Predicts a single binary foreground mask over the ENTIRE frame with NO
    instance grouping. This is a hybrid of two existing segmentation modules:

    * :class:`BottomUpSegmentationLightningModule` — whole-frame ``image`` input
      and tiling compatibility, but WITHOUT its instance-center / center-offset
      heads or the offset-based grouping (there is no per-instance separation
      here).
    * :class:`TopDownCenteredInstanceSegmentationLightningModule` — the bce-dice
      foreground loss and foreground-IoU metric, but on the whole frame rather
      than a centroid crop.

    A lone ``SegmentationHead`` (1-channel logits) is trained with
    ``compute_bce_dice_loss`` and decoded at inference by thresholding the
    foreground probability into ONE mask per frame (no grouping).

    ``forward`` mirrors :class:`BottomUpSegmentationLightningModule` — it returns
    ``{"SegmentationHead": sigmoid(logits)}`` (a dict with sigmoid ALREADY
    applied). This is load-bearing: the whole-frame ``SemanticSegmentationLayer``
    and its tiled wrapper read ``SegmentationHead`` as probabilities, and the
    tiled path Gaussian-averages the foreground map across tile overlaps (a blend
    that is only correct for probabilities, not logits). ``training_step`` /
    ``validation_step`` bypass ``forward`` and call ``self.model(X)`` directly to
    supervise the raw logits with ``compute_bce_dice_loss``.

    All-background frames are fully supported: a frame with no instances yields an
    all-zero ``foreground_mask`` (see ``generate_foreground_mask`` with an empty
    mask list), and bce-dice supervises the model to predict background
    everywhere. Enable such frames via ``data_config.use_negative_frames=True``
    (consumed in the data pipeline, not here).
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )
        seg = self.head_configs[self.model_type].segmentation
        self.seg_output_stride = seg.output_stride
        # bce-dice loss knobs (defaults preserve the symmetric unweighted loss).
        self.fg_bce_weight = getattr(seg, "bce_weight", 0.5)
        self.fg_dice_weight = getattr(seg, "dice_weight", 0.5)
        self.fg_bce_pos_weight = getattr(seg, "bce_pos_weight", None)

    def forward(self, img):
        """Forward pass returning the foreground PROBABILITY map (sigmoid applied).

        Returns a dict ``{"SegmentationHead": sigmoid(logits)}`` (mirroring
        :class:`BottomUpSegmentationLightningModule`), so the inference layers see
        probabilities. Training/validation call ``self.model`` directly instead
        (raw logits for ``compute_bce_dice_loss``).
        """
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        output = self.model(img)
        return {"SegmentationHead": torch.sigmoid(output["SegmentationHead"])}

    def training_step(self, batch, batch_idx):
        """Training step (bce-dice on the whole-frame foreground mask)."""
        X = torch.squeeze(batch["image"], dim=1)
        y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
        X = normalize_on_gpu(X)
        pred_fg = self.model(X)["SegmentationHead"]  # logits
        loss = compute_bce_dice_loss(
            pred_fg,
            y_fg,
            bce_weight=self.fg_bce_weight,
            dice_weight=self.fg_dice_weight,
            pos_weight=self.fg_bce_pos_weight,
        )

        self.log(
            "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
        )
        self._accumulate_loss(loss)
        self.log("train/fg_loss", loss, on_step=False, on_epoch=True, sync_dist=True)
        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step (val loss + whole-frame foreground IoU)."""
        X = torch.squeeze(batch["image"], dim=1)
        y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
        X = normalize_on_gpu(X)
        pred_fg = self.model(X)["SegmentationHead"]  # logits
        val_loss = compute_bce_dice_loss(
            pred_fg,
            y_fg,
            bce_weight=self.fg_bce_weight,
            dice_weight=self.fg_dice_weight,
            pos_weight=self.fg_bce_pos_weight,
        )

        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log("val/fg_loss", val_loss, on_step=False, on_epoch=True, sync_dist=True)

        # Whole-frame foreground IoU averaged PER-SAMPLE (mean of per-image IoUs)
        # rather than pooled over the batch tensor (which over-weights large /
        # foreground-heavy images). Lightning's on_epoch aggregation then yields a
        # true mean-per-image IoU.
        pred_fg_binary = (pred_fg > 0.0).float()
        dims = (1, 2, 3)
        intersection = (pred_fg_binary * y_fg).sum(dim=dims)
        union = pred_fg_binary.sum(dim=dims) + y_fg.sum(dim=dims) - intersection
        iou = (intersection / (union + 1e-6)).mean()
        self.log("val/fg_iou", iou, on_step=False, on_epoch=True, sync_dist=True)

        # Optional mask eval (enabled by SegmentationEvaluationCallback in
        # foreground mode): semantic segmentation has ONE mask per frame (no
        # grouping), so emit the binarized whole-frame prediction and the GT
        # foreground as a single-mask pair on the SAME stride grid, appended in
        # lockstep for positional pairing by the callback.
        if self._collect_val_predictions:
            for i in range(X.shape[0]):
                pm = pred_fg_binary[i, 0].detach().cpu().numpy().astype(bool)
                gm = (y_fg[i, 0] > 0.5).detach().cpu().numpy()
                self.val_predictions.append({"masks": [pm] if pm.any() else []})
                self.val_ground_truth.append({"masks": [gm] if gm.any() else []})

    def get_visualization_data(
        self, sample, include_gt_mask: bool = False
    ) -> VisualizationData:
        """Whole-frame viz: image + predicted foreground vs GT mask overlay.

        Args:
            sample: A sample dictionary from the data pipeline.
            include_gt_mask: If True, include the ground-truth foreground mask for
                a GT-vs-prediction overlay.
        """
        ex = sample.copy()
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["image"] = ex["image"].unsqueeze(dim=0)

        # Run the raw model (not self.forward, which returns a sigmoid dict) so we
        # apply sigmoid exactly once for the foreground-probability overlay.
        with torch.no_grad():
            img = ex["image"].squeeze(1).to(self.device)
            img = normalize_on_gpu(img)
            preds = self.model(img)
        fg_prob = torch.sigmoid(preds["SegmentationHead"][0]).cpu().numpy()
        fg_prob = fg_prob.transpose(1, 2, 0)  # (H, W, 1)

        img_np = ex["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        gt_mask = None
        if include_gt_mask and "foreground_mask" in ex:
            gt_mask = ex["foreground_mask"].squeeze().cpu().numpy()  # (H, W)
        return VisualizationData(
            image=img_np,
            pred_confmaps=fg_prob,
            pred_peaks=np.zeros((0, 1, 2)),
            pred_peak_values=np.zeros((0,)),
            gt_instances=np.zeros((0, 1, 2)),
            node_names=["mask"],
            output_scale=fg_prob.shape[0] / img_np.shape[0],
            is_paired=False,
            gt_mask=gt_mask,
        )

    def visualize_example(self, sample):
        """Visualize the predicted foreground mask over the frame during training."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        return fig

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )
    seg = self.head_configs[self.model_type].segmentation
    self.seg_output_stride = seg.output_stride
    # bce-dice loss knobs (defaults preserve the symmetric unweighted loss).
    self.fg_bce_weight = getattr(seg, "bce_weight", 0.5)
    self.fg_dice_weight = getattr(seg, "dice_weight", 0.5)
    self.fg_bce_pos_weight = getattr(seg, "bce_pos_weight", None)

forward(img)

Forward pass returning the foreground PROBABILITY map (sigmoid applied).

Returns a dict {"SegmentationHead": sigmoid(logits)} (mirroring :class:BottomUpSegmentationLightningModule), so the inference layers see probabilities. Training/validation call self.model directly instead (raw logits for compute_bce_dice_loss).

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass returning the foreground PROBABILITY map (sigmoid applied).

    Returns a dict ``{"SegmentationHead": sigmoid(logits)}`` (mirroring
    :class:`BottomUpSegmentationLightningModule`), so the inference layers see
    probabilities. Training/validation call ``self.model`` directly instead
    (raw logits for ``compute_bce_dice_loss``).
    """
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    output = self.model(img)
    return {"SegmentationHead": torch.sigmoid(output["SegmentationHead"])}

get_visualization_data(sample, include_gt_mask=False)

Whole-frame viz: image + predicted foreground vs GT mask overlay.

Parameters:

Name	Type	Description	Default
sample		A sample dictionary from the data pipeline.	required
include_gt_mask	bool	If True, include the ground-truth foreground mask for a GT-vs-prediction overlay.	False

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(
    self, sample, include_gt_mask: bool = False
) -> VisualizationData:
    """Whole-frame viz: image + predicted foreground vs GT mask overlay.

    Args:
        sample: A sample dictionary from the data pipeline.
        include_gt_mask: If True, include the ground-truth foreground mask for
            a GT-vs-prediction overlay.
    """
    ex = sample.copy()
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["image"] = ex["image"].unsqueeze(dim=0)

    # Run the raw model (not self.forward, which returns a sigmoid dict) so we
    # apply sigmoid exactly once for the foreground-probability overlay.
    with torch.no_grad():
        img = ex["image"].squeeze(1).to(self.device)
        img = normalize_on_gpu(img)
        preds = self.model(img)
    fg_prob = torch.sigmoid(preds["SegmentationHead"][0]).cpu().numpy()
    fg_prob = fg_prob.transpose(1, 2, 0)  # (H, W, 1)

    img_np = ex["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    gt_mask = None
    if include_gt_mask and "foreground_mask" in ex:
        gt_mask = ex["foreground_mask"].squeeze().cpu().numpy()  # (H, W)
    return VisualizationData(
        image=img_np,
        pred_confmaps=fg_prob,
        pred_peaks=np.zeros((0, 1, 2)),
        pred_peak_values=np.zeros((0,)),
        gt_instances=np.zeros((0, 1, 2)),
        node_names=["mask"],
        output_scale=fg_prob.shape[0] / img_np.shape[0],
        is_paired=False,
        gt_mask=gt_mask,
    )

training_step(batch, batch_idx)

Training step (bce-dice on the whole-frame foreground mask).

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step (bce-dice on the whole-frame foreground mask)."""
    X = torch.squeeze(batch["image"], dim=1)
    y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
    X = normalize_on_gpu(X)
    pred_fg = self.model(X)["SegmentationHead"]  # logits
    loss = compute_bce_dice_loss(
        pred_fg,
        y_fg,
        bce_weight=self.fg_bce_weight,
        dice_weight=self.fg_dice_weight,
        pos_weight=self.fg_bce_pos_weight,
    )

    self.log(
        "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
    )
    self._accumulate_loss(loss)
    self.log("train/fg_loss", loss, on_step=False, on_epoch=True, sync_dist=True)
    return loss

validation_step(batch, batch_idx)

Validation step (val loss + whole-frame foreground IoU).

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step (val loss + whole-frame foreground IoU)."""
    X = torch.squeeze(batch["image"], dim=1)
    y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
    X = normalize_on_gpu(X)
    pred_fg = self.model(X)["SegmentationHead"]  # logits
    val_loss = compute_bce_dice_loss(
        pred_fg,
        y_fg,
        bce_weight=self.fg_bce_weight,
        dice_weight=self.fg_dice_weight,
        pos_weight=self.fg_bce_pos_weight,
    )

    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log("val/fg_loss", val_loss, on_step=False, on_epoch=True, sync_dist=True)

    # Whole-frame foreground IoU averaged PER-SAMPLE (mean of per-image IoUs)
    # rather than pooled over the batch tensor (which over-weights large /
    # foreground-heavy images). Lightning's on_epoch aggregation then yields a
    # true mean-per-image IoU.
    pred_fg_binary = (pred_fg > 0.0).float()
    dims = (1, 2, 3)
    intersection = (pred_fg_binary * y_fg).sum(dim=dims)
    union = pred_fg_binary.sum(dim=dims) + y_fg.sum(dim=dims) - intersection
    iou = (intersection / (union + 1e-6)).mean()
    self.log("val/fg_iou", iou, on_step=False, on_epoch=True, sync_dist=True)

    # Optional mask eval (enabled by SegmentationEvaluationCallback in
    # foreground mode): semantic segmentation has ONE mask per frame (no
    # grouping), so emit the binarized whole-frame prediction and the GT
    # foreground as a single-mask pair on the SAME stride grid, appended in
    # lockstep for positional pairing by the callback.
    if self._collect_val_predictions:
        for i in range(X.shape[0]):
            pm = pred_fg_binary[i, 0].detach().cpu().numpy().astype(bool)
            gm = (y_fg[i, 0] > 0.5).detach().cpu().numpy()
            self.val_predictions.append({"masks": [pm] if pm.any() else []})
            self.val_ground_truth.append({"masks": [gm] if gm.any() else []})

visualize_example(sample)

Visualize the predicted foreground mask over the frame during training.

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize the predicted foreground mask over the frame during training."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    return fig

SingleInstanceLightningModule

Bases: LightningModel

Lightning Module for SingleInstance Model.

This is a subclass of the LightningModel to configure the training/ validation steps and forward pass specific to Single Instance model. Single Instance models predict keypoint locations directly from the input image without requiring a separate detection step.

Parameters:

Name	Type	Description	Default
model_type	str	Type of the model. One of single_instance, centered_instance, centroid, bottomup, multi_class_bottomup, multi_class_topdown.	required
backbone_type	str	Backbone model. One of unet, convnext and swint.	required
backbone_config	Union[str, Dict[str, Any], DictConfig]	Backbone configuration. Can be: - String: One of the preset backbone types: - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"] - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"] - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"] - Dictionary: Custom configuration with structure: { "unet": {UNetConfig parameters}, "convnext": {ConvNextConfig parameters}, "swint": {SwinTConfig parameters} } Only one backbone type should be specified in the dictionary. - DictConfig: OmegaConf DictConfig object containing backbone configuration.	required
head_configs	DictConfig	Head configuration dictionary containing model-specific parameters. For Single Instance: confmaps with part_names, sigma, output_stride. For Centroid: confmaps with anchor_part, sigma, output_stride. For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride. For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight. For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight. For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.	required
pretrained_backbone_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.	None
pretrained_head_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.	None
init_weights	Optional[str]	Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.	'xavier'
lr_scheduler	Optional[Union[str, DictConfig]]	Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.	None
online_mining	Optional[bool]	If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).	False
hard_to_easy_ratio	Optional[float]	Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.	2.0
min_hard_keypoints	Optional[int]	Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.	2
max_hard_keypoints	Optional[int]	Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.	None
loss_scale	Optional[float]	Factor to scale hard keypoint losses by. Default: 5.0.	5.0
optimizer	Optional[str]	Optimizer name. One of ["Adam", "AdamW"].	'Adam'
learning_rate	Optional[float]	Learning rate for the optimizer. Default: 1e-3.	0.001
amsgrad	Optional[bool]	Enable AMSGrad with the optimizer. Default: False.	False

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model.
get_visualization_data	Extract visualization data from a sample.
training_step	Training step.
validation_step	Validation step.
visualize_example	Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

class SingleInstanceLightningModule(LightningModel):
    """Lightning Module for SingleInstance Model.

    This is a subclass of the `LightningModel` to configure the training/ validation steps and
    forward pass specific to Single Instance model. Single Instance models predict keypoint locations
    directly from the input image without requiring a separate detection step.

    Args:
        model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
        backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
        backbone_config: Backbone configuration. Can be:
            - String: One of the preset backbone types:
                - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"]
                - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"]
                - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"]
            - Dictionary: Custom configuration with structure:
                {
                    "unet": {UNetConfig parameters},
                    "convnext": {ConvNextConfig parameters},
                    "swint": {SwinTConfig parameters}
                }
                Only one backbone type should be specified in the dictionary.
            - DictConfig: OmegaConf DictConfig object containing backbone configuration.
        head_configs: Head configuration dictionary containing model-specific parameters.
            For Single Instance: confmaps with part_names, sigma, output_stride.
            For Centroid: confmaps with anchor_part, sigma, output_stride.
            For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride.
            For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight.
            For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight.
            For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.
        pretrained_backbone_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.
        pretrained_head_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.
        init_weights: Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.
        lr_scheduler: Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.
        online_mining: If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).
        hard_to_easy_ratio: Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.
        min_hard_keypoints: Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.
        max_hard_keypoints: Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.
        loss_scale: Factor to scale hard keypoint losses by. Default: 5.0.
        optimizer: Optimizer name. One of ["Adam", "AdamW"].
        learning_rate: Learning rate for the optimizer. Default: 1e-3.
        amsgrad: Enable AMSGrad with the optimizer. Default: False.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )

        self.single_instance_inf_layer = SingleInstanceInferenceModel(
            torch_model=self.forward,
            peak_threshold=0.2,
            input_scale=1.0,
            return_confmaps=True,
            output_stride=self.head_configs.single_instance.confmaps.output_stride,
        )
        self.node_names = self.head_configs.single_instance.confmaps.part_names

    def get_visualization_data(self, sample) -> VisualizationData:
        """Extract visualization data from a sample.

        Args:
            sample: A sample dictionary from the data pipeline.

        Returns:
            VisualizationData containing image, confmaps, peaks, etc.
        """
        ex = sample.copy()
        ex["eff_scale"] = torch.tensor([1.0])
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["image"] = ex["image"].unsqueeze(dim=0)
        output = self.single_instance_inf_layer(ex)[0]

        peaks = output["pred_instance_peaks"].cpu().numpy()
        peak_values = output["pred_peak_values"].cpu().numpy()
        img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        gt_instances = ex["instances"][0].cpu().numpy()
        confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

        return VisualizationData(
            image=img,
            pred_confmaps=confmaps,
            pred_peaks=peaks,
            pred_peak_values=peak_values,
            gt_instances=gt_instances,
            node_names=list(self.node_names) if self.node_names else [],
            output_scale=confmaps.shape[0] / img.shape[0],
            is_paired=True,
        )

    def visualize_example(self, sample):
        """Visualize predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def forward(self, img):
        """Forward pass of the model."""
        # Only squeeze n_samples dim if 5D (batch, n_samples, C, H, W) -> (batch, C, H, W)
        # Avoid double-squeezing when called from validation_step which already squeezes
        if img.ndim == 5:
            img = img.squeeze(1)
        img = img.to(self.device)
        img = normalize_on_gpu(img)
        return self.model(img)["SingleInstanceConfmapsHead"]

    def training_step(self, batch, batch_idx):
        """Training step."""
        X, y = (
            torch.squeeze(batch["image"], dim=1),
            torch.squeeze(batch["confidence_maps"], dim=1),
        )
        X = normalize_on_gpu(X)

        y_preds = self.model(X)["SingleInstanceConfmapsHead"]

        loss = self._compute_negative_weighted_loss(y_preds, y, batch)
        self._log_negative_split_metrics(
            [("confmaps", y_preds, y, 1.0)], batch, stage="train"
        )
        self._log_confmap_fg_bg_loss(y_preds, y, stage="train")

        if self.online_mining is not None and self.online_mining:
            ohkm_loss = compute_ohkm_loss(
                y_gt=y,
                y_pr=y_preds,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            loss = loss + ohkm_loss

        # for part-wise loss
        if self.node_names is not None:
            batch_size, _, h, w = y.shape
            mse = (y - y_preds) ** 2
            channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
            for node_idx, name in enumerate(self.node_names):
                self.log(
                    f"train/confmaps/{name}",
                    channel_wise_loss[node_idx],
                    prog_bar=False,
                    on_step=False,
                    on_epoch=True,
                    sync_dist=True,
                )
        # Log step-level loss (every batch, uses global_step x-axis)
        self.log(
            "loss",
            loss,
            prog_bar=True,
            on_step=True,
            on_epoch=False,
            sync_dist=True,
        )

        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
        self._accumulate_loss(loss)
        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step."""
        X, y = (
            torch.squeeze(batch["image"], dim=1),
            torch.squeeze(batch["confidence_maps"], dim=1),
        )
        X = normalize_on_gpu(X)

        y_preds = self.model(X)["SingleInstanceConfmapsHead"]
        val_loss = self._compute_negative_weighted_loss(y_preds, y, batch, stage="val")
        self._log_negative_split_metrics(
            [("confmaps", y_preds, y, 1.0)], batch, stage="val"
        )
        self._log_confmap_fg_bg_loss(y_preds, y, stage="val")
        if self.online_mining is not None and self.online_mining:
            ohkm_loss = compute_ohkm_loss(
                y_gt=y,
                y_pr=y_preds,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            val_loss = val_loss + ohkm_loss
        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Collect predictions for epoch-end evaluation if enabled
        if self._collect_val_predictions:
            with torch.no_grad():
                # Squeeze n_samples dim from image for inference (batch, 1, C, H, W) -> (batch, C, H, W)
                inference_batch = {k: v for k, v in batch.items()}
                if inference_batch["image"].ndim == 5:
                    inference_batch["image"] = inference_batch["image"].squeeze(1)
                inference_output = self.single_instance_inf_layer(inference_batch)
                if isinstance(inference_output, list):
                    inference_output = inference_output[0]

            batch_size = len(batch["frame_idx"])
            for i in range(batch_size):
                eff = batch["eff_scale"][i].cpu().numpy()

                # Predictions are already in original image space (inference divides by eff_scale)
                pred_peaks = inference_output["pred_instance_peaks"][i].cpu().numpy()
                pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()

                # Transform GT from preprocessed to original image space
                # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
                gt_prep = batch["instances"][i].cpu().numpy()
                if gt_prep.ndim == 4:
                    gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
                gt_orig = gt_prep / eff
                num_inst = batch["num_instances"][i].item()
                gt_orig = gt_orig[:num_inst]  # Only valid instances

                self.val_predictions.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "pred_peaks": pred_peaks,
                        "pred_scores": pred_scores,
                    }
                )
                self.val_ground_truth.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "gt_instances": gt_orig,
                        "num_instances": num_inst,
                    }
                )

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )

    self.single_instance_inf_layer = SingleInstanceInferenceModel(
        torch_model=self.forward,
        peak_threshold=0.2,
        input_scale=1.0,
        return_confmaps=True,
        output_stride=self.head_configs.single_instance.confmaps.output_stride,
    )
    self.node_names = self.head_configs.single_instance.confmaps.part_names

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    # Only squeeze n_samples dim if 5D (batch, n_samples, C, H, W) -> (batch, C, H, W)
    # Avoid double-squeezing when called from validation_step which already squeezes
    if img.ndim == 5:
        img = img.squeeze(1)
    img = img.to(self.device)
    img = normalize_on_gpu(img)
    return self.model(img)["SingleInstanceConfmapsHead"]

get_visualization_data(sample)

Extract visualization data from a sample.

Parameters:

Name	Type	Description	Default
sample		A sample dictionary from the data pipeline.	required

Returns:

Type	Description
VisualizationData	VisualizationData containing image, confmaps, peaks, etc.

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(self, sample) -> VisualizationData:
    """Extract visualization data from a sample.

    Args:
        sample: A sample dictionary from the data pipeline.

    Returns:
        VisualizationData containing image, confmaps, peaks, etc.
    """
    ex = sample.copy()
    ex["eff_scale"] = torch.tensor([1.0])
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["image"] = ex["image"].unsqueeze(dim=0)
    output = self.single_instance_inf_layer(ex)[0]

    peaks = output["pred_instance_peaks"].cpu().numpy()
    peak_values = output["pred_peak_values"].cpu().numpy()
    img = output["image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    gt_instances = ex["instances"][0].cpu().numpy()
    confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

    return VisualizationData(
        image=img,
        pred_confmaps=confmaps,
        pred_peaks=peaks,
        pred_peak_values=peak_values,
        gt_instances=gt_instances,
        node_names=list(self.node_names) if self.node_names else [],
        output_scale=confmaps.shape[0] / img.shape[0],
        is_paired=True,
    )

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    X, y = (
        torch.squeeze(batch["image"], dim=1),
        torch.squeeze(batch["confidence_maps"], dim=1),
    )
    X = normalize_on_gpu(X)

    y_preds = self.model(X)["SingleInstanceConfmapsHead"]

    loss = self._compute_negative_weighted_loss(y_preds, y, batch)
    self._log_negative_split_metrics(
        [("confmaps", y_preds, y, 1.0)], batch, stage="train"
    )
    self._log_confmap_fg_bg_loss(y_preds, y, stage="train")

    if self.online_mining is not None and self.online_mining:
        ohkm_loss = compute_ohkm_loss(
            y_gt=y,
            y_pr=y_preds,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        loss = loss + ohkm_loss

    # for part-wise loss
    if self.node_names is not None:
        batch_size, _, h, w = y.shape
        mse = (y - y_preds) ** 2
        channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
        for node_idx, name in enumerate(self.node_names):
            self.log(
                f"train/confmaps/{name}",
                channel_wise_loss[node_idx],
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )
    # Log step-level loss (every batch, uses global_step x-axis)
    self.log(
        "loss",
        loss,
        prog_bar=True,
        on_step=True,
        on_epoch=False,
        sync_dist=True,
    )

    # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
    self._accumulate_loss(loss)
    return loss

validation_step(batch, batch_idx)

Validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step."""
    X, y = (
        torch.squeeze(batch["image"], dim=1),
        torch.squeeze(batch["confidence_maps"], dim=1),
    )
    X = normalize_on_gpu(X)

    y_preds = self.model(X)["SingleInstanceConfmapsHead"]
    val_loss = self._compute_negative_weighted_loss(y_preds, y, batch, stage="val")
    self._log_negative_split_metrics(
        [("confmaps", y_preds, y, 1.0)], batch, stage="val"
    )
    self._log_confmap_fg_bg_loss(y_preds, y, stage="val")
    if self.online_mining is not None and self.online_mining:
        ohkm_loss = compute_ohkm_loss(
            y_gt=y,
            y_pr=y_preds,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        val_loss = val_loss + ohkm_loss
    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Collect predictions for epoch-end evaluation if enabled
    if self._collect_val_predictions:
        with torch.no_grad():
            # Squeeze n_samples dim from image for inference (batch, 1, C, H, W) -> (batch, C, H, W)
            inference_batch = {k: v for k, v in batch.items()}
            if inference_batch["image"].ndim == 5:
                inference_batch["image"] = inference_batch["image"].squeeze(1)
            inference_output = self.single_instance_inf_layer(inference_batch)
            if isinstance(inference_output, list):
                inference_output = inference_output[0]

        batch_size = len(batch["frame_idx"])
        for i in range(batch_size):
            eff = batch["eff_scale"][i].cpu().numpy()

            # Predictions are already in original image space (inference divides by eff_scale)
            pred_peaks = inference_output["pred_instance_peaks"][i].cpu().numpy()
            pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()

            # Transform GT from preprocessed to original image space
            # Note: instances have shape (1, max_inst, n_nodes, 2) - squeeze n_samples dim
            gt_prep = batch["instances"][i].cpu().numpy()
            if gt_prep.ndim == 4:
                gt_prep = gt_prep.squeeze(0)  # (max_inst, n_nodes, 2)
            gt_orig = gt_prep / eff
            num_inst = batch["num_instances"][i].item()
            gt_orig = gt_orig[:num_inst]  # Only valid instances

            self.val_predictions.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "pred_peaks": pred_peaks,
                    "pred_scores": pred_scores,
                }
            )
            self.val_ground_truth.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "gt_instances": gt_orig,
                    "num_instances": num_inst,
                }
            )

visualize_example(sample)

Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

TopDownCenteredInstanceLightningModule

Bases: LightningModel

Lightning Module for TopDownCenteredInstance Model.

This is a subclass of the LightningModel to configure the training/ validation steps and forward pass specific to TopDown Centered instance model. Top-Down models use a two-stage approach: first detecting centroids, then predicting keypoints for each detected centroid.

Parameters:

Name	Type	Description	Default
model_type	str	Type of the model. One of single_instance, centered_instance, centroid, bottomup, multi_class_bottomup, multi_class_topdown.	required
backbone_type	str	Backbone model. One of unet, convnext and swint.	required
backbone_config	Union[str, Dict[str, Any], DictConfig]	Backbone configuration. Can be: - String: One of the preset backbone types: - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"] - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"] - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"] - Dictionary: Custom configuration with structure: { "unet": {UNetConfig parameters}, "convnext": {ConvNextConfig parameters}, "swint": {SwinTConfig parameters} } Only one backbone type should be specified in the dictionary. - DictConfig: OmegaConf DictConfig object containing backbone configuration.	required
head_configs	DictConfig	Head configuration dictionary containing model-specific parameters. For Single Instance: confmaps with part_names, sigma, output_stride. For Centroid: confmaps with anchor_part, sigma, output_stride. For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride. For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight. For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight. For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.	required
pretrained_backbone_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.	None
pretrained_head_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.	None
init_weights	Optional[str]	Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.	'xavier'
lr_scheduler	Optional[Union[str, DictConfig]]	Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.	None
online_mining	Optional[bool]	If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).	False
hard_to_easy_ratio	Optional[float]	Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.	2.0
min_hard_keypoints	Optional[int]	Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.	2
max_hard_keypoints	Optional[int]	Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.	None
loss_scale	Optional[float]	Factor to scale hard keypoint losses by. Default: 5.0.	5.0
optimizer	Optional[str]	Optimizer name. One of ["Adam", "AdamW"].	'Adam'
learning_rate	Optional[float]	Learning rate for the optimizer. Default: 1e-3.	0.001
amsgrad	Optional[bool]	Enable AMSGrad with the optimizer. Default: False.	False

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model.
get_visualization_data	Extract visualization data from a sample.
training_step	Training step.
validation_step	Perform validation step.
visualize_example	Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

class TopDownCenteredInstanceLightningModule(LightningModel):
    """Lightning Module for TopDownCenteredInstance Model.

    This is a subclass of the `LightningModel` to configure the training/ validation steps
    and forward pass specific to TopDown Centered instance model. Top-Down models use a two-stage
    approach: first detecting centroids, then predicting keypoints for each detected centroid.

    Args:
        model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
        backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
        backbone_config: Backbone configuration. Can be:
            - String: One of the preset backbone types:
                - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"]
                - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"]
                - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"]
            - Dictionary: Custom configuration with structure:
                {
                    "unet": {UNetConfig parameters},
                    "convnext": {ConvNextConfig parameters},
                    "swint": {SwinTConfig parameters}
                }
                Only one backbone type should be specified in the dictionary.
            - DictConfig: OmegaConf DictConfig object containing backbone configuration.
        head_configs: Head configuration dictionary containing model-specific parameters.
            For Single Instance: confmaps with part_names, sigma, output_stride.
            For Centroid: confmaps with anchor_part, sigma, output_stride.
            For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride.
            For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight.
            For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight.
            For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.
        pretrained_backbone_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.
        pretrained_head_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.
        init_weights: Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.
        lr_scheduler: Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.
        online_mining: If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).
        hard_to_easy_ratio: Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.
        min_hard_keypoints: Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.
        max_hard_keypoints: Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.
        loss_scale: Factor to scale hard keypoint losses by. Default: 5.0.
        optimizer: Optimizer name. One of ["Adam", "AdamW"].
        learning_rate: Learning rate for the optimizer. Default: 1e-3.
        amsgrad: Enable AMSGrad with the optimizer. Default: False.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )

        self.instance_peaks_inf_layer = FindInstancePeaks(
            torch_model=self.forward,
            peak_threshold=0.2,
            return_confmaps=True,
            output_stride=self.head_configs.centered_instance.confmaps.output_stride,
        )

        self.node_names = self.head_configs.centered_instance.confmaps.part_names

    def get_visualization_data(self, sample) -> VisualizationData:
        """Extract visualization data from a sample."""
        ex = sample.copy()
        ex["eff_scale"] = torch.tensor([1.0])
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["instance_image"] = ex["instance_image"].unsqueeze(dim=0)
        output = self.instance_peaks_inf_layer(ex)

        peaks = output["pred_instance_peaks"].cpu().numpy()
        peak_values = output["pred_peak_values"].cpu().numpy()
        img = output["instance_image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        gt_instances = ex["instance"].cpu().numpy()
        confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

        return VisualizationData(
            image=img,
            pred_confmaps=confmaps,
            pred_peaks=peaks,
            pred_peak_values=peak_values,
            gt_instances=gt_instances,
            node_names=list(self.node_names) if self.node_names else [],
            output_scale=confmaps.shape[0] / img.shape[0],
            is_paired=True,
        )

    def visualize_example(self, sample):
        """Visualize predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def forward(self, img):
        """Forward pass of the model."""
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        return self.model(img)["CenteredInstanceConfmapsHead"]

    def training_step(self, batch, batch_idx):
        """Training step."""
        X, y = (
            torch.squeeze(batch["instance_image"], dim=1),
            torch.squeeze(batch["confidence_maps"], dim=1),
        )
        X = normalize_on_gpu(X)

        y_preds = self.model(X)["CenteredInstanceConfmapsHead"]

        loss = nn.MSELoss()(y_preds, y)
        self._log_confmap_fg_bg_loss(y_preds, y, stage="train")

        if self.online_mining is not None and self.online_mining:
            ohkm_loss = compute_ohkm_loss(
                y_gt=y,
                y_pr=y_preds,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            loss = loss + ohkm_loss

        # for part-wise loss
        if self.node_names is not None:
            batch_size, _, h, w = y.shape
            mse = (y - y_preds) ** 2
            channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
            for node_idx, name in enumerate(self.node_names):
                self.log(
                    f"train/confmaps/{name}",
                    channel_wise_loss[node_idx],
                    prog_bar=False,
                    on_step=False,
                    on_epoch=True,
                    sync_dist=True,
                )

        # Log step-level loss (every batch, uses global_step x-axis)
        self.log(
            "loss",
            loss,
            prog_bar=True,
            on_step=True,
            on_epoch=False,
            sync_dist=True,
        )
        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
        self._accumulate_loss(loss)
        return loss

    def validation_step(self, batch, batch_idx):
        """Perform validation step."""
        X, y = (
            torch.squeeze(batch["instance_image"], dim=1),
            torch.squeeze(batch["confidence_maps"], dim=1),
        )
        X = normalize_on_gpu(X)

        y_preds = self.model(X)["CenteredInstanceConfmapsHead"]
        val_loss = nn.MSELoss()(y_preds, y)
        self._log_confmap_fg_bg_loss(y_preds, y, stage="val")
        if self.online_mining is not None and self.online_mining:
            ohkm_loss = compute_ohkm_loss(
                y_gt=y,
                y_pr=y_preds,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            val_loss = val_loss + ohkm_loss
        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Collect predictions for epoch-end evaluation if enabled
        if self._collect_val_predictions:
            # SAVE bbox BEFORE inference (it modifies in-place!)
            bbox_prep_saved = batch["instance_bbox"].clone()

            with torch.no_grad():
                inference_output = self.instance_peaks_inf_layer(batch)

            batch_size = len(batch["frame_idx"])
            for i in range(batch_size):
                eff = batch["eff_scale"][i].cpu().numpy()

                # Predictions from inference (crop-relative, original scale)
                pred_peaks_crop = (
                    inference_output["pred_instance_peaks"][i].cpu().numpy()
                )
                pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()

                # Compute bbox offset in original space from SAVED prep bbox
                # bbox has shape (n_samples=1, 4, 2) where 4 corners
                bbox_prep = bbox_prep_saved[i].squeeze(0).cpu().numpy()  # (4, 2)
                bbox_top_left_orig = (
                    bbox_prep[0] / eff
                )  # Top-left corner in original space

                # Full image coordinates (original space)
                pred_peaks_full = pred_peaks_crop + bbox_top_left_orig

                # GT transform: crop-relative preprocessed -> full image original
                gt_crop_prep = (
                    batch["instance"][i].squeeze(0).cpu().numpy()
                )  # (n_nodes, 2)
                gt_crop_orig = gt_crop_prep / eff
                gt_full_orig = gt_crop_orig + bbox_top_left_orig

                if "crop_window" in batch:
                    # Box-scaled crop: undo the input scale and the crop's own
                    # window/scale (exact on un-augmented val samples), so both
                    # land in original-frame pixels.
                    pred_peaks_full = _bbox_crop_points_to_frame(
                        pred_peaks_crop * eff, batch, i
                    )
                    gt_full_orig = _bbox_crop_points_to_frame(gt_crop_prep, batch, i)

                # One entry per crop: `crop_id` pairs this prediction with its
                # own GT instance (a frame holds several crops, each in its own
                # crop coordinates, so they must not be pooled by frame).
                crop_id = len(self.val_ground_truth)
                self.val_predictions.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "crop_id": crop_id,
                        "pred_peaks": pred_peaks_full.reshape(
                            1, -1, 2
                        ),  # (1, n_nodes, 2)
                        "pred_scores": pred_scores.reshape(1, -1),  # (1, n_nodes)
                    }
                )
                self.val_ground_truth.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "crop_id": crop_id,
                        "gt_instances": gt_full_orig.reshape(
                            1, -1, 2
                        ),  # (1, n_nodes, 2)
                        "num_instances": 1,
                    }
                )

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )

    self.instance_peaks_inf_layer = FindInstancePeaks(
        torch_model=self.forward,
        peak_threshold=0.2,
        return_confmaps=True,
        output_stride=self.head_configs.centered_instance.confmaps.output_stride,
    )

    self.node_names = self.head_configs.centered_instance.confmaps.part_names

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    return self.model(img)["CenteredInstanceConfmapsHead"]

get_visualization_data(sample)

Extract visualization data from a sample.

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(self, sample) -> VisualizationData:
    """Extract visualization data from a sample."""
    ex = sample.copy()
    ex["eff_scale"] = torch.tensor([1.0])
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["instance_image"] = ex["instance_image"].unsqueeze(dim=0)
    output = self.instance_peaks_inf_layer(ex)

    peaks = output["pred_instance_peaks"].cpu().numpy()
    peak_values = output["pred_peak_values"].cpu().numpy()
    img = output["instance_image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    gt_instances = ex["instance"].cpu().numpy()
    confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

    return VisualizationData(
        image=img,
        pred_confmaps=confmaps,
        pred_peaks=peaks,
        pred_peak_values=peak_values,
        gt_instances=gt_instances,
        node_names=list(self.node_names) if self.node_names else [],
        output_scale=confmaps.shape[0] / img.shape[0],
        is_paired=True,
    )

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    X, y = (
        torch.squeeze(batch["instance_image"], dim=1),
        torch.squeeze(batch["confidence_maps"], dim=1),
    )
    X = normalize_on_gpu(X)

    y_preds = self.model(X)["CenteredInstanceConfmapsHead"]

    loss = nn.MSELoss()(y_preds, y)
    self._log_confmap_fg_bg_loss(y_preds, y, stage="train")

    if self.online_mining is not None and self.online_mining:
        ohkm_loss = compute_ohkm_loss(
            y_gt=y,
            y_pr=y_preds,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        loss = loss + ohkm_loss

    # for part-wise loss
    if self.node_names is not None:
        batch_size, _, h, w = y.shape
        mse = (y - y_preds) ** 2
        channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
        for node_idx, name in enumerate(self.node_names):
            self.log(
                f"train/confmaps/{name}",
                channel_wise_loss[node_idx],
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )

    # Log step-level loss (every batch, uses global_step x-axis)
    self.log(
        "loss",
        loss,
        prog_bar=True,
        on_step=True,
        on_epoch=False,
        sync_dist=True,
    )
    # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
    self._accumulate_loss(loss)
    return loss

validation_step(batch, batch_idx)

Perform validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Perform validation step."""
    X, y = (
        torch.squeeze(batch["instance_image"], dim=1),
        torch.squeeze(batch["confidence_maps"], dim=1),
    )
    X = normalize_on_gpu(X)

    y_preds = self.model(X)["CenteredInstanceConfmapsHead"]
    val_loss = nn.MSELoss()(y_preds, y)
    self._log_confmap_fg_bg_loss(y_preds, y, stage="val")
    if self.online_mining is not None and self.online_mining:
        ohkm_loss = compute_ohkm_loss(
            y_gt=y,
            y_pr=y_preds,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        val_loss = val_loss + ohkm_loss
    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Collect predictions for epoch-end evaluation if enabled
    if self._collect_val_predictions:
        # SAVE bbox BEFORE inference (it modifies in-place!)
        bbox_prep_saved = batch["instance_bbox"].clone()

        with torch.no_grad():
            inference_output = self.instance_peaks_inf_layer(batch)

        batch_size = len(batch["frame_idx"])
        for i in range(batch_size):
            eff = batch["eff_scale"][i].cpu().numpy()

            # Predictions from inference (crop-relative, original scale)
            pred_peaks_crop = (
                inference_output["pred_instance_peaks"][i].cpu().numpy()
            )
            pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()

            # Compute bbox offset in original space from SAVED prep bbox
            # bbox has shape (n_samples=1, 4, 2) where 4 corners
            bbox_prep = bbox_prep_saved[i].squeeze(0).cpu().numpy()  # (4, 2)
            bbox_top_left_orig = (
                bbox_prep[0] / eff
            )  # Top-left corner in original space

            # Full image coordinates (original space)
            pred_peaks_full = pred_peaks_crop + bbox_top_left_orig

            # GT transform: crop-relative preprocessed -> full image original
            gt_crop_prep = (
                batch["instance"][i].squeeze(0).cpu().numpy()
            )  # (n_nodes, 2)
            gt_crop_orig = gt_crop_prep / eff
            gt_full_orig = gt_crop_orig + bbox_top_left_orig

            if "crop_window" in batch:
                # Box-scaled crop: undo the input scale and the crop's own
                # window/scale (exact on un-augmented val samples), so both
                # land in original-frame pixels.
                pred_peaks_full = _bbox_crop_points_to_frame(
                    pred_peaks_crop * eff, batch, i
                )
                gt_full_orig = _bbox_crop_points_to_frame(gt_crop_prep, batch, i)

            # One entry per crop: `crop_id` pairs this prediction with its
            # own GT instance (a frame holds several crops, each in its own
            # crop coordinates, so they must not be pooled by frame).
            crop_id = len(self.val_ground_truth)
            self.val_predictions.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "crop_id": crop_id,
                    "pred_peaks": pred_peaks_full.reshape(
                        1, -1, 2
                    ),  # (1, n_nodes, 2)
                    "pred_scores": pred_scores.reshape(1, -1),  # (1, n_nodes)
                }
            )
            self.val_ground_truth.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "crop_id": crop_id,
                    "gt_instances": gt_full_orig.reshape(
                        1, -1, 2
                    ),  # (1, n_nodes, 2)
                    "num_instances": 1,
                }
            )

visualize_example(sample)

Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

TopDownCenteredInstanceMultiClassLightningModule

Bases: LightningModel

Lightning Module for TopDownCenteredInstance ID Model.

This is a subclass of the LightningModel to configure the training/ validation steps and forward pass specific to TopDown Centered instance model. Multi-Class Top-Down models use a two-stage approach: first detecting centroids, then predicting keypoints and classifying instances using supervised learning with ground truth track IDs.

Parameters:

Name	Type	Description	Default
model_type	str	Type of the model. One of single_instance, centered_instance, centroid, bottomup, multi_class_bottomup, multi_class_topdown.	required
backbone_type	str	Backbone model. One of unet, convnext and swint.	required
backbone_config	Union[str, Dict[str, Any], DictConfig]	Backbone configuration. Can be: - String: One of the preset backbone types: - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"] - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"] - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"] - Dictionary: Custom configuration with structure: { "unet": {UNetConfig parameters}, "convnext": {ConvNextConfig parameters}, "swint": {SwinTConfig parameters} } Only one backbone type should be specified in the dictionary. - DictConfig: OmegaConf DictConfig object containing backbone configuration.	required
head_configs	DictConfig	Head configuration dictionary containing model-specific parameters. For Single Instance: confmaps with part_names, sigma, output_stride. For Centroid: confmaps with anchor_part, sigma, output_stride. For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride. For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight. For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight. For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.	required
pretrained_backbone_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.	None
pretrained_head_weights	Optional[str]	Path to checkpoint .ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.	None
init_weights	Optional[str]	Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.	'xavier'
lr_scheduler	Optional[Union[str, DictConfig]]	Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.	None
online_mining	Optional[bool]	If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).	False
hard_to_easy_ratio	Optional[float]	Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.	2.0
min_hard_keypoints	Optional[int]	Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.	2
max_hard_keypoints	Optional[int]	Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.	None
loss_scale	Optional[float]	Factor to scale hard keypoint losses by. Default: 5.0.	5.0
optimizer	Optional[str]	Optimizer name. One of ["Adam", "AdamW"].	'Adam'
learning_rate	Optional[float]	Learning rate for the optimizer. Default: 1e-3.	0.001
amsgrad	Optional[bool]	Enable AMSGrad with the optimizer. Default: False.	False

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model.
get_visualization_data	Extract visualization data from a sample.
training_step	Training step.
validation_step	Perform validation step.
visualize_example	Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

class TopDownCenteredInstanceMultiClassLightningModule(LightningModel):
    """Lightning Module for TopDownCenteredInstance ID Model.

    This is a subclass of the `LightningModel` to configure the training/ validation steps
    and forward pass specific to TopDown Centered instance model. Multi-Class Top-Down models
    use a two-stage approach: first detecting centroids, then predicting keypoints and
    classifying instances using supervised learning with ground truth track IDs.

    Args:
        model_type: Type of the model. One of `single_instance`, `centered_instance`, `centroid`, `bottomup`, `multi_class_bottomup`, `multi_class_topdown`.
        backbone_type: Backbone model. One of `unet`, `convnext` and `swint`.
        backbone_config: Backbone configuration. Can be:
            - String: One of the preset backbone types:
                - UNet variants: ["unet", "unet_medium_rf", "unet_large_rf"]
                - ConvNeXt variants: ["convnext", "convnext_tiny", "convnext_small", "convnext_base", "convnext_large"]
                - SwinT variants: ["swint", "swint_tiny", "swint_small", "swint_base"]
            - Dictionary: Custom configuration with structure:
                {
                    "unet": {UNetConfig parameters},
                    "convnext": {ConvNextConfig parameters},
                    "swint": {SwinTConfig parameters}
                }
                Only one backbone type should be specified in the dictionary.
            - DictConfig: OmegaConf DictConfig object containing backbone configuration.
        head_configs: Head configuration dictionary containing model-specific parameters.
            For Single Instance: confmaps with part_names, sigma, output_stride.
            For Centroid: confmaps with anchor_part, sigma, output_stride.
            For Centered Instance: confmaps with part_names, anchor_part, sigma, output_stride.
            For Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; pafs with edges, sigma, output_stride, loss_weight.
            For Multi-Class Bottom-Up: confmaps with part_names, sigma, output_stride, loss_weight; class_maps with classes, sigma, output_stride, loss_weight.
            For Multi-Class Top-Down: confmaps with part_names, anchor_part, sigma, output_stride, loss_weight; class_vectors with classes, num_fc_layers, num_fc_units, global_pool, output_stride, loss_weight.
        pretrained_backbone_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for backbone initialization. If None, random initialization is used.
        pretrained_head_weights: Path to checkpoint `.ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file for head layers initialization. If None, random initialization is used.
        init_weights: Model weights initialization method. "default" uses kaiming uniform initialization, "xavier" uses Xavier initialization.
        lr_scheduler: Learning rate scheduler configuration. Can be string ("step_lr", "reduce_lr_on_plateau") or dictionary with scheduler-specific parameters.
        online_mining: If True, online hard keypoint mining (OHKM) is enabled. Loss is computed per keypoint and sorted from lowest (easy) to highest (hard).
        hard_to_easy_ratio: Minimum ratio of individual keypoint loss to lowest keypoint loss to be considered "hard". Default: 2.0.
        min_hard_keypoints: Minimum number of keypoints considered as "hard", even if below hard_to_easy_ratio. Default: 2.
        max_hard_keypoints: Maximum number of hard keypoints to apply scaling to. If None, no limit is applied.
        loss_scale: Factor to scale hard keypoint losses by. Default: 5.0.
        optimizer: Optimizer name. One of ["Adam", "AdamW"].
        learning_rate: Learning rate for the optimizer. Default: 1e-3.
        amsgrad: Enable AMSGrad with the optimizer. Default: False.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )
        self.instance_peaks_inf_layer = TopDownMultiClassFindInstancePeaks(
            torch_model=self.forward,
            peak_threshold=0.2,
            return_confmaps=True,
            output_stride=self.head_configs.multi_class_topdown.confmaps.output_stride,
        )

        self.node_names = self.head_configs.multi_class_topdown.confmaps.part_names

    def get_visualization_data(self, sample) -> VisualizationData:
        """Extract visualization data from a sample."""
        ex = sample.copy()
        ex["eff_scale"] = torch.tensor([1.0])
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["instance_image"] = ex["instance_image"].unsqueeze(dim=0)
        output = self.instance_peaks_inf_layer(ex)

        peaks = output["pred_instance_peaks"].cpu().numpy()
        peak_values = output["pred_peak_values"].cpu().numpy()
        img = output["instance_image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        gt_instances = ex["instance"].cpu().numpy()
        confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

        return VisualizationData(
            image=img,
            pred_confmaps=confmaps,
            pred_peaks=peaks,
            pred_peak_values=peak_values,
            gt_instances=gt_instances,
            node_names=list(self.node_names) if self.node_names else [],
            output_scale=confmaps.shape[0] / img.shape[0],
            is_paired=True,
        )

    def visualize_example(self, sample):
        """Visualize predictions during training (used with callbacks)."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
        return fig

    def forward(self, img):
        """Forward pass of the model."""
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        output = self.model(img)
        return {
            "CenteredInstanceConfmapsHead": output["CenteredInstanceConfmapsHead"],
            "ClassVectorsHead": output["ClassVectorsHead"],
        }

    def training_step(self, batch, batch_idx):
        """Training step."""
        X = torch.squeeze(batch["instance_image"], dim=1)
        y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
        y_classvector = batch["class_vectors"]
        X = normalize_on_gpu(X)
        preds = self.model(X)
        classvector = preds["ClassVectorsHead"]
        confmaps = preds["CenteredInstanceConfmapsHead"]

        confmap_loss = nn.MSELoss()(confmaps, y_confmap)
        classvector_loss = nn.CrossEntropyLoss()(classvector, y_classvector)

        if self.online_mining is not None and self.online_mining:
            confmap_ohkm_loss = compute_ohkm_loss(
                y_gt=y_confmap,
                y_pr=confmaps,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            confmap_loss += confmap_ohkm_loss

        losses = {
            "CenteredInstanceConfmapsHead": confmap_loss,
            "ClassVectorsHead": classvector_loss,
        }
        loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])

        # for part-wise loss
        if self.node_names is not None:
            batch_size, _, h, w = y_confmap.shape
            mse = (y_confmap - confmaps) ** 2
            channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
            for node_idx, name in enumerate(self.node_names):
                self.log(
                    f"train/confmaps/{name}",
                    channel_wise_loss[node_idx],
                    prog_bar=False,
                    on_step=False,
                    on_epoch=True,
                    sync_dist=True,
                )

        # Log step-level loss (every batch, uses global_step x-axis)
        self.log(
            "loss",
            loss,
            prog_bar=True,
            on_step=True,
            on_epoch=False,
            sync_dist=True,
        )
        # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
        self._accumulate_loss(loss)
        self.log(
            "train/confmaps_loss",
            confmap_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "train/classvector_loss",
            classvector_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Compute classification accuracy
        with torch.no_grad():
            pred_classes = torch.argmax(classvector, dim=1)
            gt_classes = torch.argmax(y_classvector, dim=1)
            class_accuracy = (pred_classes == gt_classes).float().mean()
        self.log(
            "train/class_accuracy",
            class_accuracy,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        return loss

    def validation_step(self, batch, batch_idx):
        """Perform validation step."""
        X = torch.squeeze(batch["instance_image"], dim=1)
        y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
        y_classvector = batch["class_vectors"]
        X = normalize_on_gpu(X)
        preds = self.model(X)
        classvector = preds["ClassVectorsHead"]
        confmaps = preds["CenteredInstanceConfmapsHead"]

        confmap_loss = nn.MSELoss()(confmaps, y_confmap)
        classvector_loss = nn.CrossEntropyLoss()(classvector, y_classvector)

        if self.online_mining is not None and self.online_mining:
            confmap_ohkm_loss = compute_ohkm_loss(
                y_gt=y_confmap,
                y_pr=confmaps,
                hard_to_easy_ratio=self.hard_to_easy_ratio,
                min_hard_keypoints=self.min_hard_keypoints,
                max_hard_keypoints=self.max_hard_keypoints,
                loss_scale=self.loss_scale,
            )
            confmap_loss += confmap_ohkm_loss

        losses = {
            "CenteredInstanceConfmapsHead": confmap_loss,
            "ClassVectorsHead": classvector_loss,
        }
        val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "val/confmaps_loss",
            confmap_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log(
            "val/classvector_loss",
            classvector_loss,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Compute classification accuracy
        with torch.no_grad():
            pred_classes = torch.argmax(classvector, dim=1)
            gt_classes = torch.argmax(y_classvector, dim=1)
            class_accuracy = (pred_classes == gt_classes).float().mean()
        self.log(
            "val/class_accuracy",
            class_accuracy,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )

        # Collect predictions for epoch-end evaluation if enabled
        if self._collect_val_predictions:
            # SAVE bbox BEFORE inference (it modifies in-place!)
            bbox_prep_saved = batch["instance_bbox"].clone()

            with torch.no_grad():
                inference_output = self.instance_peaks_inf_layer(batch)

            batch_size = len(batch["frame_idx"])
            for i in range(batch_size):
                eff = batch["eff_scale"][i].cpu().numpy()

                # Predictions from inference (crop-relative, original scale)
                pred_peaks_crop = (
                    inference_output["pred_instance_peaks"][i].cpu().numpy()
                )
                pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()

                # Compute bbox offset in original space from SAVED prep bbox
                # bbox has shape (n_samples=1, 4, 2) where 4 corners
                bbox_prep = bbox_prep_saved[i].squeeze(0).cpu().numpy()  # (4, 2)
                bbox_top_left_orig = (
                    bbox_prep[0] / eff
                )  # Top-left corner in original space

                # Full image coordinates (original space)
                pred_peaks_full = pred_peaks_crop + bbox_top_left_orig

                # GT transform: crop-relative preprocessed -> full image original
                gt_crop_prep = (
                    batch["instance"][i].squeeze(0).cpu().numpy()
                )  # (n_nodes, 2)
                gt_crop_orig = gt_crop_prep / eff
                gt_full_orig = gt_crop_orig + bbox_top_left_orig

                if "crop_window" in batch:
                    # Box-scaled crop: undo the input scale and the crop's own
                    # window/scale (exact on un-augmented val samples), so both
                    # land in original-frame pixels.
                    pred_peaks_full = _bbox_crop_points_to_frame(
                        pred_peaks_crop * eff, batch, i
                    )
                    gt_full_orig = _bbox_crop_points_to_frame(gt_crop_prep, batch, i)

                # One entry per crop: `crop_id` pairs this prediction with its
                # own GT instance (a frame holds several crops, each in its own
                # crop coordinates, so they must not be pooled by frame).
                crop_id = len(self.val_ground_truth)
                self.val_predictions.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "crop_id": crop_id,
                        "pred_peaks": pred_peaks_full.reshape(
                            1, -1, 2
                        ),  # (1, n_nodes, 2)
                        "pred_scores": pred_scores.reshape(1, -1),  # (1, n_nodes)
                    }
                )
                self.val_ground_truth.append(
                    {
                        "video_idx": batch["video_idx"][i].item(),
                        "frame_idx": batch["frame_idx"][i].item(),
                        "crop_id": crop_id,
                        "gt_instances": gt_full_orig.reshape(
                            1, -1, 2
                        ),  # (1, n_nodes, 2)
                        "num_instances": 1,
                    }
                )

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )
    self.instance_peaks_inf_layer = TopDownMultiClassFindInstancePeaks(
        torch_model=self.forward,
        peak_threshold=0.2,
        return_confmaps=True,
        output_stride=self.head_configs.multi_class_topdown.confmaps.output_stride,
    )

    self.node_names = self.head_configs.multi_class_topdown.confmaps.part_names

forward(img)

Forward pass of the model.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model."""
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    output = self.model(img)
    return {
        "CenteredInstanceConfmapsHead": output["CenteredInstanceConfmapsHead"],
        "ClassVectorsHead": output["ClassVectorsHead"],
    }

get_visualization_data(sample)

Extract visualization data from a sample.

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(self, sample) -> VisualizationData:
    """Extract visualization data from a sample."""
    ex = sample.copy()
    ex["eff_scale"] = torch.tensor([1.0])
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["instance_image"] = ex["instance_image"].unsqueeze(dim=0)
    output = self.instance_peaks_inf_layer(ex)

    peaks = output["pred_instance_peaks"].cpu().numpy()
    peak_values = output["pred_peak_values"].cpu().numpy()
    img = output["instance_image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    gt_instances = ex["instance"].cpu().numpy()
    confmaps = output["pred_confmaps"][0].cpu().numpy().transpose(1, 2, 0)

    return VisualizationData(
        image=img,
        pred_confmaps=confmaps,
        pred_peaks=peaks,
        pred_peak_values=peak_values,
        gt_instances=gt_instances,
        node_names=list(self.node_names) if self.node_names else [],
        output_scale=confmaps.shape[0] / img.shape[0],
        is_paired=True,
    )

training_step(batch, batch_idx)

Training step.

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step."""
    X = torch.squeeze(batch["instance_image"], dim=1)
    y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
    y_classvector = batch["class_vectors"]
    X = normalize_on_gpu(X)
    preds = self.model(X)
    classvector = preds["ClassVectorsHead"]
    confmaps = preds["CenteredInstanceConfmapsHead"]

    confmap_loss = nn.MSELoss()(confmaps, y_confmap)
    classvector_loss = nn.CrossEntropyLoss()(classvector, y_classvector)

    if self.online_mining is not None and self.online_mining:
        confmap_ohkm_loss = compute_ohkm_loss(
            y_gt=y_confmap,
            y_pr=confmaps,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        confmap_loss += confmap_ohkm_loss

    losses = {
        "CenteredInstanceConfmapsHead": confmap_loss,
        "ClassVectorsHead": classvector_loss,
    }
    loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])

    # for part-wise loss
    if self.node_names is not None:
        batch_size, _, h, w = y_confmap.shape
        mse = (y_confmap - confmaps) ** 2
        channel_wise_loss = torch.sum(mse, dim=(0, 2, 3)) / (batch_size * h * w)
        for node_idx, name in enumerate(self.node_names):
            self.log(
                f"train/confmaps/{name}",
                channel_wise_loss[node_idx],
                prog_bar=False,
                on_step=False,
                on_epoch=True,
                sync_dist=True,
            )

    # Log step-level loss (every batch, uses global_step x-axis)
    self.log(
        "loss",
        loss,
        prog_bar=True,
        on_step=True,
        on_epoch=False,
        sync_dist=True,
    )
    # Accumulate for epoch-averaged loss (logged in on_train_epoch_end)
    self._accumulate_loss(loss)
    self.log(
        "train/confmaps_loss",
        confmap_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "train/classvector_loss",
        classvector_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Compute classification accuracy
    with torch.no_grad():
        pred_classes = torch.argmax(classvector, dim=1)
        gt_classes = torch.argmax(y_classvector, dim=1)
        class_accuracy = (pred_classes == gt_classes).float().mean()
    self.log(
        "train/class_accuracy",
        class_accuracy,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    return loss

validation_step(batch, batch_idx)

Perform validation step.

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Perform validation step."""
    X = torch.squeeze(batch["instance_image"], dim=1)
    y_confmap = torch.squeeze(batch["confidence_maps"], dim=1)
    y_classvector = batch["class_vectors"]
    X = normalize_on_gpu(X)
    preds = self.model(X)
    classvector = preds["ClassVectorsHead"]
    confmaps = preds["CenteredInstanceConfmapsHead"]

    confmap_loss = nn.MSELoss()(confmaps, y_confmap)
    classvector_loss = nn.CrossEntropyLoss()(classvector, y_classvector)

    if self.online_mining is not None and self.online_mining:
        confmap_ohkm_loss = compute_ohkm_loss(
            y_gt=y_confmap,
            y_pr=confmaps,
            hard_to_easy_ratio=self.hard_to_easy_ratio,
            min_hard_keypoints=self.min_hard_keypoints,
            max_hard_keypoints=self.max_hard_keypoints,
            loss_scale=self.loss_scale,
        )
        confmap_loss += confmap_ohkm_loss

    losses = {
        "CenteredInstanceConfmapsHead": confmap_loss,
        "ClassVectorsHead": classvector_loss,
    }
    val_loss = sum([s * losses[t] for s, t in zip(self.loss_weights, losses)])
    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "val/confmaps_loss",
        confmap_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log(
        "val/classvector_loss",
        classvector_loss,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Compute classification accuracy
    with torch.no_grad():
        pred_classes = torch.argmax(classvector, dim=1)
        gt_classes = torch.argmax(y_classvector, dim=1)
        class_accuracy = (pred_classes == gt_classes).float().mean()
    self.log(
        "val/class_accuracy",
        class_accuracy,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )

    # Collect predictions for epoch-end evaluation if enabled
    if self._collect_val_predictions:
        # SAVE bbox BEFORE inference (it modifies in-place!)
        bbox_prep_saved = batch["instance_bbox"].clone()

        with torch.no_grad():
            inference_output = self.instance_peaks_inf_layer(batch)

        batch_size = len(batch["frame_idx"])
        for i in range(batch_size):
            eff = batch["eff_scale"][i].cpu().numpy()

            # Predictions from inference (crop-relative, original scale)
            pred_peaks_crop = (
                inference_output["pred_instance_peaks"][i].cpu().numpy()
            )
            pred_scores = inference_output["pred_peak_values"][i].cpu().numpy()

            # Compute bbox offset in original space from SAVED prep bbox
            # bbox has shape (n_samples=1, 4, 2) where 4 corners
            bbox_prep = bbox_prep_saved[i].squeeze(0).cpu().numpy()  # (4, 2)
            bbox_top_left_orig = (
                bbox_prep[0] / eff
            )  # Top-left corner in original space

            # Full image coordinates (original space)
            pred_peaks_full = pred_peaks_crop + bbox_top_left_orig

            # GT transform: crop-relative preprocessed -> full image original
            gt_crop_prep = (
                batch["instance"][i].squeeze(0).cpu().numpy()
            )  # (n_nodes, 2)
            gt_crop_orig = gt_crop_prep / eff
            gt_full_orig = gt_crop_orig + bbox_top_left_orig

            if "crop_window" in batch:
                # Box-scaled crop: undo the input scale and the crop's own
                # window/scale (exact on un-augmented val samples), so both
                # land in original-frame pixels.
                pred_peaks_full = _bbox_crop_points_to_frame(
                    pred_peaks_crop * eff, batch, i
                )
                gt_full_orig = _bbox_crop_points_to_frame(gt_crop_prep, batch, i)

            # One entry per crop: `crop_id` pairs this prediction with its
            # own GT instance (a frame holds several crops, each in its own
            # crop coordinates, so they must not be pooled by frame).
            crop_id = len(self.val_ground_truth)
            self.val_predictions.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "crop_id": crop_id,
                    "pred_peaks": pred_peaks_full.reshape(
                        1, -1, 2
                    ),  # (1, n_nodes, 2)
                    "pred_scores": pred_scores.reshape(1, -1),  # (1, n_nodes)
                }
            )
            self.val_ground_truth.append(
                {
                    "video_idx": batch["video_idx"][i].item(),
                    "frame_idx": batch["frame_idx"][i].item(),
                    "crop_id": crop_id,
                    "gt_instances": gt_full_orig.reshape(
                        1, -1, 2
                    ),  # (1, n_nodes, 2)
                    "num_instances": 1,
                }
            )

visualize_example(sample)

Visualize predictions during training (used with callbacks).

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize predictions during training (used with callbacks)."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    plot_peaks(data.gt_instances, data.pred_peaks, paired=data.is_paired)
    return fig

TopDownCenteredInstanceSegmentationLightningModule

Bases: LightningModel

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

Predicts a single binary foreground mask of the centered instance on a centroid crop. Combines the crop I/O of :class:TopDownCenteredInstanceLightningModule (instance_image input) with the bce-dice foreground loss / IoU metric of :class:BottomUpSegmentationLightningModule. Composed with a centroid model for full top-down inference; the head emits logits.

Methods:

Name	Description
__init__	Initialise the configs and the model.
forward	Forward pass of the model returning foreground-mask LOGITS.
get_visualization_data	Crop-flavored viz: image + predicted foreground vs GT mask overlay.
training_step	Training step (bce-dice on the centered-instance foreground mask).
validation_step	Validation step (val loss + foreground IoU).
visualize_example	Visualize the predicted foreground mask over the crop during training.

Source code in sleap_nn/training/lightning_modules.py

class TopDownCenteredInstanceSegmentationLightningModule(LightningModel):
    """Lightning Module for top-down (crop-centered) instance segmentation (#622).

    Predicts a single binary foreground mask of the centered instance on a
    centroid crop. Combines the crop I/O of
    :class:`TopDownCenteredInstanceLightningModule` (``instance_image`` input)
    with the bce-dice foreground loss / IoU metric of
    :class:`BottomUpSegmentationLightningModule`. Composed with a ``centroid``
    model for full top-down inference; the head emits logits.
    """

    def __init__(
        self,
        model_type: str,
        backbone_type: str,
        backbone_config: Union[str, Dict[str, Any], DictConfig],
        head_configs: DictConfig,
        pretrained_backbone_weights: Optional[str] = None,
        pretrained_head_weights: Optional[str] = None,
        init_weights: Optional[str] = "xavier",
        lr_scheduler: Optional[Union[str, DictConfig]] = None,
        online_mining: Optional[bool] = False,
        hard_to_easy_ratio: Optional[float] = 2.0,
        min_hard_keypoints: Optional[int] = 2,
        max_hard_keypoints: Optional[int] = None,
        loss_scale: Optional[float] = 5.0,
        optimizer: Optional[str] = "Adam",
        learning_rate: Optional[float] = 1e-3,
        amsgrad: Optional[bool] = False,
        negative_loss_weight: Optional[float] = 1.0,
    ):
        """Initialise the configs and the model."""
        super().__init__(
            model_type=model_type,
            backbone_type=backbone_type,
            backbone_config=backbone_config,
            head_configs=head_configs,
            pretrained_backbone_weights=pretrained_backbone_weights,
            pretrained_head_weights=pretrained_head_weights,
            init_weights=init_weights,
            lr_scheduler=lr_scheduler,
            online_mining=online_mining,
            hard_to_easy_ratio=hard_to_easy_ratio,
            min_hard_keypoints=min_hard_keypoints,
            max_hard_keypoints=max_hard_keypoints,
            loss_scale=loss_scale,
            optimizer=optimizer,
            learning_rate=learning_rate,
            amsgrad=amsgrad,
            negative_loss_weight=negative_loss_weight,
        )
        self.seg_output_stride = self.head_configs[
            self.model_type
        ].segmentation.output_stride

    def forward(self, img):
        """Forward pass of the model returning foreground-mask LOGITS."""
        img = torch.squeeze(img, dim=1).to(self.device)
        img = normalize_on_gpu(img)
        return self.model(img)["SegmentationHead"]

    def training_step(self, batch, batch_idx):
        """Training step (bce-dice on the centered-instance foreground mask)."""
        X = torch.squeeze(batch["instance_image"], dim=1)
        y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
        X = normalize_on_gpu(X)
        pred_fg = self.model(X)["SegmentationHead"]  # logits
        loss = compute_bce_dice_loss(pred_fg, y_fg)

        self.log(
            "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
        )
        self._accumulate_loss(loss)
        self.log("train/fg_loss", loss, on_step=False, on_epoch=True, sync_dist=True)
        return loss

    def validation_step(self, batch, batch_idx):
        """Validation step (val loss + foreground IoU)."""
        X = torch.squeeze(batch["instance_image"], dim=1)
        y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
        X = normalize_on_gpu(X)
        pred_fg = self.model(X)["SegmentationHead"]  # logits
        val_loss = compute_bce_dice_loss(pred_fg, y_fg)

        self.log(
            "val/loss",
            val_loss,
            prog_bar=True,
            on_step=False,
            on_epoch=True,
            sync_dist=True,
        )
        self.log("val/fg_loss", val_loss, on_step=False, on_epoch=True, sync_dist=True)

        # Each crop is one centered instance, so the per-crop foreground IoU IS a
        # per-instance mask-quality metric. Average PER-CROP (mean of per-image
        # IoUs) rather than pooling over the batch tensor.
        pred_fg_binary = (pred_fg > 0.0).float()
        dims = (1, 2, 3)
        intersection = (pred_fg_binary * y_fg).sum(dim=dims)
        union = pred_fg_binary.sum(dim=dims) + y_fg.sum(dim=dims) - intersection
        iou = (intersection / (union + 1e-6)).mean()
        self.log("val/fg_iou", iou, on_step=False, on_epoch=True, sync_dist=True)

        # Optional instance-level mask eval (enabled by SegmentationEvaluationCallback):
        # one crop == one instance, so emit the binarized predicted crop mask and the
        # GT crop mask as a single-instance pair on the SAME stride grid, appended in
        # lockstep for positional pairing by the callback.
        if self._collect_val_predictions:
            for i in range(X.shape[0]):
                pm = pred_fg_binary[i, 0].detach().cpu().numpy().astype(bool)
                gm = (y_fg[i, 0] > 0.5).detach().cpu().numpy()
                self.val_predictions.append({"masks": [pm] if pm.any() else []})
                self.val_ground_truth.append({"masks": [gm] if gm.any() else []})

    def get_visualization_data(
        self, sample, include_gt_mask: bool = False
    ) -> VisualizationData:
        """Crop-flavored viz: image + predicted foreground vs GT mask overlay.

        Args:
            sample: A sample dictionary from the data pipeline.
            include_gt_mask: If True, include the ground-truth centered-instance
                mask for a GT-vs-prediction overlay.
        """
        ex = sample.copy()
        for k, v in ex.items():
            if isinstance(v, torch.Tensor):
                ex[k] = v.to(device=self.device)
        ex["instance_image"] = ex["instance_image"].unsqueeze(dim=0)
        with torch.no_grad():
            logits = self.forward(ex["instance_image"])
            fg_prob = torch.sigmoid(logits)[0].cpu().numpy().transpose(1, 2, 0)
        img_np = ex["instance_image"][0, 0].cpu().numpy().transpose(1, 2, 0)
        gt_mask = None
        if include_gt_mask and "foreground_mask" in ex:
            gt_mask = ex["foreground_mask"].squeeze().cpu().numpy()  # (H, W)
        return VisualizationData(
            image=img_np,
            pred_confmaps=fg_prob,
            pred_peaks=np.zeros((0, 1, 2)),
            pred_peak_values=np.zeros((0,)),
            gt_instances=np.zeros((0, 1, 2)),
            node_names=["mask"],
            output_scale=fg_prob.shape[0] / img_np.shape[0],
            is_paired=False,
            gt_mask=gt_mask,
        )

    def visualize_example(self, sample):
        """Visualize the predicted foreground mask over the crop during training."""
        data = self.get_visualization_data(sample)
        scale = 1.0
        if data.image.shape[0] < 512:
            scale = 2.0
        if data.image.shape[0] < 256:
            scale = 4.0
        fig = plot_img(data.image, dpi=72 * scale, scale=scale)
        plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
        return fig

__init__(model_type, backbone_type, backbone_config, head_configs, pretrained_backbone_weights=None, pretrained_head_weights=None, init_weights='xavier', lr_scheduler=None, online_mining=False, hard_to_easy_ratio=2.0, min_hard_keypoints=2, max_hard_keypoints=None, loss_scale=5.0, optimizer='Adam', learning_rate=0.001, amsgrad=False, negative_loss_weight=1.0)

Initialise the configs and the model.

Source code in sleap_nn/training/lightning_modules.py

def __init__(
    self,
    model_type: str,
    backbone_type: str,
    backbone_config: Union[str, Dict[str, Any], DictConfig],
    head_configs: DictConfig,
    pretrained_backbone_weights: Optional[str] = None,
    pretrained_head_weights: Optional[str] = None,
    init_weights: Optional[str] = "xavier",
    lr_scheduler: Optional[Union[str, DictConfig]] = None,
    online_mining: Optional[bool] = False,
    hard_to_easy_ratio: Optional[float] = 2.0,
    min_hard_keypoints: Optional[int] = 2,
    max_hard_keypoints: Optional[int] = None,
    loss_scale: Optional[float] = 5.0,
    optimizer: Optional[str] = "Adam",
    learning_rate: Optional[float] = 1e-3,
    amsgrad: Optional[bool] = False,
    negative_loss_weight: Optional[float] = 1.0,
):
    """Initialise the configs and the model."""
    super().__init__(
        model_type=model_type,
        backbone_type=backbone_type,
        backbone_config=backbone_config,
        head_configs=head_configs,
        pretrained_backbone_weights=pretrained_backbone_weights,
        pretrained_head_weights=pretrained_head_weights,
        init_weights=init_weights,
        lr_scheduler=lr_scheduler,
        online_mining=online_mining,
        hard_to_easy_ratio=hard_to_easy_ratio,
        min_hard_keypoints=min_hard_keypoints,
        max_hard_keypoints=max_hard_keypoints,
        loss_scale=loss_scale,
        optimizer=optimizer,
        learning_rate=learning_rate,
        amsgrad=amsgrad,
        negative_loss_weight=negative_loss_weight,
    )
    self.seg_output_stride = self.head_configs[
        self.model_type
    ].segmentation.output_stride

forward(img)

Forward pass of the model returning foreground-mask LOGITS.

Source code in sleap_nn/training/lightning_modules.py

def forward(self, img):
    """Forward pass of the model returning foreground-mask LOGITS."""
    img = torch.squeeze(img, dim=1).to(self.device)
    img = normalize_on_gpu(img)
    return self.model(img)["SegmentationHead"]

get_visualization_data(sample, include_gt_mask=False)

Crop-flavored viz: image + predicted foreground vs GT mask overlay.

Parameters:

Name	Type	Description	Default
sample		A sample dictionary from the data pipeline.	required
include_gt_mask	bool	If True, include the ground-truth centered-instance mask for a GT-vs-prediction overlay.	False

Source code in sleap_nn/training/lightning_modules.py

def get_visualization_data(
    self, sample, include_gt_mask: bool = False
) -> VisualizationData:
    """Crop-flavored viz: image + predicted foreground vs GT mask overlay.

    Args:
        sample: A sample dictionary from the data pipeline.
        include_gt_mask: If True, include the ground-truth centered-instance
            mask for a GT-vs-prediction overlay.
    """
    ex = sample.copy()
    for k, v in ex.items():
        if isinstance(v, torch.Tensor):
            ex[k] = v.to(device=self.device)
    ex["instance_image"] = ex["instance_image"].unsqueeze(dim=0)
    with torch.no_grad():
        logits = self.forward(ex["instance_image"])
        fg_prob = torch.sigmoid(logits)[0].cpu().numpy().transpose(1, 2, 0)
    img_np = ex["instance_image"][0, 0].cpu().numpy().transpose(1, 2, 0)
    gt_mask = None
    if include_gt_mask and "foreground_mask" in ex:
        gt_mask = ex["foreground_mask"].squeeze().cpu().numpy()  # (H, W)
    return VisualizationData(
        image=img_np,
        pred_confmaps=fg_prob,
        pred_peaks=np.zeros((0, 1, 2)),
        pred_peak_values=np.zeros((0,)),
        gt_instances=np.zeros((0, 1, 2)),
        node_names=["mask"],
        output_scale=fg_prob.shape[0] / img_np.shape[0],
        is_paired=False,
        gt_mask=gt_mask,
    )

training_step(batch, batch_idx)

Training step (bce-dice on the centered-instance foreground mask).

Source code in sleap_nn/training/lightning_modules.py

def training_step(self, batch, batch_idx):
    """Training step (bce-dice on the centered-instance foreground mask)."""
    X = torch.squeeze(batch["instance_image"], dim=1)
    y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
    X = normalize_on_gpu(X)
    pred_fg = self.model(X)["SegmentationHead"]  # logits
    loss = compute_bce_dice_loss(pred_fg, y_fg)

    self.log(
        "loss", loss, prog_bar=True, on_step=True, on_epoch=False, sync_dist=True
    )
    self._accumulate_loss(loss)
    self.log("train/fg_loss", loss, on_step=False, on_epoch=True, sync_dist=True)
    return loss

validation_step(batch, batch_idx)

Validation step (val loss + foreground IoU).

Source code in sleap_nn/training/lightning_modules.py

def validation_step(self, batch, batch_idx):
    """Validation step (val loss + foreground IoU)."""
    X = torch.squeeze(batch["instance_image"], dim=1)
    y_fg = torch.squeeze(batch["foreground_mask"], dim=1)
    X = normalize_on_gpu(X)
    pred_fg = self.model(X)["SegmentationHead"]  # logits
    val_loss = compute_bce_dice_loss(pred_fg, y_fg)

    self.log(
        "val/loss",
        val_loss,
        prog_bar=True,
        on_step=False,
        on_epoch=True,
        sync_dist=True,
    )
    self.log("val/fg_loss", val_loss, on_step=False, on_epoch=True, sync_dist=True)

    # Each crop is one centered instance, so the per-crop foreground IoU IS a
    # per-instance mask-quality metric. Average PER-CROP (mean of per-image
    # IoUs) rather than pooling over the batch tensor.
    pred_fg_binary = (pred_fg > 0.0).float()
    dims = (1, 2, 3)
    intersection = (pred_fg_binary * y_fg).sum(dim=dims)
    union = pred_fg_binary.sum(dim=dims) + y_fg.sum(dim=dims) - intersection
    iou = (intersection / (union + 1e-6)).mean()
    self.log("val/fg_iou", iou, on_step=False, on_epoch=True, sync_dist=True)

    # Optional instance-level mask eval (enabled by SegmentationEvaluationCallback):
    # one crop == one instance, so emit the binarized predicted crop mask and the
    # GT crop mask as a single-instance pair on the SAME stride grid, appended in
    # lockstep for positional pairing by the callback.
    if self._collect_val_predictions:
        for i in range(X.shape[0]):
            pm = pred_fg_binary[i, 0].detach().cpu().numpy().astype(bool)
            gm = (y_fg[i, 0] > 0.5).detach().cpu().numpy()
            self.val_predictions.append({"masks": [pm] if pm.any() else []})
            self.val_ground_truth.append({"masks": [gm] if gm.any() else []})

visualize_example(sample)

Visualize the predicted foreground mask over the crop during training.

Source code in sleap_nn/training/lightning_modules.py

def visualize_example(self, sample):
    """Visualize the predicted foreground mask over the crop during training."""
    data = self.get_visualization_data(sample)
    scale = 1.0
    if data.image.shape[0] < 512:
        scale = 2.0
    if data.image.shape[0] < 256:
        scale = 4.0
    fig = plot_img(data.image, dpi=72 * scale, scale=scale)
    plot_confmaps(data.pred_confmaps, output_scale=data.output_scale)
    return fig

set_embedding_burn_in_from_config(module, config)

Honor data_config.preprocessing.burn_in on an EmbeddingLightningModule.

burn_in is a pure runtime toggle (read in training_step / forward / validation_step), so setting it after construction is safe. The canonical default is False (matching PreprocessingConfig.burn_in and the LM __init__); BOTH the training factory (:meth:LightningModel.get_lightning_model_from_config) and the inference loader (:func:sleap_nn.inference.loaders._load_lightning_module) call this so a model trained maskless (burn_in=False, the "centroid-crop" objective) also runs maskless at inference — no train/inference mismatch.

When burn_in is False the crop is NOT multiplied by the mask and _standardize falls back to whole-crop standardize (the full square crop around the centroid / mask-COM is kept, background included).

Also honors data_config.preprocessing.background_fill (what the masked-out background is replaced with when burn_in is on): black (default, the original mask-multiply), grey (mid-grey), mean (foreground mean — neutral, equivalent to black for standardized inputs), or noise (per-pixel noise).

Parameters:

Name	Type	Description	Default
module		The constructed EmbeddingLightningModule.	required
config		The full training config (carries data_config.preprocessing).	required

Source code in sleap_nn/training/lightning_modules.py

def set_embedding_burn_in_from_config(module, config) -> None:
    """Honor ``data_config.preprocessing.burn_in`` on an ``EmbeddingLightningModule``.

    ``burn_in`` is a pure runtime toggle (read in ``training_step`` / ``forward`` /
    ``validation_step``), so setting it after construction is safe. The canonical
    default is ``False`` (matching ``PreprocessingConfig.burn_in`` and the LM
    ``__init__``); BOTH the training factory
    (:meth:`LightningModel.get_lightning_model_from_config`) and the inference loader
    (:func:`sleap_nn.inference.loaders._load_lightning_module`) call this so a model
    trained maskless (``burn_in=False``, the "centroid-crop" objective) also runs
    maskless at inference — no train/inference mismatch.

    When ``burn_in`` is ``False`` the crop is NOT multiplied by the mask and
    ``_standardize`` falls back to whole-crop standardize (the full square crop around
    the centroid / mask-COM is kept, background included).

    Also honors ``data_config.preprocessing.background_fill`` (what the masked-out
    background is replaced with when ``burn_in`` is on): ``black`` (default, the
    original mask-multiply), ``grey`` (mid-grey), ``mean`` (foreground mean — neutral,
    equivalent to ``black`` for standardized inputs), or ``noise`` (per-pixel noise).

    Args:
        module: The constructed ``EmbeddingLightningModule``.
        config: The full training config (carries ``data_config.preprocessing``).
    """
    module.burn_in = bool(
        OmegaConf.select(config, "data_config.preprocessing.burn_in", default=False)
    )
    background_fill = OmegaConf.select(
        config, "data_config.preprocessing.background_fill", default="black"
    )
    if background_fill not in _EMBEDDING_BACKGROUND_FILLS:
        message = (
            f"Unknown data_config.preprocessing.background_fill "
            f"'{background_fill}'; choose one of "
            f"{'|'.join(_EMBEDDING_BACKGROUND_FILLS)}."
        )
        logger.error(message)
        raise ValueError(message)
    module.background_fill = background_fill

validate_embedding_identity(objective, identity, has_identities=False, batch_groups=None)

Enforce the identity-equality gates for the embedding objective.

Each positive/negative source silently asserts "same / different animal"; a wrong assertion trains the appearance model on label noise (pulling two different animals together, or pushing the same animal apart). The data_config.identity block DECLARES the data's semantics so the objective can be validated against it:

• pos aug_view: two views = same id (same crop) -> no assumption.
• pos tracklet: same (video, track) = same animal -> warn if identity.tracks_are_proofread is False (tracker swaps poison training).
• pos global_id: same GLOBAL identity = same animal. Grounded by a real sio.Identity when the data carries one (has_identities); otherwise falls back to the track name, so it errors when the data has no identities AND identity.track_names_are_global is False (names must be globally consistent).
• neg same_frame: two detections / frame = different animals -> warn if identity.detections_deduplicated is False (a double / over-segmented detection becomes a hard negative against itself).
• negatives at all: with batch_groups (the training crops' groups), error if no batch the sampler can draw holds a single negative pair (one group, or one group per video under restrict_same_video / within_video, or groups_per_batch=1): the contrastive loss would have nothing to train on.

Absent objective fields are filled by :func:~sleap_nn.config.model_config.resolve_embedding_objective (the same resolution :class:EmbeddingLightningModule uses); absent identity fields take the conservative :class:IdentityConfig defaults.

Parameters:

Name	Type	Description	Default
objective		The head_configs.embedding.embedding.objective node (DictConfig or None — defaults assumed when absent). Reads positives.scope, negatives.sources / restrict_same_video and sampler.kind.	required
identity		The data_config.identity node (DictConfig or None — defaults assumed when absent).	required
has_identities	bool	Whether the training labels carry global sio.Identity annotations. When True, global_id grouping is grounded by the real identities and needs no track_names_are_global promise.	False
batch_groups	Optional[Tuple[ndarray, ndarray, ndarray]]	Optional (group_ids, video_ids, frame_ids) of the training crops (:func:~sleap_nn.data.custom_datasets.embedding_training_groups). When given, the objective is checked for being able to draw negatives.	None

Raises:

Type	Description
ValueError	if positives.scope='global_id' but the data carries no sio.Identity annotations AND does not declare identity.track_names_are_global=True; if no training batch can hold a negative pair; or if the objective has an unknown option value (see resolve_embedding_objective).

Source code in sleap_nn/training/lightning_modules.py

def validate_embedding_identity(
    objective,
    identity,
    has_identities: bool = False,
    batch_groups: Optional[Tuple[np.ndarray, np.ndarray, np.ndarray]] = None,
):
    """Enforce the identity-equality gates for the embedding objective.

    Each positive/negative source silently asserts "same / different animal"; a wrong
    assertion trains the appearance model on label noise (pulling two *different*
    animals together, or pushing the *same* animal apart). The ``data_config.identity``
    block DECLARES the data's semantics so the objective can be validated against it:

    - pos ``aug_view``: two views = same id (same crop) -> no assumption.
    - pos ``tracklet``: same ``(video, track)`` = same animal -> **warn** if
      ``identity.tracks_are_proofread`` is False (tracker swaps poison training).
    - pos ``global_id``: same GLOBAL identity = same animal. Grounded by a real
      ``sio.Identity`` when the data carries one (``has_identities``); otherwise falls
      back to the track name, so it **errors** when the data has no identities AND
      ``identity.track_names_are_global`` is False (names must be globally consistent).
    - neg ``same_frame``: two detections / frame = different animals -> **warn** if
      ``identity.detections_deduplicated`` is False (a double / over-segmented
      detection becomes a hard negative against itself).
    - negatives at all: with ``batch_groups`` (the training crops' groups), **error**
      if no batch the sampler can draw holds a single negative pair (one group, or
      one group per video under ``restrict_same_video`` / ``within_video``, or
      ``groups_per_batch=1``): the contrastive loss would have nothing to train on.

    Absent objective fields are filled by
    :func:`~sleap_nn.config.model_config.resolve_embedding_objective` (the same
    resolution :class:`EmbeddingLightningModule` uses); absent identity fields take
    the conservative :class:`IdentityConfig` defaults.

    Args:
        objective: The ``head_configs.embedding.embedding.objective`` node (``DictConfig``
            or ``None`` — defaults assumed when absent). Reads ``positives.scope``,
            ``negatives.sources`` / ``restrict_same_video`` and ``sampler.kind``.
        identity: The ``data_config.identity`` node (``DictConfig`` or ``None`` —
            defaults assumed when absent).
        has_identities: Whether the training labels carry global ``sio.Identity``
            annotations. When ``True``, ``global_id`` grouping is grounded by the real
            identities and needs no ``track_names_are_global`` promise.
        batch_groups: Optional ``(group_ids, video_ids, frame_ids)`` of the training
            crops (:func:`~sleap_nn.data.custom_datasets.embedding_training_groups`).
            When given, the objective is checked for being able to draw negatives.

    Raises:
        ValueError: if ``positives.scope='global_id'`` but the data carries no
            ``sio.Identity`` annotations AND does not declare
            ``identity.track_names_are_global=True``; if no training batch can hold
            a negative pair; or if the objective has an unknown option value (see
            ``resolve_embedding_objective``).
    """
    # Resolve objective semantics exactly as the LightningModule does (one resolver,
    # so a default cannot mean one thing here and another there).
    resolved = resolve_embedding_objective(objective)
    scope = resolved.positives.scope
    neg_sources = list(resolved.negatives.sources)

    # Declared data semantics (default to the conservative IdentityConfig defaults).
    if identity is not None:
        tracks_are_proofread = bool(
            OmegaConf.select(identity, "tracks_are_proofread", default=False)
        )
        track_names_are_global = bool(
            OmegaConf.select(identity, "track_names_are_global", default=False)
        )
        detections_deduplicated = bool(
            OmegaConf.select(identity, "detections_deduplicated", default=True)
        )
    else:
        tracks_are_proofread = False
        track_names_are_global = False
        detections_deduplicated = True

    if scope == "global_id" and not (track_names_are_global or has_identities):
        raise ValueError(
            "head_configs.embedding.embedding.objective.positives.scope='global_id' "
            "requires either (a) the labels to carry global `sio.Identity` annotations "
            "(the ground-truth cross-video animal identity), or (b) "
            "data_config.identity.track_names_are_global=True (the same track name must "
            "mean the same animal across videos). Add identities to your labels, set "
            "track_names_are_global=True if your track names are globally consistent, "
            "or use scope='tracklet' (video-local) instead."
        )
    if scope == "tracklet" and not tracks_are_proofread:
        logger.warning(
            "embedding objective positives.scope='tracklet' but "
            "data_config.identity.tracks_are_proofread=False: tracker swaps will pull "
            "DIFFERENT animals together (training on label noise)."
        )
    if scope == "tracklet" and "in_batch" in neg_sources:
        # Under tracklet scope identity is per-(video, track), so the same animal in two
        # different videos is NOT a positive. With in-batch negatives and
        # restrict_same_video=False those cross-video same-animal pairs become hard
        # negatives — training the model to push the same animal apart across videos,
        # the exact opposite of cross-session re-ID. The invariant is documented on both
        # NegativesConfig.restrict_same_video and build_contrastive_masks; enforce it.
        if not resolved.negatives.restrict_same_video:
            raise ValueError(
                "embedding objective positives.scope='tracklet' with 'in_batch' "
                "negatives REQUIRES "
                "head_configs.embedding.embedding.objective.negatives."
                "restrict_same_video=True. Without it, two crops of the SAME animal in "
                "different videos become in-batch hard negatives and the model is "
                "trained to push that animal apart across videos. Set "
                "restrict_same_video=True, or use a video-global scope='global_id' with "
                "globally-consistent track names."
            )
    if scope == "tracklet":
        # `restrict_same_video=True` (required above) makes every cross-video pair a
        # non-negative, so the batch must actually CONTAIN same-video crops for an
        # anchor to have any negative at all. The `pk` sampler draws its P groups
        # globally, so most of a batch is cross-video and a large share of anchors
        # end up with zero negatives -- measured at 29-96% depending on how many
        # videos the data spans, against 0% for `within_video`, which draws each
        # batch from one video. Not a correctness invariant (the loss masks out
        # zero-negative anchors), so a warning rather than an error.
        if resolved.sampler.kind == "pk":
            logger.warning(
                "embedding objective positives.scope='tracklet' with "
                "sampler.kind='pk': PK draws groups across ALL videos while "
                "restrict_same_video=True discards cross-video negatives, so many "
                "anchors will see NO negatives and contribute nothing to the loss "
                "(it averages only the anchors that have one). Use "
                "sampler.kind='within_video' for tracklet-scope training."
            )

    if scope == "aug_view" and resolved.sampler.kind != "random":
        # Under aug_view every detection is its own group, so a group-aware sampler
        # draws P crops and repeats each K times: a batch of P distinct crops.
        logger.warning(
            "embedding objective positives.scope='aug_view' makes every detection its "
            f"own group, so sampler.kind='{resolved.sampler.kind}' fills a batch with "
            f"{resolved.sampler.samples_per_group} views of each of "
            f"{resolved.sampler.groups_per_batch} crops. Use sampler.kind='random' for "
            "batches of distinct crops."
        )

    if "same_frame" in neg_sources and not detections_deduplicated:
        logger.warning(
            "embedding objective negatives include 'same_frame' but "
            "data_config.identity.detections_deduplicated=False: a double / "
            "over-segmented detection will be treated as a hard negative against itself."
        )

    if batch_groups is not None:
        reason = _why_no_batch_has_negatives(resolved, *batch_groups)
        if reason is not None:
            raise ValueError(
                "The embedding objective can never see a negative pair: " + reason
            )

validate_embedding_views(objective, use_augmentations_train, augmentation_config)

Check that the two contrastive views of a crop can differ.

The training step contrasts two views of every crop; they are two independent draws of the configured augmentation (EmbeddingDataset._pack_sample). With use_augmentations_train=False, or no augmentation that can fire, both views are the same tensor (EmbeddingLightningModule._two_views reuses view 1).

• positives.scope='aug_view': the other view is the ONLY positive, so each crop's positive is itself and the objective reduces to telling crops apart pixel for pixel -> error.
• other scopes: identity positives still carry the signal, but the model learns no invariance to the augmentations -> warning.

Parameters:

Name	Type	Description	Default
objective		The head_configs.embedding.embedding.objective node.	required
use_augmentations_train	bool	data_config.use_augmentations_train.	required
augmentation_config		data_config.augmentation_config.	required

Raises:

Type	Description
ValueError	For aug_view with identical views.

Source code in sleap_nn/training/lightning_modules.py

def validate_embedding_views(
    objective, use_augmentations_train: bool, augmentation_config
) -> None:
    """Check that the two contrastive views of a crop can differ.

    The training step contrasts two views of every crop; they are two independent
    draws of the configured augmentation (``EmbeddingDataset._pack_sample``). With
    ``use_augmentations_train=False``, or no augmentation that can fire, both views
    are the same tensor (``EmbeddingLightningModule._two_views`` reuses view 1).

    - ``positives.scope='aug_view'``: the other view is the ONLY positive, so each
      crop's positive is itself and the objective reduces to telling crops apart
      pixel for pixel -> **error**.
    - other scopes: identity positives still carry the signal, but the model learns no
      invariance to the augmentations -> **warning**.

    Args:
        objective: The ``head_configs.embedding.embedding.objective`` node.
        use_augmentations_train: ``data_config.use_augmentations_train``.
        augmentation_config: ``data_config.augmentation_config``.

    Raises:
        ValueError: For ``aug_view`` with identical views.
    """
    if use_augmentations_train and not _augmentation_is_noop(augmentation_config):
        return
    scope = resolve_embedding_objective(objective).positives.scope
    cause = (
        "data_config.use_augmentations_train=False"
        if not use_augmentations_train
        else "data_config.augmentation_config has no augmentation that can fire"
    )
    if scope == "aug_view":
        raise ValueError(
            f"positives.scope='aug_view' needs two different views of each crop, but "
            f"{cause}, so both views are the same tensor and each crop's only "
            "positive is itself. Enable use_augmentations_train with an intensity "
            "and/or geometric augmentation."
        )
    logger.warning(
        f"embedding training: {cause}, so the two contrastive views of each crop are "
        "identical. Identity positives still train the model, but it learns no "
        "invariance to the augmentations."
    )