Source code for sleap.nn.config.model

import attr
from typing import Optional, Text, List, Sequence, Tuple
from sleap.nn.config.utils import oneof


[docs]@attr.s(auto_attribs=True) class SingleInstanceConfmapsHeadConfig: """Configurations for single instance confidence map heads. These heads are used in single instance models that make the assumption that only one of each body part is present in the image. These heads produce confidence maps with a single peak for each part type which can be detected via global peak finding. Do not use this head if there is more than one animal present in the image. Attributes: 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. If not specified, all body parts in the skeleton will be used. sigma: 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: 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[Text]] = None sigma: float = 5.0 output_stride: int = 1
[docs]@attr.s(auto_attribs=True) class CentroidsHeadConfig: """Configurations for centroid confidence map heads. These heads are used in topdown models that rely on centroid detection to detect instances for cropping before predicting the remaining body parts. Multiple centroids can be present (one per instance), so their coordinates can be recovered in infernece via local peak finding. Attributes: anchor_part: Text name of a body part (node) to use as the anchor point. If None, the midpoint of the bounding box of all visible instance points will be used as the anchor. The bounding box midpoint will also be used if 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. sigma: 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: 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. """ anchor_part: Optional[Text] = None sigma: float = 5.0 output_stride: int = 1
[docs]@attr.s(auto_attribs=True) class CenteredInstanceConfmapsHeadConfig: """Configurations for centered instance confidence map heads. These heads are used in topdown multi-instance models that make the assumption that there is an instance reliably centered in the cropped input image. These heads are useful when centroids are easy to detect as they learn complex relationships between the geometry of body parts, even when animals are occluded. This comes at the cost of a strong reliance on the accuracy of the instance-centered cropping, i.e., it is heavily limited by the accuracy of the centroid model. Additionally, since one image crop is evaluated per instance, topdown models scale linearly with the number of animals in the frame, which can result in poor performance when many instances are present. Use this head when centroids are easy to detect, preferably using a consistent body part as an anchor, and when there are few animals that cover a small region of the full frame. Attributes: anchor_part: Text name of a body part (node) to use as the anchor point. If None, the midpoint of the bounding box of all visible instance points will be used as the anchor. The bounding box midpoint will also be used if 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. 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. If not specified, all body parts in the skeleton will be used. sigma: 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: 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. """ anchor_part: Optional[Text] = None part_names: Optional[List[Text]] = None sigma: float = 5.0 output_stride: int = 1
[docs]@attr.s(auto_attribs=True) class MultiInstanceConfmapsHeadConfig: """Configurations for multi-instance confidence map heads. These heads are used in bottom-up multi-instance models that do not make any assumption about the connectivity of the body parts. These heads will generate multiple local peaks for each body part type and must be detected using local peak finding. Although this head alone is sufficient to detect multiple copies of each body part type, it provides no information as to which sets of points should be grouped together to the same instance. If this is required, a head that provides connectivity or grouping information is required, e.g., part affinity fields. Use this head when multiple instances of each body part are present and do not need to be grouped or will be grouped using additional information. This head type has the advantage that it only needs to evaluate each frame once to find all peaks, in contrast to topdown models that must be evaluated for each crop. This constant scaling with the number of instances can be especially beneficial when there are many animals present in the frame. Attributes: 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. If not specified, all body parts in the skeleton will be used. sigma: 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: 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: 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[Text]] = None sigma: float = 5.0 output_stride: int = 1 loss_weight: float = 1.0
[docs]@attr.s(auto_attribs=True) class PartAffinityFieldsHeadConfig: """Configurations for multi-instance part affinity field heads. These heads are used in bottom-up multi-instance models that require information about body part connectivity in order to group multiple detections of each body part type into distinct instances. Part affinity fields are an image-space representation of the directed graph that defines the skeleton. Pixels that are close to the line (directed edge) formed between pairs of nodes of the same instance will contain unit vectors pointing along the direction of the the connection. The similarity between this line and the average of the unit vectors at the pixels underneath the line can be used as a matching score to associate candidate pairs of body part detections. Use this head when multiple instances of each body part are present and need to be grouped to coherent instances. This head type has the advantage that it only needs to evaluate each frame once to find all peaks, in contrast to topdown models that must be evaluated for each crop. This constant scaling with the number of instances can be especially beneficial when there are many animals present in the frame. Attributes: 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. If not set, all edges in the skeleton will be used. sigma: Spread of the Gaussian distribution that weigh the part affinity fields as a function of their distance from the edge they represent. 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 edge distance, so can be less useful in disambiguating between edges that are nearby and parallel in direction. 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: The stride of the output part affinity fields relative to the input image. This is the reciprocal of the resolution, e.g., an output stride of 2 results in PAFs 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: 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[Sequence[Tuple[Text, Text]]] = None sigma: float = 15.0 output_stride: int = 1 loss_weight: float = 1.0
[docs]@attr.s(auto_attribs=True) class MultiInstanceConfig: """Configuration for combined multi-instance confidence map and PAF model heads. This configuration specifies a multi-head model that outputs both multi-instance confidence maps and part affinity fields, which together enable multi-instance pose estimation in a bottom-up fashion, i.e., no instance cropping or centroids are required. Attributes: confmaps: Part confidence map configuration (see the description in `MultiInstanceConfmapsHeadConfig`). pafs: Part affinity fields configuration (see the description in `PartAffinityFieldsHeadConfig`). """ confmaps: MultiInstanceConfmapsHeadConfig = attr.ib( factory=MultiInstanceConfmapsHeadConfig ) pafs: PartAffinityFieldsHeadConfig = attr.ib(factory=PartAffinityFieldsHeadConfig)
[docs]@oneof @attr.s(auto_attribs=True) class HeadsConfig: """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`. multi_instance: An instance of `MultiInstanceConfig`. """ single_instance: Optional[SingleInstanceConfmapsHeadConfig] = None centroid: Optional[CentroidsHeadConfig] = None centered_instance: Optional[CenteredInstanceConfmapsHeadConfig] = None multi_instance: Optional[MultiInstanceConfig] = None
[docs]@attr.s(auto_attribs=True) class LEAPConfig: """LEAP backbone configuration. Attributes: max_stride: Determines the number of downsampling blocks in the network, increasing receptive field size at the cost of network size. output_stride: Determines the number of upsampling blocks in the network. filters: Base number of filters in the network. filters_rate: Factor to scale the number of filters by at each block. up_interpolate: If True, use bilinear upsampling instead of transposed convolutions for upsampling. This can save computations but may lower overall accuracy. stacks: Number of repeated stacks of the network (excluding the stem). """ max_stride: int = 8 # determines down blocks output_stride: int = 1 # determines up blocks filters: int = 64 filters_rate: float = 2 up_interpolate: bool = False stacks: int = 1
[docs]@attr.s(auto_attribs=True) class UNetConfig: """UNet backbone configuration. Attributes: stem_stride: If not None, controls how many stem blocks to use for initial downsampling. These are useful for learned downsampling that is able to retain spatial information while reducing large input image sizes. max_stride: Determines the number of downsampling blocks in the network, increasing receptive field size at the cost of network size. output_stride: Determines the number of upsampling blocks in the network. filters: Base number of filters in the network. filters_rate: Factor to scale the number of filters by at each block. middle_block: If True, add an intermediate block between the downsampling and upsampling branch for additional processing for features at the largest receptive field size. This will not introduce an extra pooling step. up_interpolate: If True, use bilinear upsampling instead of transposed convolutions for upsampling. This can save computations but may lower overall accuracy. stacks: Number of repeated stacks of the network (excluding the stem). """ stem_stride: Optional[int] = None max_stride: int = 16 output_stride: int = 1 filters: int = 64 filters_rate: float = 2 middle_block: bool = True up_interpolate: bool = False stacks: int = 1
[docs]@attr.s(auto_attribs=True) class HourglassConfig: """Hourglass backbone configuration. Attributes: stem_stride: Controls how many stem blocks to use for initial downsampling. These are useful for learned downsampling that is able to retain spatial information while reducing large input image sizes. max_stride: Determines the number of downsampling blocks in the network, increasing receptive field size at the cost of network size. output_stride: Determines the number of upsampling blocks in the network. filters: Base number of filters in the network. filters_increase: Constant to increase the number of filters by at each block. stacks: Number of repeated stacks of the network (excluding the stem). """ stem_stride: int = 4 max_stride: int = 64 output_stride: int = 4 stem_filters: int = 128 filters: int = 256 filter_increase: int = 128 stacks: int = 3
[docs]@attr.s(auto_attribs=True) class UpsamplingConfig: """Upsampling stack configuration. Attributes: method: If "transposed_conv", use a strided transposed convolution to perform learnable upsampling. If "interpolation", bilinear upsampling will be used instead. skip_connections: If "add", incoming feature tensors form skip connection with upsampled features via element-wise addition. Height/width are matched via stride and a 1x1 linear conv is applied if the channel counts do no match up. If "concatenate", the skip connection is formed via channel-wise concatenation. If None, skip connections will not be formed. block_stride: The striding of the upsampling *layer* (not tensor). This is typically set to 2, such that the tensor doubles in size with each upsampling step, but can be set higher to upsample to the desired `output_stride` directly in fewer steps. filters: Integer that specifies the base number of filters in each convolution layer. This will be scaled by the `filters_rate` at every upsampling step. filters_rate: Factor to scale the number of filters in the convolution layers after each upsampling step. If set to 1, the number of filters won't change. refine_convs: If greater than 0, specifies the number of 3x3 convolutions that will be applied after the upsampling step for refinement. These layers can serve the purpose of "mixing" the skip connection fused features, or to refine the current feature map after upsampling, which can help to prevent aliasing and checkerboard effects. If 0, no additional convolutions will be applied. conv_batchnorm: Specifies whether batch norm should be applied after each convolution (and before the ReLU activation). transposed_conv_kernel_size: Size of the kernel for the transposed convolution. No effect if bilinear upsampling is used. """ method: Text = attr.ib( default="interpolation", validator=attr.validators.in_(["interpolation", "transposed_conv"]), ) skip_connections: Optional[Text] = attr.ib( default=None, validator=attr.validators.optional(attr.validators.in_(["add", "concatenate"])), ) block_stride: int = 2 filters: int = 64 filters_rate: float = 1 refine_convs: int = 2 batch_norm: bool = True transposed_conv_kernel_size: int = 4
[docs]@attr.s(auto_attribs=True) class ResNetConfig: """ResNet backbone configuration. Attributes: version: Name of the ResNetV1 variant. Can be one of: "ResNet50", "ResNet101", or "ResNet152". weights: Controls how the network weights are initialized. If "random", the network is not pretrained. If "frozen", the network uses pretrained weights and keeps them fixed. If "tunable", the network uses pretrained weights and allows them to be trainable. upsampling: A `UpsamplingConfig` that defines an upsampling branch if not None. max_stride: Stride of the backbone feature activations. These should be <= 32. output_stride: Stride of the final output. If the upsampling branch is not defined, the output stride is controlled via dilated convolutions or reduced pooling in the backbone. """ version: Text = attr.ib( default="ResNet50", validator=attr.validators.in_(["ResNet50", "ResNet101", "ResNet152"]), ) weights: Text = attr.ib( default="frozen", validator=attr.validators.in_(["random", "frozen", "tunable"]) ) upsampling: Optional[UpsamplingConfig] = None max_stride: int = 32 output_stride: int = 4
[docs]@oneof @attr.s(auto_attribs=True) class BackboneConfig: """Configurations related to the model backbone. Only one field can be set and will determine which backbone architecture to use. Attributes: leap: A `LEAPConfig` instance. unet: A `UNetConfig` instance. hourglass: A `HourglassConfig` instance. resnet: A `ResNetConfig` instance. """ leap: Optional[LEAPConfig] = None unet: Optional[UNetConfig] = None hourglass: Optional[HourglassConfig] = None resnet: Optional[ResNetConfig] = None
[docs]@attr.s(auto_attribs=True) class ModelConfig: """Configurations related to model architecture. Attributes: backbone: Configurations related to the main network architecture. heads: Configurations related to the output heads. """ backbone: BackboneConfig = attr.ib(factory=BackboneConfig) heads: HeadsConfig = attr.ib(factory=HeadsConfig)