Skip to content

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
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).

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
model_mapper

Map the legacy model configuration to the new model configuration.

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 str

Pooling over the encoder feature map. One of gem (generalized-mean, learnable exponent), max, avg.

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.

loss_weight float

Scalar loss weight.

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.

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``.
        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``.
        loss_weight: Scalar loss weight.
        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.
        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: str = "gem"
    normalize: bool = True
    output_stride: int = 32
    loss_weight: float = 1.0
    freeze_backbone: bool = False
    anchor_part: Optional[str] = None
    centroid_method: Optional[str] = None
    centroid_fallback: Optional[str] = None
    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.

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`.
    """

    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

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 = "supcon"
    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 Optional[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).

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).

proximity_filter_px Optional[float]

Reserved (P2); unused in P1.

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``).
        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).
        proximity_filter_px: Reserved (P2); unused in P1.
    """

    sources: Optional[List[str]] = None  # default ["same_frame", "in_batch"]
    exclude_same_track: bool = True
    restrict_same_video: bool = False
    proximity_filter_px: Optional[float] = None

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, no identity labels), tracklet (same (video, track) — video-local identity; cross-video pairs are UNKNOWN and excluded from the loss), or global_id (same track name across all videos — requires globally consistent names; gated 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,
            no identity labels), ``tracklet`` (same ``(video, track)`` — video-local
            identity; cross-video pairs are UNKNOWN and excluded from the loss), or
            ``global_id`` (same track name across all videos — requires globally
            consistent names; gated 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 = "global_id"
    aug_views: int = 2

PretrainedConfig

Configuration for an external pretrained backbone (HuggingFace).

Reuses an external pretrained image encoder (ConvNeXtV2, ResNet, Swinv2, DINOv2, ...) as the model backbone via transformers AutoBackbone. See sleap_nn.architectures.pretrained.PretrainedBackbone for the two integration surfaces (hierarchical decoder vs. encoder-only pooled).

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 for spatial heads), "encoder" (Case B: encoder-only pooled bottleneck for class-vectors/embedding heads). 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. 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; patch size for a ViT). Default: 32.

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 (ConvNeXtV2, ResNet, Swinv2,
    DINOv2, ...) as the model backbone via `transformers` `AutoBackbone`. See
    ``sleap_nn.architectures.pretrained.PretrainedBackbone`` for the two
    integration surfaces (hierarchical decoder vs. encoder-only pooled).

    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 for spatial heads), ``"encoder"`` (Case B: encoder-only
            pooled bottleneck for class-vectors/embedding heads). *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.
            *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; patch size for a ViT). *Default*: ``32``.
    """

    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

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) | 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 (then doubled by two-view aug).

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) |
            ``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`` (then doubled by two-view aug).
    """

    kind: str = "pk"
    groups_per_batch: int = 8
    samples_per_group: int = 16

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

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,
    )