model_config

sleap_nn.config.model_config

Serializable configuration classes for specifying all model config parameters.

These configuration classes are intended to specify all the parameters required to initialize the model config.

Classes:

Name	Description
BBoxCenterHeadConfig	Center-heatmap head of a bbox_centernet model.
BBoxCenterNetConfig	Head config for the bbox_centernet (CenterNet-style) model type.
BBoxOffsetHeadConfig	Sub-cell center offset head of a bbox_centernet model.
BBoxSSDConfig	Head config for the bbox_ssd (SSD-style, anchor-based) model type.
BBoxSizeHeadConfig	Box-size regression head of a bbox_centernet model.
BackboneConfig	Configurations related to model backbone configuration.
BottomUpConfMapsConfig	Bottomup configuration map.
BottomUpConfig	bottomup head_config.
BottomUpMultiClassConfig	Head config for BottomUp Id models.
BottomUpSegmentationConfig	Head config for bottom-up instance segmentation models.
CenterOffsetConfig	Configuration for the center offset head.
CenteredInstanceConfMapsConfig	Centered Instance configuration map.
CenteredInstanceConfig	centered_instance head_config.
CenteredInstanceSegmentationConfig	Head config for top-down crop-centered instance segmentation models.
CenteredInstanceSegmentationHeadConfig	Foreground-mask head config for top-down crop-centered segmentation.
CentroidConfMapsConfig	Centroid configuration map.
CentroidConfig	centroid head_config.
ClassMapConfig	Class map head config.
ClassVectorsConfig	Configurations for class vectors heads.
ConvNextBaseConfig	Convnext configuration for backbone.
ConvNextConfig	Convnext configuration for backbone.
ConvNextLargeConfig	Convnext configuration for backbone.
ConvNextSmallConfig	Convnext configuration for backbone.
EmbeddingConfig	Head config for the embedding (crop -> vector, re-ID) model type.
EmbeddingHeadConfig	Configuration for the embedding head (the adapter on a pooled encoder feature).
HeadConfig	Configurations related to the model output head type.
InstanceCenterConfig	Configuration for the instance center heatmap head.
LossConfig	Contrastive loss for the embedding objective (the loss axis).
ModelConfig	Configurations related to model architecture.
NegativesConfig	Negative-pair eligibility for the embedding objective.
ObjectiveConfig	Pluggable training objective = positives x negatives x loss.
PAFConfig	PAF configuration map.
PositivesConfig	Positive-pair sampling for the embedding objective.
PretrainedConfig	Configuration for an external pretrained backbone (HuggingFace).
SSDHeadConfig	Multi-scale anchor head of a bbox_ssd (SSD-style) model.
SamplerConfig	Group-aware batch sampler that realizes the objective.
SegmentationHeadConfig	Configuration for the foreground segmentation head.
SemanticSegmentationConfig	Head config for whole-frame semantic (foreground) segmentation models.
SingleInstanceConfMapsConfig	Single Instance configuration map.
SingleInstanceConfig	single instance head_config.
SwinTBaseConfig	SwinT configuration for backbone.
SwinTConfig	SwinT configuration (tiny) for backbone.
SwinTSmallConfig	SwinT configuration (small) for backbone.
TopDownCenteredInstanceMultiClassConfig	Head config for TopDown centered instance ID models.
UNetConfig	UNet config for backbone.
UNetLargeRFConfig	UNet config for backbone with large receptive field.
UNetMediumRFConfig	UNet config for backbone with medium receptive field.

Functions:

Name	Description
drop_removed_embedding_keys	Remove the removed embedding keys from an embedding head leaf, in place.
model_mapper	Map the legacy model configuration to the new model configuration.
resolve_embedding_objective	Return the embedding objective with every sub-config filled in and validated.
ssd_level_strides	Pyramid level strides of an SSDHeadConfig (or its DictConfig), finest first.

BBoxCenterHeadConfig

Center-heatmap head of a bbox_centernet model.

One sigmoid heatmap channel per category with a Gaussian peak at each box center (CenterNet, Zhou et al. 2019), trained with the penalty-reduced pixelwise focal loss of CornerNet/CenterNet.

Attributes:

Name	Type	Description
categories	Optional[List[str]]	(List[str]) Ordered class names; one heatmap channel per entry. None (default) is filled in at training setup from the categories of the training boxes, in first-seen order; boxes without a category belong to the single class "object".
sigma	Optional[float]	(float) Fixed Gaussian sigma in input pixels. None (default) uses CenterNet's size-adaptive radius instead: the largest radius at which a box whose corners move by that much still has IoU >= min_overlap with the GT box, and sigma = (2 * radius + 1) / 6.
min_overlap	float	(float) IoU used by the size-adaptive radius. Default: 0.7.
output_stride	int	(int) Stride of the heatmap relative to the input image. Must equal the size and offset strides. Default: 4.
loss_weight	float	(float) Weight of the heatmap focal loss. Default: 1.0.
focal_alpha	float	(float) Focal exponent on the prediction error. Default: 2.0.
focal_beta	float	(float) Exponent of the penalty reduction around GT peaks. Default: 4.0.

Source code in sleap_nn/config/model_config.py

@define
class BBoxCenterHeadConfig:
    """Center-heatmap head of a ``bbox_centernet`` model.

    One sigmoid heatmap channel per category with a Gaussian peak at each box
    center (CenterNet, Zhou et al. 2019), trained with the penalty-reduced pixelwise
    focal loss of CornerNet/CenterNet.

    Attributes:
        categories: (List[str]) Ordered class names; one heatmap channel per entry.
            ``None`` (default) is filled in at training setup from the categories of
            the training boxes, in first-seen order; boxes without a category belong
            to the single class ``"object"``.
        sigma: (float) Fixed Gaussian sigma in input pixels. ``None`` (default) uses
            CenterNet's size-adaptive radius instead: the largest radius at which a
            box whose corners move by that much still has IoU >= ``min_overlap``
            with the GT box, and ``sigma = (2 * radius + 1) / 6``.
        min_overlap: (float) IoU used by the size-adaptive radius. *Default*: ``0.7``.
        output_stride: (int) Stride of the heatmap relative to the input image.
            Must equal the ``size`` and ``offset`` strides. *Default*: ``4``.
        loss_weight: (float) Weight of the heatmap focal loss. *Default*: ``1.0``.
        focal_alpha: (float) Focal exponent on the prediction error. *Default*:
            ``2.0``.
        focal_beta: (float) Exponent of the penalty reduction around GT peaks.
            *Default*: ``4.0``.
    """

    categories: Optional[List[str]] = None
    sigma: Optional[float] = None
    min_overlap: float = 0.7
    output_stride: int = 4
    loss_weight: float = 1.0
    focal_alpha: float = 2.0
    focal_beta: float = 4.0

BBoxCenterNetConfig

Head config for the bbox_centernet (CenterNet-style) model type.

A whole-frame, anchor-free detector: a center heatmap plus per-center size and offset regression. Decoding takes the heatmap's local maxima as detections and reads each one's box from the size/offset maps (no anchors, no grouping).

Source code in sleap_nn/config/model_config.py

@define
class BBoxCenterNetConfig:
    """Head config for the ``bbox_centernet`` (CenterNet-style) model type.

    A whole-frame, anchor-free detector: a center heatmap plus per-center size and
    offset regression. Decoding takes the heatmap's local maxima as detections and
    reads each one's box from the size/offset maps (no anchors, no grouping).
    """

    center: Optional[BBoxCenterHeadConfig] = None
    size: Optional[BBoxSizeHeadConfig] = None
    offset: Optional[BBoxOffsetHeadConfig] = None

BBoxOffsetHeadConfig

Sub-cell center offset head of a bbox_centernet model.

Two channels recovering the center's position inside its output-grid cell, i.e. the quantization error of the stride, with an L1 loss masked to center cells.

Attributes:

Name	Type	Description
output_stride	int	(int) Must equal the center head's stride. Default: 4.
loss_weight	float	(float) Weight of the offset L1 loss. Default: 1.0.

Source code in sleap_nn/config/model_config.py

@define
class BBoxOffsetHeadConfig:
    """Sub-cell center offset head of a ``bbox_centernet`` model.

    Two channels recovering the center's position inside its output-grid cell,
    i.e. the quantization error of the stride, with an L1 loss masked to center
    cells.

    Attributes:
        output_stride: (int) Must equal the center head's stride. *Default*: ``4``.
        loss_weight: (float) Weight of the offset L1 loss. *Default*: ``1.0``.
    """

    output_stride: int = 4
    loss_weight: float = 1.0

BBoxSSDConfig

Head config for the bbox_ssd (SSD-style, anchor-based) model type.

Wraps the single multi-scale ssd leaf (mirrors SemanticSegmentationConfig).

Source code in sleap_nn/config/model_config.py

@define
class BBoxSSDConfig:
    """Head config for the ``bbox_ssd`` (SSD-style, anchor-based) model type.

    Wraps the single multi-scale ``ssd`` leaf (mirrors ``SemanticSegmentationConfig``).
    """

    ssd: Optional[SSDHeadConfig] = None

BBoxSizeHeadConfig

Box-size regression head of a bbox_centernet model.

Two channels (width, height) regressed at each box's center cell, in units of output-grid cells, with an L1 loss masked to center cells.

Attributes:

Name	Type	Description
output_stride	int	(int) Must equal the center head's stride. Default: 4.
loss_weight	float	(float) Weight of the size L1 loss. Default: 0.1 (CenterNet's value; sizes are large numbers next to the focal loss).

Source code in sleap_nn/config/model_config.py

@define
class BBoxSizeHeadConfig:
    """Box-size regression head of a ``bbox_centernet`` model.

    Two channels (width, height) regressed at each box's center cell, in units of
    output-grid cells, with an L1 loss masked to center cells.

    Attributes:
        output_stride: (int) Must equal the center head's stride. *Default*: ``4``.
        loss_weight: (float) Weight of the size L1 loss. *Default*: ``0.1``
            (CenterNet's value; sizes are large numbers next to the focal loss).
    """

    output_stride: int = 4
    loss_weight: float = 0.1

BackboneConfig

Configurations related to model backbone configuration.

Attributes:

Name	Type	Description
unet	Optional[UNetConfig]	An instance of UNetConfig.
convnext	Optional[ConvNextConfig]	An instance of ConvNextConfig.
swint	Optional[SwinTConfig]	An instance of SwinTConfig.
pretrained	Optional[PretrainedConfig]	An instance of PretrainedConfig (external HuggingFace pretrained backbone).

Source code in sleap_nn/config/model_config.py

@oneof
@define
class BackboneConfig:
    """Configurations related to model backbone configuration.

    Attributes:
        unet: An instance of `UNetConfig`.
        convnext: An instance of `ConvNextConfig`.
        swint: An instance of `SwinTConfig`.
        pretrained: An instance of `PretrainedConfig` (external HuggingFace
            pretrained backbone).
    """

    unet: Optional[UNetConfig] = None
    convnext: Optional[ConvNextConfig] = None
    swint: Optional[SwinTConfig] = None
    pretrained: Optional[PretrainedConfig] = None

BottomUpConfMapsConfig

Bottomup configuration map.

Attributes:

Name	Type	Description
part_names	Optional[List[str]]	(List[str]) None if nodes from sio.Labels file can be used directly. Else provide text name of the body parts (nodes) that the head will be configured to produce. The number of parts determines the number of channels in the output. If not specified, all body parts in the skeleton will be used. This config does not apply for 'PartAffinityFieldsHead'.
sigma	float	(float) Spread of the Gaussian distribution of the confidence maps as a scalar float. Smaller values are more precise but may be difficult to learn as they have a lower density within the image space. Larger values are easier to learn but are less precise with respect to the peak coordinate. This spread is in units of pixels of the model input image, i.e., the image resolution after any input scaling is applied.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
loss_weight	Optional[float]	(float) Scalar float used to weigh the loss term for this head during training. Increase this to encourage the optimization to focus on improving this specific output in multi-head models.

Source code in sleap_nn/config/model_config.py

@define
class BottomUpConfMapsConfig:
    """Bottomup configuration map.

    Attributes:
        part_names: (List[str]) None if nodes from sio.Labels file can be used directly.
            Else provide text name of the body parts (nodes) that the head will be
            configured to produce. The number of parts determines the number of channels
            in the output. If not specified, all body parts in the skeleton will be used.
            This config does not apply for 'PartAffinityFieldsHead'.
        sigma: (float) Spread of the Gaussian distribution of the confidence maps as a
            scalar float. Smaller values are more precise but may be difficult to learn
            as they have a lower density within the image space. Larger values are easier
            to learn but are less precise with respect to the peak coordinate. This spread
            is in units of pixels of the model input image, i.e., the image resolution
            after any input scaling is applied.
        output_stride: (int) The stride of the output confidence maps relative to the
            input image. This is the reciprocal of the resolution, e.g., an output stride
            of 2 results in confidence maps that are 0.5x the size of the input.
            Increasing this value can considerably speed up model performance and
            decrease memory requirements, at the cost of decreased spatial resolution.
        loss_weight: (float) Scalar float used to weigh the loss term for this head
            during training. Increase this to encourage the optimization to focus on
            improving this specific output in multi-head models.
    """

    part_names: Optional[List[str]] = None
    sigma: float = 5.0
    output_stride: int = 1
    loss_weight: Optional[float] = None

BottomUpConfig

bottomup head_config.

Source code in sleap_nn/config/model_config.py

@define
class BottomUpConfig:
    """bottomup head_config."""

    confmaps: Optional[BottomUpConfMapsConfig] = None
    pafs: Optional[PAFConfig] = None

BottomUpMultiClassConfig

Head config for BottomUp Id models.

Source code in sleap_nn/config/model_config.py

@define
class BottomUpMultiClassConfig:
    """Head config for BottomUp Id models."""

    confmaps: Optional[BottomUpConfMapsConfig] = None
    class_maps: Optional[ClassMapConfig] = None

BottomUpSegmentationConfig

Head config for bottom-up instance segmentation models.

Source code in sleap_nn/config/model_config.py

@define
class BottomUpSegmentationConfig:
    """Head config for bottom-up instance segmentation models."""

    segmentation: Optional[SegmentationHeadConfig] = None
    center: Optional[InstanceCenterConfig] = None
    offsets: Optional[CenterOffsetConfig] = None

CenterOffsetConfig

Configuration for the center offset head.

Attributes:

Name	Type	Description
output_stride	int	(int) The stride of the output offset maps relative to the input image. Default: 2.
loss_weight	float	(float) Scalar float used to weigh the loss term for this head during training. Default: 0.1.

Source code in sleap_nn/config/model_config.py

@define
class CenterOffsetConfig:
    """Configuration for the center offset head.

    Attributes:
        output_stride: (int) The stride of the output offset maps relative to the
            input image. Default: 2.
        loss_weight: (float) Scalar float used to weigh the loss term for this head
            during training. Default: 0.1.
    """

    output_stride: int = 2
    loss_weight: float = 0.1

CenteredInstanceConfMapsConfig

Centered Instance configuration map.

Attributes:

Name	Type	Description
part_names	Optional[List[str]]	(List[str]) None if nodes from sio.Labels file can be used directly. Else provide text name of the body parts (nodes) that the head will be configured to produce. The number of parts determines the number of channels in the output. If not specified, all body parts in the skeleton will be used. This config does not apply for 'PartAffinityFieldsHead'.
anchor_part	Optional[str]	(str) Node name to use as the anchor point. If None, the NaN-ignoring mean of all visible instance nodes will be used as the anchor. The same mean-of-visible-nodes fallback is used when the anchor part is specified but not visible in the instance. Setting a reliable anchor point can significantly improve topdown model accuracy as they benefit from a consistent geometry of the body parts relative to the center of the image. Default is None.
centroid_method	Optional[str]	(str) How the crop center is derived from the instance's points, spelled as in sio.Instance.to_centroid: "center_of_mass" (mean of visible nodes), "bbox_center" (midpoint of the visible nodes' bounding box), "geometric_median" (Weiszfeld median — the least affected by a MISLOCALIZED node; measured on real pose data, one node off by a body length moves it ~1.7x less than the mean and ~5x less than the bbox midpoint. Not more stable than the mean under node dropout), or "anchor" (the anchor_part node). None (default) infers it: "anchor" when anchor_part is set, else "center_of_mass" — i.e. exactly the historical behavior, so existing configs are unchanged. Setting both anchor_part and a non-anchor centroid_method is an error (they name different centroids); use centroid_fallback for that. Default is None.
centroid_fallback	Optional[str]	(str) The reduce method used when anchor_part is configured but that node is not visible: "center_of_mass" (default), "bbox_center" or "geometric_median". Only meaningful for the anchor method. Unlike sio's fallback=None, sleap-nn always falls back rather than emitting a NaN centroid. Default is None (= "center_of_mass").
sigma	float	(float) Spread of the Gaussian distribution of the confidence maps as a scalar float. Smaller values are more precise but may be difficult to learn as they have a lower density within the image space. Larger values are easier to learn but are less precise with respect to the peak coordinate. This spread is in units of pixels of the model input image, i.e., the image resolution after any input scaling is applied.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
loss_weight	float	(float) Scalar float used to weigh the loss term for this head during training. Increase this to encourage the optimization to focus on improving this specific output in multi-head models.

Source code in sleap_nn/config/model_config.py

@define
class CenteredInstanceConfMapsConfig:
    """Centered Instance configuration map.

    Attributes:
        part_names: (List[str]) None if nodes from sio.Labels file can be used directly.
            Else provide text name of the body parts (nodes) that the head will be
            configured to produce. The number of parts determines the number of channels
            in the output. If not specified, all body parts in the skeleton will be used.
            This config does not apply for 'PartAffinityFieldsHead'.
        anchor_part: (str) Node name to use as the anchor point. If None, the
            NaN-ignoring mean of all visible instance nodes will be used as the
            anchor. The same mean-of-visible-nodes fallback is used when the
            anchor part is specified but not visible in the instance. Setting a
            reliable anchor point can significantly improve topdown model
            accuracy as they benefit from a consistent geometry of the body parts
            relative to the center of the image. Default is None.
        centroid_method: (str) How the crop center is derived from the instance's
            points, spelled as in ``sio.Instance.to_centroid``:
            ``"center_of_mass"`` (mean of visible nodes), ``"bbox_center"``
            (midpoint of the visible nodes' bounding box), ``"geometric_median"``
            (Weiszfeld median — the least affected by a MISLOCALIZED node; measured
            on real pose data, one node off by a body length moves it ~1.7x less
            than the mean and ~5x less than the bbox midpoint. Not more stable
            than the mean under node dropout), or ``"anchor"``
            (the ``anchor_part`` node). ``None`` (default) infers it: ``"anchor"``
            when ``anchor_part`` is set, else ``"center_of_mass"`` — i.e. exactly
            the historical behavior, so existing configs are unchanged. Setting
            both ``anchor_part`` and a non-anchor ``centroid_method`` is an error
            (they name different centroids); use ``centroid_fallback`` for that.
            Default is None.
        centroid_fallback: (str) The reduce method used when ``anchor_part`` is
            configured but that node is not visible: ``"center_of_mass"``
            (default), ``"bbox_center"`` or ``"geometric_median"``. Only
            meaningful for the anchor method. Unlike ``sio``'s ``fallback=None``,
            sleap-nn always falls back rather than emitting a NaN centroid.
            Default is None (= ``"center_of_mass"``).
        sigma: (float) Spread of the Gaussian distribution of the confidence maps as a
            scalar float. Smaller values are more precise but may be difficult to learn
            as they have a lower density within the image space. Larger values are
            easier to learn but are less precise with respect to the peak coordinate.
            This spread is in units of pixels of the model input image, i.e., the image
            resolution after any input scaling is applied.
        output_stride: (int) The stride of the output confidence maps relative to the
            input image. This is the reciprocal of the resolution, e.g., an output
            stride of 2 results in confidence maps that are 0.5x the size of the input.
            Increasing this value can considerably speed up model performance and
            decrease memory requirements, at the cost of decreased spatial resolution.
        loss_weight: (float) Scalar float used to weigh the loss term for this head
            during training. Increase this to encourage the optimization to focus on
            improving this specific output in multi-head models.
    """

    part_names: Optional[List[str]] = None
    anchor_part: Optional[str] = None
    centroid_method: Optional[str] = None
    centroid_fallback: Optional[str] = None
    sigma: float = 5.0
    output_stride: int = 1
    loss_weight: float = 1.0

CenteredInstanceConfig

centered_instance head_config.

Source code in sleap_nn/config/model_config.py

@define
class CenteredInstanceConfig:
    """centered_instance head_config."""

    confmaps: Optional[CenteredInstanceConfMapsConfig] = None

CenteredInstanceSegmentationConfig

Head config for top-down crop-centered instance segmentation models.

A single foreground-mask head predicting the centered instance's mask on a centroid crop (the segmentation analog of centered_instance; composed with a centroid model for full top-down inference).

Source code in sleap_nn/config/model_config.py

@define
class CenteredInstanceSegmentationConfig:
    """Head config for top-down crop-centered instance segmentation models.

    A single foreground-mask head predicting the *centered* instance's mask on a
    centroid crop (the segmentation analog of ``centered_instance``; composed
    with a ``centroid`` model for full top-down inference).
    """

    segmentation: Optional[CenteredInstanceSegmentationHeadConfig] = None

CenteredInstanceSegmentationHeadConfig

Foreground-mask head config for top-down crop-centered segmentation.

The bottom-up SegmentationHeadConfig fields plus anchor_part. Keeping anchor_part INSIDE the head leaf (rather than as a sibling of segmentation) matches centered_instance's confmaps.anchor_part and the codebase invariant that every per-type head config is a dict of head-leaf configs each carrying output_stride — so model/config code that iterates head leaves (loss weights, output strides, etc.) never trips over it.

Attributes:

Name	Type	Description
output_stride	int	(int) Stride of the output mask relative to the input crop. Default: 2.
loss_weight	float	(float) Scalar weight for the bce-dice loss term. Default: 1.0.
anchor_part	Optional[str]	(str) Optional node name used to center crops during training. None (default) centers on the mean of each instance's visible nodes.
centroid_method	Optional[str]	(str) How the centroid is derived from the instance's points, spelled as in sio.Instance.to_centroid: "center_of_mass" (mean of visible nodes), "bbox_center" (midpoint of the visible nodes' bounding box), "geometric_median" (Weiszfeld median — the least affected by a MISLOCALIZED node; measured on real pose data, one node off by a body length moves it ~1.7x less than the mean and ~5x less than the bbox midpoint. Not more stable than the mean under node dropout), or "anchor" (the anchor_part node). None (default) infers it: "anchor" when anchor_part is set, else "center_of_mass" — i.e. exactly the historical behavior, so existing configs are unchanged. Setting both anchor_part and a non-anchor centroid_method is an error (they name different centroids); use centroid_fallback for that. Default is None.
centroid_fallback	Optional[str]	(str) The reduce method used when anchor_part is configured but that node is not visible: "center_of_mass" (default), "bbox_center" or "geometric_median". Only meaningful for the anchor method. Unlike sio's fallback=None, sleap-nn always falls back rather than emitting a NaN centroid. Default is None (= "center_of_mass").

Source code in sleap_nn/config/model_config.py

@define
class CenteredInstanceSegmentationHeadConfig:
    """Foreground-mask head config for top-down crop-centered segmentation.

    The bottom-up ``SegmentationHeadConfig`` fields plus ``anchor_part``. Keeping
    ``anchor_part`` INSIDE the head leaf (rather than as a sibling of
    ``segmentation``) matches ``centered_instance``'s ``confmaps.anchor_part`` and
    the codebase invariant that every per-type head config is a dict of head-leaf
    configs each carrying ``output_stride`` — so model/config code that iterates
    head leaves (loss weights, output strides, etc.) never trips over it.

    Attributes:
        output_stride: (int) Stride of the output mask relative to the input
            crop. Default: 2.
        loss_weight: (float) Scalar weight for the bce-dice loss term. Default: 1.0.
        anchor_part: (str) Optional node name used to center crops during
            training. ``None`` (default) centers on the mean of each instance's
            visible nodes.
        centroid_method: (str) How the centroid is derived from the instance's
            points, spelled as in ``sio.Instance.to_centroid``:
            ``"center_of_mass"`` (mean of visible nodes), ``"bbox_center"``
            (midpoint of the visible nodes' bounding box), ``"geometric_median"``
            (Weiszfeld median — the least affected by a MISLOCALIZED node; measured
            on real pose data, one node off by a body length moves it ~1.7x less
            than the mean and ~5x less than the bbox midpoint. Not more stable
            than the mean under node dropout), or ``"anchor"``
            (the ``anchor_part`` node). ``None`` (default) infers it: ``"anchor"``
            when ``anchor_part`` is set, else ``"center_of_mass"`` — i.e. exactly
            the historical behavior, so existing configs are unchanged. Setting
            both ``anchor_part`` and a non-anchor ``centroid_method`` is an error
            (they name different centroids); use ``centroid_fallback`` for that.
            Default is None.
        centroid_fallback: (str) The reduce method used when ``anchor_part`` is
            configured but that node is not visible: ``"center_of_mass"``
            (default), ``"bbox_center"`` or ``"geometric_median"``. Only
            meaningful for the anchor method. Unlike ``sio``'s ``fallback=None``,
            sleap-nn always falls back rather than emitting a NaN centroid.
            Default is None (= ``"center_of_mass"``).
    """

    output_stride: int = 2
    loss_weight: float = 1.0
    anchor_part: Optional[str] = None
    centroid_method: Optional[str] = None
    centroid_fallback: Optional[str] = None

CentroidConfMapsConfig

Centroid configuration map.

Attributes:

Name	Type	Description
anchor_part	Optional[str]	(str) Node name to use as the anchor point. If None, the NaN-ignoring mean of all visible instance nodes will be used as the anchor. The same mean-of-visible-nodes fallback is used when the anchor part is specified but not visible in the instance. Setting a reliable anchor point can significantly improve topdown model accuracy as they benefit from a consistent geometry of the body parts relative to the center of the image. Default is None.
centroid_method	Optional[str]	(str) How the centroid is derived from the instance's points, spelled as in sio.Instance.to_centroid: "center_of_mass" (mean of visible nodes), "bbox_center" (midpoint of the visible nodes' bounding box), "geometric_median" (Weiszfeld median — the least affected by a MISLOCALIZED node; measured on real pose data, one node off by a body length moves it ~1.7x less than the mean and ~5x less than the bbox midpoint. Not more stable than the mean under node dropout), or "anchor" (the anchor_part node). None (default) infers it: "anchor" when anchor_part is set, else "center_of_mass" — i.e. exactly the historical behavior, so existing configs are unchanged. Setting both anchor_part and a non-anchor centroid_method is an error (they name different centroids); use centroid_fallback for that. Default is None.
centroid_fallback	Optional[str]	(str) The reduce method used when anchor_part is configured but that node is not visible: "center_of_mass" (default), "bbox_center" or "geometric_median". Only meaningful for the anchor method. Unlike sio's fallback=None, sleap-nn always falls back rather than emitting a NaN centroid. Default is None (= "center_of_mass").
centroid_source	Optional[str]	(str) Which centroid the model is trained to predict. The centroid head must use ONE source for the whole dataset — mixing user-annotated and computed centroids trains the head against two conflicting definitions of "centroid". Options: - "user": train on first-class UserCentroid annotations; frames with pose instances but no user centroid are dropped. - "computed": derive every centroid from instance keypoints (the anchor_part node, else the mean of visible nodes); any UserCentroid annotations are ignored. - None (default): infer the source from the training labels (user centroids present -> "user", else "computed") and emit a loud warning, since a silently-chosen target is a training footgun. Set this explicitly to silence the warning.
sigma	float	(float) Spread of the Gaussian distribution of the confidence maps as a scalar float. Smaller values are more precise but may be difficult to learn as they have a lower density within the image space. Larger values are easier to learn but are less precise with respect to the peak coordinate. This spread is in units of pixels of the model input image, i.e., the image resolution after any input scaling is applied.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
use_sigmoid_activation	bool	(bool) If True, applies a sigmoid activation to the head's output so it is a calibrated (0, 1) probability rather than raw/unbounded regression output. Required when training with focal_loss_alpha (a focal-style loss needs Ŷ ∈ (0, 1) for its log terms, and peak_threshold-based inference stays meaningful only if the raw output is already a probability). Default False -- no change from existing behavior for plain MSE training.
focal_loss_alpha	float	(float) If nonzero, replaces the plain MSE train loss with a CenterNet/CornerNet-style penalty-reduced pixelwise focal loss (see sleap_nn.training.losses.compute_centroid_focal_loss) -- down-weights already-confident pixels on both the positive (near a true peak) and negative side, focusing training on ambiguous pixels (e.g. where two animals' peaks are close together). Requires use_sigmoid_activation: true so the head's output is a calibrated (0, 1) probability. 0 disables this (plain MSE, matching every other config). Default: 0.0.
focal_loss_beta	float	(float) Penalty-reduction exponent for negative pixels near a true peak. Only has an effect when focal_loss_alpha != 0. Default: 4.0 (standard CenterNet value).
focal_loss_pos_threshold	float	(float) Minimum target confmap value for a pixel to count as "positive" (near a true peak) in the focal loss. sleap-nn's Gaussian confmap targets are continuous (sub-pixel), so the peak pixel's value is rarely exactly 1.0 -- unlike the original CenterNet formulation's integer-snapped peaks -- hence a threshold rather than exact equality. Only has an effect when focal_loss_alpha != 0. Default: 0.5.

Source code in sleap_nn/config/model_config.py

@define
class CentroidConfMapsConfig:
    """Centroid configuration map.

    Attributes:
        anchor_part: (str) Node name to use as the anchor point. If None, the
            NaN-ignoring mean of all visible instance nodes will be used as the
            anchor. The same mean-of-visible-nodes fallback is used when the
            anchor part is specified but not visible in the instance. Setting a
            reliable anchor point can significantly improve topdown model
            accuracy as they benefit from a consistent geometry of the body parts
            relative to the center of the image. Default is None.
        centroid_method: (str) How the centroid is derived from the instance's
            points, spelled as in ``sio.Instance.to_centroid``:
            ``"center_of_mass"`` (mean of visible nodes), ``"bbox_center"``
            (midpoint of the visible nodes' bounding box), ``"geometric_median"``
            (Weiszfeld median — the least affected by a MISLOCALIZED node; measured
            on real pose data, one node off by a body length moves it ~1.7x less
            than the mean and ~5x less than the bbox midpoint. Not more stable
            than the mean under node dropout), or ``"anchor"``
            (the ``anchor_part`` node). ``None`` (default) infers it: ``"anchor"``
            when ``anchor_part`` is set, else ``"center_of_mass"`` — i.e. exactly
            the historical behavior, so existing configs are unchanged. Setting
            both ``anchor_part`` and a non-anchor ``centroid_method`` is an error
            (they name different centroids); use ``centroid_fallback`` for that.
            Default is None.
        centroid_fallback: (str) The reduce method used when ``anchor_part`` is
            configured but that node is not visible: ``"center_of_mass"``
            (default), ``"bbox_center"`` or ``"geometric_median"``. Only
            meaningful for the anchor method. Unlike ``sio``'s ``fallback=None``,
            sleap-nn always falls back rather than emitting a NaN centroid.
            Default is None (= ``"center_of_mass"``).
        centroid_source: (str) Which centroid the model is trained to predict.
            The centroid head must use ONE source for the whole dataset —
            mixing user-annotated and computed centroids trains the head against
            two conflicting definitions of "centroid". Options:
            - ``"user"``: train on first-class ``UserCentroid`` annotations;
              frames with pose instances but no user centroid are dropped.
            - ``"computed"``: derive every centroid from instance keypoints
              (the ``anchor_part`` node, else the mean of visible nodes); any
              ``UserCentroid`` annotations are ignored.
            - ``None`` (default): infer the source from the training labels
              (user centroids present -> ``"user"``, else ``"computed"``) and
              emit a loud warning, since a silently-chosen target is a training
              footgun. Set this explicitly to silence the warning.
        sigma: (float) Spread of the Gaussian distribution of the confidence maps as a
            scalar float. Smaller values are more precise but may be difficult to learn as
            they have a lower density within the image space. Larger values are easier to
            learn but are less precise with respect to the peak coordinate. This spread is
            in units of pixels of the model input image, i.e., the image resolution after
            any input scaling is applied.
        output_stride: (int) The stride of the output confidence maps relative to the
            input image. This is the reciprocal of the resolution, e.g., an output
            stride of 2 results in confidence maps that are 0.5x the size of the input.
            Increasing this value can considerably speed up model performance and
            decrease memory requirements, at the cost of decreased spatial resolution.
        use_sigmoid_activation: (bool) If `True`, applies a sigmoid activation to the
            head's output so it is a calibrated `(0, 1)` probability rather than
            raw/unbounded regression output. Required when training with
            `focal_loss_alpha` (a focal-style loss needs `Ŷ ∈ (0, 1)` for its log
            terms, and `peak_threshold`-based inference stays meaningful only if the
            raw output is already a probability). Default `False` -- no change from
            existing behavior for plain MSE training.
        focal_loss_alpha: (float) If nonzero, replaces the plain MSE train loss with a
            CenterNet/CornerNet-style penalty-reduced pixelwise focal loss (see
            `sleap_nn.training.losses.compute_centroid_focal_loss`) -- down-weights
            already-confident pixels on both the positive (near a true peak) and
            negative side, focusing training on ambiguous pixels (e.g. where two
            animals' peaks are close together). Requires `use_sigmoid_activation:
            true` so the head's output is a calibrated `(0, 1)` probability. `0`
            disables this (plain MSE, matching every other config). *Default*: `0.0`.
        focal_loss_beta: (float) Penalty-reduction exponent for negative pixels near a
            true peak. Only has an effect when `focal_loss_alpha != 0`. *Default*:
            `4.0` (standard CenterNet value).
        focal_loss_pos_threshold: (float) Minimum target confmap value for a pixel to
            count as "positive" (near a true peak) in the focal loss. sleap-nn's
            Gaussian confmap targets are continuous (sub-pixel), so the peak pixel's
            value is rarely exactly `1.0` -- unlike the original CenterNet
            formulation's integer-snapped peaks -- hence a threshold rather than exact
            equality. Only has an effect when `focal_loss_alpha != 0`. *Default*:
            `0.5`.
    """

    anchor_part: Optional[str] = None
    centroid_method: Optional[str] = None
    centroid_fallback: Optional[str] = None
    centroid_source: Optional[str] = None
    sigma: float = 5.0
    output_stride: int = 1
    use_sigmoid_activation: bool = False
    focal_loss_alpha: float = field(default=0.0, validator=validators.ge(0))
    focal_loss_beta: float = field(default=4.0, validator=validators.ge(0))
    focal_loss_pos_threshold: float = field(default=0.5, validator=validators.ge(0))

CentroidConfig

centroid head_config.

Source code in sleap_nn/config/model_config.py

@define
class CentroidConfig:
    """centroid head_config."""

    confmaps: Optional[CentroidConfMapsConfig] = None

ClassMapConfig

Class map head config.

Attributes:

Name	Type	Description
classes	Optional[List[str]]	(List[str]) List of class (track) names. Default is None. When None, these are inferred from the track names in the labels file.
class_output	str	(str) How a predicted class is interpreted as an sleap_io object at inference. One of "track" (default) — emit only a video-local sio.Track per class (the classification-as-tracking output) — or "identity" — ALSO stamp a global sio.Identity(name=<class name>), for classes that enumerate unique individuals (e.g. named animals for re-ID). The sio.Track is always emitted regardless. Only set "identity" when each class is a distinct animal; otherwise the shared identity name would falsely claim all instances of a class are the same animal. The simplified sleap-io Identity matches by NAME, so the class name is the canonical cross-file key and nothing per-class is frozen at train time.
sigma	float	(float) Spread of the Gaussian distribution of the confidence maps as a scalar float. Smaller values are more precise but may be difficult to learn as they have a lower density within the image space. Larger values are easier to learn but are less precise with respect to the peak coordinate. This spread is in units of pixels of the model input image, i.e., the image resolution after any input scaling is applied.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
loss_weight	Optional[float]	(float) Scalar float used to weigh the loss term for this head during training. Increase this to encourage the optimization to focus on improving this specific output in multi-head models.

Source code in sleap_nn/config/model_config.py

@define
class ClassMapConfig:
    """Class map head config.

    Attributes:
        classes: (List[str]) List of class (track) names. Default is `None`. When `None`, these are inferred from the track names in the labels file.
        class_output: (str) How a predicted class is interpreted as an `sleap_io`
            object at inference. One of ``"track"`` (default) — emit only a
            video-local `sio.Track` per class (the classification-as-tracking
            output) — or ``"identity"`` — ALSO stamp a global
            `sio.Identity(name=<class name>)`, for classes that enumerate unique
            individuals (e.g. named animals for re-ID). The `sio.Track` is always
            emitted regardless. Only set ``"identity"`` when each class is a
            distinct animal; otherwise the shared identity name would falsely
            claim all instances of a class are the same animal. The simplified
            sleap-io `Identity` matches by NAME, so the class name is the
            canonical cross-file key and nothing per-class is frozen at train time.
        sigma: (float) Spread of the Gaussian distribution of the confidence maps as
            a scalar float. Smaller values are more precise but may be difficult to
            learn as they have a lower density within the image space. Larger values
            are easier to learn but are less precise with respect to the peak
            coordinate. This spread is in units of pixels of the model input image,
            i.e., the image resolution after any input scaling is applied.
        output_stride: (int) The stride of the output confidence maps relative to
            the input image. This is the reciprocal of the resolution, e.g., an output
            stride of 2 results in confidence maps that are 0.5x the size of the
            input. Increasing this value can considerably speed up model performance
            and decrease memory requirements, at the cost of decreased spatial
            resolution.
        loss_weight: (float) Scalar float used to weigh the loss term for this head
            during training. Increase this to encourage the optimization to focus on
            improving this specific output in multi-head models.
    """

    classes: Optional[List[str]] = None
    class_output: str = field(
        default="track", validator=validators.in_(("track", "identity"))
    )
    sigma: float = 5.0
    output_stride: int = 1
    loss_weight: Optional[float] = None

ClassVectorsConfig

Configurations for class vectors heads.

These heads are used in top-down multi-instance models that classify detected points using a fixed set of learned classes (e.g., animal identities).

Attributes:

Name	Type	Description
classes	Optional[List[str]]	List of string names of the classes that this head will predict.
class_output	str	How a predicted class is interpreted as an sleap_io object at inference. One of "track" (default) — emit only a video-local sio.Track per class (classification-as-tracking) — or "identity" — ALSO stamp a global sio.Identity(name=<class name>), for classes that enumerate unique individuals (e.g. named animals for re-ID). The sio.Track is always emitted. Only set "identity" when each class is a distinct animal; a shared identity name otherwise falsely claims all instances of a class are the same animal. The simplified sleap-io Identity matches by NAME, so the class name is the canonical cross-file key and nothing per-class is frozen at train time.
num_fc_layers	int	Number of fully-connected layers before the classification output layer. These can help in transforming general image features into classification-specific features.
num_fc_units	int	Number of units (dimensions) in the fully-connected layers before classification. Increasing this can improve the representational capacity in the pre-classification layers.
output_stride	int	(Ideally this should be same as the backbone's maxstride). The stride of the output class maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in maps that are 0.5x the size of the input. This should be the same size as the confidence maps they are associated with.
loss_weight	float	Scalar float used to weigh the loss term for this head during training. Increase this to encourage the optimization to focus on improving this specific output in multi-head models.

Source code in sleap_nn/config/model_config.py

@define
class ClassVectorsConfig:
    """Configurations for class vectors heads.

    These heads are used in top-down multi-instance models that classify detected
    points using a fixed set of learned classes (e.g., animal identities).

    Attributes:
        classes: List of string names of the classes that this head will predict.
        class_output: How a predicted class is interpreted as an ``sleap_io`` object
            at inference. One of ``"track"`` (default) — emit only a video-local
            ``sio.Track`` per class (classification-as-tracking) — or
            ``"identity"`` — ALSO stamp a global ``sio.Identity(name=<class
            name>)``, for classes that enumerate unique individuals (e.g. named
            animals for re-ID). The ``sio.Track`` is always emitted. Only set
            ``"identity"`` when each class is a distinct animal; a shared identity
            name otherwise falsely claims all instances of a class are the same
            animal. The simplified sleap-io ``Identity`` matches by NAME, so the
            class name is the canonical cross-file key and nothing per-class is
            frozen at train time.
        num_fc_layers: Number of fully-connected layers before the classification output
            layer. These can help in transforming general image features into
            classification-specific features.
        num_fc_units: Number of units (dimensions) in the fully-connected layers before
            classification. Increasing this can improve the representational capacity in
            the pre-classification layers.
        output_stride: (Ideally this should be same as the backbone's maxstride).
            The stride of the output class maps relative to the input image.
            This is the reciprocal of the resolution, e.g., an output stride of 2
            results in maps that are 0.5x the size of the input. This should be the same
            size as the confidence maps they are associated with.
        loss_weight: Scalar float used to weigh the loss term for this head during
            training. Increase this to encourage the optimization to focus on improving
            this specific output in multi-head models.
    """

    classes: Optional[List[str]] = None
    class_output: str = field(
        default="track", validator=validators.in_(("track", "identity"))
    )
    num_fc_layers: int = 1
    num_fc_units: int = 64
    global_pool: bool = True
    output_stride: int = 1
    loss_weight: float = 1.0

ConvNextBaseConfig

Bases: ConvNextConfig

Convnext configuration for backbone.

Attributes:

Name	Type	Description
pre_trained_weights	Optional[str]	(str) Pretrained weights file name supported only for ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
arch	Optional[dict]	(Default is Tiny architecture config. No need to provide if model_type is provided) depths: (List(int)) Number of layers in each block. Default: [3, 3, 9, 3]. channels: (List(int)) Number of channels in each block. Default: [96, 192, 384, 768].
model_type	str	(str) One of the ConvNext architecture types: ["tiny", "small", "base", "large"]. Default: "tiny".
stem_patch_kernel	int	(int) Size of the convolutional kernels in the stem layer. Default is 4.
stem_patch_stride	int	(int) Convolutional stride in the stem layer. Default is 2.
in_channels	int	(int) Number of input channels. Default is 1.
kernel_size	int	(int) Size of the convolutional kernels. Default is 3.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default is 2.
convs_per_block	int	(int) Number of convolutional layers per block. Default is 2.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
max_stride	int	Factor by which input image size is reduced through the layers. This is always 32 for all convnext architectures.

Methods:

Name	Description
validate_pre_trained_weights	Validate pre_trained_weights.

Source code in sleap_nn/config/model_config.py

@define
class ConvNextBaseConfig(ConvNextConfig):
    """Convnext configuration for backbone.

    Attributes:
        pre_trained_weights: (str) Pretrained weights file name supported only for
            ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
        arch: (Default is Tiny architecture config. No need to provide if model_type
            is provided)
            depths: (List(int)) Number of layers in each block. Default: [3, 3, 9, 3].
            channels: (List(int)) Number of channels in each block. Default:
                [96, 192, 384, 768].
        model_type: (str) One of the ConvNext architecture types:
            ["tiny", "small", "base", "large"]. Default: "tiny".
        stem_patch_kernel: (int) Size of the convolutional kernels in the stem layer.
            Default is 4.
        stem_patch_stride: (int) Convolutional stride in the stem layer. Default is 2.
        in_channels: (int) Number of input channels. Default is 1.
        kernel_size: (int) Size of the convolutional kernels. Default is 3.
        filters_rate: (float) Factor to adjust the number of filters per block.
            Default is 2.
        convs_per_block: (int) Number of convolutional layers per block. Default is 2.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed
            convolutions for upsampling. Interpolation is faster but transposed
            convolutions may be able to learn richer or more complex upsampling to
            recover details from higher scales. Default: True.
        output_stride: (int) The stride of the output confidence maps relative to the
            input image. This is the reciprocal of the resolution, e.g., an output stride
            of 2 results in confidence maps that are 0.5x the size of the input.
            Increasing this value can considerably speed up model performance and
            decrease memory requirements, at the cost of decreased spatial resolution.
        max_stride: Factor by which input image size is reduced through the layers.
            This is always `32` for all convnext architectures.
    """

    pre_trained_weights: Optional[str] = field(
        default=None,
        validator=lambda instance, attr, value: instance.validate_pre_trained_weights(
            value
        ),
    )
    model_type: str = "base"  # Options: tiny, small, base, large
    arch: Optional[dict] = None
    stem_patch_kernel: int = 4
    stem_patch_stride: int = 2
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 2
    convs_per_block: int = 2
    up_interpolate: bool = True
    output_stride: int = 1
    max_stride: int = 32

    def validate_pre_trained_weights(self, value):
        """Validate pre_trained_weights.

        Check:
        convnext_weights are one of
        (
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        )
        """
        if value is None:
            return

        convnext_weights = [
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        ]

        if value not in convnext_weights:
            message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
            logger.error(message)
            raise ValueError(message)

validate_pre_trained_weights(value)

Validate pre_trained_weights.

Check: convnext_weights are one of ( "ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights", )

Source code in sleap_nn/config/model_config.py

def validate_pre_trained_weights(self, value):
    """Validate pre_trained_weights.

    Check:
    convnext_weights are one of
    (
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    )
    """
    if value is None:
        return

    convnext_weights = [
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    ]

    if value not in convnext_weights:
        message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
        logger.error(message)
        raise ValueError(message)

ConvNextConfig

Convnext configuration for backbone.

Attributes:

Name	Type	Description
pre_trained_weights	Optional[str]	(str) Pretrained weights file name supported only for ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
arch	Optional[dict]	(Default is Tiny architecture config. No need to provide if model_type is provided) depths: (List[int]) Number of layers in each block. Default: [3, 3, 9, 3]. channels: (List[int]) Number of channels in each block. Default: [96, 192, 384, 768].
model_type	str	(str) One of the ConvNext architecture types: ["tiny", "small", "base", "large"]. Default: "tiny".
stem_patch_kernel	int	(int) Size of the convolutional kernels in the stem layer. Default: 4.
stem_patch_stride	int	(int) Convolutional stride in the stem layer. Default: 2.
in_channels	int	(int) Number of input channels. Default: 1.
kernel_size	int	(int) Size of the convolutional kernels. Default: 3.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default: 2.
convs_per_block	int	(int) Number of convolutional layers per block. Default: 2.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. Default: 1.
max_stride	int	(int) Factor by which input image size is reduced through the layers. This is always 32 for all convnext architectures. Default: 32.

Methods:

Name	Description
validate_pre_trained_weights	Validate pre_trained_weights.

Source code in sleap_nn/config/model_config.py

@define
class ConvNextConfig:
    """Convnext configuration for backbone.

    Attributes:
        pre_trained_weights: (str) Pretrained weights file name supported only for
            ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
        arch: (Default is Tiny architecture config. No need to provide if model_type is provided)
            depths: (List[int]) Number of layers in each block. *Default*: `[3, 3, 9, 3]`.
            channels: (List[int]) Number of channels in each block. *Default*: `[96, 192, 384, 768]`.
        model_type: (str) One of the ConvNext architecture types: ["tiny", "small", "base", "large"]. *Default*: `"tiny"`.
        stem_patch_kernel: (int) Size of the convolutional kernels in the stem layer. *Default*: `4`.
        stem_patch_stride: (int) Convolutional stride in the stem layer. *Default*: `2`.
        in_channels: (int) Number of input channels. *Default*: `1`.
        kernel_size: (int) Size of the convolutional kernels. *Default*: `3`.
        filters_rate: (float) Factor to adjust the number of filters per block. *Default*: `2`.
        convs_per_block: (int) Number of convolutional layers per block. *Default*: `2`.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. *Default*: `True`.
        output_stride: (int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. *Default*: `1`.
        max_stride: (int) Factor by which input image size is reduced through the layers. This is always `32` for all convnext architectures. *Default*: `32`.
    """

    pre_trained_weights: Optional[str] = field(
        default=None,
        validator=lambda instance, attr, value: instance.validate_pre_trained_weights(
            value
        ),
    )
    model_type: str = "tiny"  # Options: tiny, small, base, large
    arch: Optional[dict] = None
    stem_patch_kernel: int = 4
    stem_patch_stride: int = 2
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 2
    convs_per_block: int = 2
    up_interpolate: bool = True
    output_stride: int = 1
    max_stride: int = 32

    def validate_pre_trained_weights(self, value):
        """Validate pre_trained_weights.

        Check:
        convnext_weights are one of
        (
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        )
        """
        if value is None:
            return

        convnext_weights = [
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        ]

        if value not in convnext_weights:
            message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
            logger.error(message)
            raise ValueError(message)

validate_pre_trained_weights(value)

Validate pre_trained_weights.

Check: convnext_weights are one of ( "ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights", )

Source code in sleap_nn/config/model_config.py

def validate_pre_trained_weights(self, value):
    """Validate pre_trained_weights.

    Check:
    convnext_weights are one of
    (
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    )
    """
    if value is None:
        return

    convnext_weights = [
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    ]

    if value not in convnext_weights:
        message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
        logger.error(message)
        raise ValueError(message)

ConvNextLargeConfig

Bases: ConvNextConfig

Convnext configuration for backbone.

Attributes:

Name	Type	Description
pre_trained_weights	Optional[str]	(str) Pretrained weights file name supported only for ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
arch	Optional[dict]	(Default is Tiny architecture config. No need to provide if model_type is provided) depths: (List(int)) Number of layers in each block. Default: [3, 3, 9, 3]. channels: (List(int)) Number of channels in each block. Default: [96, 192, 384, 768].
model_type	str	(str) One of the ConvNext architecture types: ["tiny", "small", "base", "large"]. Default: "tiny".
stem_patch_kernel	int	(int) Size of the convolutional kernels in the stem layer. Default is 4.
stem_patch_stride	int	(int) Convolutional stride in the stem layer. Default is 2.
in_channels	int	(int) Number of input channels. Default is 1.
kernel_size	int	(int) Size of the convolutional kernels. Default is 3.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default is 2.
convs_per_block	int	(int) Number of convolutional layers per block. Default is 2.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
max_stride	int	Factor by which input image size is reduced through the layers. This is always 32 for all convnext architectures.

Methods:

Name	Description
validate_pre_trained_weights	Validate pre_trained_weights.

Source code in sleap_nn/config/model_config.py

@define
class ConvNextLargeConfig(ConvNextConfig):
    """Convnext configuration for backbone.

    Attributes:
        pre_trained_weights: (str) Pretrained weights file name supported only for
            ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
        arch: (Default is Tiny architecture config. No need to provide if model_type
            is provided)
            depths: (List(int)) Number of layers in each block. Default: [3, 3, 9, 3].
            channels: (List(int)) Number of channels in each block. Default:
                [96, 192, 384, 768].
        model_type: (str) One of the ConvNext architecture types:
            ["tiny", "small", "base", "large"]. Default: "tiny".
        stem_patch_kernel: (int) Size of the convolutional kernels in the stem layer.
            Default is 4.
        stem_patch_stride: (int) Convolutional stride in the stem layer. Default is 2.
        in_channels: (int) Number of input channels. Default is 1.
        kernel_size: (int) Size of the convolutional kernels. Default is 3.
        filters_rate: (float) Factor to adjust the number of filters per block.
            Default is 2.
        convs_per_block: (int) Number of convolutional layers per block. Default is 2.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed
            convolutions for upsampling. Interpolation is faster but transposed
            convolutions may be able to learn richer or more complex upsampling to
            recover details from higher scales. Default: True.
        output_stride: (int) The stride of the output confidence maps relative to the
            input image. This is the reciprocal of the resolution, e.g., an output stride
            of 2 results in confidence maps that are 0.5x the size of the input.
            Increasing this value can considerably speed up model performance and
            decrease memory requirements, at the cost of decreased spatial resolution.
        max_stride: Factor by which input image size is reduced through the layers.
            This is always `32` for all convnext architectures.
    """

    pre_trained_weights: Optional[str] = field(
        default=None,
        validator=lambda instance, attr, value: instance.validate_pre_trained_weights(
            value
        ),
    )
    model_type: str = "large"  # Options: tiny, small, base, large
    arch: Optional[dict] = None
    stem_patch_kernel: int = 4
    stem_patch_stride: int = 2
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 2
    convs_per_block: int = 2
    up_interpolate: bool = True
    output_stride: int = 1
    max_stride: int = 32

    def validate_pre_trained_weights(self, value):
        """Validate pre_trained_weights.

        Check:
        convnext_weights are one of
        (
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        )
        """
        if value is None:
            return

        convnext_weights = [
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        ]

        if value not in convnext_weights:
            message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
            logger.error(message)
            raise ValueError(message)

validate_pre_trained_weights(value)

Validate pre_trained_weights.

Check: convnext_weights are one of ( "ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights", )

Source code in sleap_nn/config/model_config.py

def validate_pre_trained_weights(self, value):
    """Validate pre_trained_weights.

    Check:
    convnext_weights are one of
    (
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    )
    """
    if value is None:
        return

    convnext_weights = [
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    ]

    if value not in convnext_weights:
        message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
        logger.error(message)
        raise ValueError(message)

ConvNextSmallConfig

Bases: ConvNextConfig

Convnext configuration for backbone.

Attributes:

Name	Type	Description
pre_trained_weights	Optional[str]	(str) Pretrained weights file name supported only for ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
arch	Optional[dict]	(Default is Tiny architecture config. No need to provide if model_type is provided) depths: (List(int)) Number of layers in each block. Default: [3, 3, 9, 3]. channels: (List(int)) Number of channels in each block. Default: [96, 192, 384, 768].
model_type	str	(str) One of the ConvNext architecture types: ["tiny", "small", "base", "large"]. Default: "tiny".
stem_patch_kernel	int	(int) Size of the convolutional kernels in the stem layer. Default is 4.
stem_patch_stride	int	(int) Convolutional stride in the stem layer. Default is 2.
in_channels	int	(int) Number of input channels. Default is 1.
kernel_size	int	(int) Size of the convolutional kernels. Default is 3.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default is 2.
convs_per_block	int	(int) Number of convolutional layers per block. Default is 2.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
max_stride	int	Factor by which input image size is reduced through the layers. This is always 32 for all convnext architectures.

Methods:

Name	Description
validate_pre_trained_weights	Validate pre_trained_weights.

Source code in sleap_nn/config/model_config.py

@define
class ConvNextSmallConfig(ConvNextConfig):
    """Convnext configuration for backbone.

    Attributes:
        pre_trained_weights: (str) Pretrained weights file name supported only for
            ConvNext backbones. For ConvNext, one of ["ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights"].
        arch: (Default is Tiny architecture config. No need to provide if model_type
            is provided)
            depths: (List(int)) Number of layers in each block. Default: [3, 3, 9, 3].
            channels: (List(int)) Number of channels in each block. Default:
                [96, 192, 384, 768].
        model_type: (str) One of the ConvNext architecture types:
            ["tiny", "small", "base", "large"]. Default: "tiny".
        stem_patch_kernel: (int) Size of the convolutional kernels in the stem layer.
            Default is 4.
        stem_patch_stride: (int) Convolutional stride in the stem layer. Default is 2.
        in_channels: (int) Number of input channels. Default is 1.
        kernel_size: (int) Size of the convolutional kernels. Default is 3.
        filters_rate: (float) Factor to adjust the number of filters per block.
            Default is 2.
        convs_per_block: (int) Number of convolutional layers per block. Default is 2.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed
            convolutions for upsampling. Interpolation is faster but transposed
            convolutions may be able to learn richer or more complex upsampling to
            recover details from higher scales. Default: True.
        output_stride: (int) The stride of the output confidence maps relative to the
            input image. This is the reciprocal of the resolution, e.g., an output stride
            of 2 results in confidence maps that are 0.5x the size of the input.
            Increasing this value can considerably speed up model performance and
            decrease memory requirements, at the cost of decreased spatial resolution.
        max_stride: Factor by which input image size is reduced through the layers.
            This is always `32` for all convnext architectures.
    """

    pre_trained_weights: Optional[str] = field(
        default=None,
        validator=lambda instance, attr, value: instance.validate_pre_trained_weights(
            value
        ),
    )
    model_type: str = "small"  # Options: tiny, small, base, large
    arch: Optional[dict] = None
    stem_patch_kernel: int = 4
    stem_patch_stride: int = 2
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 2
    convs_per_block: int = 2
    up_interpolate: bool = True
    output_stride: int = 1
    max_stride: int = 32

    def validate_pre_trained_weights(self, value):
        """Validate pre_trained_weights.

        Check:
        convnext_weights are one of
        (
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        )
        """
        if value is None:
            return

        convnext_weights = [
            "ConvNeXt_Base_Weights",
            "ConvNeXt_Tiny_Weights",
            "ConvNeXt_Small_Weights",
            "ConvNeXt_Large_Weights",
        ]

        if value not in convnext_weights:
            message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
            logger.error(message)
            raise ValueError(message)

validate_pre_trained_weights(value)

Validate pre_trained_weights.

Check: convnext_weights are one of ( "ConvNeXt_Base_Weights", "ConvNeXt_Tiny_Weights", "ConvNeXt_Small_Weights", "ConvNeXt_Large_Weights", )

Source code in sleap_nn/config/model_config.py

def validate_pre_trained_weights(self, value):
    """Validate pre_trained_weights.

    Check:
    convnext_weights are one of
    (
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    )
    """
    if value is None:
        return

    convnext_weights = [
        "ConvNeXt_Base_Weights",
        "ConvNeXt_Tiny_Weights",
        "ConvNeXt_Small_Weights",
        "ConvNeXt_Large_Weights",
    ]

    if value not in convnext_weights:
        message = f"Invalid pre-trained weights for ConvNext. Must be one of {convnext_weights}"
        logger.error(message)
        raise ValueError(message)

EmbeddingConfig

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

A single pooled-encoder head producing a per-instance embedding vector. Wraps one leaf head config (mirrors CenteredInstanceSegmentationConfig).

Source code in sleap_nn/config/model_config.py

@define
class EmbeddingConfig:
    """Head config for the ``embedding`` (crop -> vector, re-ID) model type.

    A single pooled-encoder head producing a per-instance embedding vector. Wraps
    one leaf head config (mirrors ``CenteredInstanceSegmentationConfig``).
    """

    embedding: Optional[EmbeddingHeadConfig] = None

EmbeddingHeadConfig

Configuration for the embedding head (the adapter on a pooled encoder feature).

Shape mirrors ClassVectorsConfig: [pool] -> Flatten -> N x (Linear+ReLU) -> Linear(embedding_dim) -> [L2Norm]. The training objective is nested here.

Attributes:

Name	Type	Description
embedding_dim	int	Output embedding dimensionality.
num_fc_layers	int	Number of FC layers before the embedding output.
num_fc_units	int	Units in the pre-embedding FC layers.
pool	Optional[str]	Pooling over the encoder feature map. One of gem (generalized-mean, learnable exponent), max, avg, or None (default) for the backbone's default, resolved at training setup and written into the saved config: avg for a FROZEN pretrained encoder (pretrained.freeze or freeze_backbone), gem otherwise. GeM pools only the positive part of each channel, and a frozen pretrained encoder's LayerNorm output is about half negative: on the gerbil re-ID set avg beat GeM by +0.04 val rank-1 on frozen DINOv2 and ConvNeXtV2 (8 seeds each, p < 0.02) and tied when the encoder was fine-tuned (see sleap_nn.architectures.heads.default_embedding_pool). Configs saved before this default existed carry an explicit gem and are unchanged.
normalize	bool	L2-normalize the embedding (applied identically train + inference).
output_stride	int	Stride of the pooled feature. Should equal the backbone max_stride so the decoder is empty and the head taps middle_output.
freeze_backbone	bool	(bool) Freeze the backbone's PRETRAINED ENCODER and train only what sits on top of it. The encoder is the part pretrained weights load into: the HuggingFace model of a pretrained backbone (the same thing backbone_config.pretrained.freeze freezes), the ImageNet encoder of a native convnext / swint with pre_trained_weights, the stem + encoder blocks of a unet. The frozen encoder gets no gradient AND stays in eval mode for the whole run, so BatchNorm keeps its pretrained running statistics and dropout / stochastic depth are off. Everything else stays trainable: the randomly initialized middle blocks of a native convnext / swint, any decoder, the embedding head and the train-only projection head. Setting it without pretrained weights (no pretrained.weights, no pre_trained_weights, no model_config.pretrained_backbone_weights) logs a warning: the encoder would then stay at its random initialization. Default is False.
anchor_part	Optional[str]	(str) Node name used to center the re-ID crop. None (default) centers on the mean of each instance's visible nodes.
centroid_method	Optional[str]	(str) How the crop center is derived from the instance's points, spelled as in sio.Instance.to_centroid: "center_of_mass" (mean of visible nodes), "bbox_center" (midpoint of the visible nodes' bounding box), "geometric_median" (Weiszfeld median — the least affected by a MISLOCALIZED node; measured on real pose data, one node off by a body length moves it ~1.7x less than the mean and ~5x less than the bbox midpoint. Not more stable than the mean under node dropout), or "anchor" (the anchor_part node). None (default) infers it: "anchor" when anchor_part is set, else "center_of_mass" — i.e. exactly the historical behavior, so existing configs are unchanged. Setting both anchor_part and a non-anchor centroid_method is an error (they name different centroids); use centroid_fallback for that. Only used in the pose detection mode — a mask-driven embedding dataset crops on the mask's own center of mass. Default is None.
detection_mode	Optional[str]	(str) The detection carrier the model crops: "pose" (lf.instances, crop on the pose centroid) or "mask" (lf.masks, crop on the mask center, carrying the mask for burn_in). None (default) lets training pick: "mask" under data_config.preprocessing.burn_in, else the carrier holding more training samples (ties go to masks). Training WRITES the carrier it used here, and inference embeds that carrier whenever a file holds it -- a top-down segmentation output holds both. Set it to force a carrier.
centroid_fallback	Optional[str]	(str) The reduce method used when anchor_part is configured but that node is not visible: "center_of_mass" (default), "bbox_center" or "geometric_median". Only meaningful for the anchor method. Unlike sio's fallback=None, sleap-nn always falls back rather than emitting a NaN centroid. Default is None (= "center_of_mass").
objective	Optional[ObjectiveConfig]	The pluggable training objective (positives x negatives x loss).

Source code in sleap_nn/config/model_config.py

@define
class EmbeddingHeadConfig:
    """Configuration for the embedding head (the adapter on a pooled encoder feature).

    Shape mirrors ``ClassVectorsConfig``: ``[pool] -> Flatten -> N x (Linear+ReLU) ->
    Linear(embedding_dim) -> [L2Norm]``. The training objective is nested here.

    Attributes:
        embedding_dim: Output embedding dimensionality.
        num_fc_layers: Number of FC layers before the embedding output.
        num_fc_units: Units in the pre-embedding FC layers.
        pool: Pooling over the encoder feature map. One of ``gem`` (generalized-mean,
            learnable exponent), ``max``, ``avg``, or ``None`` (default) for the
            backbone's default, resolved at training setup and written into the
            saved config: ``avg`` for a FROZEN ``pretrained`` encoder
            (``pretrained.freeze`` or ``freeze_backbone``), ``gem`` otherwise. GeM
            pools only the positive part of each channel, and a frozen pretrained
            encoder's LayerNorm output is about half negative: on the gerbil re-ID
            set avg beat GeM by +0.04 val rank-1 on frozen DINOv2 and ConvNeXtV2
            (8 seeds each, p < 0.02) and tied when the encoder was fine-tuned (see
            ``sleap_nn.architectures.heads.default_embedding_pool``). Configs saved
            before this default existed carry an explicit ``gem`` and are
            unchanged.
        normalize: L2-normalize the embedding (applied identically train + inference).
        output_stride: Stride of the pooled feature. Should equal the backbone
            ``max_stride`` so the decoder is empty and the head taps ``middle_output``.
        freeze_backbone: (bool) Freeze the backbone's PRETRAINED ENCODER and train
            only what sits on top of it. The encoder is the part pretrained weights
            load into: the HuggingFace model of a ``pretrained`` backbone (the same
            thing ``backbone_config.pretrained.freeze`` freezes), the ImageNet
            encoder of a native ``convnext`` / ``swint`` with
            ``pre_trained_weights``, the stem + encoder blocks of a ``unet``. The
            frozen encoder gets no gradient AND stays in eval mode for the whole run,
            so BatchNorm keeps its pretrained running statistics and dropout /
            stochastic depth are off. Everything else stays trainable: the randomly
            initialized middle blocks of a native ``convnext`` / ``swint``, any
            decoder, the embedding head and the train-only projection head. Setting
            it without pretrained weights (no ``pretrained.weights``, no
            ``pre_trained_weights``, no ``model_config.pretrained_backbone_weights``)
            logs a warning: the encoder would then stay at its random
            initialization. Default is False.
        anchor_part: (str) Node name used to center the re-ID crop. ``None``
            (default) centers on the mean of each instance's visible nodes.
        centroid_method: (str) How the crop center is derived from the instance's
            points, spelled as in ``sio.Instance.to_centroid``:
            ``"center_of_mass"`` (mean of visible nodes), ``"bbox_center"``
            (midpoint of the visible nodes' bounding box), ``"geometric_median"``
            (Weiszfeld median — the least affected by a MISLOCALIZED node; measured
            on real pose data, one node off by a body length moves it ~1.7x less
            than the mean and ~5x less than the bbox midpoint. Not more stable
            than the mean under node dropout), or ``"anchor"``
            (the ``anchor_part`` node). ``None`` (default) infers it: ``"anchor"``
            when ``anchor_part`` is set, else ``"center_of_mass"`` — i.e. exactly
            the historical behavior, so existing configs are unchanged. Setting
            both ``anchor_part`` and a non-anchor ``centroid_method`` is an error
            (they name different centroids); use ``centroid_fallback`` for that.
            Only used in the pose detection mode — a mask-driven embedding dataset
            crops on the mask's own center of mass. Default is None.
        detection_mode: (str) The detection carrier the model crops: ``"pose"``
            (``lf.instances``, crop on the pose centroid) or ``"mask"`` (``lf.masks``,
            crop on the mask center, carrying the mask for ``burn_in``). ``None``
            (default) lets training pick: ``"mask"`` under
            ``data_config.preprocessing.burn_in``, else the carrier holding more
            training samples (ties go to masks). Training WRITES the carrier it used
            here, and inference embeds that carrier whenever a file holds it -- a
            top-down segmentation output holds both. Set it to force a carrier.
        centroid_fallback: (str) The reduce method used when ``anchor_part`` is
            configured but that node is not visible: ``"center_of_mass"``
            (default), ``"bbox_center"`` or ``"geometric_median"``. Only
            meaningful for the anchor method. Unlike ``sio``'s ``fallback=None``,
            sleap-nn always falls back rather than emitting a NaN centroid.
            Default is None (= ``"center_of_mass"``).
        objective: The pluggable training objective (positives x negatives x loss).
    """

    embedding_dim: int = 128
    num_fc_layers: int = 1
    num_fc_units: int = 256
    pool: Optional[str] = None
    normalize: bool = True
    output_stride: int = 32
    freeze_backbone: bool = False
    anchor_part: Optional[str] = None
    centroid_method: Optional[str] = None
    centroid_fallback: Optional[str] = None
    detection_mode: Optional[str] = field(
        default=None, validator=validators.optional(validators.in_(("pose", "mask")))
    )
    objective: Optional[ObjectiveConfig] = None

HeadConfig

Configurations related to the model output head type.

Only one attribute of this class can be set, which defines the model output type.

Attributes:

Name	Type	Description
single_instance	Optional[SingleInstanceConfig]	An instance of SingleInstanceConfmapsHeadConfig.
centroid	Optional[CentroidConfig]	An instance of CentroidsHeadConfig.
centered_instance	Optional[CenteredInstanceConfig]	An instance of CenteredInstanceConfmapsHeadConfig.
bottomup	Optional[BottomUpConfig]	An instance of BottomUpConfig.
multi_class_bottomup	Optional[BottomUpMultiClassConfig]	An instance of BottomUpMultiClassConfig.
multi_class_topdown	Optional[TopDownCenteredInstanceMultiClassConfig]	An instance of TopDownCenteredInstanceMultiClassConfig.
bottomup_segmentation	Optional[BottomUpSegmentationConfig]	An instance of BottomUpSegmentationConfig.
centered_instance_segmentation	Optional[CenteredInstanceSegmentationConfig]	An instance of CenteredInstanceSegmentationConfig.
semantic_segmentation	Optional[SemanticSegmentationConfig]	An instance of SemanticSegmentationConfig.
embedding	Optional[EmbeddingConfig]	An instance of EmbeddingConfig.
bbox_centernet	Optional[BBoxCenterNetConfig]	An instance of BBoxCenterNetConfig.
bbox_ssd	Optional[BBoxSSDConfig]	An instance of BBoxSSDConfig.

Source code in sleap_nn/config/model_config.py

@oneof
@define
class HeadConfig:
    """Configurations related to the model output head type.

    Only one attribute of this class can be set, which defines the model output type.

    Attributes:
        single_instance: An instance of `SingleInstanceConfmapsHeadConfig`.
        centroid: An instance of `CentroidsHeadConfig`.
        centered_instance: An instance of `CenteredInstanceConfmapsHeadConfig`.
        bottomup: An instance of `BottomUpConfig`.
        multi_class_bottomup: An instance of `BottomUpMultiClassConfig`.
        multi_class_topdown: An instance of `TopDownCenteredInstanceMultiClassConfig`.
        bottomup_segmentation: An instance of `BottomUpSegmentationConfig`.
        centered_instance_segmentation: An instance of
            `CenteredInstanceSegmentationConfig`.
        semantic_segmentation: An instance of `SemanticSegmentationConfig`.
        embedding: An instance of `EmbeddingConfig`.
        bbox_centernet: An instance of `BBoxCenterNetConfig`.
        bbox_ssd: An instance of `BBoxSSDConfig`.
    """

    single_instance: Optional[SingleInstanceConfig] = None
    centroid: Optional[CentroidConfig] = None
    centered_instance: Optional[CenteredInstanceConfig] = None
    bottomup: Optional[BottomUpConfig] = None
    multi_class_bottomup: Optional[BottomUpMultiClassConfig] = None
    multi_class_topdown: Optional[TopDownCenteredInstanceMultiClassConfig] = None
    bottomup_segmentation: Optional[BottomUpSegmentationConfig] = None
    centered_instance_segmentation: Optional[CenteredInstanceSegmentationConfig] = None
    semantic_segmentation: Optional[SemanticSegmentationConfig] = None
    embedding: Optional[EmbeddingConfig] = None
    bbox_centernet: Optional[BBoxCenterNetConfig] = None
    bbox_ssd: Optional[BBoxSSDConfig] = None

InstanceCenterConfig

Configuration for the instance center heatmap head.

This config is used exclusively by the bottomup_segmentation head; it does not affect centroid (CentroidConfMapsConfig.sigma) or bottom-up pose (BottomUpConfMapsConfig) models.

Attributes:

Name	Type	Description
sigma	float	(float) Standard deviation of the Gaussian distribution used to generate center heatmaps, in pixels at original image resolution. Default: 4.0. Caveat: 4.0 was validated empirically on mice only (compact bodies; it cut learned center over-detection from 60.75 to 14.2 peaks/frame relative to the old default of 10.0, which over-fires badly on elongated bodies). It is the better general default than 10.0, but the optimal value is dataset-dependent: elongated/large animals may want a larger sigma, tiny animals smaller. The target Gaussian is always isotropic (anisotropic targets were tested and performed worse). Because this is a per-head config field, users can tune it per dataset without affecting any other model type.
output_stride	int	(int) The stride of the output center heatmaps relative to the input image. Default: 2.
loss_weight	float	(float) Scalar float used to weigh the loss term for this head during training. Default: 1.0.

Source code in sleap_nn/config/model_config.py

@define
class InstanceCenterConfig:
    """Configuration for the instance center heatmap head.

    This config is used exclusively by the ``bottomup_segmentation`` head; it does
    not affect centroid (``CentroidConfMapsConfig.sigma``) or bottom-up pose
    (``BottomUpConfMapsConfig``) models.

    Attributes:
        sigma: (float) Standard deviation of the Gaussian distribution used to generate
            center heatmaps, in pixels at original image resolution. Default: 4.0.

            Caveat: 4.0 was validated empirically on mice only (compact bodies; it cut
            learned center over-detection from 60.75 to 14.2 peaks/frame relative to the
            old default of 10.0, which over-fires badly on elongated bodies). It is the
            better general default than 10.0, but the optimal value is dataset-dependent:
            elongated/large animals may want a larger sigma, tiny animals smaller. The
            target Gaussian is always isotropic (anisotropic targets were tested and
            performed worse). Because this is a per-head config field, users can tune it
            per dataset without affecting any other model type.
        output_stride: (int) The stride of the output center heatmaps relative to the
            input image. Default: 2.
        loss_weight: (float) Scalar float used to weigh the loss term for this head
            during training. Default: 1.0.
    """

    sigma: float = 4.0
    output_stride: int = 2
    loss_weight: float = 1.0

LossConfig

Contrastive loss for the embedding objective (the loss axis).

Attributes:

Name	Type	Description
name	str	One of supcon | infonce | triplet.
temperature	float	Softmax temperature for supcon / infonce.
margin	float	Margin for triplet.

Source code in sleap_nn/config/model_config.py

@define
class LossConfig:
    """Contrastive loss for the embedding objective (the loss axis).

    Attributes:
        name: One of ``supcon`` | ``infonce`` | ``triplet``.
        temperature: Softmax temperature for ``supcon`` / ``infonce``.
        margin: Margin for ``triplet``.
    """

    name: str = field(
        default="supcon", validator=_objective_option("loss", EMBEDDING_LOSSES)
    )
    temperature: float = 0.1
    margin: float = 0.2

ModelConfig

Configurations related to model architecture.

Attributes:

Name	Type	Description
init_weights	str	(str) model weights initialization method. "default" uses kaiming uniform initialization and "xavier" uses Xavier initialization method.
pretrained_backbone_weights	Optional[str]	Path of the ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file with which the backbone is initialized. If None, random init is used.
pretrained_head_weights	Optional[str]	Path of the ckpt (or .h5 file from SLEAP - only UNet backbone is supported) file with which the head layers are initialized. If None, random init is used.
backbone_config	BackboneConfig	initialize either UNetConfig, ConvNextConfig, or SwinTConfig based on input from backbone_type
head_configs	HeadConfig	(Dict) Dictionary with the following keys having head configs for the model to be trained. Note: Configs should be provided only for the model to train and others should be None
total_params	Optional[int]	(int) Total number of parameters in the model. This is automatically computed when the training starts.

Source code in sleap_nn/config/model_config.py

@define
class ModelConfig:
    """Configurations related to model architecture.

    Attributes:
        init_weights: (str) model weights initialization method. "default" uses kaiming
            uniform initialization and "xavier" uses Xavier initialization method.
        pretrained_backbone_weights: Path of the `ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file with which the backbone
            is initialized. If `None`, random init is used.
        pretrained_head_weights: Path of the `ckpt` (or `.h5` file from SLEAP - only UNet backbone is supported) file with which the head layers
            are initialized. If `None`, random init is used.
        backbone_config: initialize either UNetConfig, ConvNextConfig, or SwinTConfig
            based on input from backbone_type
        head_configs: (Dict) Dictionary with the following keys having head configs for
            the model to be trained. Note: Configs should be provided only for the model
            to train and others should be None
        total_params: (int) Total number of parameters in the model. This is automatically
            computed when the training starts.
    """

    init_weights: str = "default"
    pretrained_backbone_weights: Optional[str] = None
    pretrained_head_weights: Optional[str] = None
    backbone_config: BackboneConfig = field(factory=BackboneConfig)
    head_configs: HeadConfig = field(factory=HeadConfig)
    total_params: Optional[int] = None

NegativesConfig

Negative-pair eligibility for the embedding objective.

A negative must be a KNOWN-different pair, never merely "not known-positive".

Attributes:

Name	Type	Description
sources	List[str]	List of negative sources. in_batch = any other crop in the batch; same_frame = crops in the same frame (hard negatives; asserts detections_deduplicated). Default is both, ["same_frame", "in_batch"].
exclude_same_track	bool	Drop same-(video, track) pairs from the negatives.
restrict_same_video	bool	Restrict negatives to same-video pairs. REQUIRED for scope=tracklet (video-local ids): cross-video pairs are unknown and must not be used as negatives (avoids false negatives that push the same animal apart across videos).

Source code in sleap_nn/config/model_config.py

@define
class NegativesConfig:
    """Negative-pair eligibility for the embedding objective.

    A negative must be a KNOWN-different pair, never merely "not known-positive".

    Attributes:
        sources: List of negative sources. ``in_batch`` = any other crop in the
            batch; ``same_frame`` = crops in the same frame (hard negatives;
            asserts ``detections_deduplicated``). Default is both,
            ``["same_frame", "in_batch"]``.
        exclude_same_track: Drop same-``(video, track)`` pairs from the negatives.
        restrict_same_video: Restrict negatives to same-video pairs. REQUIRED for
            ``scope=tracklet`` (video-local ids): cross-video pairs are unknown and
            must not be used as negatives (avoids false negatives that push the same
            animal apart across videos).
    """

    # A tuple, not `field(factory=...)`: OmegaConf builds this class from its type
    # when merging a partial `negatives:` block and cannot call an attrs Factory.
    # It lands as a list in every config (OmegaConf copies it into a ListConfig).
    sources: List[str] = field(
        default=EMBEDDING_NEGATIVE_SOURCES,
        validator=_objective_option("negatives", EMBEDDING_NEGATIVE_SOURCES, many=True),
    )
    exclude_same_track: bool = True
    restrict_same_video: bool = False

ObjectiveConfig

Pluggable training objective = positives x negatives x loss.

The sampler composes the batch, a mask-builder turns each item's (video, frame, group, item_id) into (pos_mask, neg_mask), and the loss consumes (embeddings, pos_mask, neg_mask).

Attributes:

Name	Type	Description
positives	Optional[PositivesConfig]	Positive-pair sampling config.
negatives	Optional[NegativesConfig]	Negative-pair eligibility config.
loss	Optional[LossConfig]	Contrastive loss config.
sampler	Optional[SamplerConfig]	Group-aware batch sampler config.
use_projection	bool	Add a train-only projection head (discarded at inference) for supcon / infonce.
projection_dim	int	Width of the projection head.

Source code in sleap_nn/config/model_config.py

@define
class ObjectiveConfig:
    """Pluggable training objective = positives x negatives x loss.

    The sampler composes the batch, a mask-builder turns each item's
    ``(video, frame, group, item_id)`` into ``(pos_mask, neg_mask)``, and the loss
    consumes ``(embeddings, pos_mask, neg_mask)``.

    Attributes:
        positives: Positive-pair sampling config.
        negatives: Negative-pair eligibility config.
        loss: Contrastive loss config.
        sampler: Group-aware batch sampler config.
        use_projection: Add a train-only projection head (discarded at inference)
            for ``supcon`` / ``infonce``.
        projection_dim: Width of the projection head.
    """

    # Nested configs default to None (idiomatic for OmegaConf merge from a plain dict;
    # `field(factory=...)` defaults break structured-schema instantiation). The string
    # factory + the LightningModule fill them with defaults when absent.
    positives: Optional[PositivesConfig] = None
    negatives: Optional[NegativesConfig] = None
    loss: Optional[LossConfig] = None
    sampler: Optional[SamplerConfig] = None
    use_projection: bool = True
    projection_dim: int = 128

PAFConfig

PAF configuration map.

Attributes:

Name	Type	Description
edges	Optional[List[List[str]]]	(List[str]) None if edges from sio.Labels file can be used directly. Note: Only for 'PartAffinityFieldsHead'. List of indices (src, dest) that form an edge.
sigma	float	(float) Spread of the Gaussian distribution of the confidence maps as a scalar float. Smaller values are more precise but may be difficult to learn as they have a lower density within the image space. Larger values are easier to learn but are less precise with respect to the peak coordinate. This spread is in units of pixels of the model input image, i.e., the image resolution after any input scaling is applied.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.
loss_weight	Optional[float]	(float) Scalar float used to weigh the loss term for this head during training. Increase this to encourage the optimization to focus on improving this specific output in multi-head models.

Source code in sleap_nn/config/model_config.py

@define
class PAFConfig:
    """PAF configuration map.

    Attributes:
        edges: (List[str]) None if edges from sio.Labels file can be used directly.
            Note: Only for 'PartAffinityFieldsHead'. List of indices (src, dest) that
            form an edge.
        sigma: (float) Spread of the Gaussian distribution of the confidence maps as
            a scalar float. Smaller values are more precise but may be difficult to
            learn as they have a lower density within the image space. Larger values
            are easier to learn but are less precise with respect to the peak
            coordinate. This spread is in units of pixels of the model input image,
            i.e., the image resolution after any input scaling is applied.
        output_stride: (int) The stride of the output confidence maps relative to
            the input image. This is the reciprocal of the resolution, e.g., an output
            stride of 2 results in confidence maps that are 0.5x the size of the
            input. Increasing this value can considerably speed up model performance
            and decrease memory requirements, at the cost of decreased spatial
            resolution.
        loss_weight: (float) Scalar float used to weigh the loss term for this head
            during training. Increase this to encourage the optimization to focus on
            improving this specific output in multi-head models.
    """

    edges: Optional[List[List[str]]] = None
    sigma: float = 15.0
    output_stride: int = 1
    loss_weight: Optional[float] = None

PositivesConfig

Positive-pair sampling for the embedding objective.

Attributes:

Name	Type	Description
scope	str	Which crops are positives of an anchor. One of aug_view (only the anchor's own augmented views — self-supervised: every detection is a sample and its own group, so unlabeled data trains; needs two DIFFERENT views, i.e. use_augmentations_train with some augmentation configured, and pairs best with sampler.kind='random'), tracklet (same (video, track) — video-local identity; cross-video pairs are UNKNOWN and excluded from the loss), or global_id (same global identity across all videos: a sio.Identity when the detection carries one, else its track name, which then must be globally consistent -- declared by data_config.identity.track_names_are_global).
aug_views	int	Number of augmented views of each anchor (always positives). Fixed at 2 (the standard two-view contrastive setup the training_step implements); any other value is unsupported in P1 and raises a ValueError at model construction rather than being silently ignored.

Source code in sleap_nn/config/model_config.py

@define
class PositivesConfig:
    """Positive-pair sampling for the embedding objective.

    Attributes:
        scope: Which crops are positives of an anchor. One of
            ``aug_view`` (only the anchor's own augmented views — self-supervised:
            every detection is a sample and its own group, so unlabeled data trains;
            needs two DIFFERENT views, i.e. ``use_augmentations_train`` with some
            augmentation configured, and pairs best with ``sampler.kind='random'``),
            ``tracklet`` (same ``(video, track)`` — video-local identity; cross-video
            pairs are UNKNOWN and excluded from the loss), or ``global_id`` (same
            global identity across all videos: a ``sio.Identity`` when the detection
            carries one, else its track name, which then must be globally consistent
            -- declared by ``data_config.identity.track_names_are_global``).
        aug_views: Number of augmented views of each anchor (always positives).
            Fixed at 2 (the standard two-view contrastive setup the ``training_step``
            implements); any other value is unsupported in P1 and raises a
            ``ValueError`` at model construction rather than being silently ignored.
    """

    scope: str = field(
        default="global_id",
        validator=_objective_option("positives", EMBEDDING_POSITIVE_SCOPES),
    )
    aug_views: int = 2

PretrainedConfig

Configuration for an external pretrained backbone (HuggingFace).

Reuses an external pretrained image encoder from the HuggingFace Hub as the model backbone: hierarchical encoders (ConvNeXtV2, ResNet, Swinv2, DINOv3 ConvNeXt, ...) via transformers AutoBackbone, and plain ViTs (DINOv2/v3, CLIP, SigLIP, MAE, ...) via AutoModel hidden states; timm/... ids load through transformers' TimmWrapperModel. See sleap_nn.architectures.pretrained.PretrainedBackbone for the three integration surfaces (hierarchical decoder, encoder-only pooled, plain ViT + pyramid neck + decoder).

Attributes:

Name	Type	Description
source	str	(str) Backbone source. Only "hf" (HuggingFace transformers) is currently supported. Default: "hf".
model_name	str	(str) HuggingFace model id, e.g. "facebook/convnextv2-nano-22k-224", "microsoft/resnet-50", "facebook/dinov2-with-registers-base".
weights	bool	(bool) If True, download and load the pretrained weights; if False, build the architecture with random init from the model config (no network) — useful for tests and cold-start training. Default: True.
mode	str	(str) One of "auto", "decoder" (Case A: hierarchical encoder + sleap decoder; Case C for a plain ViT: pyramid neck + decoder), "encoder" (Case B: encoder-only pooled bottleneck for class-vectors/embedding heads). "auto" picks the decoder whenever the model has a spatial head. Default: "auto".
freeze	bool	(bool) If True, freeze the pretrained encoder (feature extraction; only the decoder/head train). Default: False.
revision	Optional[str]	(str) Optional HuggingFace revision (commit sha / tag) to pin for reproducibility. Default: None.
normalize	bool	(bool) If True, apply model-specific per-channel mean/std normalization inside the backbone forward (the data pipeline only rescales to [0, 1]). Default: True.
image_mean	Optional[List[float]]	(List[float]) Optional explicit per-channel mean (length 3). If None and normalize is set, read from the model's AutoImageProcessor, falling back to ImageNet stats. Default: None.
image_std	Optional[List[float]]	(List[float]) Optional explicit per-channel std (length 3). Default: None.
out_indices	Optional[List[int]]	(List[int]) Optional explicit stage indices to tap (Case A). If None, all stages are requested and deduplicated by stride. Indices follow the source: an HF backbone's stage_names (0 is the stem, so a ConvNeXt's stages are [1, 2, 3, 4]) or a timm/... checkpoint's feature_info (0 is the first stage). Default: None.
in_channels	int	(int) Number of input channels the stem expects. Pretrained stems are 3-channel; grayscale is replicated upstream. Default: 3.
filters_rate	float	(float) Decoder filter growth factor (Case A). Default: 2.0.
convs_per_block	int	(int) Refinement convs per decoder block (Case A). Default: 2.
kernel_size	int	(int) Decoder conv kernel size (Case A). Default: 3.
up_interpolate	bool	(bool) Bilinear upsampling (vs. transposed conv) in the decoder (Case A). Default: True.
output_stride	int	(int) Stride of the finest decoder output (Case A). Default: 2.
max_stride	int	(int) Deepest stride the encoder reaches (32 for hierarchical CNN/Swin; twice the token stride for a plain ViT, i.e. 32 at the default token stride of 16). Default: 32.
layers	Optional[List[int]]	(List[int]) Plain-ViT layers to tap for the feature-pyramid neck, 1-indexed (negative counts from the last). None taps 4 evenly spaced layers ending at the last (vit_neck: dpt) or the last layer (sfp). Default: None.
vit_neck	str	(str) How a plain ViT's single-stride token maps become the stride s/4..2s pyramid the decoder needs: "dpt" (one tapped layer per level) or "sfp" (ViTDet Simple Feature Pyramid from the last tapped layer). Default: "dpt".
neck_channels	Optional[List[int]]	(List[int]) Channels of the 4 neck levels, finest first. None uses [h/8, h/4, h/2, h] for ViT hidden size h. Default: None.
token_stride	Optional[int]	(int) Effective stride of a plain ViT's token grid; the input is resampled so the grid is H / token_stride tokens tall whatever the patch size (power of two). None uses the power of two nearest the patch size (16 for a patch-14 DINOv2). Use 8 for a finer grid at ~4x encoder cost (then max_stride: 16 suffices). Default: None.
gradient_checkpointing	bool	(bool) Enable gradient checkpointing in the encoder while training (less memory, slower). Default: False.
encoder_lr_scale	float	(float) Learning-rate multiplier for the pretrained encoder relative to the neck/decoder/heads. Default: 1.0.
layer_decay	Optional[float]	(float) Optional layer-wise LR decay for a plain ViT encoder: block i of L trains at encoder_lr_scale * layer_decay**(L + 1 - i) times the base LR. Default: None (off).
detail_stem	bool	(bool) Add a light conv stream on the image whose stride 2/4/8 features join a plain ViT's pyramid as decoder skips (real high-resolution detail below the token stride). Default: False.
detail_channels	Optional[List[int]]	(List[int]) Detail-stem channels at strides 2, 4, 8 (exactly three). None uses [32, 64, 128]. Default: None.

Source code in sleap_nn/config/model_config.py

@define
class PretrainedConfig:
    """Configuration for an external pretrained backbone (HuggingFace).

    Reuses an external pretrained image encoder from the HuggingFace Hub as the
    model backbone: hierarchical encoders (ConvNeXtV2, ResNet, Swinv2, DINOv3
    ConvNeXt, ...) via `transformers` `AutoBackbone`, and plain ViTs (DINOv2/v3,
    CLIP, SigLIP, MAE, ...) via `AutoModel` hidden states; ``timm/...`` ids load
    through transformers' `TimmWrapperModel`. See
    ``sleap_nn.architectures.pretrained.PretrainedBackbone`` for the three
    integration surfaces (hierarchical decoder, encoder-only pooled, plain ViT +
    pyramid neck + decoder).

    Attributes:
        source: (str) Backbone source. Only ``"hf"`` (HuggingFace `transformers`)
            is currently supported. *Default*: ``"hf"``.
        model_name: (str) HuggingFace model id, e.g.
            ``"facebook/convnextv2-nano-22k-224"``, ``"microsoft/resnet-50"``,
            ``"facebook/dinov2-with-registers-base"``.
        weights: (bool) If ``True``, download and load the pretrained weights; if
            ``False``, build the architecture with random init from the model
            config (no network) — useful for tests and cold-start training.
            *Default*: ``True``.
        mode: (str) One of ``"auto"``, ``"decoder"`` (Case A: hierarchical encoder
            + sleap decoder; Case C for a plain ViT: pyramid neck + decoder),
            ``"encoder"`` (Case B: encoder-only pooled bottleneck for
            class-vectors/embedding heads). ``"auto"`` picks the decoder whenever
            the model has a spatial head. *Default*: ``"auto"``.
        freeze: (bool) If ``True``, freeze the pretrained encoder (feature
            extraction; only the decoder/head train). *Default*: ``False``.
        revision: (str) Optional HuggingFace revision (commit sha / tag) to pin for
            reproducibility. *Default*: ``None``.
        normalize: (bool) If ``True``, apply model-specific per-channel mean/std
            normalization inside the backbone forward (the data pipeline only
            rescales to ``[0, 1]``). *Default*: ``True``.
        image_mean: (List[float]) Optional explicit per-channel mean (length 3). If
            ``None`` and ``normalize`` is set, read from the model's
            `AutoImageProcessor`, falling back to ImageNet stats. *Default*: ``None``.
        image_std: (List[float]) Optional explicit per-channel std (length 3).
            *Default*: ``None``.
        out_indices: (List[int]) Optional explicit stage indices to tap (Case A).
            If ``None``, all stages are requested and deduplicated by stride.
            Indices follow the source: an HF backbone's ``stage_names`` (``0`` is
            the stem, so a ConvNeXt's stages are ``[1, 2, 3, 4]``) or a
            ``timm/...`` checkpoint's ``feature_info`` (``0`` is the first stage).
            *Default*: ``None``.
        in_channels: (int) Number of input channels the stem expects. Pretrained
            stems are 3-channel; grayscale is replicated upstream. *Default*: ``3``.
        filters_rate: (float) Decoder filter growth factor (Case A). *Default*: ``2.0``.
        convs_per_block: (int) Refinement convs per decoder block (Case A).
            *Default*: ``2``.
        kernel_size: (int) Decoder conv kernel size (Case A). *Default*: ``3``.
        up_interpolate: (bool) Bilinear upsampling (vs. transposed conv) in the
            decoder (Case A). *Default*: ``True``.
        output_stride: (int) Stride of the finest decoder output (Case A).
            *Default*: ``2``.
        max_stride: (int) Deepest stride the encoder reaches (``32`` for
            hierarchical CNN/Swin; twice the token stride for a plain ViT, i.e.
            ``32`` at the default token stride of 16). *Default*: ``32``.
        layers: (List[int]) Plain-ViT layers to tap for the feature-pyramid neck,
            1-indexed (negative counts from the last). ``None`` taps 4 evenly
            spaced layers ending at the last (``vit_neck: dpt``) or the last layer
            (``sfp``). *Default*: ``None``.
        vit_neck: (str) How a plain ViT's single-stride token maps become the
            stride ``s/4..2s`` pyramid the decoder needs: ``"dpt"`` (one tapped
            layer per level) or ``"sfp"`` (ViTDet Simple Feature Pyramid from the
            last tapped layer). *Default*: ``"dpt"``.
        neck_channels: (List[int]) Channels of the 4 neck levels, finest first.
            ``None`` uses ``[h/8, h/4, h/2, h]`` for ViT hidden size ``h``.
            *Default*: ``None``.
        token_stride: (int) Effective stride of a plain ViT's token grid; the
            input is resampled so the grid is ``H / token_stride`` tokens tall
            whatever the patch size (power of two). ``None`` uses the power of two
            nearest the patch size (16 for a patch-14 DINOv2). Use ``8`` for a
            finer grid at ~4x encoder cost (then ``max_stride: 16`` suffices).
            *Default*: ``None``.
        gradient_checkpointing: (bool) Enable gradient checkpointing in the
            encoder while training (less memory, slower). *Default*: ``False``.
        encoder_lr_scale: (float) Learning-rate multiplier for the pretrained
            encoder relative to the neck/decoder/heads. *Default*: ``1.0``.
        layer_decay: (float) Optional layer-wise LR decay for a plain ViT
            encoder: block ``i`` of ``L`` trains at
            ``encoder_lr_scale * layer_decay**(L + 1 - i)`` times the base LR.
            *Default*: ``None`` (off).
        detail_stem: (bool) Add a light conv stream on the image whose stride
            2/4/8 features join a plain ViT's pyramid as decoder skips (real
            high-resolution detail below the token stride). *Default*: ``False``.
        detail_channels: (List[int]) Detail-stem channels at strides 2, 4, 8
            (exactly three). ``None`` uses ``[32, 64, 128]``. *Default*: ``None``.
    """

    source: str = "hf"
    model_name: str = "facebook/convnextv2-nano-22k-224"
    weights: bool = True
    mode: str = "auto"
    freeze: bool = False
    revision: Optional[str] = None
    normalize: bool = True
    image_mean: Optional[List[float]] = None
    image_std: Optional[List[float]] = None
    out_indices: Optional[List[int]] = None
    in_channels: int = 3
    filters_rate: float = 2.0
    convs_per_block: int = 2
    kernel_size: int = 3
    up_interpolate: bool = True
    output_stride: int = 2
    max_stride: int = 32
    layers: Optional[List[int]] = None
    vit_neck: str = "dpt"
    neck_channels: Optional[List[int]] = None
    token_stride: Optional[int] = None
    gradient_checkpointing: bool = False
    encoder_lr_scale: float = 1.0
    layer_decay: Optional[float] = None
    detail_stem: bool = False
    detail_channels: Optional[List[int]] = None

SSDHeadConfig

Multi-scale anchor head of a bbox_ssd (SSD-style) model.

Binds to num_levels feature maps of the backbone at strides output_stride * 2**i (the decoder's top-down path is the feature pyramid; the bottleneck supplies the deepest level). Each level runs a small conv tower and predicts, for each of its anchors, len(categories) + 1 class logits (index 0 is background) and 4 box deltas in SSD's (cx, cy, w, h) encoding. Training matches anchors to boxes by IoU and uses softmax cross-entropy with hard negative mining plus smooth-L1 on the positives (Liu et al. 2016).

Attributes:

Name	Type	Description
categories	Optional[List[str]]	(List[str]) Ordered class names (background excluded). None (default) is filled in at training setup, as for bbox_centernet.
output_stride	int	(int) Stride of the finest pyramid level. Default: 8.
num_levels	int	(int) Number of pyramid levels. The deepest one, output_stride * 2**(num_levels - 1), must not exceed the backbone's max_stride (it is raised to fit, as for any head stride). Default: 3.
anchor_sizes	Optional[List[float]]	(List[float]) Base anchor size (sqrt of the area, in input pixels) per level, finest first. None (default) sizes them at training setup from the training boxes: geometrically spaced between the 5th and 95th percentile of sqrt(w * h). The resolved sizes are written into the saved config so inference uses the same anchors.
anchor_scales	List[float]	(List[float]) Multipliers on each level's base size. Default: (1.0, 1.26, 1.587) (2 ** (0, 1/3, 2/3)).
aspect_ratios	List[float]	(List[float]) Anchor h / w ratios. Default: (0.5, 1.0, 2.0).
pos_iou_threshold	float	(float) An anchor with IoU >= this with some box is a positive for it. Each box's best anchor is always a positive too. Default: 0.5.
neg_iou_threshold	float	(float) Anchors whose best IoU is below this are background; those in between are ignored. Default: 0.4.
neg_pos_ratio	float	(float) Hard negative mining keeps this many of the highest-loss background anchors per positive. Default: 3.0.
box_variances	List[float]	(List[float]) SSD's (cx, cy, w, h) encoding variances. Default: (0.1, 0.1, 0.2, 0.2).
head_channels	int	(int) Channels of each level's conv tower. Default: 128.
head_convs	int	(int) 3x3 conv + ReLU layers in each level's tower before the predictors. Default: 1.
loss_weight	float	(float) Weight of the box regression loss relative to the classification loss. Default: 1.0.

Source code in sleap_nn/config/model_config.py

@define
class SSDHeadConfig:
    """Multi-scale anchor head of a ``bbox_ssd`` (SSD-style) model.

    Binds to ``num_levels`` feature maps of the backbone at strides
    ``output_stride * 2**i`` (the decoder's top-down path is the feature pyramid;
    the bottleneck supplies the deepest level). Each level runs a small conv tower
    and predicts, for each of its anchors, ``len(categories) + 1`` class logits
    (index 0 is background) and 4 box deltas in SSD's ``(cx, cy, w, h)`` encoding.
    Training matches anchors to boxes by IoU and uses softmax cross-entropy with
    hard negative mining plus smooth-L1 on the positives (Liu et al. 2016).

    Attributes:
        categories: (List[str]) Ordered class names (background excluded). ``None``
            (default) is filled in at training setup, as for ``bbox_centernet``.
        output_stride: (int) Stride of the finest pyramid level. *Default*: ``8``.
        num_levels: (int) Number of pyramid levels. The deepest one,
            ``output_stride * 2**(num_levels - 1)``, must not exceed the backbone's
            ``max_stride`` (it is raised to fit, as for any head stride). *Default*:
            ``3``.
        anchor_sizes: (List[float]) Base anchor size (sqrt of the area, in input
            pixels) per level, finest first. ``None`` (default) sizes them at
            training setup from the training boxes: geometrically spaced between the
            5th and 95th percentile of ``sqrt(w * h)``. The resolved sizes are
            written into the saved config so inference uses the same anchors.
        anchor_scales: (List[float]) Multipliers on each level's base size.
            *Default*: ``(1.0, 1.26, 1.587)`` (``2 ** (0, 1/3, 2/3)``).
        aspect_ratios: (List[float]) Anchor ``h / w`` ratios. *Default*:
            ``(0.5, 1.0, 2.0)``.
        pos_iou_threshold: (float) An anchor with IoU >= this with some box is a
            positive for it. Each box's best anchor is always a positive too.
            *Default*: ``0.5``.
        neg_iou_threshold: (float) Anchors whose best IoU is below this are
            background; those in between are ignored. *Default*: ``0.4``.
        neg_pos_ratio: (float) Hard negative mining keeps this many of the
            highest-loss background anchors per positive. *Default*: ``3.0``.
        box_variances: (List[float]) SSD's ``(cx, cy, w, h)`` encoding variances.
            *Default*: ``(0.1, 0.1, 0.2, 0.2)``.
        head_channels: (int) Channels of each level's conv tower. *Default*:
            ``128``.
        head_convs: (int) 3x3 conv + ReLU layers in each level's tower before the
            predictors. *Default*: ``1``.
        loss_weight: (float) Weight of the box regression loss relative to the
            classification loss. *Default*: ``1.0``.
    """

    categories: Optional[List[str]] = None
    output_stride: int = 8
    num_levels: int = field(default=3, validator=validators.ge(1))
    anchor_sizes: Optional[List[float]] = None
    # Tuples, not `field(factory=...)`: OmegaConf cannot call an attrs Factory when
    # it builds this class from its type (see `NegativesConfig.sources`).
    anchor_scales: List[float] = (1.0, 1.26, 1.587)
    aspect_ratios: List[float] = (0.5, 1.0, 2.0)
    pos_iou_threshold: float = 0.5
    neg_iou_threshold: float = 0.4
    neg_pos_ratio: float = 3.0
    box_variances: List[float] = (0.1, 0.1, 0.2, 0.2)
    head_channels: int = 128
    head_convs: int = 1
    loss_weight: float = 1.0

SamplerConfig

Group-aware batch sampler that realizes the objective.

Attributes:

Name	Type	Description
kind	str	pk (P groups x K crops) | within_video (one video per batch, so cross-video pairs never co-occur — the correct video-local sampler; a batch holds min(P, groups in that video) x K crops) | random (aug-view-only / self-supervised).
groups_per_batch	int	P — number of groups (identities/tracklets) per batch.
samples_per_group	int	K — crops per group per batch. The effective batch size is P x K (min(P, groups in the video) x K for within_video), then doubled by two-view aug. The steps in an epoch are counted with that size, so an epoch is one pass over the crops.

Source code in sleap_nn/config/model_config.py

@define
class SamplerConfig:
    """Group-aware batch sampler that realizes the objective.

    Attributes:
        kind: ``pk`` (P groups x K crops) | ``within_video`` (one video per batch, so
            cross-video pairs never co-occur — the correct video-local sampler; a
            batch holds ``min(P, groups in that video) x K`` crops) |
            ``random`` (aug-view-only / self-supervised).
        groups_per_batch: P — number of groups (identities/tracklets) per batch.
        samples_per_group: K — crops per group per batch. The effective batch size is
            ``P x K`` (``min(P, groups in the video) x K`` for ``within_video``), then
            doubled by two-view aug. The steps in an epoch are counted with that size,
            so an epoch is one pass over the crops.
    """

    kind: str = field(
        default="pk", validator=_objective_option("sampler", EMBEDDING_SAMPLER_KINDS)
    )
    groups_per_batch: int = field(default=8, validator=validators.ge(1))
    samples_per_group: int = field(default=16, validator=validators.ge(1))

SegmentationHeadConfig

Configuration for the foreground segmentation head.

Shared by bottomup_segmentation and semantic_segmentation (both a plain foreground head). The loss / target knobs default to the historical behavior, so an unset config trains exactly as before.

Attributes:

Name	Type	Description
output_stride	int	(int) The stride of the output segmentation maps relative to the input image. Default: 2. Setting 1 (dense, full-resolution) avoids the stride-downsample of the foreground target and lets the decoder draw sharp thin structures — recommended for thin/high-res objects (e.g. plant roots).
loss_weight	float	(float) Scalar float used to weigh the loss term for this head during training. Default: 1.0.
bce_weight	float	(float) Weight of the BCE term in the bce-dice foreground loss. Default: 0.5. Tilt toward Dice (e.g. 0.3 BCE / 0.7 Dice) to reduce the easy-background/thick-object dominance for thin foreground.
dice_weight	float	(float) Weight of the Dice term in the bce-dice loss. Default: 0.5.
bce_pos_weight	Optional[float]	(Optional[float]) Positive-class weight for the BCE term. For thin/rare foreground (<1% of pixels), a value >1 (e.g. ~5-20) up-weights the foreground so the head stays confident on faint thin structures. None (default) leaves BCE unweighted.
target_maxpool	bool	(bool) Downsample the foreground target with max-pool semantics (any foreground pixel in a stride cell -> foreground) instead of area-average + 0.5 threshold. Default: False. Set True when output_stride > 1 to keep thin structures that would otherwise erode below 50% cell coverage. Inert at output_stride=1 (no downsample).

Source code in sleap_nn/config/model_config.py

@define
class SegmentationHeadConfig:
    """Configuration for the foreground segmentation head.

    Shared by ``bottomup_segmentation`` and ``semantic_segmentation`` (both a
    plain foreground head). The loss / target knobs default to the historical
    behavior, so an unset config trains exactly as before.

    Attributes:
        output_stride: (int) The stride of the output segmentation maps relative to the
            input image. Default: 2. Setting ``1`` (dense, full-resolution) avoids the
            stride-downsample of the foreground target and lets the decoder draw sharp
            thin structures — recommended for thin/high-res objects (e.g. plant roots).
        loss_weight: (float) Scalar float used to weigh the loss term for this head
            during training. Default: 1.0.
        bce_weight: (float) Weight of the BCE term in the bce-dice foreground loss.
            Default: 0.5. Tilt toward Dice (e.g. 0.3 BCE / 0.7 Dice) to reduce the
            easy-background/thick-object dominance for thin foreground.
        dice_weight: (float) Weight of the Dice term in the bce-dice loss. Default: 0.5.
        bce_pos_weight: (Optional[float]) Positive-class weight for the BCE term. For
            thin/rare foreground (<1% of pixels), a value >1 (e.g. ~5-20) up-weights
            the foreground so the head stays confident on faint thin structures.
            ``None`` (default) leaves BCE unweighted.
        target_maxpool: (bool) Downsample the foreground target with max-pool
            semantics (any foreground pixel in a stride cell -> foreground) instead of
            area-average + 0.5 threshold. Default: ``False``. Set ``True`` when
            ``output_stride`` > 1 to keep thin structures that would otherwise erode
            below 50% cell coverage. Inert at ``output_stride=1`` (no downsample).
    """

    output_stride: int = 2
    loss_weight: float = 1.0
    bce_weight: float = 0.5
    dice_weight: float = 0.5
    bce_pos_weight: Optional[float] = None
    target_maxpool: bool = False

SemanticSegmentationConfig

Head config for whole-frame semantic (foreground) segmentation models.

A single foreground-mask head predicting one binary foreground/background mask over the whole frame — no instance grouping, no center/offset heads. It is the whole-frame analog of centered_instance_segmentation (which runs the same fg-only SegmentationHead on a centroid crop) and the group-free sibling of bottomup_segmentation (which adds center + offset heads to group the foreground into instances). Tiling-compatible.

Reuses the bottom-up SegmentationHeadConfig leaf (output_stride + loss_weight); it carries no anchor_part because there is no crop.

Source code in sleap_nn/config/model_config.py

@define
class SemanticSegmentationConfig:
    """Head config for whole-frame semantic (foreground) segmentation models.

    A single foreground-mask head predicting one binary foreground/background mask
    over the *whole frame* — no instance grouping, no center/offset heads. It is
    the whole-frame analog of ``centered_instance_segmentation`` (which runs the
    same fg-only ``SegmentationHead`` on a centroid crop) and the group-free
    sibling of ``bottomup_segmentation`` (which adds center + offset heads to group
    the foreground into instances). Tiling-compatible.

    Reuses the bottom-up ``SegmentationHeadConfig`` leaf (``output_stride`` +
    ``loss_weight``); it carries no ``anchor_part`` because there is no crop.
    """

    segmentation: Optional[SegmentationHeadConfig] = None

SingleInstanceConfMapsConfig

Single Instance configuration map.

Attributes:

Name	Type	Description
part_names	Optional[List[str]]	(List[str]) None if nodes from sio.Labels file can be used directly. Else provide text name of the body parts (nodes) that the head will be configured to produce. The number of parts determines the number of channels in the output. If not specified, all body parts in the skeleton will be used. This config does not apply for 'PartAffinityFieldsHead'.
sigma	float	(float) Spread of the Gaussian distribution of the confidence maps as a scalar float. Smaller values are more precise but may be difficult to learn as they have a lower density within the image space. Larger values are easier to learn but are less precise with respect to the peak coordinate. This spread is in units of pixels of the model input image, i.e., the image resolution after any input scaling is applied.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution.

Source code in sleap_nn/config/model_config.py

@define
class SingleInstanceConfMapsConfig:
    """Single Instance configuration map.

    Attributes:
        part_names: (List[str]) None if nodes from sio.Labels file can be used directly.
            Else provide text name of the body parts (nodes) that the head will be
            configured to produce. The number of parts determines the number of channels
            in the output. If not specified, all body parts in the skeleton will be used.
            This config does not apply for 'PartAffinityFieldsHead'.
        sigma: (float) Spread of the Gaussian distribution of the confidence maps as a
            scalar float. Smaller values are more precise but may be difficult to learn
            as they have a lower density within the image space. Larger values are
            easier to learn but are less precise with respect to the peak coordinate.
            This spread is in units of pixels of the model input image,
            i.e., the image resolution after any input scaling is applied.
        output_stride: (int) The stride of the output confidence maps relative to the
            input image. This is the reciprocal of the resolution, e.g., an output
            stride of 2 results in confidence maps that are 0.5x the size of the input.
            Increasing this value can considerably speed up model performance and
            decrease memory requirements, at the cost of decreased spatial resolution.
    """

    part_names: Optional[List[str]] = None
    sigma: float = 5.0
    output_stride: int = 1

SingleInstanceConfig

single instance head_config.

Source code in sleap_nn/config/model_config.py

@define
class SingleInstanceConfig:
    """single instance head_config."""

    confmaps: Optional[SingleInstanceConfMapsConfig] = None

SwinTBaseConfig

Bases: SwinTConfig

SwinT configuration for backbone.

Attributes:

Name	Type	Description
pre_trained_weights	Optional[str]	(str) Pretrained weights file name supported only for SwinT backbones. For SwinT, one of ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"].
model_type	str	(str) One of the SwinT architecture types: ["tiny", "small", "base"]. Default: "base".
arch	Optional[dict]	Dictionary of embed dimension, depths and number of heads in each layer. Default is "Tiny architecture". {'embed': 96, 'depths': [2,2,6,2], 'channels':[3, 6, 12, 24]}. Default: None.
max_stride	int	(int) Factor by which input image size is reduced through the layers. This is always 32 for all swint architectures. Default: 32.
patch_size	int	(int) Patch size for the stem layer of SwinT. Default: 4.
stem_patch_stride	int	(int) Stride for the patch. Default: 2.
window_size	int	(int) Window size. Default: 7.
in_channels	int	(int) Number of input channels. Default: 1.
kernel_size	int	(int) Size of the convolutional kernels. Default: 3.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default: 2.
convs_per_block	int	(int) Number of convolutional layers per block. Default: 2.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. Default: 1.

Methods:

Name	Description
validate_model_type	Validate model_type.
validate_pre_trained_weights	Validate pre_trained_weights.

Source code in sleap_nn/config/model_config.py

@define
class SwinTBaseConfig(SwinTConfig):
    """SwinT configuration for backbone.

    Attributes:
        pre_trained_weights: (str) Pretrained weights file name supported only for
            SwinT backbones. For SwinT, one of ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"].
        model_type: (str) One of the SwinT architecture types: ["tiny", "small", "base"]. *Default*: `"base"`.
        arch: Dictionary of embed dimension, depths and number of heads in each layer. Default is "Tiny architecture". {'embed': 96, 'depths': [2,2,6,2], 'channels':[3, 6, 12, 24]}. *Default*: `None`.
        max_stride: (int) Factor by which input image size is reduced through the layers. This is always `32` for all swint architectures. *Default*: `32`.
        patch_size: (int) Patch size for the stem layer of SwinT. *Default*: `4`.
        stem_patch_stride: (int) Stride for the patch. *Default*: `2`.
        window_size: (int) Window size. *Default*: `7`.
        in_channels: (int) Number of input channels. *Default*: `1`.
        kernel_size: (int) Size of the convolutional kernels. *Default*: `3`.
        filters_rate: (float) Factor to adjust the number of filters per block. *Default*: `2`.
        convs_per_block: (int) Number of convolutional layers per block. *Default*: `2`.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. *Default*: `True`.
        output_stride: (int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. *Default*: `1`.

    """

    pre_trained_weights: Optional[str] = field(
        default=None,
        validator=lambda instance, attr, value: instance.validate_pre_trained_weights(
            value
        ),
    )
    model_type: str = field(
        default="base",
        validator=lambda instance, attr, value: instance.validate_model_type(value),
    )
    arch: Optional[dict] = None
    max_stride: int = 32
    patch_size: int = 4
    stem_patch_stride: int = 2
    window_size: int = 7
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 2
    convs_per_block: int = 2
    up_interpolate: bool = True
    output_stride: int = 1

    def validate_model_type(self, value):
        """Validate model_type.

        Ensure model_type is one of "tiny", "small", or "base".
        """
        valid_types = ["tiny", "small", "base"]
        if value not in valid_types:
            message = f"Invalid model_type. Must be one of {valid_types}"
            logger.error(message)
            raise ValueError(message)

    def validate_pre_trained_weights(self, value):
        """Validate pre_trained_weights.

        Check:
        swint_weights are one of
        (
            "Swin_T_Weights",
            "Swin_S_Weights",
            "Swin_B_Weights"
        )
        """
        if value is None:
            return

        swint_weights = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]

        if value not in swint_weights:
            message = (
                f"Invalid pre-trained weights for SwinT. Must be one of {swint_weights}"
            )
            logger.error(message)
            raise ValueError(message)

validate_model_type(value)

Validate model_type.

Ensure model_type is one of "tiny", "small", or "base".

Source code in sleap_nn/config/model_config.py

def validate_model_type(self, value):
    """Validate model_type.

    Ensure model_type is one of "tiny", "small", or "base".
    """
    valid_types = ["tiny", "small", "base"]
    if value not in valid_types:
        message = f"Invalid model_type. Must be one of {valid_types}"
        logger.error(message)
        raise ValueError(message)

validate_pre_trained_weights(value)

Validate pre_trained_weights.

Check: swint_weights are one of ( "Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights" )

Source code in sleap_nn/config/model_config.py

def validate_pre_trained_weights(self, value):
    """Validate pre_trained_weights.

    Check:
    swint_weights are one of
    (
        "Swin_T_Weights",
        "Swin_S_Weights",
        "Swin_B_Weights"
    )
    """
    if value is None:
        return

    swint_weights = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]

    if value not in swint_weights:
        message = (
            f"Invalid pre-trained weights for SwinT. Must be one of {swint_weights}"
        )
        logger.error(message)
        raise ValueError(message)

SwinTConfig

SwinT configuration (tiny) for backbone.

Attributes:

Name	Type	Description
pre_trained_weights	Optional[str]	(str) Pretrained weights file name supported only for SwinT backbones. For SwinT, one of ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"].
model_type	str	(str) One of the SwinT architecture types: ["tiny", "small", "base"]. Default: "tiny".
arch	Optional[dict]	Dictionary of embed dimension, depths and number of heads in each layer. Default is "Tiny architecture". {'embed': 96, 'depths': [2,2,6,2], 'channels':[3, 6, 12, 24]}. Default: None.
max_stride	int	(int) Factor by which input image size is reduced through the layers. This is always 32 for all swint architectures. Default: 32.
patch_size	int	(int) Patch size for the stem layer of SwinT. Default: 4.
stem_patch_stride	int	(int) Stride for the patch. Default: 2.
window_size	int	(int) Window size. Default: 7.
in_channels	int	(int) Number of input channels. Default: 1.
kernel_size	int	(int) Size of the convolutional kernels. Default: 3.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default: 2.
convs_per_block	int	(int) Number of convolutional layers per block. Default: 2.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. Default: 1.

Methods:

Name	Description
validate_model_type	Validate model_type.
validate_pre_trained_weights	Validate pre_trained_weights.

Source code in sleap_nn/config/model_config.py

@define
class SwinTConfig:
    """SwinT configuration (tiny) for backbone.

    Attributes:
        pre_trained_weights: (str) Pretrained weights file name supported only for
            SwinT backbones. For SwinT, one of ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"].
        model_type: (str) One of the SwinT architecture types: ["tiny", "small", "base"]. *Default*: `"tiny"`.
        arch: Dictionary of embed dimension, depths and number of heads in each layer. Default is "Tiny architecture". {'embed': 96, 'depths': [2,2,6,2], 'channels':[3, 6, 12, 24]}. *Default*: `None`.
        max_stride: (int) Factor by which input image size is reduced through the layers. This is always `32` for all swint architectures. *Default*: `32`.
        patch_size: (int) Patch size for the stem layer of SwinT. *Default*: `4`.
        stem_patch_stride: (int) Stride for the patch. *Default*: `2`.
        window_size: (int) Window size. *Default*: `7`.
        in_channels: (int) Number of input channels. *Default*: `1`.
        kernel_size: (int) Size of the convolutional kernels. *Default*: `3`.
        filters_rate: (float) Factor to adjust the number of filters per block. *Default*: `2`.
        convs_per_block: (int) Number of convolutional layers per block. *Default*: `2`.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. *Default*: `True`.
        output_stride: (int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. *Default*: `1`.
    """

    pre_trained_weights: Optional[str] = field(
        default=None,
        validator=lambda instance, attr, value: instance.validate_pre_trained_weights(
            value
        ),
    )
    model_type: str = field(
        default="tiny",
        validator=lambda instance, attr, value: instance.validate_model_type(value),
    )
    arch: Optional[dict] = None
    max_stride: int = 32
    patch_size: int = 4
    stem_patch_stride: int = 2
    window_size: int = 7
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 2
    convs_per_block: int = 2
    up_interpolate: bool = True
    output_stride: int = 1

    def validate_model_type(self, value):
        """Validate model_type.

        Ensure model_type is one of "tiny", "small", or "base".
        """
        valid_types = ["tiny", "small", "base"]
        if value not in valid_types:
            message = f"Invalid model_type. Must be one of {valid_types}"
            logger.error(message)
            raise ValueError(message)

    def validate_pre_trained_weights(self, value):
        """Validate pre_trained_weights.

        Check:
        swint_weights are one of
        (
            "Swin_T_Weights",
            "Swin_S_Weights",
            "Swin_B_Weights"
        )
        """
        if value is None:
            return

        swint_weights = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]

        if value not in swint_weights:
            message = (
                f"Invalid pre-trained weights for SwinT. Must be one of {swint_weights}"
            )
            logger.error(message)
            raise ValueError(message)

validate_model_type(value)

Validate model_type.

Ensure model_type is one of "tiny", "small", or "base".

Source code in sleap_nn/config/model_config.py

def validate_model_type(self, value):
    """Validate model_type.

    Ensure model_type is one of "tiny", "small", or "base".
    """
    valid_types = ["tiny", "small", "base"]
    if value not in valid_types:
        message = f"Invalid model_type. Must be one of {valid_types}"
        logger.error(message)
        raise ValueError(message)

validate_pre_trained_weights(value)

Validate pre_trained_weights.

Check: swint_weights are one of ( "Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights" )

Source code in sleap_nn/config/model_config.py

def validate_pre_trained_weights(self, value):
    """Validate pre_trained_weights.

    Check:
    swint_weights are one of
    (
        "Swin_T_Weights",
        "Swin_S_Weights",
        "Swin_B_Weights"
    )
    """
    if value is None:
        return

    swint_weights = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]

    if value not in swint_weights:
        message = (
            f"Invalid pre-trained weights for SwinT. Must be one of {swint_weights}"
        )
        logger.error(message)
        raise ValueError(message)

SwinTSmallConfig

Bases: SwinTConfig

SwinT configuration (small) for backbone.

Attributes:

Name	Type	Description
pre_trained_weights	Optional[str]	(str) Pretrained weights file name supported only for SwinT backbones. For SwinT, one of ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"].
model_type	str	(str) One of the SwinT architecture types: ["tiny", "small", "base"]. Default: "small".
arch	Optional[dict]	Dictionary of embed dimension, depths and number of heads in each layer. Default is "Tiny architecture". {'embed': 96, 'depths': [2,2,6,2], 'channels':[3, 6, 12, 24]}. Default: None.
max_stride	int	(int) Factor by which input image size is reduced through the layers. This is always 32 for all swint architectures. Default: 32.
patch_size	int	(int) Patch size for the stem layer of SwinT. Default: 4.
stem_patch_stride	int	(int) Stride for the patch. Default: 2.
window_size	int	(int) Window size. Default: 7.
in_channels	int	(int) Number of input channels. Default: 1.
kernel_size	int	(int) Size of the convolutional kernels. Default: 3.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default: 2.
convs_per_block	int	(int) Number of convolutional layers per block. Default: 2.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. Default: 1.

Methods:

Name	Description
validate_model_type	Validate model_type.
validate_pre_trained_weights	Validate pre_trained_weights.

Source code in sleap_nn/config/model_config.py

@define
class SwinTSmallConfig(SwinTConfig):
    """SwinT configuration (small) for backbone.

    Attributes:
        pre_trained_weights: (str) Pretrained weights file name supported only for
            SwinT backbones. For SwinT, one of ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"].
        model_type: (str) One of the SwinT architecture types: ["tiny", "small", "base"]. *Default*: `"small"`.
        arch: Dictionary of embed dimension, depths and number of heads in each layer. Default is "Tiny architecture". {'embed': 96, 'depths': [2,2,6,2], 'channels':[3, 6, 12, 24]}. *Default*: `None`.
        max_stride: (int) Factor by which input image size is reduced through the layers. This is always `32` for all swint architectures. *Default*: `32`.
        patch_size: (int) Patch size for the stem layer of SwinT. *Default*: `4`.
        stem_patch_stride: (int) Stride for the patch. *Default*: `2`.
        window_size: (int) Window size. *Default*: `7`.
        in_channels: (int) Number of input channels. *Default*: `1`.
        kernel_size: (int) Size of the convolutional kernels. *Default*: `3`.
        filters_rate: (float) Factor to adjust the number of filters per block. *Default*: `2`.
        convs_per_block: (int) Number of convolutional layers per block. *Default*: `2`.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. *Default*: `True`.
        output_stride: (int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. *Default*: `1`.
    """

    pre_trained_weights: Optional[str] = field(
        default=None,
        validator=lambda instance, attr, value: instance.validate_pre_trained_weights(
            value
        ),
    )
    model_type: str = field(
        default="small",
        validator=lambda instance, attr, value: instance.validate_model_type(value),
    )
    arch: Optional[dict] = None
    max_stride: int = 32
    patch_size: int = 4
    stem_patch_stride: int = 2
    window_size: int = 7
    in_channels: int = 1
    kernel_size: int = 3
    filters_rate: float = 2
    convs_per_block: int = 2
    up_interpolate: bool = True
    output_stride: int = 1

    def validate_model_type(self, value):
        """Validate model_type.

        Ensure model_type is one of "tiny", "small", or "base".
        """
        valid_types = ["tiny", "small", "base"]
        if value not in valid_types:
            message = f"Invalid model_type. Must be one of {valid_types}"
            logger.error(message)
            raise ValueError(message)

    def validate_pre_trained_weights(self, value):
        """Validate pre_trained_weights.

        Check:
        swint_weights are one of
        (
            "Swin_T_Weights",
            "Swin_S_Weights",
            "Swin_B_Weights"
        )
        """
        if value is None:
            return

        swint_weights = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]

        if value not in swint_weights:
            message = (
                f"Invalid pre-trained weights for SwinT. Must be one of {swint_weights}"
            )
            logger.error(message)
            raise ValueError(message)

validate_model_type(value)

Validate model_type.

Ensure model_type is one of "tiny", "small", or "base".

Source code in sleap_nn/config/model_config.py

def validate_model_type(self, value):
    """Validate model_type.

    Ensure model_type is one of "tiny", "small", or "base".
    """
    valid_types = ["tiny", "small", "base"]
    if value not in valid_types:
        message = f"Invalid model_type. Must be one of {valid_types}"
        logger.error(message)
        raise ValueError(message)

validate_pre_trained_weights(value)

Validate pre_trained_weights.

Check: swint_weights are one of ( "Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights" )

Source code in sleap_nn/config/model_config.py

def validate_pre_trained_weights(self, value):
    """Validate pre_trained_weights.

    Check:
    swint_weights are one of
    (
        "Swin_T_Weights",
        "Swin_S_Weights",
        "Swin_B_Weights"
    )
    """
    if value is None:
        return

    swint_weights = ["Swin_T_Weights", "Swin_S_Weights", "Swin_B_Weights"]

    if value not in swint_weights:
        message = (
            f"Invalid pre-trained weights for SwinT. Must be one of {swint_weights}"
        )
        logger.error(message)
        raise ValueError(message)

TopDownCenteredInstanceMultiClassConfig

Head config for TopDown centered instance ID models.

Source code in sleap_nn/config/model_config.py

@define
class TopDownCenteredInstanceMultiClassConfig:
    """Head config for TopDown centered instance ID models."""

    confmaps: Optional[CenteredInstanceConfMapsConfig] = None
    class_vectors: Optional[ClassVectorsConfig] = None

UNetConfig

UNet config for backbone.

Attributes:

Name	Type	Description
in_channels	int	(int) Number of input channels. Default: 1.
kernel_size	int	(int) Size of the convolutional kernels. Default: 3.
filters	int	(int) Base number of filters in the network. Default: 32.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default: 1.5.
max_stride	int	(int) Scalar integer specifying the maximum stride that the image must be divisible by. Default: 16.
stem_stride	Optional[int]	(int) If not None, will create additional "down" blocks for initial downsampling based on the stride. These will be configured identically to the down blocks below. Default: None.
middle_block	bool	(bool) If True, add an additional block at the end of the encoder. Default: True.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
stacks	int	(int) Number of upsampling blocks in the decoder. Default: 1.
convs_per_block	int	(int) Number of convolutional layers per block. Default: 2.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. Default: 1.

Source code in sleap_nn/config/model_config.py

@define
class UNetConfig:
    """UNet config for backbone.

    Attributes:
        in_channels: (int) Number of input channels. *Default*: `1`.
        kernel_size: (int) Size of the convolutional kernels. *Default*: `3`.
        filters: (int) Base number of filters in the network. *Default*: `32`.
        filters_rate: (float) Factor to adjust the number of filters per block. *Default*: `1.5`.
        max_stride: (int) Scalar integer specifying the maximum stride that the image must be divisible by. *Default*: `16`.
        stem_stride: (int) If not None, will create additional "down" blocks for initial downsampling based on the stride. These will be configured identically to the down blocks below. *Default*: `None`.
        middle_block: (bool) If True, add an additional block at the end of the encoder. *Default*: `True`.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. *Default*: `True`.
        stacks: (int) Number of upsampling blocks in the decoder. *Default*: `1`.
        convs_per_block: (int) Number of convolutional layers per block. *Default*: `2`.
        output_stride: (int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. *Default*: `1`.
    """

    in_channels: int = 1
    kernel_size: int = 3
    filters: int = 32
    filters_rate: float = 1.5
    max_stride: int = 16
    stem_stride: Optional[int] = None
    middle_block: bool = True
    up_interpolate: bool = True
    stacks: int = 1
    convs_per_block: int = 2
    output_stride: int = 1

UNetLargeRFConfig

Bases: UNetConfig

UNet config for backbone with large receptive field.

Attributes:

Name	Type	Description
in_channels	int	(int) Number of input channels. Default: 1.
kernel_size	int	(int) Size of the convolutional kernels. Default: 3.
filters	int	(int) Base number of filters in the network. Default: 24.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default: 1.5.
max_stride	int	(int) Scalar integer specifying the maximum stride that the image must be divisible by. Default: 32.
stem_stride	Optional[int]	(int) If not None, will create additional "down" blocks for initial downsampling based on the stride. These will be configured identically to the down blocks below. Default: None.
middle_block	bool	(bool) If True, add an additional block at the end of the encoder. Default: True.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
stacks	int	(int) Number of upsampling blocks in the decoder. Default: 1.
convs_per_block	int	(int) Number of convolutional layers per block. Default: 2.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. Default: 1.

Source code in sleap_nn/config/model_config.py

@define
class UNetLargeRFConfig(UNetConfig):
    """UNet config for backbone with large receptive field.

    Attributes:
        in_channels: (int) Number of input channels. *Default*: `1`.
        kernel_size: (int) Size of the convolutional kernels. *Default*: `3`.
        filters: (int) Base number of filters in the network. *Default*: `24`.
        filters_rate: (float) Factor to adjust the number of filters per block. *Default*: `1.5`.
        max_stride: (int) Scalar integer specifying the maximum stride that the image must be divisible by. *Default*: `32`.
        stem_stride: (int) If not None, will create additional "down" blocks for initial downsampling based on the stride. These will be configured identically to the down blocks below. *Default*: `None`.
        middle_block: (bool) If True, add an additional block at the end of the encoder. *Default*: `True`.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. *Default*: `True`.
        stacks: (int) Number of upsampling blocks in the decoder. *Default*: `1`.
        convs_per_block: (int) Number of convolutional layers per block. *Default*: `2`.
        output_stride: (int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. *Default*: `1`.
    """

    in_channels: int = 1
    kernel_size: int = 3
    filters: int = 24
    filters_rate: float = 1.5
    max_stride: int = 32
    stem_stride: Optional[int] = None
    middle_block: bool = True
    up_interpolate: bool = True
    stacks: int = 1
    convs_per_block: int = 2
    output_stride: int = 1

UNetMediumRFConfig

Bases: UNetConfig

UNet config for backbone with medium receptive field.

Attributes:

Name	Type	Description
in_channels	int	(int) Number of input channels. Default: 1.
kernel_size	int	(int) Size of the convolutional kernels. Default: 3.
filters	int	(int) Base number of filters in the network. Default: 32.
filters_rate	float	(float) Factor to adjust the number of filters per block. Default: 2.
max_stride	int	(int) Scalar integer specifying the maximum stride that the image must be divisible by. Default: 16.
stem_stride	Optional[int]	(int) If not None, will create additional "down" blocks for initial downsampling based on the stride. These will be configured identically to the down blocks below. Default: None.
middle_block	bool	(bool) If True, add an additional block at the end of the encoder. Default: True.
up_interpolate	bool	(bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. Default: True.
stacks	int	(int) Number of upsampling blocks in the decoder. Default: 1.
convs_per_block	int	(int) Number of convolutional layers per block. Default: 2.
output_stride	int	(int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. Default: 1.

Source code in sleap_nn/config/model_config.py

@define
class UNetMediumRFConfig(UNetConfig):
    """UNet config for backbone with medium receptive field.

    Attributes:
        in_channels: (int) Number of input channels. *Default*: `1`.
        kernel_size: (int) Size of the convolutional kernels. *Default*: `3`.
        filters: (int) Base number of filters in the network. *Default*: `32`.
        filters_rate: (float) Factor to adjust the number of filters per block. *Default*: `2`.
        max_stride: (int) Scalar integer specifying the maximum stride that the image must be divisible by. *Default*: `16`.
        stem_stride: (int) If not None, will create additional "down" blocks for initial downsampling based on the stride. These will be configured identically to the down blocks below. *Default*: `None`.
        middle_block: (bool) If True, add an additional block at the end of the encoder. *Default*: `True`.
        up_interpolate: (bool) If True, use bilinear interpolation instead of transposed convolutions for upsampling. Interpolation is faster but transposed convolutions may be able to learn richer or more complex upsampling to recover details from higher scales. *Default*: `True`.
        stacks: (int) Number of upsampling blocks in the decoder. *Default*: `1`.
        convs_per_block: (int) Number of convolutional layers per block. *Default*: `2`.
        output_stride: (int) The stride of the output confidence maps relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in confidence maps that are 0.5x the size of the input. Increasing this value can considerably speed up model performance and decrease memory requirements, at the cost of decreased spatial resolution. *Default*: `1`.
    """

    in_channels: int = 1
    kernel_size: int = 3
    filters: int = 32
    filters_rate: float = 2
    max_stride: int = 16
    stem_stride: Optional[int] = None
    middle_block: bool = True
    up_interpolate: bool = True
    stacks: int = 1
    convs_per_block: int = 2
    output_stride: int = 1

drop_removed_embedding_keys(leaf, where=_EMBEDDING_LEAF)

Remove the removed embedding keys from an embedding head leaf, in place.

Logs a warning per key removed.

Parameters:

Name	Type	Description	Default
leaf		The head_configs.embedding.embedding node (a dict or DictConfig), or None.	required
where	str	The leaf's config path, for the warning.	_EMBEDDING_LEAF

Returns:

Type	Description
	The names of the keys removed (full config paths).

Source code in sleap_nn/config/model_config.py

def drop_removed_embedding_keys(leaf, where: str = _EMBEDDING_LEAF):
    """Remove the removed embedding keys from an embedding head leaf, in place.

    Logs a warning per key removed.

    Args:
        leaf: The ``head_configs.embedding.embedding`` node (a ``dict`` or
            ``DictConfig``), or ``None``.
        where: The leaf's config path, for the warning.

    Returns:
        The names of the keys removed (full config paths).
    """
    from omegaconf import DictConfig, open_dict

    if leaf is None:
        return []
    removed = []

    def pop(node, keys, path):
        if node is None:
            return
        for key in keys:
            if key in node:
                if isinstance(node, DictConfig):
                    with open_dict(node):
                        node.pop(key)
                else:
                    node.pop(key)
                removed.append(f"{path}.{key}")

    pop(leaf, REMOVED_EMBEDDING_HEAD_KEYS, where)
    objective = leaf.get("objective") if hasattr(leaf, "get") else None
    negatives = objective.get("negatives") if hasattr(objective, "get") else None
    pop(negatives, REMOVED_EMBEDDING_NEGATIVES_KEYS, f"{where}.objective.negatives")
    for key in removed:
        logger.warning(
            f"Ignoring `{key}`: this option was never used and has been removed "
            "(an embedding config saved by an older sleap-nn still carries it). "
            "Delete it from the config to silence this warning."
        )
    return removed

model_mapper(legacy_config)

Map the legacy model configuration to the new model configuration.

Parameters:

Name	Type	Description	Default
legacy_config	dict	A dictionary containing the legacy model configuration.	required

Returns:

Type	Description
ModelConfig	An instance of ModelConfig with the mapped configuration.

Source code in sleap_nn/config/model_config.py

def model_mapper(legacy_config: dict) -> ModelConfig:
    """Map the legacy model configuration to the new model configuration.

    Args:
        legacy_config: A dictionary containing the legacy model configuration.

    Returns:
        An instance of `ModelConfig` with the mapped configuration.
    """
    legacy_config_model = legacy_config.get("model", {})
    backbone_cfg_args = {}
    head_cfg_args = {}
    if legacy_config_model.get("backbone", {}).get("unet", None) is not None:
        backbone_cfg_args["unet"] = UNetConfig(
            filters=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("filters", 32),
            filters_rate=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("filters_rate", 1.5),
            max_stride=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("max_stride", 16),
            stem_stride=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("stem_stride", 16),
            middle_block=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("middle_block", True),
            up_interpolate=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("up_interpolate", True),
            stacks=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("stacks", 1),
            output_stride=legacy_config_model.get("backbone", {})
            .get("unet", {})
            .get("output_stride", 1),
        )

    backbone_cfg = BackboneConfig(**backbone_cfg_args)

    if legacy_config_model.get("heads", {}).get("single_instance", None) is not None:
        head_cfg_args["single_instance"] = SingleInstanceConfig(
            confmaps=SingleInstanceConfMapsConfig(
                part_names=legacy_config_model.get("heads", {})
                .get("single_instance", {})
                .get("part_names", None),
                sigma=legacy_config_model.get("heads", {})
                .get("single_instance", {})
                .get("sigma", 5.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("single_instance", {})
                .get("output_stride", 1),
            )
        )
    if legacy_config_model.get("heads", {}).get("centroid", None) is not None:
        head_cfg_args["centroid"] = CentroidConfig(
            confmaps=CentroidConfMapsConfig(
                anchor_part=legacy_config_model.get("heads", {})
                .get("centroid", {})
                .get("anchor_part", None),
                sigma=legacy_config_model.get("heads", {})
                .get("centroid", {})
                .get("sigma", 5.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("centroid", {})
                .get("output_stride", 1),
            )
        )
    if legacy_config_model.get("heads", {}).get("centered_instance", None) is not None:
        head_cfg_args["centered_instance"] = CenteredInstanceConfig(
            confmaps=CenteredInstanceConfMapsConfig(
                anchor_part=legacy_config_model.get("heads", {})
                .get("centered_instance", {})
                .get("anchor_part", None),
                sigma=legacy_config_model.get("heads", {})
                .get("centered_instance", {})
                .get("sigma", 5.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("centered_instance", {})
                .get("output_stride", 1),
                part_names=legacy_config_model.get("heads", {})
                .get("centered_instance", {})
                .get("part_names", None),
            )
        )
    if legacy_config_model.get("heads", {}).get("multi_instance", None) is not None:
        head_cfg_args["bottomup"] = BottomUpConfig(
            confmaps=BottomUpConfMapsConfig(
                loss_weight=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("confmaps", {})
                .get("loss_weight", 1.0),
                sigma=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("confmaps", {})
                .get("sigma", 5.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("confmaps", {})
                .get("output_stride", 1),
                part_names=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("confmaps", {})
                .get("part_names", None),
            ),
            pafs=PAFConfig(
                edges=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("pafs", {})
                .get("edges", None),
                sigma=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("pafs", {})
                .get("sigma", 15.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("pafs", {})
                .get("output_stride", 1),
                loss_weight=legacy_config_model.get("heads", {})
                .get("multi_instance", {})
                .get("pafs", {})
                .get("loss_weight", 1.0),
            ),
        )
    if (
        legacy_config_model.get("heads", {}).get("multi_class_bottomup", None)
        is not None
    ):
        head_cfg_args["multi_class_bottomup"] = BottomUpMultiClassConfig(
            confmaps=BottomUpConfMapsConfig(
                loss_weight=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("confmaps", {})
                .get("loss_weight", 1.0),
                sigma=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("confmaps", {})
                .get("sigma", 5.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("confmaps", {})
                .get("output_stride", 1),
                part_names=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("confmaps", {})
                .get("part_names", None),
            ),
            class_maps=ClassMapConfig(
                sigma=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("class_maps", {})
                .get("sigma", 15.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("class_maps", {})
                .get("output_stride", 1),
                loss_weight=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("class_maps", {})
                .get("loss_weight", 1.0),
                classes=legacy_config_model.get("heads", {})
                .get("multi_class_bottomup", {})
                .get("class_maps", {})
                .get("classes", None),
            ),
        )

    if (
        legacy_config_model.get("heads", {}).get("multi_class_topdown", None)
        is not None
    ):
        head_cfg_args["multi_class_topdown"] = TopDownCenteredInstanceMultiClassConfig(
            confmaps=CenteredInstanceConfMapsConfig(
                loss_weight=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("confmaps", {})
                .get("loss_weight", 1.0),
                sigma=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("confmaps", {})
                .get("sigma", 5.0),
                output_stride=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("confmaps", {})
                .get("output_stride", 1),
                anchor_part=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("confmaps", {})
                .get("anchor_part", None),
                part_names=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("confmaps", {})
                .get("part_names", None),
            ),
            class_vectors=ClassVectorsConfig(
                classes=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("class_vectors", {})
                .get("classes", None),
                num_fc_layers=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("class_vectors", {})
                .get("num_fc_layers", 2),
                num_fc_units=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("class_vectors", {})
                .get("num_fc_units", 1024),
                global_pool=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("class_vectors", {})
                .get("global_pool", True),
                output_stride=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("class_vectors", {})
                .get("output_stride", 1),
                loss_weight=legacy_config_model.get("heads", {})
                .get("multi_class_topdown", {})
                .get("class_vectors", {})
                .get("loss_weight", 1.0),
            ),
        )

    head_cfg = HeadConfig(**head_cfg_args)

    trained_weights_path = legacy_config_model.get("base_checkpoint", None)

    return ModelConfig(
        backbone_config=backbone_cfg,
        head_configs=head_cfg,
        pretrained_backbone_weights=trained_weights_path,
        pretrained_head_weights=trained_weights_path,
    )

resolve_embedding_objective(objective)

Return the embedding objective with every sub-config filled in and validated.

The single reading of head_configs.embedding.embedding.objective shared by the EmbeddingLightningModule, the identity validator and the batch sampler, so they cannot disagree about a default (negatives.sources defaulting to None once meant "both sources" to one reader and "no negatives" to another). The defaults are the attrs defaults above.

Parameters:

Name	Type	Description	Default
objective		The objective node: a DictConfig, a plain dict, an ObjectiveConfig, or None (all defaults). Absent or None sub-configs and fields take their defaults.	required

Returns:

Type	Description
ObjectiveConfig	A fully-populated ObjectiveConfig.

Raises:

Type	Description
ValueError	If a field has an unknown value (e.g. a misspelled positives.scope, negatives.sources entry, loss.name or sampler.kind) or the node has an unknown key.

Source code in sleap_nn/config/model_config.py

def resolve_embedding_objective(objective) -> ObjectiveConfig:
    """Return the embedding objective with every sub-config filled in and validated.

    The single reading of ``head_configs.embedding.embedding.objective`` shared by
    the ``EmbeddingLightningModule``, the identity validator and the batch sampler, so
    they cannot disagree about a default (``negatives.sources`` defaulting to ``None``
    once meant "both sources" to one reader and "no negatives" to another). The
    defaults are the attrs defaults above.

    Args:
        objective: The ``objective`` node: a ``DictConfig``, a plain ``dict``, an
            ``ObjectiveConfig``, or ``None`` (all defaults). Absent or ``None``
            sub-configs and fields take their defaults.

    Returns:
        A fully-populated ``ObjectiveConfig``.

    Raises:
        ValueError: If a field has an unknown value (e.g. a misspelled
            ``positives.scope``, ``negatives.sources`` entry, ``loss.name`` or
            ``sampler.kind``) or the node has an unknown key.
    """
    from omegaconf import DictConfig, OmegaConf

    if isinstance(objective, ObjectiveConfig):
        objective = asdict(objective)
    elif isinstance(objective, DictConfig):
        objective = OmegaConf.to_container(objective, resolve=True)
    objective = dict(objective or {})

    def build(cls, value):
        if value is None:
            return cls()
        if not isinstance(value, dict):
            value = asdict(value)
        return cls(**{k: v for k, v in value.items() if v is not None})

    negatives = objective.pop("negatives", None)
    if isinstance(negatives, dict):
        # A config saved before the removal still carries the removed keys.
        negatives = {
            k: v
            for k, v in negatives.items()
            if k not in REMOVED_EMBEDDING_NEGATIVES_KEYS
        }
    try:
        return ObjectiveConfig(
            positives=build(PositivesConfig, objective.pop("positives", None)),
            negatives=build(NegativesConfig, negatives),
            loss=build(LossConfig, objective.pop("loss", None)),
            sampler=build(SamplerConfig, objective.pop("sampler", None)),
            **{k: v for k, v in objective.items() if v is not None},
        )
    except (TypeError, ValueError) as e:
        message = str(e)
        if not message.startswith("head_configs."):  # not already a named key
            message = f"Invalid head_configs.embedding.embedding.objective: {message}"
        raise ValueError(message) from e

ssd_level_strides(ssd_config)

Pyramid level strides of an SSDHeadConfig (or its DictConfig), finest first.

Source code in sleap_nn/config/model_config.py

def ssd_level_strides(ssd_config) -> List[int]:
    """Pyramid level strides of an ``SSDHeadConfig`` (or its DictConfig), finest first."""
    return [
        int(ssd_config.output_stride) * 2**i for i in range(int(ssd_config.num_levels))
    ]