heads

sleap_nn.architectures.heads

Model head definitions for defining model output types.

Classes:

Name	Description
CenterOffsetHead	Head for predicting per-pixel offset vectors to instance centers.
CenteredInstanceConfmapsHead	Head for specifying centered instance confidence maps.
CentroidConfmapsHead	Head for specifying instance centroid confidence maps.
ClassMapsHead	Head for specifying class identity maps.
ClassVectorsHead	Head for specifying classification heads.
EmbeddingHead	Head for crop -> embedding-vector (re-ID) models.
GeM	Generalized-mean pooling: (mean(x.clamp(min=eps)^p))^(1/p) over HxW.
Head	Base class for model output heads.
InstanceCenterHead	Head for predicting instance center heatmaps.
L2Norm	L2-normalize along dim (so embeddings live on the unit hypersphere).
MultiInstanceConfmapsHead	Head for specifying multi-instance confidence maps.
OffsetRefinementHead	Head for specifying offset refinement maps.
PartAffinityFieldsHead	Head for specifying multi-instance part affinity fields.
SegmentationHead	Head for predicting binary foreground segmentation masks.
SingleInstanceConfmapsHead	Head for specifying single instance confidence maps.

CenterOffsetHead

Bases: Head

Head for predicting per-pixel offset vectors to instance centers.

Outputs a 2-channel map where each pixel's value is (dx, dy) pointing from the pixel to its instance's center. Only meaningful on foreground pixels.

Attributes:

Name	Type	Description
output_stride		Stride of the output head tensor.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

class CenterOffsetHead(Head):
    """Head for predicting per-pixel offset vectors to instance centers.

    Outputs a 2-channel map where each pixel's value is (dx, dy) pointing
    from the pixel to its instance's center. Only meaningful on foreground pixels.

    Attributes:
        output_stride: Stride of the output head tensor.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        output_stride: int = 2,
        loss_weight: float = 0.1,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 2

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "smooth_l1"

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(output_stride=2, loss_weight=0.1)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    output_stride: int = 2,
    loss_weight: float = 0.1,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)

CenteredInstanceConfmapsHead

Bases: Head

Head for specifying centered instance confidence maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
anchor_part		Name of the part to use as an anchor node. If not specified, the bounding box centroid will be used.
centroid_method		How the centroid is derived from the instance's points -- "center_of_mass", "bbox_center", "geometric_median" or "anchor"; None infers it from anchor_part. Data-pipeline metadata only (does not affect the head tensor); stored here so it rides the checkpoint and inference can reproduce the training geometry. See sleap_nn.data.instance_centroids.
centroid_fallback		Reduce method used when anchor_part is not visible. Passthrough metadata, as centroid_method.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class CenteredInstanceConfmapsHead(Head):
    """Head for specifying centered instance confidence maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        anchor_part: Name of the part to use as an anchor node. If not specified, the
            bounding box centroid will be used.
        centroid_method: How the centroid is derived from the instance's points --
            ``"center_of_mass"``, ``"bbox_center"``, ``"geometric_median"`` or
            ``"anchor"``; ``None`` infers it from ``anchor_part``. Data-pipeline
            metadata only (does not affect the head tensor); stored here so it
            rides the checkpoint and inference can reproduce the training
            geometry. See ``sleap_nn.data.instance_centroids``.
        centroid_fallback: Reduce method used when ``anchor_part`` is not visible.
            Passthrough metadata, as ``centroid_method``.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        anchor_part: Optional[Text] = None,
        centroid_method: Optional[Text] = None,
        centroid_fallback: Optional[Text] = None,
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.anchor_part = anchor_part
        # Data-pipeline metadata only (does not affect the head tensor); see
        # `CentroidConfmapsHead` and `sleap_nn.data.instance_centroids`.
        self.centroid_method = centroid_method
        self.centroid_fallback = centroid_fallback
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.part_names)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
    ) -> "CenteredInstanceConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            part_names: 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. This must be provided if the `part_names`
                attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.part_names is not None:
            part_names = config.part_names
        elif part_names is None:
            message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            anchor_part=config.anchor_part,
            centroid_method=config.get("centroid_method", None),
            centroid_fallback=config.get("centroid_fallback", None),
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, anchor_part=None, centroid_method=None, centroid_fallback=None, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    anchor_part: Optional[Text] = None,
    centroid_method: Optional[Text] = None,
    centroid_fallback: Optional[Text] = None,
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.anchor_part = anchor_part
    # Data-pipeline metadata only (does not affect the head tensor); see
    # `CentroidConfmapsHead` and `sleap_nn.data.instance_centroids`.
    self.centroid_method = centroid_method
    self.centroid_fallback = centroid_fallback
    self.sigma = sigma

from_config(config, part_names=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		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. This must be provided if the part_names attribute of the configuration is not set.

Returns:

Type	Description
CenteredInstanceConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
) -> "CenteredInstanceConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        part_names: 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. This must be provided if the `part_names`
            attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.part_names is not None:
        part_names = config.part_names
    elif part_names is None:
        message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        anchor_part=config.anchor_part,
        centroid_method=config.get("centroid_method", None),
        centroid_fallback=config.get("centroid_fallback", None),
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

CentroidConfmapsHead

Bases: Head

Head for specifying instance centroid confidence maps.

Attributes:

Name	Type	Description
anchor_part		Name of the part to use as an anchor node. If not specified, the bounding box centroid will be used.
centroid_source		Data-pipeline setting for which centroid the model is trained to predict ("user" / "computed" / None to infer). The head output is identical regardless; it is stored here only so the head can be built from the confmaps config, which co-locates it with anchor_part. See CentroidConfMapsConfig and resolve_centroid_source.
centroid_method		How the centroid is derived from the instance's points -- "center_of_mass", "bbox_center", "geometric_median" or "anchor"; None infers it from anchor_part. Data-pipeline metadata only (does not affect the head tensor); stored here so it rides the checkpoint and inference can reproduce the training geometry. See sleap_nn.data.instance_centroids.
centroid_fallback		Reduce method used when anchor_part is not visible. Passthrough metadata, as centroid_method.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.
use_sigmoid_activation		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 a focal-style loss (see CentroidConfMapsConfig.focal_loss_alpha), which needs Ŷ ∈ (0, 1) for its log(Ŷ)/log(1-Ŷ) terms, and so that peak_threshold-based inference (which compares raw confmap values against a fixed cutoff) stays meaningful. Default False (plain "identity" activation, i.e. no change from existing behavior) -- same opt-in pattern already used by SegmentationHead.
focal_loss_alpha		Training-loss metadata only (does not affect the head tensor); kept so the head is constructible from **head_config.confmaps. See CentroidConfMapsConfig.focal_loss_alpha -- actually consumed by CentroidLightningModule, not this head.
focal_loss_beta		Same as focal_loss_alpha -- passthrough only.
focal_loss_pos_threshold		Same as focal_loss_alpha -- passthrough only.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class CentroidConfmapsHead(Head):
    """Head for specifying instance centroid confidence maps.

    Attributes:
        anchor_part: Name of the part to use as an anchor node. If not specified, the
            bounding box centroid will be used.
        centroid_source: Data-pipeline setting for which centroid the model is
            trained to predict (``"user"`` / ``"computed"`` / ``None`` to infer).
            The head output is identical regardless; it is stored here only so
            the head can be built from the confmaps config, which co-locates it
            with ``anchor_part``. See ``CentroidConfMapsConfig`` and
            ``resolve_centroid_source``.
        centroid_method: How the centroid is derived from the instance's points --
            ``"center_of_mass"``, ``"bbox_center"``, ``"geometric_median"`` or
            ``"anchor"``; ``None`` infers it from ``anchor_part``. Data-pipeline
            metadata only (does not affect the head tensor); stored here so it
            rides the checkpoint and inference can reproduce the training
            geometry. See ``sleap_nn.data.instance_centroids``.
        centroid_fallback: Reduce method used when ``anchor_part`` is not visible.
            Passthrough metadata, as ``centroid_method``.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
        use_sigmoid_activation: 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 a focal-style loss (see
            `CentroidConfMapsConfig.focal_loss_alpha`), which needs `Ŷ ∈ (0, 1)` for its
            `log(Ŷ)`/`log(1-Ŷ)` terms, and so that `peak_threshold`-based inference
            (which compares raw confmap values against a fixed cutoff) stays meaningful.
            Default `False` (plain "identity" activation, i.e. no change from existing
            behavior) -- same opt-in pattern already used by `SegmentationHead`.
        focal_loss_alpha: Training-loss metadata only (does not affect the head
            tensor); kept so the head is constructible from
            ``**head_config.confmaps``. See ``CentroidConfMapsConfig.focal_loss_alpha``
            -- actually consumed by ``CentroidLightningModule``, not this head.
        focal_loss_beta: Same as ``focal_loss_alpha`` -- passthrough only.
        focal_loss_pos_threshold: Same as ``focal_loss_alpha`` -- passthrough only.
    """

    def __init__(
        self,
        anchor_part: Optional[Text] = None,
        centroid_source: Optional[Text] = None,
        centroid_method: Optional[Text] = None,
        centroid_fallback: Optional[Text] = None,
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
        use_sigmoid_activation: bool = False,
        focal_loss_alpha: float = 0.0,
        focal_loss_beta: float = 4.0,
        focal_loss_pos_threshold: float = 0.5,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.anchor_part = anchor_part
        # Data-pipeline metadata only (does not affect the head tensor); kept so
        # the head is constructible from ``**head_config.confmaps``.
        self.centroid_source = centroid_source
        self.centroid_method = centroid_method
        self.centroid_fallback = centroid_fallback
        self.sigma = sigma
        self.use_sigmoid_activation = use_sigmoid_activation
        # Training-loss metadata only (consumed by `CentroidLightningModule`, not
        # this head) -- kept here purely so the head is constructible from
        # `**head_config.confmaps`, which co-locates them with `use_sigmoid_activation`.
        self.focal_loss_alpha = focal_loss_alpha
        self.focal_loss_beta = focal_loss_beta
        self.focal_loss_pos_threshold = focal_loss_pos_threshold

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 1

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "sigmoid" if self.use_sigmoid_activation else "identity"

    @classmethod
    def from_config(cls, config: DictConfig) -> "CentroidConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.

        Returns:
            The instantiated head with the specified configuration options.
        """
        return cls(
            anchor_part=config.anchor_part,
            centroid_source=config.get("centroid_source", None),
            centroid_method=config.get("centroid_method", None),
            centroid_fallback=config.get("centroid_fallback", None),
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
            use_sigmoid_activation=getattr(config, "use_sigmoid_activation", False),
            focal_loss_alpha=getattr(config, "focal_loss_alpha", 0.0),
            focal_loss_beta=getattr(config, "focal_loss_beta", 4.0),
            focal_loss_pos_threshold=getattr(config, "focal_loss_pos_threshold", 0.5),
        )

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

__init__(anchor_part=None, centroid_source=None, centroid_method=None, centroid_fallback=None, sigma=5.0, output_stride=1, loss_weight=1.0, use_sigmoid_activation=False, focal_loss_alpha=0.0, focal_loss_beta=4.0, focal_loss_pos_threshold=0.5)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    anchor_part: Optional[Text] = None,
    centroid_source: Optional[Text] = None,
    centroid_method: Optional[Text] = None,
    centroid_fallback: Optional[Text] = None,
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
    use_sigmoid_activation: bool = False,
    focal_loss_alpha: float = 0.0,
    focal_loss_beta: float = 4.0,
    focal_loss_pos_threshold: float = 0.5,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.anchor_part = anchor_part
    # Data-pipeline metadata only (does not affect the head tensor); kept so
    # the head is constructible from ``**head_config.confmaps``.
    self.centroid_source = centroid_source
    self.centroid_method = centroid_method
    self.centroid_fallback = centroid_fallback
    self.sigma = sigma
    self.use_sigmoid_activation = use_sigmoid_activation
    # Training-loss metadata only (consumed by `CentroidLightningModule`, not
    # this head) -- kept here purely so the head is constructible from
    # `**head_config.confmaps`, which co-locates them with `use_sigmoid_activation`.
    self.focal_loss_alpha = focal_loss_alpha
    self.focal_loss_beta = focal_loss_beta
    self.focal_loss_pos_threshold = focal_loss_pos_threshold

from_config(config) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.

Returns:

Type	Description
CentroidConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(cls, config: DictConfig) -> "CentroidConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.

    Returns:
        The instantiated head with the specified configuration options.
    """
    return cls(
        anchor_part=config.anchor_part,
        centroid_source=config.get("centroid_source", None),
        centroid_method=config.get("centroid_method", None),
        centroid_fallback=config.get("centroid_fallback", None),
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
        use_sigmoid_activation=getattr(config, "use_sigmoid_activation", False),
        focal_loss_alpha=getattr(config, "focal_loss_alpha", 0.0),
        focal_loss_beta=getattr(config, "focal_loss_beta", 4.0),
        focal_loss_pos_threshold=getattr(config, "focal_loss_pos_threshold", 0.5),
    )

ClassMapsHead

Bases: Head

Head for specifying class identity maps.

Attributes:

Name	Type	Description
classes		List of string names of the classes.
sigma		Spread of the class maps around each node.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.
class_output		How the classes map to sleap-io objects ("track" / "identity"). Carried through verbatim from the head config; unused by the architecture/loss. Accepted so the config can be splatted into the constructor. Appended AFTER the pre-existing arguments so a positional ClassMapsHead(classes, 5.0) still binds sigma.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class ClassMapsHead(Head):
    """Head for specifying class identity maps.

    Attributes:
        classes: List of string names of the classes.
        sigma: Spread of the class maps around each node.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
        class_output: How the classes map to sleap-io objects (``"track"`` /
            ``"identity"``). Carried through verbatim from the head config; unused
            by the architecture/loss. Accepted so the config can be splatted into
            the constructor. Appended AFTER the pre-existing arguments so a
            positional ``ClassMapsHead(classes, 5.0)`` still binds ``sigma``.
    """

    def __init__(
        self,
        classes: List[Text],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
        class_output: str = "track",
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.classes = classes
        self.class_output = class_output
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.classes)

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "sigmoid"

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        classes: Optional[List[Text]] = None,
    ) -> "ClassMapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.
            classes: List of string names of the classes that this head will predict.
                This must be set if the `classes` attribute of the configuration is not
                set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.classes is not None:
            classes = config.classes
        return cls(
            classes=classes,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

__init__(classes, sigma=5.0, output_stride=1, loss_weight=1.0, class_output='track')

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    classes: List[Text],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
    class_output: str = "track",
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.classes = classes
    self.class_output = class_output
    self.sigma = sigma

from_config(config, classes=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
classes		List of string names of the classes that this head will predict. This must be set if the classes attribute of the configuration is not set.

Returns:

Type	Description
ClassMapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    classes: Optional[List[Text]] = None,
) -> "ClassMapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.
        classes: List of string names of the classes that this head will predict.
            This must be set if the `classes` attribute of the configuration is not
            set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.classes is not None:
        classes = config.classes
    return cls(
        classes=classes,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

ClassVectorsHead

Bases: Head

Head for specifying classification heads.

Attributes:

Name	Type	Description
classes		List of string names of the classes.
num_fc_layers		Number of fully connected layers after flattening input features.
num_fc_units		Number of units (dimensions) in fully connected layers prior to classification output.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.
class_output		How the classes map to sleap-io objects ("track" / "identity"). Carried through verbatim from the head config; unused by the architecture/loss. Accepted so the config can be splatted into the constructor. Appended AFTER the pre-existing arguments so a positional ClassVectorsHead(classes, 2) still binds num_fc_layers.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.
make_head	Make head output tensor from input feature tensor.

Source code in sleap_nn/architectures/heads.py

class ClassVectorsHead(Head):
    """Head for specifying classification heads.

    Attributes:
        classes: List of string names of the classes.
        num_fc_layers: Number of fully connected layers after flattening input features.
        num_fc_units: Number of units (dimensions) in fully connected layers prior to
            classification output.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
        class_output: How the classes map to sleap-io objects (``"track"`` /
            ``"identity"``). Carried through verbatim from the head config; unused
            by the architecture/loss. Accepted so the config can be splatted into
            the constructor. Appended AFTER the pre-existing arguments so a
            positional ``ClassVectorsHead(classes, 2)`` still binds
            ``num_fc_layers``.
    """

    def __init__(
        self,
        classes: List[Text],
        num_fc_layers: int = 1,
        num_fc_units: int = 64,
        global_pool: bool = True,
        output_stride: int = 1,
        loss_weight: float = 1.0,
        class_output: str = "track",
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.classes = classes
        self.class_output = class_output
        self.num_fc_layers = num_fc_layers
        self.num_fc_units = num_fc_units
        self.global_pool = global_pool

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.classes)

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "softmax"

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "categorical_crossentropy"

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        classes: Optional[List[Text]] = None,
    ) -> "ClassVectorsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.
            classes: List of string names of the classes that this head will predict.
                This must be set if the `classes` attribute of the configuration is not
                set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.classes is not None:
            classes = config.classes
        return cls(
            classes=classes,
            num_fc_layers=config.num_fc_layers,
            num_fc_units=config.num_fc_units,
            global_pool=config.global_pool,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

    def make_head(self, x_in: int) -> nn.Sequential:
        """Make head output tensor from input feature tensor.

        Args:
            x_in: An int for the input shape after applying AdaptiveMaxPool2d on dim=1, assuming inputs of shape (B, C, H, W).

        Returns:
            A `nn.Sequential` with the correct shape for the head.
        """
        from collections import OrderedDict

        module_dict = OrderedDict()

        if self.global_pool:
            module_dict[f"pre_classification_global_pool"] = nn.AdaptiveMaxPool2d(1)

        module_dict[f"pre_classification_flatten"] = nn.Flatten(start_dim=1)

        for i in range(self.num_fc_layers):
            if i == 0:
                module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                    x_in, self.num_fc_units
                )
            else:
                module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                    self.num_fc_units, self.num_fc_units
                )
            module_dict[f"pre_classification{i}_relu"] = get_act_fn("relu")

        module_dict[f"ClassVectorsHead"] = nn.Linear(self.num_fc_units, self.channels)
        module_dict[f"softmax"] = get_act_fn("softmax")

        return nn.Sequential(module_dict)

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(classes, num_fc_layers=1, num_fc_units=64, global_pool=True, output_stride=1, loss_weight=1.0, class_output='track')

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    classes: List[Text],
    num_fc_layers: int = 1,
    num_fc_units: int = 64,
    global_pool: bool = True,
    output_stride: int = 1,
    loss_weight: float = 1.0,
    class_output: str = "track",
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.classes = classes
    self.class_output = class_output
    self.num_fc_layers = num_fc_layers
    self.num_fc_units = num_fc_units
    self.global_pool = global_pool

from_config(config, classes=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
classes		List of string names of the classes that this head will predict. This must be set if the classes attribute of the configuration is not set.

Returns:

Type	Description
ClassVectorsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    classes: Optional[List[Text]] = None,
) -> "ClassVectorsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.
        classes: List of string names of the classes that this head will predict.
            This must be set if the `classes` attribute of the configuration is not
            set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.classes is not None:
        classes = config.classes
    return cls(
        classes=classes,
        num_fc_layers=config.num_fc_layers,
        num_fc_units=config.num_fc_units,
        global_pool=config.global_pool,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

make_head(x_in)

Make head output tensor from input feature tensor.

Parameters:

Name	Type	Description	Default
x_in	int	An int for the input shape after applying AdaptiveMaxPool2d on dim=1, assuming inputs of shape (B, C, H, W).	required

Returns:

Type	Description
Sequential	A nn.Sequential with the correct shape for the head.

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: int) -> nn.Sequential:
    """Make head output tensor from input feature tensor.

    Args:
        x_in: An int for the input shape after applying AdaptiveMaxPool2d on dim=1, assuming inputs of shape (B, C, H, W).

    Returns:
        A `nn.Sequential` with the correct shape for the head.
    """
    from collections import OrderedDict

    module_dict = OrderedDict()

    if self.global_pool:
        module_dict[f"pre_classification_global_pool"] = nn.AdaptiveMaxPool2d(1)

    module_dict[f"pre_classification_flatten"] = nn.Flatten(start_dim=1)

    for i in range(self.num_fc_layers):
        if i == 0:
            module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                x_in, self.num_fc_units
            )
        else:
            module_dict[f"pre_classification{i}_fc"] = nn.Linear(
                self.num_fc_units, self.num_fc_units
            )
        module_dict[f"pre_classification{i}_relu"] = get_act_fn("relu")

    module_dict[f"ClassVectorsHead"] = nn.Linear(self.num_fc_units, self.channels)
    module_dict[f"softmax"] = get_act_fn("softmax")

    return nn.Sequential(module_dict)

EmbeddingHead

Bases: Head

Head for crop -> embedding-vector (re-ID) models.

Mirrors ClassVectorsHead (a pooled, non-spatial head): [pool] -> Flatten -> num_fc_layers x (Linear+ReLU) -> Linear(embedding_dim) -> [L2Norm]. The pooled feature comes from the backbone's middle_output (lone head, empty decoder).

Attributes:

Name	Type	Description
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: gem | max | avg.
normalize		L2-normalize the output embedding.
output_stride		Should equal the backbone max_stride (so the decoder is empty and the head taps middle_output).
loss_weight		Weight of the loss term for this head.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a head-leaf configuration.
make_head	Make the head output module from the pooled-feature input channels.

Source code in sleap_nn/architectures/heads.py

class EmbeddingHead(Head):
    """Head for crop -> embedding-vector (re-ID) models.

    Mirrors ``ClassVectorsHead`` (a pooled, non-spatial head): ``[pool] -> Flatten ->
    num_fc_layers x (Linear+ReLU) -> Linear(embedding_dim) -> [L2Norm]``. The pooled
    feature comes from the backbone's ``middle_output`` (lone head, empty decoder).

    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: ``gem`` | ``max`` | ``avg``.
        normalize: L2-normalize the output embedding.
        output_stride: Should equal the backbone max_stride (so the decoder is empty
            and the head taps ``middle_output``).
        loss_weight: Weight of the loss term for this head.
    """

    def __init__(
        self,
        embedding_dim: int = 128,
        num_fc_layers: int = 1,
        num_fc_units: int = 256,
        pool: str = "gem",
        normalize: bool = True,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.embedding_dim = embedding_dim
        self.num_fc_layers = num_fc_layers
        self.num_fc_units = num_fc_units
        self.pool = pool
        self.normalize = normalize

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return self.embedding_dim

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "identity"

    @property
    def loss_function(self) -> str:
        """Return the loss-function name (informational).

        The contrastive loss is batch-level and is driven by the embedding
        LightningModule's ``objective``, not by this string.
        """
        return "supcon"

    @classmethod
    def from_config(cls, config: DictConfig) -> "EmbeddingHead":
        """Create this head from a head-leaf configuration."""
        return cls(
            embedding_dim=config.embedding_dim,
            num_fc_layers=config.num_fc_layers,
            num_fc_units=config.num_fc_units,
            pool=config.pool,
            normalize=config.normalize,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

    def make_head(self, x_in: int) -> nn.Sequential:
        """Make the head output module from the pooled-feature input channels.

        Args:
            x_in: Number of channels of the encoder feature map (its channel dim).

        Returns:
            An ``nn.Sequential`` mapping ``[B, x_in, H, W] -> [B, embedding_dim]``.
        """
        module_dict = OrderedDict()
        if self.pool == "gem":
            module_dict["pre_embedding_pool"] = GeM()
        elif self.pool == "max":
            module_dict["pre_embedding_pool"] = nn.AdaptiveMaxPool2d(1)
        elif self.pool == "avg":
            module_dict["pre_embedding_pool"] = nn.AdaptiveAvgPool2d(1)
        else:
            message = f"Unknown pool '{self.pool}'; choose one of gem|max|avg."
            logger.error(message)
            raise ValueError(message)

        module_dict["pre_embedding_flatten"] = nn.Flatten(start_dim=1)

        d_in = x_in
        for i in range(self.num_fc_layers):
            module_dict[f"pre_embedding{i}_fc"] = nn.Linear(d_in, self.num_fc_units)
            module_dict[f"pre_embedding{i}_relu"] = get_act_fn("relu")
            d_in = self.num_fc_units

        module_dict["EmbeddingHead"] = nn.Linear(d_in, self.embedding_dim)
        if self.normalize:
            module_dict["l2norm"] = L2Norm(dim=1)

        return nn.Sequential(module_dict)

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the loss-function name (informational).

The contrastive loss is batch-level and is driven by the embedding LightningModule's objective, not by this string.

__init__(embedding_dim=128, num_fc_layers=1, num_fc_units=256, pool='gem', normalize=True, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    embedding_dim: int = 128,
    num_fc_layers: int = 1,
    num_fc_units: int = 256,
    pool: str = "gem",
    normalize: bool = True,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.embedding_dim = embedding_dim
    self.num_fc_layers = num_fc_layers
    self.num_fc_units = num_fc_units
    self.pool = pool
    self.normalize = normalize

from_config(config) classmethod

Create this head from a head-leaf configuration.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(cls, config: DictConfig) -> "EmbeddingHead":
    """Create this head from a head-leaf configuration."""
    return cls(
        embedding_dim=config.embedding_dim,
        num_fc_layers=config.num_fc_layers,
        num_fc_units=config.num_fc_units,
        pool=config.pool,
        normalize=config.normalize,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

make_head(x_in)

Make the head output module from the pooled-feature input channels.

Parameters:

Name	Type	Description	Default
x_in	int	Number of channels of the encoder feature map (its channel dim).	required

Returns:

Type	Description
Sequential	An nn.Sequential mapping [B, x_in, H, W] -> [B, embedding_dim].

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: int) -> nn.Sequential:
    """Make the head output module from the pooled-feature input channels.

    Args:
        x_in: Number of channels of the encoder feature map (its channel dim).

    Returns:
        An ``nn.Sequential`` mapping ``[B, x_in, H, W] -> [B, embedding_dim]``.
    """
    module_dict = OrderedDict()
    if self.pool == "gem":
        module_dict["pre_embedding_pool"] = GeM()
    elif self.pool == "max":
        module_dict["pre_embedding_pool"] = nn.AdaptiveMaxPool2d(1)
    elif self.pool == "avg":
        module_dict["pre_embedding_pool"] = nn.AdaptiveAvgPool2d(1)
    else:
        message = f"Unknown pool '{self.pool}'; choose one of gem|max|avg."
        logger.error(message)
        raise ValueError(message)

    module_dict["pre_embedding_flatten"] = nn.Flatten(start_dim=1)

    d_in = x_in
    for i in range(self.num_fc_layers):
        module_dict[f"pre_embedding{i}_fc"] = nn.Linear(d_in, self.num_fc_units)
        module_dict[f"pre_embedding{i}_relu"] = get_act_fn("relu")
        d_in = self.num_fc_units

    module_dict["EmbeddingHead"] = nn.Linear(d_in, self.embedding_dim)
    if self.normalize:
        module_dict["l2norm"] = L2Norm(dim=1)

    return nn.Sequential(module_dict)

GeM

Bases: Module

Generalized-mean pooling: (mean(x.clamp(min=eps)^p))^(1/p) over HxW.

The exponent p is learnable (init 3.0). The clamp(min=eps) BEFORE the fractional power guards against NaNs (a fractional power of a negative/zero base). Returns a flattened [B, C] tensor.

Methods:

Name	Description
__init__	Initialize the pooling layer.
forward	Pool [B, C, H, W] to [B, C] by the generalized mean over HxW.

Source code in sleap_nn/architectures/heads.py

class GeM(nn.Module):
    """Generalized-mean pooling: ``(mean(x.clamp(min=eps)^p))^(1/p)`` over HxW.

    The exponent ``p`` is learnable (init 3.0). The ``clamp(min=eps)`` BEFORE the
    fractional power guards against NaNs (a fractional power of a negative/zero base).
    Returns a flattened ``[B, C]`` tensor.
    """

    def __init__(self, p: float = 3.0, eps: float = 1e-6, learnable: bool = True):
        """Initialize the pooling layer.

        Args:
            p: Initial generalized-mean exponent (``p=1`` averages, larger ``p``
                approaches max-pooling).
            eps: Small floor applied to the activation before the power to avoid a
                fractional power of a non-positive value.
            learnable: If ``True``, ``p`` is a trainable parameter; otherwise it is a
                fixed buffer.
        """
        super().__init__()
        if learnable:
            self.p = nn.Parameter(torch.tensor(float(p)))
        else:
            self.register_buffer("p", torch.tensor(float(p)))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """Pool ``[B, C, H, W]`` to ``[B, C]`` by the generalized mean over ``HxW``."""
        # Clamp the (learnable) exponent to a sane floor: p -> 0 makes ``1/p`` explode
        # and a negative p inverts the mean — either can yield inf/NaN embeddings that
        # corrupt the whole batch's contrastive loss. The activation eps-clamp guards
        # the base; this guards the exponent.
        p = self.p.clamp(min=1.0)
        xp = x.clamp(min=self.eps).pow(p)
        return F.adaptive_avg_pool2d(xp, 1).pow(1.0 / p).flatten(1)

__init__(p=3.0, eps=1e-06, learnable=True)

Initialize the pooling layer.

Parameters:

Name	Type	Description	Default
p	float	Initial generalized-mean exponent (p=1 averages, larger p approaches max-pooling).	3.0
eps	float	Small floor applied to the activation before the power to avoid a fractional power of a non-positive value.	1e-06
learnable	bool	If True, p is a trainable parameter; otherwise it is a fixed buffer.	True

Source code in sleap_nn/architectures/heads.py

def __init__(self, p: float = 3.0, eps: float = 1e-6, learnable: bool = True):
    """Initialize the pooling layer.

    Args:
        p: Initial generalized-mean exponent (``p=1`` averages, larger ``p``
            approaches max-pooling).
        eps: Small floor applied to the activation before the power to avoid a
            fractional power of a non-positive value.
        learnable: If ``True``, ``p`` is a trainable parameter; otherwise it is a
            fixed buffer.
    """
    super().__init__()
    if learnable:
        self.p = nn.Parameter(torch.tensor(float(p)))
    else:
        self.register_buffer("p", torch.tensor(float(p)))
    self.eps = eps

forward(x)

Pool [B, C, H, W] to [B, C] by the generalized mean over HxW.

Source code in sleap_nn/architectures/heads.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """Pool ``[B, C, H, W]`` to ``[B, C]`` by the generalized mean over ``HxW``."""
    # Clamp the (learnable) exponent to a sane floor: p -> 0 makes ``1/p`` explode
    # and a negative p inverts the mean — either can yield inf/NaN embeddings that
    # corrupt the whole batch's contrastive loss. The activation eps-clamp guards
    # the base; this guards the exponent.
    p = self.p.clamp(min=1.0)
    xp = x.clamp(min=self.eps).pow(p)
    return F.adaptive_avg_pool2d(xp, 1).pow(1.0 / p).flatten(1)

Head

Base class for model output heads.

Attributes:

Name	Type	Description
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
make_head	Make head output tensor from input feature tensor.

Source code in sleap_nn/architectures/heads.py

class Head:
    """Base class for model output heads.

    Attributes:
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(self, output_stride: int = 1, loss_weight: float = 1.0) -> None:
        """Initialize the object with the specified attributes."""
        self.output_stride = output_stride
        self.loss_weight = loss_weight

    @property
    def name(self) -> str:
        """Name of the head."""
        return type(self).__name__

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        message = "Subclasses must implement this method."
        logger.error(message)
        raise NotImplementedError(message)

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "identity"

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "mse"

    def make_head(self, x_in: int) -> nn.Sequential:
        """Make head output tensor from input feature tensor.

        Args:
            x_in: An int input for the input channels.

        Returns:
            A `nn.Sequential` with the correct shape for the head.
        """
        module_dict = OrderedDict()
        module_dict[self.name] = nn.Sequential(
            nn.Conv2d(
                in_channels=x_in,
                out_channels=self.channels,
                kernel_size=1,
                stride=1,
                padding="same",
            ),
            get_act_fn(self.activation),
        )

        return nn.Sequential(module_dict)

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

name property

Name of the head.

__init__(output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(self, output_stride: int = 1, loss_weight: float = 1.0) -> None:
    """Initialize the object with the specified attributes."""
    self.output_stride = output_stride
    self.loss_weight = loss_weight

make_head(x_in)

Make head output tensor from input feature tensor.

Parameters:

Name	Type	Description	Default
x_in	int	An int input for the input channels.	required

Returns:

Type	Description
Sequential	A nn.Sequential with the correct shape for the head.

Source code in sleap_nn/architectures/heads.py

def make_head(self, x_in: int) -> nn.Sequential:
    """Make head output tensor from input feature tensor.

    Args:
        x_in: An int input for the input channels.

    Returns:
        A `nn.Sequential` with the correct shape for the head.
    """
    module_dict = OrderedDict()
    module_dict[self.name] = nn.Sequential(
        nn.Conv2d(
            in_channels=x_in,
            out_channels=self.channels,
            kernel_size=1,
            stride=1,
            padding="same",
        ),
        get_act_fn(self.activation),
    )

    return nn.Sequential(module_dict)

InstanceCenterHead

Bases: Head

Head for predicting instance center heatmaps.

Outputs a single-channel Gaussian heatmap with peaks at each instance's mask centroid. Similar to CentroidConfmapsHead but for mask-derived centers.

Attributes:

Name	Type	Description
sigma		Standard deviation of the Gaussian in pixels.
output_stride		Stride of the output head tensor.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

class InstanceCenterHead(Head):
    """Head for predicting instance center heatmaps.

    Outputs a single-channel Gaussian heatmap with peaks at each instance's
    mask centroid. Similar to CentroidConfmapsHead but for mask-derived centers.

    Attributes:
        sigma: Standard deviation of the Gaussian in pixels.
        output_stride: Stride of the output head tensor.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        sigma: float = 4.0,
        output_stride: int = 2,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 1

channels property

Return the number of channels in the tensor output by this head.

__init__(sigma=4.0, output_stride=2, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    sigma: float = 4.0,
    output_stride: int = 2,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.sigma = sigma

L2Norm

Bases: Module

L2-normalize along dim (so embeddings live on the unit hypersphere).

Methods:

Name	Description
__init__	Initialize the layer.
forward	L2-normalize x along dim.

Source code in sleap_nn/architectures/heads.py

class L2Norm(nn.Module):
    """L2-normalize along ``dim`` (so embeddings live on the unit hypersphere)."""

    def __init__(self, dim: int = 1):
        """Initialize the layer.

        Args:
            dim: Dimension along which to L2-normalize.
        """
        super().__init__()
        self.dim = dim

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """L2-normalize ``x`` along ``dim``."""
        return F.normalize(x, dim=self.dim)

__init__(dim=1)

Initialize the layer.

Parameters:

Name	Type	Description	Default
dim	int	Dimension along which to L2-normalize.	1

Source code in sleap_nn/architectures/heads.py

def __init__(self, dim: int = 1):
    """Initialize the layer.

    Args:
        dim: Dimension along which to L2-normalize.
    """
    super().__init__()
    self.dim = dim

forward(x)

L2-normalize x along dim.

Source code in sleap_nn/architectures/heads.py

def forward(self, x: torch.Tensor) -> torch.Tensor:
    """L2-normalize ``x`` along ``dim``."""
    return F.normalize(x, dim=self.dim)

MultiInstanceConfmapsHead

Bases: Head

Head for specifying multi-instance confidence maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class MultiInstanceConfmapsHead(Head):
    """Head for specifying multi-instance confidence maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.part_names)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
    ) -> "MultiInstanceConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            part_names: 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. This must be provided if the `part_names`
                attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.part_names is not None:
            part_names = config.part_names
        elif part_names is None:
            message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.sigma = sigma

from_config(config, part_names=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		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. This must be provided if the part_names attribute of the configuration is not set.

Returns:

Type	Description
MultiInstanceConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
) -> "MultiInstanceConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        part_names: 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. This must be provided if the `part_names`
            attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.part_names is not None:
        part_names = config.part_names
    elif part_names is None:
        message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

OffsetRefinementHead

Bases: Head

Head for specifying offset refinement maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
sigma_threshold		Threshold of confidence map values to use for defining the boundary of the offset maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class OffsetRefinementHead(Head):
    """Head for specifying offset refinement maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        sigma_threshold: Threshold of confidence map values to use for defining the
            boundary of the offset maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        sigma_threshold: float = 0.2,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.sigma_threshold = sigma_threshold

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return int(len(self.part_names) * 2)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
        sigma_threshold: float = 0.2,
        loss_weight: float = 1.0,
    ) -> "OffsetRefinementHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head parameters.
            part_names: 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. This must be provided if the `part_names`
                attribute of the configuration is not set.
            sigma_threshold: Minimum confidence map value below which offsets will be
                replaced with zeros.
            loss_weight: Weight of the loss associated with this head.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if hasattr(config, "part_names"):
            if config.part_names is not None:
                part_names = config.part_names
        elif hasattr(config, "anchor_part"):
            part_names = [config.anchor_part]
        else:
            message = "Required attribute 'part_names' is missing in the configuration."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            output_stride=config.output_stride,
            sigma_threshold=sigma_threshold,
            loss_weight=loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, sigma_threshold=0.2, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    sigma_threshold: float = 0.2,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.sigma_threshold = sigma_threshold

from_config(config, part_names=None, sigma_threshold=0.2, loss_weight=1.0) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		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. This must be provided if the part_names attribute of the configuration is not set.
sigma_threshold		Minimum confidence map value below which offsets will be replaced with zeros.
loss_weight		Weight of the loss associated with this head.

Returns:

Type	Description
OffsetRefinementHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
    sigma_threshold: float = 0.2,
    loss_weight: float = 1.0,
) -> "OffsetRefinementHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head parameters.
        part_names: 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. This must be provided if the `part_names`
            attribute of the configuration is not set.
        sigma_threshold: Minimum confidence map value below which offsets will be
            replaced with zeros.
        loss_weight: Weight of the loss associated with this head.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if hasattr(config, "part_names"):
        if config.part_names is not None:
            part_names = config.part_names
    elif hasattr(config, "anchor_part"):
        part_names = [config.anchor_part]
    else:
        message = "Required attribute 'part_names' is missing in the configuration."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        output_stride=config.output_stride,
        sigma_threshold=sigma_threshold,
        loss_weight=loss_weight,
    )

PartAffinityFieldsHead

Bases: Head

Head for specifying multi-instance part affinity fields.

Attributes:

Name	Type	Description
edges		List of tuples of (source, destination) node names.
sigma		Spread of the part affinity fields.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class PartAffinityFieldsHead(Head):
    """Head for specifying multi-instance part affinity fields.

    Attributes:
        edges: List of tuples of `(source, destination)` node names.
        sigma: Spread of the part affinity fields.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        edges: Sequence[Tuple[Text, Text]],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.edges = edges
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return int(len(self.edges) * 2)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        edges: Optional[Sequence[Tuple[Text, Text]]] = None,
    ) -> "PartAffinityFieldsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            edges: List of 2-tuples of the form `(source_node, destination_node)` that
                define pairs of text names of the directed edges of the graph. This must
                be set if the `edges` attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.edges is not None:
            edges = config.edges
        return cls(
            edges=edges,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(edges, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    edges: Sequence[Tuple[Text, Text]],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.edges = edges
    self.sigma = sigma

from_config(config, edges=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
edges		List of 2-tuples of the form (source_node, destination_node) that define pairs of text names of the directed edges of the graph. This must be set if the edges attribute of the configuration is not set.

Returns:

Type	Description
PartAffinityFieldsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    edges: Optional[Sequence[Tuple[Text, Text]]] = None,
) -> "PartAffinityFieldsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        edges: List of 2-tuples of the form `(source_node, destination_node)` that
            define pairs of text names of the directed edges of the graph. This must
            be set if the `edges` attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.edges is not None:
        edges = config.edges
    return cls(
        edges=edges,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )

SegmentationHead

Bases: Head

Head for predicting binary foreground segmentation masks.

Outputs a single-channel map with sigmoid activation representing the probability that each pixel belongs to any instance (foreground).

Attributes:

Name	Type	Description
output_stride		Stride of the output head tensor.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

class SegmentationHead(Head):
    """Head for predicting binary foreground segmentation masks.

    Outputs a single-channel map with sigmoid activation representing the
    probability that each pixel belongs to any instance (foreground).

    Attributes:
        output_stride: Stride of the output head tensor.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        output_stride: int = 2,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return 1

    @property
    def activation(self) -> str:
        """Return the activation function of the head output layer."""
        return "identity"

    @property
    def loss_function(self) -> str:
        """Return the name of the loss function to use for this head."""
        return "bce_dice"

activation property

Return the activation function of the head output layer.

channels property

Return the number of channels in the tensor output by this head.

loss_function property

Return the name of the loss function to use for this head.

__init__(output_stride=2, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    output_stride: int = 2,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)

SingleInstanceConfmapsHead

Bases: Head

Head for specifying single instance confidence maps.

Attributes:

Name	Type	Description
part_names		List of strings specifying the part names associated with channels.
sigma		Spread of the confidence maps.
output_stride		Stride of the output head tensor. The input tensor is expected to be at the same stride.
loss_weight		Weight of the loss term for this head during optimization.

Methods:

Name	Description
__init__	Initialize the object with the specified attributes.
from_config	Create this head from a set of configurations.

Source code in sleap_nn/architectures/heads.py

class SingleInstanceConfmapsHead(Head):
    """Head for specifying single instance confidence maps.

    Attributes:
        part_names: List of strings specifying the part names associated with channels.
        sigma: Spread of the confidence maps.
        output_stride: Stride of the output head tensor. The input tensor is expected to
            be at the same stride.
        loss_weight: Weight of the loss term for this head during optimization.
    """

    def __init__(
        self,
        part_names: List[Text],
        sigma: float = 5.0,
        output_stride: int = 1,
        loss_weight: float = 1.0,
    ) -> None:
        """Initialize the object with the specified attributes."""
        super().__init__(output_stride, loss_weight)
        self.part_names = part_names
        self.sigma = sigma

    @property
    def channels(self) -> int:
        """Return the number of channels in the tensor output by this head."""
        return len(self.part_names)

    @classmethod
    def from_config(
        cls,
        config: DictConfig,
        part_names: Optional[List[Text]] = None,
    ) -> "SingleInstanceConfmapsHead":
        """Create this head from a set of configurations.

        Attributes:
            config: A `DictConfig` instance specifying the head
                parameters.
            part_names: 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. This must be provided if the `part_names`
                attribute of the configuration is not set.

        Returns:
            The instantiated head with the specified configuration options.
        """
        if config.part_names is not None:
            part_names = config.part_names
        elif part_names is None:
            message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
            logger.error(message)
            raise ValueError(message)
        return cls(
            part_names=part_names,
            sigma=config.sigma,
            output_stride=config.output_stride,
            loss_weight=config.loss_weight,
        )

channels property

Return the number of channels in the tensor output by this head.

__init__(part_names, sigma=5.0, output_stride=1, loss_weight=1.0)

Initialize the object with the specified attributes.

Source code in sleap_nn/architectures/heads.py

def __init__(
    self,
    part_names: List[Text],
    sigma: float = 5.0,
    output_stride: int = 1,
    loss_weight: float = 1.0,
) -> None:
    """Initialize the object with the specified attributes."""
    super().__init__(output_stride, loss_weight)
    self.part_names = part_names
    self.sigma = sigma

from_config(config, part_names=None) classmethod

Create this head from a set of configurations.

Attributes:

Name	Type	Description
config		A DictConfig instance specifying the head parameters.
part_names		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. This must be provided if the part_names attribute of the configuration is not set.

Returns:

Type	Description
SingleInstanceConfmapsHead	The instantiated head with the specified configuration options.

Source code in sleap_nn/architectures/heads.py

@classmethod
def from_config(
    cls,
    config: DictConfig,
    part_names: Optional[List[Text]] = None,
) -> "SingleInstanceConfmapsHead":
    """Create this head from a set of configurations.

    Attributes:
        config: A `DictConfig` instance specifying the head
            parameters.
        part_names: 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. This must be provided if the `part_names`
            attribute of the configuration is not set.

    Returns:
        The instantiated head with the specified configuration options.
    """
    if config.part_names is not None:
        part_names = config.part_names
    elif part_names is None:
        message = "Required attribute 'part_names' is missing in the configuration or in `from_config` input."
        logger.error(message)
        raise ValueError(message)
    return cls(
        part_names=part_names,
        sigma=config.sigma,
        output_stride=config.output_stride,
        loss_weight=config.loss_weight,
    )