layers

sleap_nn.inference.layers

Inference layers — model-type-aware wrappers around a runtime backend.

Layers are model-type-aware (peak finding, NMS, multi-class identity grouping). Backends are runtime-aware (PyTorch, ONNX, TensorRT). Crossing the two gives 6 × 3 = 18 conceptual variants — but with this protocol-based split we only ship 6 + 3 = 9 classes total, with zero duplication.

Modules:

Name	Description
backends	Runtime backends for inference layers.
base	InferenceLayer — abstract base for every model-type layer.
bottomup	BottomUpLayer — single-stage multi-instance inference via PAF grouping.
bottomup_multiclass	BottomUpMultiClassLayer — multi-class variant of bottom-up inference.
centered_instance	CenteredInstanceLayer — predicts keypoints from instance-centered crops.
centroid	CentroidLayer — predicts instance centroids from a confmap model.
configs	PreprocessConfig / PostprocessConfig — value types parameterizing layers.
embedding	Embedding (crop -> appearance vector, re-ID) inference layers.
exported	Export-adapter layers — thin translators around exported ONNX/TRT models.
segmentation	SegmentationLayer — bottom-up instance segmentation inference.
single_instance	SingleInstanceLayer — single-pose-per-frame inference.
tiled	Sliding-window tiled inference wrappers.
topdown	TopDownLayer — composes CentroidLayer + CenteredInstanceLayer.
topdown_multiclass	TopDownMultiClassLayer — multi-class variant of top-down inference.
topdown_segmentation	Top-down (crop-centered) instance-segmentation inference layers (#622).

Classes:

Name	Description
InferenceLayer	Abstract base for model-type-specific inference layers.
PostprocessConfig	Knobs that govern how raw model outputs become keypoints.
PreprocessConfig	Preprocessing knobs applied before the model forward pass.
SingleInstanceLayer	Single-pose-per-frame inference layer.

InferenceLayer

Bases: ABC

Abstract base for model-type-specific inference layers.

Subclasses implement preprocess (image → tensor + PreprocInfo), postprocess (raw backend output + PreprocInfo → Outputs), and may override predict for composed layers (top-down). The default predict is preprocess → backend → postprocess.

Attributes:

Name	Type	Description
backend		The runtime backend (TorchBackend etc.).
preprocess_config		Knobs governing input transformation.
postprocess_config		Knobs governing peak decoding and what intermediate tensors to keep.
output_stride		Confmap → input-pixel stride. Read from the model's head config at construction.

Methods:

Name	Description
__call__	Alias for :meth:predict.
__init__	Validate the backend protocol and stash configs.
postprocess	Turn the backend's raw dict into a structured Outputs.
predict	Run the full preprocess → backend → postprocess pipeline.
preprocess	Run the full preprocessing chain on a raw frame.
warmup	Prime the backend by running predict() on a synthesized frame.

Source code in sleap_nn/inference/layers/base.py

class InferenceLayer(ABC):
    """Abstract base for model-type-specific inference layers.

    Subclasses implement ``preprocess`` (image → tensor + ``PreprocInfo``),
    ``postprocess`` (raw backend output + ``PreprocInfo`` → ``Outputs``),
    and may override ``predict`` for composed layers (top-down). The
    default ``predict`` is preprocess → backend → postprocess.

    Attributes:
        backend: The runtime backend (``TorchBackend`` etc.).
        preprocess_config: Knobs governing input transformation.
        postprocess_config: Knobs governing peak decoding and what
            intermediate tensors to keep.
        output_stride: Confmap → input-pixel stride. Read from the model's
            head config at construction.
    """

    def __init__(
        self,
        backend: ModelBackend,
        preprocess_config: PreprocessConfig,
        postprocess_config: PostprocessConfig,
        output_stride: int,
        max_stride: int = 1,
    ) -> None:
        """Validate the backend protocol and stash configs."""
        if not isinstance(backend, ModelBackend):
            raise TypeError(
                f"backend must satisfy ModelBackend, got {type(backend).__name__}"
            )
        self.backend = backend
        self.preprocess_config = preprocess_config
        self.postprocess_config = postprocess_config
        self.output_stride = output_stride
        self.max_stride = max_stride

    # Class-level attribute for ``_extract_confmaps``.  Subclasses that
    # use confmap-based postprocessing should set this to the model's
    # canonical head key (e.g. ``"SingleInstanceConfmapsHead"``).
    _HEAD_OUTPUT_KEY: str = ""

    # ──────────────────────────────────────────────────────────────────
    # Subclass contract
    # ──────────────────────────────────────────────────────────────────

    def preprocess(self, image: ImageInput) -> Tuple[torch.Tensor, PreprocInfo]:
        """Run the full preprocessing chain on a raw frame.

        Delegates to :meth:`_apply_full_preprocess`:
        ensure_rgb/grayscale -> per-sample sizematcher (records eff_scale) ->
        input_scale -> pad_to_stride -> ``n_samples`` wrap.

        Subclasses that need non-standard behaviour (e.g. a different
        ``output_stride`` attribute or extra logic) can override this.
        """
        x = self._to_4d_tensor(image)
        scaled_5d, eff_scale, orig_hw = self._apply_full_preprocess(
            x, max_stride=self.max_stride, unsqueeze_n_samples=True
        )

        info = PreprocInfo(
            original_size=orig_hw,
            processed_size=tuple(scaled_5d.shape[-2:]),
            eff_scale=eff_scale,
            input_scale=self.preprocess_config.scale,
            output_stride=self.output_stride,
        )
        return scaled_5d, info

    @abstractmethod
    def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
        """Turn the backend's raw dict into a structured ``Outputs``."""

    # ──────────────────────────────────────────────────────────────────
    # Default forward — subclasses override for composed layers
    # ──────────────────────────────────────────────────────────────────

    def predict(self, image: ImageInput) -> Outputs:
        """Run the full preprocess → backend → postprocess pipeline."""
        x, info = self.preprocess(image)
        raw = self.backend(x)
        return self.postprocess(raw, info)

    def __call__(self, image: ImageInput) -> Outputs:
        """Alias for :meth:`predict`."""
        return self.predict(image)

    # ──────────────────────────────────────────────────────────────────
    # Warmup helper — subclasses define ``warmup_input_shape``
    # ──────────────────────────────────────────────────────────────────

    def warmup(self, sample_shape: Tuple[int, ...] | None = None) -> None:
        """Prime the backend by running ``predict()`` on a synthesized frame.

        The synthesized frame goes through the layer's full ``preprocess``
        chain (sizematcher → input_scale → ensure_rgb/grayscale → pad →
        n_samples wrap) so the model receives an input with the same
        rank / channel-count / device contract as real inference, and
        cuDNN's algorithm cache is primed for the right shape.

        When ``sample_shape`` is ``None`` (the default), a tiny raw frame
        is synthesized and routed through the layer's full ``preprocess``
        chain so cuDNN's algorithm cache is primed for the correct input
        shape. This avoids shape-mismatch crashes that can occur when a
        bare ``backend.warmup`` bypasses ``preprocess`` and cuDNN caches
        an algorithm for a degenerate dummy shape.

        Args:
            sample_shape: Escape hatch. When provided, dispatches straight
                to ``backend.warmup``. Prefer the default (synthesized
                real frame) on cuda / mps.
        """
        if sample_shape is not None:
            self.backend.warmup(sample_shape)
            return
        if self.backend.device == "cpu":
            return  # warmup is a no-op on CPU; first forward is already cold-start
        # Synthesize a tiny 3-channel uint8 frame in raw-video shape
        # (H, W, C). ``preprocess`` will route it through sizematcher (when
        # ``max_height``/``max_width`` are set), channel coercion, input
        # scale, stride pad, and the n_samples wrap — producing the exact
        # post-preprocess shape real inference uses.
        cfg = self.preprocess_config
        h = min(cfg.max_height or 96, 256)
        w = min(cfg.max_width or 96, 256)
        dummy = np.zeros((h, w, 3), dtype=np.uint8)
        try:
            self.predict(dummy)
        except Exception:  # noqa: BLE001 — warmup is best-effort
            pass
        if self.backend.device.startswith("cuda"):
            torch.cuda.synchronize()
        elif self.backend.device == "mps":
            torch.mps.synchronize()

    @property
    def warmup_input_shape(self) -> Tuple[int, ...]:
        """Warmup shape -- only used when ``sample_shape`` is passed.

        The default ``warmup()`` path ignores this and synthesizes a real
        raw frame instead.
        """
        return (1, 1, 64, 64)

    # ──────────────────────────────────────────────────────────────────
    # Shared confmap extraction
    # ──────────────────────────────────────────────────────────────────

    # Key used by ``TorchBackend`` when the Lightning forward returns a
    # bare ``Tensor`` (wrapped as ``{"output": tensor}``).
    _TORCH_OUTPUT_KEY: str = "output"

    def _extract_confmaps(self, raw_out: dict) -> torch.Tensor:
        """Pull the confmap tensor out of the backend's dict.

        ``TorchBackend`` wraps a tensor-returning Lightning forward under
        ``"output"``; if the model returned a dict directly, we look for
        the canonical head name stored in ``_HEAD_OUTPUT_KEY``.

        Subclasses set ``_HEAD_OUTPUT_KEY`` to their model's canonical
        head output key (e.g. ``"SingleInstanceConfmapsHead"``).
        """
        if self._TORCH_OUTPUT_KEY in raw_out:
            return raw_out[self._TORCH_OUTPUT_KEY]
        if self._HEAD_OUTPUT_KEY and self._HEAD_OUTPUT_KEY in raw_out:
            return raw_out[self._HEAD_OUTPUT_KEY]
        # Fall back to the single tensor in the dict, if there's exactly one.
        tensors = [v for v in raw_out.values() if isinstance(v, torch.Tensor)]
        if len(tensors) == 1:
            return tensors[0]
        head = self._HEAD_OUTPUT_KEY or "(not set)"
        raise KeyError(
            f"{type(self).__name__}.postprocess could not find confmaps in "
            f"raw_out keys={list(raw_out.keys())}; expected "
            f"'{self._TORCH_OUTPUT_KEY}' or '{head}'."
        )

    # ──────────────────────────────────────────────────────────────────
    # Helpers shared by every subclass
    # ──────────────────────────────────────────────────────────────────

    @staticmethod
    def _to_4d_tensor(image: ImageInput) -> torch.Tensor:
        """Coerce an image input to ``(B, C, H, W)``, preserving dtype.

        Accepts:

        - ``(H, W)`` grayscale numpy/torch
        - ``(H, W, C)`` channel-last numpy/torch
        - ``(B, H, W, C)`` channel-last
        - ``(C, H, W)`` channel-first
        - ``(B, C, H, W)`` channel-first

        Returns ``(B, C, H, W)`` with the same dtype as the input. uint8
        inputs stay uint8 so subsequent ``tvf.resize`` calls produce
        clean integer outputs (eager float conversion produces
        255.00006... values that diverge from the clean uint8 path).
        """
        if isinstance(image, np.ndarray):
            t = torch.from_numpy(image)
        elif isinstance(image, torch.Tensor):
            t = image
        else:
            raise TypeError(
                f"image must be np.ndarray or torch.Tensor, got {type(image).__name__}"
            )

        if t.ndim == 2:  # (H, W)
            t = t.unsqueeze(0).unsqueeze(0)  # (1, 1, H, W)
        elif t.ndim == 3:
            # Heuristic: smaller-trailing-dim → channel-last; otherwise
            # already channel-first single sample.
            if t.shape[-1] <= 4 and t.shape[0] > 4:  # (H, W, C)
                t = t.permute(2, 0, 1).unsqueeze(0)  # (1, C, H, W)
            else:  # (C, H, W)
                t = t.unsqueeze(0)
        elif t.ndim == 4:
            # Same heuristic for batched: trailing channel-last → permute.
            if t.shape[-1] <= 4 and t.shape[1] > 4:
                t = t.permute(0, 3, 1, 2)
        else:
            raise ValueError(f"unexpected image rank {t.ndim}: shape {tuple(t.shape)}")

        return t

    @classmethod
    def _to_4d_float_tensor(cls, image: ImageInput) -> torch.Tensor:
        """Coerce to ``(B, C, H, W)`` ``torch.float32``.

        Thin wrapper over :meth:`_to_4d_tensor` for callers that
        explicitly want float32 (ONNX backends that don't accept uint8,
        GT-path helpers, etc.). Layer ``preprocess()`` methods use
        ``_to_4d_tensor`` to preserve the uint8 ``tvf.resize`` path.
        """
        return cls._to_4d_tensor(image).float()

    # ──────────────────────────────────────────────────────────────────
    # Shared raw-frame preprocessing chain
    # ──────────────────────────────────────────────────────────────────

    def _apply_full_preprocess(
        self,
        x: torch.Tensor,
        *,
        max_stride: int = 1,
        unsqueeze_n_samples: bool = True,
        skip_sizematcher: bool = False,
    ) -> Tuple[torch.Tensor, torch.Tensor, Tuple[int, int]]:
        """Run the standard preprocessing chain on a (B, C, H, W) tensor.

        Each step short-circuits when its config field is the identity
        (``None``/``False``/``1.0``), so a raw-frame layer running on a
        properly-sized batch sees zero extra ops.

        Stages applied in order:

        1. ``ensure_rgb`` / ``ensure_grayscale`` -- channel coercion.
        2. Per-sample ``apply_sizematcher`` to
           ``(preprocess_config.max_height, preprocess_config.max_width)``,
           returning a per-sample ``eff_scale`` for the coord-undo ladder.
        3. ``resize_image`` by ``preprocess_config.scale`` -- global input
           scale.
        4. ``apply_pad_to_stride`` to ``max_stride``. Use the model's
           max_stride; ``1`` is a no-op.
        5. ``unsqueeze(dim=1)`` to add the ``n_samples`` axis so the
           Lightning forward's unconditional ``squeeze(dim=1)`` resolves
           to the expected rank. Skip when the layer's forward accepts
           4D directly.

        Args:
            x: ``(B, C, H, W)`` float32 tensor from :meth:`_to_4d_float_tensor`.
            max_stride: Model's required input stride; the input is padded
                bottom-right to a multiple of this. ``1`` is the identity.
            unsqueeze_n_samples: When ``True`` (the default for
                multi-instance layers) wraps with a ``(B, 1, C, H, W)``
                ``n_samples`` axis. Top-down crops feed
                :class:`CenteredInstanceLayer` post-crop and don't need
                sizematcher — those callers pass ``False``.
            skip_sizematcher: When ``True``, bypass the per-sample
                ``apply_sizematcher`` step entirely and return an all-ones
                ``eff_scale``. Used by the tiled-inference path
                (:class:`~sleap_nn.inference.layers.tiled.TiledLayer`), which
                processes each frame at native resolution (only ``input_scale``
                applies) instead of shrinking it to
                ``(max_height, max_width)``. The default (``False``) is
                byte-identical to the pre-existing behaviour.

        Returns:
            ``(processed_tensor, eff_scale, original_HW)``:

            * ``processed_tensor``: ``(B, 1, C, H', W')`` if
              ``unsqueeze_n_samples`` else ``(B, C, H', W')``.
            * ``eff_scale``: ``(B,)`` per-sample sizematcher scale factor.
              All ones when no sizematcher is configured.
            * ``original_HW``: ``(H, W)`` of the input before any resize.
        """
        # Local imports avoid a circular base.py → data.* → ... → base.py path.
        from sleap_nn.data.normalization import convert_to_grayscale, convert_to_rgb
        from sleap_nn.data.resizing import (
            apply_pad_to_stride,
            apply_sizematcher,
            resize_image,
        )

        cfg = self.preprocess_config
        B, _C, H, W = x.shape
        orig_hw = (H, W)

        # 1. Channel coercion. Check ensure_rgb first to match legacy precedence
        # when both are set (a misconfiguration PreprocessConfig now rejects). #584.
        if cfg.ensure_rgb and x.shape[-3] != 3:
            x = convert_to_rgb(x)
        elif cfg.ensure_grayscale and x.shape[-3] != 1:
            x = convert_to_grayscale(x)

        # 2. Per-sample sizematcher → eff_scale. Skipped entirely when
        # ``skip_sizematcher`` (tiled inference runs frames at native res).
        if not skip_sizematcher and (
            cfg.max_height is not None or cfg.max_width is not None
        ):
            resized_frames: list = []
            eff_scales: list = []
            for b in range(B):
                # apply_sizematcher accepts (C, H, W); preserves device.
                r, scale = apply_sizematcher(x[b], cfg.max_height, cfg.max_width)
                resized_frames.append(r)
                eff_scales.append(float(scale))
            x = torch.stack(resized_frames, dim=0)
            eff_scale = torch.tensor(eff_scales, dtype=torch.float32, device=x.device)
        else:
            eff_scale = torch.ones(B, dtype=torch.float32, device=x.device)

        # 3. Input scale.
        if cfg.scale != 1.0:
            x = resize_image(x, cfg.scale)

        # 4. Pad to stride.
        if max_stride != 1:
            x = apply_pad_to_stride(x, max_stride)

        # 5. n_samples wrap.
        if unsqueeze_n_samples:
            x = x.unsqueeze(1)

        return x, eff_scale, orig_hw

warmup_input_shape property

Warmup shape -- only used when sample_shape is passed.

The default warmup() path ignores this and synthesizes a real raw frame instead.

__call__(image)

Alias for :meth:predict.

Source code in sleap_nn/inference/layers/base.py

def __call__(self, image: ImageInput) -> Outputs:
    """Alias for :meth:`predict`."""
    return self.predict(image)

__init__(backend, preprocess_config, postprocess_config, output_stride, max_stride=1)

Validate the backend protocol and stash configs.

Source code in sleap_nn/inference/layers/base.py

def __init__(
    self,
    backend: ModelBackend,
    preprocess_config: PreprocessConfig,
    postprocess_config: PostprocessConfig,
    output_stride: int,
    max_stride: int = 1,
) -> None:
    """Validate the backend protocol and stash configs."""
    if not isinstance(backend, ModelBackend):
        raise TypeError(
            f"backend must satisfy ModelBackend, got {type(backend).__name__}"
        )
    self.backend = backend
    self.preprocess_config = preprocess_config
    self.postprocess_config = postprocess_config
    self.output_stride = output_stride
    self.max_stride = max_stride

postprocess(raw_out, info) abstractmethod

Turn the backend's raw dict into a structured Outputs.

Source code in sleap_nn/inference/layers/base.py

@abstractmethod
def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
    """Turn the backend's raw dict into a structured ``Outputs``."""

predict(image)

Run the full preprocess → backend → postprocess pipeline.

Source code in sleap_nn/inference/layers/base.py

def predict(self, image: ImageInput) -> Outputs:
    """Run the full preprocess → backend → postprocess pipeline."""
    x, info = self.preprocess(image)
    raw = self.backend(x)
    return self.postprocess(raw, info)

preprocess(image)

Run the full preprocessing chain on a raw frame.

Delegates to :meth:_apply_full_preprocess: ensure_rgb/grayscale -> per-sample sizematcher (records eff_scale) -> input_scale -> pad_to_stride -> n_samples wrap.

Subclasses that need non-standard behaviour (e.g. a different output_stride attribute or extra logic) can override this.

Source code in sleap_nn/inference/layers/base.py

def preprocess(self, image: ImageInput) -> Tuple[torch.Tensor, PreprocInfo]:
    """Run the full preprocessing chain on a raw frame.

    Delegates to :meth:`_apply_full_preprocess`:
    ensure_rgb/grayscale -> per-sample sizematcher (records eff_scale) ->
    input_scale -> pad_to_stride -> ``n_samples`` wrap.

    Subclasses that need non-standard behaviour (e.g. a different
    ``output_stride`` attribute or extra logic) can override this.
    """
    x = self._to_4d_tensor(image)
    scaled_5d, eff_scale, orig_hw = self._apply_full_preprocess(
        x, max_stride=self.max_stride, unsqueeze_n_samples=True
    )

    info = PreprocInfo(
        original_size=orig_hw,
        processed_size=tuple(scaled_5d.shape[-2:]),
        eff_scale=eff_scale,
        input_scale=self.preprocess_config.scale,
        output_stride=self.output_stride,
    )
    return scaled_5d, info

warmup(sample_shape=None)

Prime the backend by running predict() on a synthesized frame.

The synthesized frame goes through the layer's full preprocess chain (sizematcher → input_scale → ensure_rgb/grayscale → pad → n_samples wrap) so the model receives an input with the same rank / channel-count / device contract as real inference, and cuDNN's algorithm cache is primed for the right shape.

When sample_shape is None (the default), a tiny raw frame is synthesized and routed through the layer's full preprocess chain so cuDNN's algorithm cache is primed for the correct input shape. This avoids shape-mismatch crashes that can occur when a bare backend.warmup bypasses preprocess and cuDNN caches an algorithm for a degenerate dummy shape.

Parameters:

Name	Type	Description	Default
sample_shape	Tuple[int, ...] | None	Escape hatch. When provided, dispatches straight to backend.warmup. Prefer the default (synthesized real frame) on cuda / mps.	None

Source code in sleap_nn/inference/layers/base.py

def warmup(self, sample_shape: Tuple[int, ...] | None = None) -> None:
    """Prime the backend by running ``predict()`` on a synthesized frame.

    The synthesized frame goes through the layer's full ``preprocess``
    chain (sizematcher → input_scale → ensure_rgb/grayscale → pad →
    n_samples wrap) so the model receives an input with the same
    rank / channel-count / device contract as real inference, and
    cuDNN's algorithm cache is primed for the right shape.

    When ``sample_shape`` is ``None`` (the default), a tiny raw frame
    is synthesized and routed through the layer's full ``preprocess``
    chain so cuDNN's algorithm cache is primed for the correct input
    shape. This avoids shape-mismatch crashes that can occur when a
    bare ``backend.warmup`` bypasses ``preprocess`` and cuDNN caches
    an algorithm for a degenerate dummy shape.

    Args:
        sample_shape: Escape hatch. When provided, dispatches straight
            to ``backend.warmup``. Prefer the default (synthesized
            real frame) on cuda / mps.
    """
    if sample_shape is not None:
        self.backend.warmup(sample_shape)
        return
    if self.backend.device == "cpu":
        return  # warmup is a no-op on CPU; first forward is already cold-start
    # Synthesize a tiny 3-channel uint8 frame in raw-video shape
    # (H, W, C). ``preprocess`` will route it through sizematcher (when
    # ``max_height``/``max_width`` are set), channel coercion, input
    # scale, stride pad, and the n_samples wrap — producing the exact
    # post-preprocess shape real inference uses.
    cfg = self.preprocess_config
    h = min(cfg.max_height or 96, 256)
    w = min(cfg.max_width or 96, 256)
    dummy = np.zeros((h, w, 3), dtype=np.uint8)
    try:
        self.predict(dummy)
    except Exception:  # noqa: BLE001 — warmup is best-effort
        pass
    if self.backend.device.startswith("cuda"):
        torch.cuda.synchronize()
    elif self.backend.device == "mps":
        torch.mps.synchronize()

PostprocessConfig

Knobs that govern how raw model outputs become keypoints.

Distinct from the post-inference FilterConfig: this struct governs the decoding step (peak finding, integral refinement, NMS), while FilterConfig filters the keypoints that come out the other side. peak_threshold here decides which confmap pixels become peaks; min_peak_value in FilterConfig filters peaks the decoder already returned.

Attributes:

Name	Type	Description
peak_threshold	float	Minimum confmap activation to consider a peak.
refinement	Literal['integral', 'none']	"integral" runs sub-pixel integral regression around each rough peak. "none" returns grid-aligned peaks.
integral_patch_size	int	Side length of the refinement patch.
max_instances	Optional[int]	Cap on instances per frame (centroid layer only).
return_confmaps	bool	Keep (B, N, H, W) confmaps on the Outputs (heavy — opt-in for visualization / debugging).
return_pafs	bool	Keep (B, 2E, H, W) PAFs (bottom-up only; heavy).
return_paf_graph	bool	Keep the bottom-up PAF graph tuple (opt-in).
return_class_maps	bool	Keep (B, C, H, W) class maps (multi-class bottom-up; heavy).
return_class_vectors	bool	Keep (B, I, N, C) class logits (multi-class top-down).

Source code in sleap_nn/inference/layers/configs.py

@attrs.frozen
class PostprocessConfig:
    """Knobs that govern how raw model outputs become keypoints.

    Distinct from the post-inference ``FilterConfig``: this struct governs
    the *decoding* step (peak finding, integral refinement, NMS), while
    ``FilterConfig`` filters the keypoints that come out the other side.
    ``peak_threshold`` here decides which confmap pixels become peaks;
    ``min_peak_value`` in ``FilterConfig`` filters peaks the decoder
    already returned.

    Attributes:
        peak_threshold: Minimum confmap activation to consider a peak.
        refinement: ``"integral"`` runs sub-pixel integral regression around
            each rough peak. ``"none"`` returns grid-aligned peaks.
        integral_patch_size: Side length of the refinement patch.
        max_instances: Cap on instances per frame (centroid layer only).
        return_confmaps: Keep ``(B, N, H, W)`` confmaps on the ``Outputs``
            (heavy — opt-in for visualization / debugging).
        return_pafs: Keep ``(B, 2E, H, W)`` PAFs (bottom-up only; heavy).
        return_paf_graph: Keep the bottom-up PAF graph tuple (opt-in).
        return_class_maps: Keep ``(B, C, H, W)`` class maps (multi-class
            bottom-up; heavy).
        return_class_vectors: Keep ``(B, I, N, C)`` class logits
            (multi-class top-down).
    """

    peak_threshold: float = 0.2
    refinement: Literal["integral", "none"] = "integral"
    integral_patch_size: int = 5
    max_instances: Optional[int] = None

    return_confmaps: bool = False
    return_pafs: bool = False
    return_paf_graph: bool = False
    return_class_maps: bool = False
    return_class_vectors: bool = False

    @property
    def effective_refinement(self) -> Optional[str]:
        """Return the refinement string or ``None`` when ``"none"``.

        Every postprocess site needs ``refinement=None`` (not the string
        ``"none"``) to disable refinement. This property centralises that
        coercion.
        """
        return self.refinement if self.refinement != "none" else None

effective_refinement property

Return the refinement string or None when "none".

Every postprocess site needs refinement=None (not the string "none") to disable refinement. This property centralises that coercion.

PreprocessConfig

Preprocessing knobs applied before the model forward pass.

Defaults are the no-op identity for every field — calling the layer on an already-correctly-shaped batch produces zero extra work.

Attributes:

Name	Type	Description
ensure_rgb	Optional[bool]	True forces 3-channel RGB; False forces 1-channel grayscale; None leaves channels untouched.
ensure_grayscale	Optional[bool]	Inverse of ensure_rgb. Mutually exclusive.
max_height	Optional[int]	Resize so height ≤ this (preserves aspect ratio). None means no max.
max_width	Optional[int]	Same for width.
scale	float	Multiplicative input-scale factor applied (after size matching) via :func:sleap_nn.data.resizing.resize_image (tvf.resize) on the live ckpt path. 1.0 is identity.
crop_size	Optional[Tuple[int, int]]	Top-down stage 2 only — square crop side length.

Methods:

Name	Description
__attrs_post_init__	Reject the contradictory ensure_rgb=True + ensure_grayscale=True.

Source code in sleap_nn/inference/layers/configs.py

@attrs.frozen
class PreprocessConfig:
    """Preprocessing knobs applied before the model forward pass.

    Defaults are the no-op identity for every field — calling the layer
    on an already-correctly-shaped batch produces zero extra work.

    Attributes:
        ensure_rgb: ``True`` forces 3-channel RGB; ``False`` forces 1-channel
            grayscale; ``None`` leaves channels untouched.
        ensure_grayscale: Inverse of ``ensure_rgb``. Mutually exclusive.
        max_height: Resize so height ≤ this (preserves aspect ratio). ``None``
            means no max.
        max_width: Same for width.
        scale: Multiplicative input-scale factor applied (after size matching)
            via :func:`sleap_nn.data.resizing.resize_image` (``tvf.resize``) on
            the live ckpt path. ``1.0`` is identity.
        crop_size: Top-down stage 2 only — square crop side length.
    """

    ensure_rgb: Optional[bool] = None
    ensure_grayscale: Optional[bool] = None
    max_height: Optional[int] = None
    max_width: Optional[int] = None
    scale: float = 1.0
    crop_size: Optional[Tuple[int, int]] = None

    def __attrs_post_init__(self) -> None:
        """Reject the contradictory ``ensure_rgb=True`` + ``ensure_grayscale=True``."""
        if self.ensure_rgb and self.ensure_grayscale:
            raise ValueError(
                "ensure_rgb and ensure_grayscale cannot both be True; choose one "
                "(or leave both None to keep the source channel count)."
            )

__attrs_post_init__()

Reject the contradictory ensure_rgb=True + ensure_grayscale=True.

Source code in sleap_nn/inference/layers/configs.py

def __attrs_post_init__(self) -> None:
    """Reject the contradictory ``ensure_rgb=True`` + ``ensure_grayscale=True``."""
    if self.ensure_rgb and self.ensure_grayscale:
        raise ValueError(
            "ensure_rgb and ensure_grayscale cannot both be True; choose one "
            "(or leave both None to keep the source channel count)."
        )

SingleInstanceLayer

Bases: InferenceLayer

Single-pose-per-frame inference layer.

Parameters:

Name	Type	Description	Default
backend	ModelBackend	Runtime backend (e.g. TorchBackend(model=lightning_module)).	required
preprocess_config	Optional[PreprocessConfig]	Pre-forward transformation knobs.	None
postprocess_config	Optional[PostprocessConfig]	Peak decoding + intermediate-return knobs.	None
output_stride	int	Stride between confmap and input pixels (read from the head config at construction).	required
max_stride	int	Backbone-network stride; inputs are padded bottom-right to a multiple of this in preprocess. Default 1 (no pad).	1

Methods:

Name	Description
__init__	Compose the layer with default empty configs when omitted.
postprocess	Decode confmaps → keypoints, reverse coord ladder, build Outputs.

Source code in sleap_nn/inference/layers/single_instance.py

class SingleInstanceLayer(InferenceLayer):
    """Single-pose-per-frame inference layer.

    Args:
        backend: Runtime backend (e.g. ``TorchBackend(model=lightning_module)``).
        preprocess_config: Pre-forward transformation knobs.
        postprocess_config: Peak decoding + intermediate-return knobs.
        output_stride: Stride between confmap and input pixels (read from
            the head config at construction).
        max_stride: Backbone-network stride; inputs are padded bottom-right
            to a multiple of this in ``preprocess``. Default ``1`` (no pad).
    """

    _HEAD_OUTPUT_KEY: str = "SingleInstanceConfmapsHead"

    def __init__(
        self,
        backend: ModelBackend,
        output_stride: int,
        max_stride: int = 1,
        preprocess_config: Optional[PreprocessConfig] = None,
        postprocess_config: Optional[PostprocessConfig] = None,
    ) -> None:
        """Compose the layer with default empty configs when omitted."""
        super().__init__(
            backend=backend,
            preprocess_config=preprocess_config or PreprocessConfig(),
            postprocess_config=postprocess_config or PostprocessConfig(),
            output_stride=output_stride,
            max_stride=max_stride,
        )

    # ──────────────────────────────────────────────────────────────────
    # Postprocess
    # ──────────────────────────────────────────────────────────────────

    def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
        """Decode confmaps → keypoints, reverse coord ladder, build ``Outputs``.

        This layer always runs the torch decode path: it is only ever built
        with a ``TorchBackend`` (``does_baked_postproc=False``). The exported
        ONNX/TRT path uses the separate :class:`ExportedSingleInstanceLayer`
        adapter, which returns already-final peaks without the coord ladder —
        so this method must NOT special-case a baked backend (doing so would
        double-apply the ladder; #584).
        """
        confmaps = self._extract_confmaps(raw_out)
        peaks, vals = find_global_peaks(
            confmaps.detach(),
            threshold=self.postprocess_config.peak_threshold,
            refinement=self.postprocess_config.effective_refinement,
            integral_patch_size=self.postprocess_config.integral_patch_size,
        )

        # Coord ladder: confmap pixels → input pixels → original-image pixels.
        peaks = undo_stride(peaks, info.output_stride)
        peaks = undo_input_scale(peaks, info.input_scale)
        peaks = undo_eff_scale(peaks, info.eff_scale)

        # ``find_global_peaks`` returns (B, N, 2) — single-instance has I=1.
        # Reshape to the canonical (B, I=1, N, 2) / (B, I=1, N) Outputs shape.
        peaks_BIN2 = peaks.unsqueeze(1)
        vals_BIN = vals.unsqueeze(1)

        outputs = Outputs(
            pred_keypoints=peaks_BIN2,
            pred_peak_values=vals_BIN,
            preprocess_info=info,
        )
        if self.postprocess_config.return_confmaps and confmaps is not None:
            outputs = attrs.evolve(outputs, pred_confmaps=confmaps.detach())
        return outputs

__init__(backend, output_stride, max_stride=1, preprocess_config=None, postprocess_config=None)

Compose the layer with default empty configs when omitted.

Source code in sleap_nn/inference/layers/single_instance.py

def __init__(
    self,
    backend: ModelBackend,
    output_stride: int,
    max_stride: int = 1,
    preprocess_config: Optional[PreprocessConfig] = None,
    postprocess_config: Optional[PostprocessConfig] = None,
) -> None:
    """Compose the layer with default empty configs when omitted."""
    super().__init__(
        backend=backend,
        preprocess_config=preprocess_config or PreprocessConfig(),
        postprocess_config=postprocess_config or PostprocessConfig(),
        output_stride=output_stride,
        max_stride=max_stride,
    )

postprocess(raw_out, info)

Decode confmaps → keypoints, reverse coord ladder, build Outputs.

This layer always runs the torch decode path: it is only ever built with a TorchBackend (does_baked_postproc=False). The exported ONNX/TRT path uses the separate :class:ExportedSingleInstanceLayer adapter, which returns already-final peaks without the coord ladder — so this method must NOT special-case a baked backend (doing so would double-apply the ladder; #584).

Source code in sleap_nn/inference/layers/single_instance.py

def postprocess(self, raw_out: dict, info: PreprocInfo) -> Outputs:
    """Decode confmaps → keypoints, reverse coord ladder, build ``Outputs``.

    This layer always runs the torch decode path: it is only ever built
    with a ``TorchBackend`` (``does_baked_postproc=False``). The exported
    ONNX/TRT path uses the separate :class:`ExportedSingleInstanceLayer`
    adapter, which returns already-final peaks without the coord ladder —
    so this method must NOT special-case a baked backend (doing so would
    double-apply the ladder; #584).
    """
    confmaps = self._extract_confmaps(raw_out)
    peaks, vals = find_global_peaks(
        confmaps.detach(),
        threshold=self.postprocess_config.peak_threshold,
        refinement=self.postprocess_config.effective_refinement,
        integral_patch_size=self.postprocess_config.integral_patch_size,
    )

    # Coord ladder: confmap pixels → input pixels → original-image pixels.
    peaks = undo_stride(peaks, info.output_stride)
    peaks = undo_input_scale(peaks, info.input_scale)
    peaks = undo_eff_scale(peaks, info.eff_scale)

    # ``find_global_peaks`` returns (B, N, 2) — single-instance has I=1.
    # Reshape to the canonical (B, I=1, N, 2) / (B, I=1, N) Outputs shape.
    peaks_BIN2 = peaks.unsqueeze(1)
    vals_BIN = vals.unsqueeze(1)

    outputs = Outputs(
        pred_keypoints=peaks_BIN2,
        pred_peak_values=vals_BIN,
        preprocess_info=info,
    )
    if self.postprocess_config.return_confmaps and confmaps is not None:
        outputs = attrs.evolve(outputs, pred_confmaps=confmaps.detach())
    return outputs