sleap.nn.heads

Model head definitions for defining model output types.

class sleap.nn.heads.CenteredInstanceConfmapsHead(part_names: List[str], anchor_part: str | None = None, sigma: float = 5.0, output_stride: int = 1, loss_weight: float = 1.0)

Head for specifying centered instance confidence maps.

part_names

List of strings specifying the part names associated with channels.

Type:

List[str]

anchor_part

Name of the part to use as an anchor node. If not specified, the bounding box centroid will be used.

Type:

str | None

sigma

Spread of the confidence maps.

Type:

float

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property channels: int

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

classmethod from_config(config: CenteredInstanceConfmapsHeadConfig, part_names: List[str] | None = None) → CenteredInstanceConfmapsHead

Create this head from a set of configurations.

config

A CenteredInstanceConfmapsHeadConfig 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.

class sleap.nn.heads.CentroidConfmapsHead(anchor_part: str | None = None, sigma: float = 5.0, output_stride: int = 1, loss_weight: float = 1.0)

Head for specifying instance centroid confidence maps.

anchor_part

Name of the part to use as an anchor node. If not specified, the bounding box centroid will be used.

Type:

str | None

sigma

Spread of the confidence maps.

Type:

float

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property channels: int

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

classmethod from_config(config: CentroidsHeadConfig) → CentroidConfmapsHead

Create this head from a set of configurations.

config

A CentroidsHeadConfig instance specifying the head parameters.

Returns:

The instantiated head with the specified configuration options.

class sleap.nn.heads.ClassMapsHead(classes: List[str], sigma: float = 5.0, output_stride: int = 1, loss_weight: float = 1.0)

Head for specifying class identity maps.

classes

List of string names of the classes.

Type:

List[str]

sigma

Spread of the class maps around each node.

Type:

float

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property activation: str

Return the activation function of the head output layer.

property channels: int

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

classmethod from_config(config: ClassMapsHeadConfig, classes: List[str] | None = None) → ClassMapsHead

Create this head from a set of configurations.

config

A ClassMapsHeadConfig 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.

class sleap.nn.heads.ClassVectorsHead(classes: List[str], num_fc_layers: int = 1, num_fc_units: int = 64, global_pool: bool = True, output_stride: int = 1, loss_weight: float = 1.0)

Head for specifying classification heads.

classes

List of string names of the classes.

Type:

List[str]

num_fc_layers

Number of fully connected layers after flattening input features.

Type:

int

num_fc_units

Number of units (dimensions) in fully connected layers prior to classification output.

Type:

int

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property activation: str

Return the activation function of the head output layer.

property channels: int

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

classmethod from_config(config: ClassVectorsHeadConfig, classes: List[str] | None = None) → ClassVectorsHead

Create this head from a set of configurations.

config

A ClassVectorsHeadConfig 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.

property loss_function: str

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

make_head(x_in: Tensor, name: str | None = None) → Tensor

Make head output tensor from input feature tensor.

Parameters:

• x_in – An input tf.Tensor.

• name – If provided, specifies the name of the output layer. If not (the default), uses the name of the head as the layer name.

Returns:

A tf.Tensor with the correct shape for the head.

class sleap.nn.heads.Head(output_stride: int = 1, loss_weight: float = 1.0)

Base class for model output heads.

property activation: str

Return the activation function of the head output layer.

abstract property channels: int

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

property loss_function: str

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

make_head(x_in: Tensor, name: str | None = None) → Tensor

Make head output tensor from input feature tensor.

Parameters:

• x_in – An input tf.Tensor.

• name – If provided, specifies the name of the output layer. If not (the default), uses the name of the head as the layer name.

Returns:

A tf.Tensor with the correct shape for the head.

class sleap.nn.heads.MultiInstanceConfmapsHead(part_names: List[str], sigma: float = 5.0, output_stride: int = 1, loss_weight: float = 1.0)

Head for specifying multi-instance confidence maps.

part_names

List of strings specifying the part names associated with channels.

Type:

List[str]

sigma

Spread of the confidence maps.

Type:

float

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property channels: int

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

classmethod from_config(config: MultiInstanceConfmapsHeadConfig, part_names: List[str] | None = None) → MultiInstanceConfmapsHead

Create this head from a set of configurations.

config

A MultiInstanceConfmapsHeadConfig 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.

class sleap.nn.heads.OffsetRefinementHead(part_names: List[str], output_stride: int = 1, sigma_threshold: float = 0.2, loss_weight: float = 1.0)

Head for specifying offset refinement maps.

part_names

List of strings specifying the part names associated with channels.

Type:

List[str]

sigma_threshold

Threshold of confidence map values to use for defining the boundary of the offset maps.

Type:

float

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property channels: int

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

classmethod from_config(config: CentroidsHeadConfig | SingleInstanceConfmapsHeadConfig | CenteredInstanceConfmapsHeadConfig | MultiInstanceConfmapsHeadConfig, part_names: List[str] | None = None, sigma_threshold: float = 0.2, loss_weight: float = 1.0) → OffsetRefinementHead

Create this head from a set of configurations.

config

A ConfmapConfig 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.

class sleap.nn.heads.PartAffinityFieldsHead(edges: Sequence[Tuple[str, str]], sigma: float = 15.0, output_stride: int = 1, loss_weight: float = 1.0)

Head for specifying multi-instance part affinity fields.

edges

List of tuples of (source, destination) node names.

Type:

Sequence[Tuple[str, str]]

sigma

Spread of the part affinity fields.

Type:

float

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property channels: int

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

classmethod from_config(config: PartAffinityFieldsHeadConfig, edges: Sequence[Tuple[str, str]] | None = None) → PartAffinityFieldsHead

Create this head from a set of configurations.

config

A PartAffinityFieldsHeadConfig 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.

class sleap.nn.heads.SingleInstanceConfmapsHead(part_names: List[str], sigma: float = 5.0, output_stride: int = 1, loss_weight: float = 1.0)

Head for specifying single instance confidence maps.

part_names

List of strings specifying the part names associated with channels.

Type:

List[str]

sigma

Spread of the confidence maps.

Type:

float

output_stride

Stride of the output head tensor. The input tensor is expected to be at the same stride.

Type:

int

loss_weight

Weight of the loss term for this head during optimization.

Type:

float

property channels: int

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

classmethod from_config(config: SingleInstanceConfmapsHeadConfig, part_names: List[str] | None = None) → SingleInstanceConfmapsHead

Create this head from a set of configurations.

config

A SingleInstanceConfmapsHeadConfig 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.