centered_instance

sleap_nn.export.wrappers.centered_instance

Centered-instance ONNX wrapper.

Classes:

Name	Description
CenteredInstanceONNXWrapper	ONNX-exportable wrapper for centered-instance models.

CenteredInstanceONNXWrapper

Bases: BaseExportWrapper

ONNX-exportable wrapper for centered-instance models.

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

Methods:

Name	Description
__init__	Initialize centered instance ONNX wrapper.
forward	Run centered-instance inference on crops.

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

class CenteredInstanceONNXWrapper(BaseExportWrapper):
    """ONNX-exportable wrapper for centered-instance models.

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

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

        Args:
            model: Centered instance model for pose estimation.
            output_stride: Output stride for confidence maps.
            input_scale: Input scaling factor.
            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 centered-instance inference on crops."""
        image = self._normalize_uint8(image)
        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
            )

        confmaps = self._extract_tensor(
            self.model(image), ["centered", "instance", "confmap"]
        )
        peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)
        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 centered instance ONNX wrapper.

Parameters:

Name	Type	Description	Default
model	Module	Centered instance model for pose estimation.	required
output_stride	int	Output stride for confidence maps.	4
input_scale	float	Input scaling factor.	1.0
peak_threshold	float	Minimum confidence for a peak to be considered valid.	0.2

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

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

    Args:
        model: Centered instance model for pose estimation.
        output_stride: Output stride for confidence maps.
        input_scale: Input scaling factor.
        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 centered-instance inference on crops.

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

def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
    """Run centered-instance inference on crops."""
    image = self._normalize_uint8(image)
    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
        )

    confmaps = self._extract_tensor(
        self.model(image), ["centered", "instance", "confmap"]
    )
    peaks, values = self._find_global_peaks(confmaps, self.peak_threshold)
    peaks = peaks * (self.output_stride / self.input_scale)

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