architecture_estimates

sleap_nn.config_generator.architecture_estimates

Shared architecture-related estimates for UNet config recommendations.

Single source of truth for the formulas used by the config generator (TUI), recommender, generator, and YAML emitter. Mirrors the formulas used by the config picker web app at docs/configuration/config-picker/app.html so the two surfaces produce equivalent recommendations and estimates.

References: - Canonical receptive-field formula: example_notebooks/receptive_field_guide.py (https://distill.pub/2019/computing-receptive-fields/, Eq. 2) - UNet implementation (ground truth for parameter count): sleap_nn/architectures/unet.py - Web-app counterparts: computeReceptiveField, estimateParamsAccurate, computeAugmentationPadding, computeSuggestedCropSize in app.html.

Functions:

Name	Description
compute_augmentation_padding	Pixels of padding required so a rotated/scaled bbox stays in bounds.
compute_backbone_context_margin	Half the surrounding context (px) a tile edge needs to keep seams valid.
compute_max_stride_for_animal_size	Smallest max_stride whose receptive field covers the animal.
compute_pad_to_stride	Round (height, width) up so each is a multiple of max_stride.
compute_receptive_field	Compute the receptive field of the deepest encoder layer of a UNet.
compute_suggested_crop_size	Suggest a crop size that fits the largest instance with optional padding.
compute_suggested_tile_overlap	Overlap (px) large enough that a seam-straddling object is whole in one tile.
compute_suggested_tile_size	Square tile side that fits an object plus context on both sides.
decoder_blocks	Number of decoder upsampling blocks needed to reach output_stride.
encoder_blocks	Number of encoder downsampling blocks for the given max stride.
estimate_unet_params	Estimate trainable parameter count of a UNet head + body.
recommend_default_max_stride	Bucket-based default max_stride recommendation.

compute_augmentation_padding(bbox_size, rotation_max=0.0, scale_max=1.0)

Pixels of padding required so a rotated/scaled bbox stays in bounds.

For a square bbox rotated by angle theta, the worst-case bounding-box expansion is |cos(theta)| + |sin(theta)|, which peaks at sqrt(2) at 45°. Scaling expands the bbox by max(scale_max, 1.0).

Parameters:

Name	Type	Description	Default
bbox_size	float	Original bbox dimension in pixels.	required
rotation_max	float	Max absolute rotation in degrees.	0.0
scale_max	float	Max scale factor (1.0 == no scaling).	1.0

Returns:

Type	Description
int	Padding in pixels (ceiling), 0 if no augmentation expansion needed.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_augmentation_padding(
    bbox_size: float,
    rotation_max: float = 0.0,
    scale_max: float = 1.0,
) -> int:
    """Pixels of padding required so a rotated/scaled bbox stays in bounds.

    For a square bbox rotated by angle theta, the worst-case bounding-box
    expansion is ``|cos(theta)| + |sin(theta)|``, which peaks at sqrt(2) at 45°.
    Scaling expands the bbox by ``max(scale_max, 1.0)``.

    Args:
        bbox_size: Original bbox dimension in pixels.
        rotation_max: Max absolute rotation in degrees.
        scale_max: Max scale factor (1.0 == no scaling).

    Returns:
        Padding in pixels (ceiling), 0 if no augmentation expansion needed.
    """
    if rotation_max == 0 and scale_max <= 1.0:
        return 0

    rotation_factor = 1.0
    if rotation_max > 0:
        if abs(rotation_max) >= 45:
            rotation_factor = math.sqrt(2)
        else:
            rad = math.radians(min(abs(rotation_max), 90))
            rotation_factor = abs(math.cos(rad)) + abs(math.sin(rad))

    expansion = rotation_factor * max(scale_max, 1.0)
    expanded = bbox_size * expansion
    return math.ceil(expanded - bbox_size)

compute_backbone_context_margin(backbone_type, max_stride, convs_per_block=2, kernel_size=3)

Half the surrounding context (px) a tile edge needs to keep seams valid.

A tile-edge output pixel has part of its receptive field off-tile; sizing the overlap by this margin ensures an adjacent tile's center (full RF) covers that region.

• UNet: half the deepest-encoder receptive field (compute_receptive_field).
• ConvNext / SwinT: a fixed per-family constant (windowed attention / patch-merging make an analytic RF invalid; see design DQ4).
• Any other backbone (pretrained / unsupported): raises, since tiling is not supported there.

Parameters:

Name	Type	Description	Default
backbone_type	str	One of "unet", "convnext", "swint".	required
max_stride	int	Backbone total downsampling factor (UNet only).	required
convs_per_block	int	Convs per UNet down block (UNet only).	2
kernel_size	int	UNet conv kernel size (UNet only).	3

Returns:

Type	Description
int	Context margin in input pixels.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_backbone_context_margin(
    backbone_type: str,
    max_stride: int,
    convs_per_block: int = 2,
    kernel_size: int = 3,
) -> int:
    """Half the surrounding context (px) a tile edge needs to keep seams valid.

    A tile-edge output pixel has part of its receptive field off-tile; sizing
    the overlap by this margin ensures an adjacent tile's *center* (full RF)
    covers that region.

    - UNet: half the deepest-encoder receptive field (``compute_receptive_field``).
    - ConvNext / SwinT: a fixed per-family constant (windowed attention /
      patch-merging make an analytic RF invalid; see design DQ4).
    - Any other backbone (pretrained / unsupported): raises, since tiling is
      not supported there.

    Args:
        backbone_type: One of ``"unet"``, ``"convnext"``, ``"swint"``.
        max_stride: Backbone total downsampling factor (UNet only).
        convs_per_block: Convs per UNet down block (UNet only).
        kernel_size: UNet conv kernel size (UNet only).

    Returns:
        Context margin in input pixels.
    """
    if backbone_type == "unet":
        rf = compute_receptive_field(max_stride, convs_per_block, kernel_size)
        return int(math.ceil(rf / 2))
    if backbone_type in _BACKBONE_CONTEXT_MARGIN_PX:
        return _BACKBONE_CONTEXT_MARGIN_PX[backbone_type]
    raise ValueError(
        f"Tiling context margin is undefined for backbone {backbone_type!r} "
        "(pretrained / unsupported backbones cannot be tiled)."
    )

compute_max_stride_for_animal_size(animal_size, candidates=SUPPORTED_MAX_STRIDES)

Smallest max_stride whose receptive field covers the animal.

Parameters:

Name	Type	Description	Default
animal_size	float	Maximum animal bounding-box dimension in input pixels (already scaled by input_scale if applicable).	required
candidates	Tuple[int, ...]	Strides to consider, ascending.	SUPPORTED_MAX_STRIDES

Returns:

Type	Description
int	Smallest stride in candidates whose RF >= animal_size. Falls back to the largest candidate if none cover the animal.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_max_stride_for_animal_size(
    animal_size: float,
    candidates: Tuple[int, ...] = SUPPORTED_MAX_STRIDES,
) -> int:
    """Smallest max_stride whose receptive field covers the animal.

    Args:
        animal_size: Maximum animal bounding-box dimension in input pixels
            (already scaled by ``input_scale`` if applicable).
        candidates: Strides to consider, ascending.

    Returns:
        Smallest stride in ``candidates`` whose RF >= ``animal_size``. Falls
        back to the largest candidate if none cover the animal.
    """
    for stride in sorted(candidates):
        if compute_receptive_field(stride) >= animal_size:
            return stride
    return max(candidates)

compute_pad_to_stride(height, width, max_stride)

Round (height, width) up so each is a multiple of max_stride.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_pad_to_stride(height: int, width: int, max_stride: int) -> Tuple[int, int]:
    """Round (height, width) up so each is a multiple of ``max_stride``."""
    h_padded = math.ceil(height / max_stride) * max_stride
    w_padded = math.ceil(width / max_stride) * max_stride
    return h_padded, w_padded

compute_receptive_field(max_stride, convs_per_block=2, kernel_size=3)

Compute the receptive field of the deepest encoder layer of a UNet.

Each downsampling block has convs_per_block convolutions (stride 1, kernel kernel_size) followed by a 2x2 stride-2 pool. RF is built up layer-by-layer with the canonical formula::

RF = 1 + sum((kernel[l] - 1) * prod(strides[:l])) for l in 0..L-1

Parameters:

Name	Type	Description	Default
max_stride	int	Total downsampling factor of the encoder (must be a positive power of 2).	required
convs_per_block	int	Number of conv layers per down block.	2
kernel_size	int	Kernel size of the conv layers.	3

Returns:

Type	Description
int	Receptive field in input pixels.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_receptive_field(
    max_stride: int,
    convs_per_block: int = 2,
    kernel_size: int = 3,
) -> int:
    """Compute the receptive field of the deepest encoder layer of a UNet.

    Each downsampling block has ``convs_per_block`` convolutions (stride 1,
    kernel ``kernel_size``) followed by a 2x2 stride-2 pool. RF is built up
    layer-by-layer with the canonical formula::

        RF = 1 + sum((kernel[l] - 1) * prod(strides[:l])) for l in 0..L-1

    Args:
        max_stride: Total downsampling factor of the encoder (must be a
            positive power of 2).
        convs_per_block: Number of conv layers per down block.
        kernel_size: Kernel size of the conv layers.

    Returns:
        Receptive field in input pixels.
    """
    down_blocks = int(math.log2(max_stride))
    if 2**down_blocks != max_stride or max_stride < 1:
        raise ValueError(f"max_stride must be a positive power of 2, got {max_stride}")

    block_strides = [1] * convs_per_block + [2]
    block_kernels = [kernel_size] * convs_per_block + [2]

    strides = block_strides * down_blocks
    kernels = block_kernels * down_blocks

    rf = 1
    prod = 1
    for stride, kernel in zip(strides, kernels):
        rf += (kernel - 1) * prod
        prod *= stride
    return rf

compute_suggested_crop_size(max_bbox_dim, max_stride, use_augmentation=False, user_padding=None, rotation_max=0.0, scale_max=1.0)

Suggest a crop size that fits the largest instance with optional padding.

Mirrors the web app's computeSuggestedCropSize (app.html:3402).

• If user_padding is provided, it overrides any auto-computed padding (including 0, which means "no padding").
• Else if use_augmentation, padding is computed from rotation_max / scale_max.
• Result is rounded UP to the next multiple of max_stride.

Parameters:

Name	Type	Description	Default
max_bbox_dim	float	Largest instance bbox dimension (height or width).	required
max_stride	int	Network max stride; result will be a multiple of this.	required
use_augmentation	bool	Whether to add padding for rotation/scale aug.	False
user_padding	Optional[int]	Explicit padding override.	None
rotation_max	float	Max rotation in degrees (used when use_augmentation).	0.0
scale_max	float	Max scale factor (used when use_augmentation).	1.0

Returns:

Type	Description
int	Suggested crop size in pixels, divisible by max_stride.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_suggested_crop_size(
    max_bbox_dim: float,
    max_stride: int,
    use_augmentation: bool = False,
    user_padding: Optional[int] = None,
    rotation_max: float = 0.0,
    scale_max: float = 1.0,
) -> int:
    """Suggest a crop size that fits the largest instance with optional padding.

    Mirrors the web app's ``computeSuggestedCropSize`` (app.html:3402).

    - If ``user_padding`` is provided, it overrides any auto-computed padding
      (including 0, which means "no padding").
    - Else if ``use_augmentation``, padding is computed from
      ``rotation_max`` / ``scale_max``.
    - Result is rounded UP to the next multiple of ``max_stride``.

    Args:
        max_bbox_dim: Largest instance bbox dimension (height or width).
        max_stride: Network max stride; result will be a multiple of this.
        use_augmentation: Whether to add padding for rotation/scale aug.
        user_padding: Explicit padding override.
        rotation_max: Max rotation in degrees (used when use_augmentation).
        scale_max: Max scale factor (used when use_augmentation).

    Returns:
        Suggested crop size in pixels, divisible by ``max_stride``.
    """
    if user_padding is not None and user_padding >= 0:
        padding = user_padding
    elif use_augmentation:
        padding = compute_augmentation_padding(max_bbox_dim, rotation_max, scale_max)
    else:
        padding = 0

    size_with_padding = max_bbox_dim + padding
    return math.ceil(size_with_padding / max_stride) * max_stride

compute_suggested_tile_overlap(tile_size, max_bbox_dim, confmap_sigma, output_stride, backbone_margin, min_overlap_fraction=0.25, sigma_multiple=3.0)

Overlap (px) large enough that a seam-straddling object is whole in one tile.

Covers half the object extent + a few confmap sigmas + backbone context, is at least min_overlap_fraction of the tile, rounded UP to a multiple of output_stride, and clamped to leave a positive stride (>= output_stride).

Parameters:

Name	Type	Description	Default
tile_size	int	Square tile side (from compute_suggested_tile_size).	required
max_bbox_dim	float	Largest instance bbox dimension.	required
confmap_sigma	float	Confidence-map Gaussian sigma (input pixels).	required
output_stride	int	Head output stride.	required
backbone_margin	int	Per-side context margin.	required
min_overlap_fraction	float	Minimum overlap as a fraction of tile_size.	0.25
sigma_multiple	float	How many sigmas of blob to keep whole across a seam.	3.0

Returns:

Type	Description
int	Suggested overlap in pixels, divisible by output_stride.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_suggested_tile_overlap(
    tile_size: int,
    max_bbox_dim: float,
    confmap_sigma: float,
    output_stride: int,
    backbone_margin: int,
    min_overlap_fraction: float = 0.25,
    sigma_multiple: float = 3.0,
) -> int:
    """Overlap (px) large enough that a seam-straddling object is whole in one tile.

    Covers half the object extent + a few confmap sigmas + backbone context, is
    at least ``min_overlap_fraction`` of the tile, rounded UP to a multiple of
    ``output_stride``, and clamped to leave a positive stride (``>= output_stride``).

    Args:
        tile_size: Square tile side (from ``compute_suggested_tile_size``).
        max_bbox_dim: Largest instance bbox dimension.
        confmap_sigma: Confidence-map Gaussian sigma (input pixels).
        output_stride: Head output stride.
        backbone_margin: Per-side context margin.
        min_overlap_fraction: Minimum overlap as a fraction of ``tile_size``.
        sigma_multiple: How many sigmas of blob to keep whole across a seam.

    Returns:
        Suggested overlap in pixels, divisible by ``output_stride``.
    """
    object_overlap = (
        0.5 * float(max_bbox_dim)
        + sigma_multiple * float(confmap_sigma)
        + int(backbone_margin)
    )
    frac_floor = float(min_overlap_fraction) * int(tile_size)
    overlap = math.ceil(max(object_overlap, frac_floor) / output_stride) * output_stride
    max_overlap = int(tile_size) - int(output_stride)
    return int(min(overlap, max_overlap))

compute_suggested_tile_size(max_bbox_dim, max_stride, output_stride, backbone_margin, object_multiple=2.0, min_tile_multiples=2)

Square tile side that fits an object plus context on both sides.

Rounded UP to a multiple of lcm(max_stride, output_stride) so both the confmap target grid (subsamples by output_stride) and the backbone (divides by max_stride) stay exact. Depends ONLY on object extent + margin (no overlap) to avoid a cycle with compute_suggested_tile_overlap.

Parameters:

Name	Type	Description	Default
max_bbox_dim	float	Largest instance bbox dimension (height or width).	required
max_stride	int	Backbone total downsampling factor.	required
output_stride	int	Head output stride.	required
backbone_margin	int	Per-side context margin (compute_backbone_context_margin).	required
object_multiple	float	Multiple of the object extent to span.	2.0
min_tile_multiples	int	Floor on the tile side, in units of the divisor.	2

Returns:

Type	Description
int	Suggested square tile side in pixels, divisible by both strides.

Source code in sleap_nn/config_generator/architecture_estimates.py

def compute_suggested_tile_size(
    max_bbox_dim: float,
    max_stride: int,
    output_stride: int,
    backbone_margin: int,
    object_multiple: float = 2.0,
    min_tile_multiples: int = 2,
) -> int:
    """Square tile side that fits an object plus context on both sides.

    Rounded UP to a multiple of ``lcm(max_stride, output_stride)`` so both the
    confmap target grid (subsamples by ``output_stride``) and the backbone
    (divides by ``max_stride``) stay exact. Depends ONLY on object extent +
    margin (no overlap) to avoid a cycle with ``compute_suggested_tile_overlap``.

    Args:
        max_bbox_dim: Largest instance bbox dimension (height or width).
        max_stride: Backbone total downsampling factor.
        output_stride: Head output stride.
        backbone_margin: Per-side context margin (``compute_backbone_context_margin``).
        object_multiple: Multiple of the object extent to span.
        min_tile_multiples: Floor on the tile side, in units of the divisor.

    Returns:
        Suggested square tile side in pixels, divisible by both strides.
    """
    divisor = math.lcm(int(max_stride), int(output_stride))
    raw = object_multiple * float(max_bbox_dim) + 2 * int(backbone_margin)
    tile = math.ceil(raw / divisor) * divisor
    return int(max(tile, min_tile_multiples * divisor))

decoder_blocks(max_stride, output_stride)

Number of decoder upsampling blocks needed to reach output_stride.

Source code in sleap_nn/config_generator/architecture_estimates.py

def decoder_blocks(max_stride: int, output_stride: int) -> int:
    """Number of decoder upsampling blocks needed to reach ``output_stride``."""
    if output_stride <= 0:
        return encoder_blocks(max_stride)
    return int(math.log2(max_stride / output_stride))

encoder_blocks(max_stride)

Number of encoder downsampling blocks for the given max stride.

Source code in sleap_nn/config_generator/architecture_estimates.py

def encoder_blocks(max_stride: int) -> int:
    """Number of encoder downsampling blocks for the given max stride."""
    return int(math.log2(max_stride))

estimate_unet_params(filters, max_stride, output_stride, in_channels, num_keypoints, filters_rate=1.5)

Estimate trainable parameter count of a UNet head + body.

Mirrors the web app's estimateParamsAccurate (app.html:3446) and matches the structure of the real UNet (architectures/unet.py): encoder + middle/bottleneck block + decoder + 1x1 head.

Each encoder block is 2x (kxk conv + bias) with k=3. Decoder blocks take a skip connection from the matching encoder level so their input channel count is f + skip_f.

Parameters:

Name	Type	Description	Default
filters	int	Base filter count in the first encoder block.	required
max_stride	int	Determines encoder depth (log2(max_stride) blocks).	required
output_stride	int	Determines decoder depth (log2(max_stride/output_stride)).	required
in_channels	int	Network input channels (1 grayscale, 3 RGB).	required
num_keypoints	int	Number of output channels in the head.	required
filters_rate	float	Multiplier applied to filter count per encoder block.	1.5

Returns:

Type	Description
int	Estimated parameter count (weights + biases).

Source code in sleap_nn/config_generator/architecture_estimates.py

def estimate_unet_params(
    filters: int,
    max_stride: int,
    output_stride: int,
    in_channels: int,
    num_keypoints: int,
    filters_rate: float = 1.5,
) -> int:
    """Estimate trainable parameter count of a UNet head + body.

    Mirrors the web app's ``estimateParamsAccurate`` (app.html:3446) and
    matches the structure of the real UNet (``architectures/unet.py``):
    encoder + middle/bottleneck block + decoder + 1x1 head.

    Each encoder block is ``2x (kxk conv + bias)`` with k=3. Decoder blocks
    take a skip connection from the matching encoder level so their input
    channel count is ``f + skip_f``.

    Args:
        filters: Base filter count in the first encoder block.
        max_stride: Determines encoder depth (``log2(max_stride)`` blocks).
        output_stride: Determines decoder depth (``log2(max_stride/output_stride)``).
        in_channels: Network input channels (1 grayscale, 3 RGB).
        num_keypoints: Number of output channels in the head.
        filters_rate: Multiplier applied to filter count per encoder block.

    Returns:
        Estimated parameter count (weights + biases).
    """
    down_blocks = encoder_blocks(max_stride)
    up_blocks = decoder_blocks(max_stride, output_stride)

    total = 0
    ch = in_channels
    f = filters

    # Encoder
    for _ in range(down_blocks):
        total += ch * f * 9 + f
        total += f * f * 9 + f
        ch = f
        f = int(f * filters_rate)

    # Middle / bottleneck
    total += ch * f * 9 + f
    total += f * f * 9 + f
    middle_filters = f

    # Decoder
    f = middle_filters
    for i in range(up_blocks):
        next_f = int(f / filters_rate)
        skip_f = (
            int(filters * (filters_rate ** (down_blocks - 1 - i)))
            if i < down_blocks
            else 0
        )
        decoder_input = f + skip_f
        total += decoder_input * next_f * 9 + next_f
        total += next_f * next_f * 9 + next_f
        f = next_f

    # 1x1 head
    total += f * num_keypoints + num_keypoints

    return total

recommend_default_max_stride(avg_animal_size, scale=1.0)

Bucket-based default max_stride recommendation.

Mirrors setDefaultParameters in docs/configuration/config-picker/app.html (lines 5371–5375): pick the stride based on the average animal bbox size after input scaling.

Parameters:

Name	Type	Description	Default
avg_animal_size	float	Average instance bbox diagonal in original pixels.	required
scale	float	Input scale factor (multiplier applied before pickup).	1.0

Returns:

Type	Description
int	Recommended max_stride: 8 if effective size < 40, 32 if > 100, else 16.

Source code in sleap_nn/config_generator/architecture_estimates.py

def recommend_default_max_stride(avg_animal_size: float, scale: float = 1.0) -> int:
    """Bucket-based default ``max_stride`` recommendation.

    Mirrors ``setDefaultParameters`` in
    ``docs/configuration/config-picker/app.html`` (lines 5371–5375): pick
    the stride based on the *average* animal bbox size after input scaling.

    Args:
        avg_animal_size: Average instance bbox diagonal in original pixels.
        scale: Input scale factor (multiplier applied before pickup).

    Returns:
        Recommended max_stride: 8 if effective size < 40, 32 if > 100, else 16.
    """
    effective = avg_animal_size * scale
    if effective < 40:
        return 8
    if effective > 100:
        return 32
    return 16