single_instance

sleap_nn.export.wrappers.single_instance

Single-instance ONNX wrapper.

Classes:

Name	Description
SingleInstanceONNXWrapper	ONNX-exportable wrapper for single-instance models.

SingleInstanceONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for single-instance models.

This wrapper handles full-frame inference assuming a single instance per frame. For each body part (channel), it finds the global maximum in the confidence map.

Expects input images as uint8 tensors in [0, 255].

Attributes:

Name	Type	Description
model		The trained backbone model that outputs confidence maps.
output_stride		Output stride of the model (e.g., 4 means confmaps are ¼ the input resolution).
input_scale		Factor to scale input images before inference.

Methods:

Name	Description
__init__	Initialize the single-instance wrapper.
forward	Run single-instance inference.

Source code in sleap_nn/export/wrappers/single_instance.py

class SingleInstanceONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for single-instance models.

    This wrapper handles full-frame inference assuming a single instance per frame.
    For each body part (channel), it finds the global maximum in the confidence map.

    Expects input images as uint8 tensors in [0, 255].

    Attributes:
        model: The trained backbone model that outputs confidence maps.
        output_stride: Output stride of the model (e.g., 4 means confmaps are 1/4 the
            input resolution).
        input_scale: Factor to scale input images before inference.
    """

    def __init__(
        self,
        model: nn.Module,
        output_stride: int = 4,
        input_scale: float = 1.0,
        peak_threshold: float = 0.2,
    ):
        """Initialize the single-instance wrapper.

        Args:
            model: The trained backbone model.
            output_stride: Output stride of the model. Default: 4.
            input_scale: Factor to scale input images. Default: 1.0.
            peak_threshold: Minimum confidence for a peak to be considered valid.
        """
        super().__init__(model)
        self.output_stride = output_stride
        self.input_scale = input_scale
        self.peak_threshold = peak_threshold

    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
        """Run single-instance inference.

        Args:
            image: Input image tensor of shape (batch, channels, height, width).
                Expected as uint8 [0, 255] values.

        Returns:
            Dictionary with:
                peaks: Peak coordinates of shape (batch, n_nodes, 2) in (x, y) format.
                peak_vals: Peak confidence values of shape (batch, n_nodes).
        """
        # Normalize uint8 [0, 255] to float32 [0, 1]
        image = self._normalize_uint8(image)

        # Apply input scaling if needed. ``antialias=True`` matches the PyTorch inference
        # path, which resizes via ``torchvision.transforms.functional.resize``
        # (antialiased by default). Without it a downscaling resize (input_scale < 1)
        # produces visibly different pixels, which shifts confidence-map peaks and —
        # once scaled back by ``output_stride / input_scale`` — the exported keypoints.
        if self.input_scale != 1.0:
            height = int(image.shape[-2] * self.input_scale)
            width = int(image.shape[-1] * self.input_scale)
            image = F.interpolate(
                image,
                size=(height, width),
                mode="bilinear",
                align_corners=False,
                antialias=True,
            )

        # Run model to get confidence maps: (batch, n_nodes, height, width)
        confmaps = self._extract_tensor(
            self.model(image), ["single", "instance", "confmap"]
        )

        # Find global peak for each channel: (batch, n_nodes, 2), (batch, n_nodes)
        peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)

        # Scale peaks from confmap coordinates to image coordinates
        peaks = peaks * (self.output_stride / self.input_scale)

        return {
            "peaks": peaks,
            "peak_vals": values,
        }

__init__(model, output_stride=4, input_scale=1.0, peak_threshold=0.2)

Initialize the single-instance wrapper.

Parameters:

Name	Type	Description	Default
model	Module	The trained backbone model.	required
output_stride	int	Output stride of the model. Default: 4.	4
input_scale	float	Factor to scale input images. Default: 1.0.	1.0
peak_threshold	float	Minimum confidence for a peak to be considered valid.	0.2

Source code in sleap_nn/export/wrappers/single_instance.py

def __init__(
    self,
    model: nn.Module,
    output_stride: int = 4,
    input_scale: float = 1.0,
    peak_threshold: float = 0.2,
):
    """Initialize the single-instance wrapper.

    Args:
        model: The trained backbone model.
        output_stride: Output stride of the model. Default: 4.
        input_scale: Factor to scale input images. Default: 1.0.
        peak_threshold: Minimum confidence for a peak to be considered valid.
    """
    super().__init__(model)
    self.output_stride = output_stride
    self.input_scale = input_scale
    self.peak_threshold = peak_threshold

forward(image)

Run single-instance inference.

Parameters:

Name	Type	Description	Default
image	Tensor	Input image tensor of shape (batch, channels, height, width). Expected as uint8 [0, 255] values.	required

Returns:

Type	Description
Dict[str, Tensor]	Dictionary with: peaks: Peak coordinates of shape (batch, n_nodes, 2) in (x, y) format. peak_vals: Peak confidence values of shape (batch, n_nodes).

Source code in sleap_nn/export/wrappers/single_instance.py

def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run single-instance inference.

    Args:
        image: Input image tensor of shape (batch, channels, height, width).
            Expected as uint8 [0, 255] values.

    Returns:
        Dictionary with:
            peaks: Peak coordinates of shape (batch, n_nodes, 2) in (x, y) format.
            peak_vals: Peak confidence values of shape (batch, n_nodes).
    """
    # Normalize uint8 [0, 255] to float32 [0, 1]
    image = self._normalize_uint8(image)

    # Apply input scaling if needed. ``antialias=True`` matches the PyTorch inference
    # path, which resizes via ``torchvision.transforms.functional.resize``
    # (antialiased by default). Without it a downscaling resize (input_scale < 1)
    # produces visibly different pixels, which shifts confidence-map peaks and —
    # once scaled back by ``output_stride / input_scale`` — the exported keypoints.
    if self.input_scale != 1.0:
        height = int(image.shape[-2] * self.input_scale)
        width = int(image.shape[-1] * self.input_scale)
        image = F.interpolate(
            image,
            size=(height, width),
            mode="bilinear",
            align_corners=False,
            antialias=True,
        )

    # Run model to get confidence maps: (batch, n_nodes, height, width)
    confmaps = self._extract_tensor(
        self.model(image), ["single", "instance", "confmap"]
    )

    # Find global peak for each channel: (batch, n_nodes, 2), (batch, n_nodes)
    peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)

    # Scale peaks from confmap coordinates to image coordinates
    peaks = peaks * (self.output_stride / self.input_scale)

    return {
        "peaks": peaks,
        "peak_vals": values,
    }