sleap-nn 0.1.0a1__py3-none-any.whl → 0.1.0a3__py3-none-any.whl

sleap_nn/export/wrappers/bottomup_multiclass.py

@@ -0,0 +1,195 @@
+"""ONNX wrapper for bottom-up multiclass (supervised ID) models."""
+
+from typing import Dict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sleap_nn.export.wrappers.base import BaseExportWrapper
+
+
+class BottomUpMultiClassONNXWrapper(BaseExportWrapper):
+    """ONNX-exportable wrapper for bottom-up multiclass (supervised ID) models.
+
+    This wrapper handles models that output both confidence maps for keypoint
+    detection and class maps for identity classification. Unlike PAF-based
+    bottom-up models, multiclass models use class maps to assign identity to
+    each detected peak, then group peaks by identity.
+
+    The wrapper performs:
+    1. Peak detection in confidence maps (GPU)
+    2. Class probability sampling at peak locations (GPU)
+    3. Returns fixed-size tensors for CPU-side grouping
+
+    Expects input images as uint8 tensors in [0, 255].
+
+    Attributes:
+        model: The underlying PyTorch model.
+        n_nodes: Number of keypoint nodes in the skeleton.
+        n_classes: Number of identity classes.
+        max_peaks_per_node: Maximum number of peaks to detect per node.
+        cms_output_stride: Output stride of the confidence map head.
+        class_maps_output_stride: Output stride of the class maps head.
+        input_scale: Scale factor applied to input images before inference.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        n_nodes: int,
+        n_classes: int = 2,
+        max_peaks_per_node: int = 20,
+        cms_output_stride: int = 4,
+        class_maps_output_stride: int = 8,
+        input_scale: float = 1.0,
+    ):
+        """Initialize the wrapper.
+
+        Args:
+            model: The underlying PyTorch model.
+            n_nodes: Number of keypoint nodes.
+            n_classes: Number of identity classes (e.g., 2 for male/female).
+            max_peaks_per_node: Maximum peaks per node to detect.
+            cms_output_stride: Output stride of confidence maps.
+            class_maps_output_stride: Output stride of class maps.
+            input_scale: Scale factor for input images.
+        """
+        super().__init__(model)
+        self.n_nodes = n_nodes
+        self.n_classes = n_classes
+        self.max_peaks_per_node = max_peaks_per_node
+        self.cms_output_stride = cms_output_stride
+        self.class_maps_output_stride = class_maps_output_stride
+        self.input_scale = input_scale
+
+    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Run bottom-up multiclass inference.
+
+        Args:
+            image: Input image tensor of shape (batch, channels, height, width).
+                   Expected to be uint8 in [0, 255].
+
+        Returns:
+            Dictionary with keys:
+                - "peaks": Detected peak coordinates (batch, n_nodes, max_peaks, 2).
+                    Coordinates are in input image space (x, y).
+                - "peak_vals": Peak confidence values (batch, n_nodes, max_peaks).
+                - "peak_mask": Boolean mask for valid peaks (batch, n_nodes, max_peaks).
+                - "class_probs": Class probabilities at each peak location
+                    (batch, n_nodes, max_peaks, n_classes).
+
+            Postprocessing on CPU uses `classify_peaks_from_maps()` to group
+            peaks by identity using Hungarian matching.
+        """
+        # Normalize uint8 [0, 255] to float32 [0, 1]
+        image = self._normalize_uint8(image)
+
+        # Apply input scaling if needed
+        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
+            )
+
+        batch_size = image.shape[0]
+
+        # Forward pass
+        out = self.model(image)
+
+        # Extract outputs
+        # Note: Use "classmaps" as a single hint to avoid "map" matching "confmaps"
+        confmaps = self._extract_tensor(out, ["confmap", "multiinstance"])
+        class_maps = self._extract_tensor(out, ["classmaps", "classmapshead"])
+
+        # Find top-k peaks per node
+        peaks, peak_vals, peak_mask = self._find_topk_peaks_per_node(
+            confmaps, self.max_peaks_per_node
+        )
+
+        # Scale peaks to input image space
+        peaks = peaks * self.cms_output_stride
+
+        # Sample class maps at peak locations
+        class_probs = self._sample_class_maps_at_peaks(class_maps, peaks, peak_mask)
+
+        # Scale peaks for output (accounting for input scale)
+        if self.input_scale != 1.0:
+            peaks = peaks / self.input_scale
+
+        return {
+            "peaks": peaks,
+            "peak_vals": peak_vals,
+            "peak_mask": peak_mask,
+            "class_probs": class_probs,
+        }
+
+    def _sample_class_maps_at_peaks(
+        self,
+        class_maps: torch.Tensor,
+        peaks: torch.Tensor,
+        peak_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        """Sample class map values at peak locations.
+
+        Args:
+            class_maps: Class maps of shape (batch, n_classes, height, width).
+            peaks: Peak coordinates in cms_output_stride space,
+                shape (batch, n_nodes, max_peaks, 2) in (x, y) order.
+            peak_mask: Boolean mask for valid peaks (batch, n_nodes, max_peaks).
+
+        Returns:
+            Class probabilities at each peak location,
+            shape (batch, n_nodes, max_peaks, n_classes).
+        """
+        batch_size, n_classes, cm_height, cm_width = class_maps.shape
+        _, n_nodes, max_peaks, _ = peaks.shape
+        device = peaks.device
+
+        # Initialize output tensor
+        class_probs = torch.zeros(
+            (batch_size, n_nodes, max_peaks, n_classes),
+            device=device,
+            dtype=class_maps.dtype,
+        )
+
+        # Convert peak coordinates to class map space
+        # peaks are in full image space (after cms_output_stride scaling)
+        peaks_cm = peaks / self.class_maps_output_stride
+
+        # Clamp coordinates to valid range
+        peaks_cm_x = peaks_cm[..., 0].clamp(0, cm_width - 1)
+        peaks_cm_y = peaks_cm[..., 1].clamp(0, cm_height - 1)
+
+        # Use grid_sample for bilinear interpolation
+        # Normalize coordinates to [-1, 1] for grid_sample
+        grid_x = (peaks_cm_x / (cm_width - 1)) * 2 - 1
+        grid_y = (peaks_cm_y / (cm_height - 1)) * 2 - 1
+
+        # Reshape for grid_sample: (batch, n_nodes * max_peaks, 1, 2)
+        grid = torch.stack([grid_x, grid_y], dim=-1)
+        grid_flat = grid.reshape(batch_size, n_nodes * max_peaks, 1, 2)
+
+        # Sample class maps: (batch, n_classes, n_nodes * max_peaks, 1)
+        sampled = F.grid_sample(
+            class_maps,
+            grid_flat,
+            mode="bilinear",
+            padding_mode="zeros",
+            align_corners=True,
+        )
+
+        # Reshape to (batch, n_nodes, max_peaks, n_classes)
+        sampled = sampled.squeeze(-1)  # (batch, n_classes, n_nodes * max_peaks)
+        sampled = sampled.permute(0, 2, 1)  # (batch, n_nodes * max_peaks, n_classes)
+        sampled = sampled.reshape(batch_size, n_nodes, max_peaks, n_classes)
+
+        # Apply softmax to get probabilities (optional - depends on training)
+        # For now, return raw values as the grouping function expects logits
+        class_probs = sampled
+
+        # Mask invalid peaks
+        class_probs = class_probs * peak_mask.unsqueeze(-1).float()
+
+        return class_probs

sleap_nn/export/wrappers/centered_instance.py

@@ -0,0 +1,56 @@
+"""Centered-instance ONNX wrapper."""
+
+from __future__ import annotations
+
+from typing import Dict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sleap_nn.export.wrappers.base import BaseExportWrapper
+
+
+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,
+    ):
+        """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.
+        """
+        super().__init__(model)
+        self.output_stride = output_stride
+        self.input_scale = input_scale
+
+    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)
+        peaks = peaks * (self.output_stride / self.input_scale)
+
+        return {
+            "peaks": peaks,
+            "peak_vals": values,
+        }

sleap_nn/export/wrappers/centroid.py

@@ -0,0 +1,58 @@
+"""Centroid ONNX wrapper."""
+
+from __future__ import annotations
+
+from typing import Dict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sleap_nn.export.wrappers.base import BaseExportWrapper
+
+
+class CentroidONNXWrapper(BaseExportWrapper):
+    """ONNX-exportable wrapper for centroid models.
+
+    Expects input images as uint8 tensors in [0, 255].
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        max_instances: int = 20,
+        output_stride: int = 2,
+        input_scale: float = 1.0,
+    ):
+        """Initialize centroid ONNX wrapper.
+
+        Args:
+            model: Centroid detection model.
+            max_instances: Maximum number of instances to detect.
+            output_stride: Output stride for confidence maps.
+            input_scale: Input scaling factor.
+        """
+        super().__init__(model)
+        self.max_instances = max_instances
+        self.output_stride = output_stride
+        self.input_scale = input_scale
+
+    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Run centroid inference and return fixed-size outputs."""
+        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), ["centroid", "confmap"])
+        peaks, values, valid = self._find_topk_peaks(confmaps, self.max_instances)
+        peaks = peaks * (self.output_stride / self.input_scale)
+
+        return {
+            "centroids": peaks,
+            "centroid_vals": values,
+            "instance_valid": valid,
+        }

sleap_nn/export/wrappers/single_instance.py

@@ -0,0 +1,83 @@
+"""Single-instance ONNX wrapper."""
+
+from __future__ import annotations
+
+from typing import Dict
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sleap_nn.export.wrappers.base import BaseExportWrapper
+
+
+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,
+    ):
+        """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.
+        """
+        super().__init__(model)
+        self.output_stride = output_stride
+        self.input_scale = input_scale
+
+    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
+        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
+            )
+
+        # 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)
+
+        # Scale peaks from confmap coordinates to image coordinates
+        peaks = peaks * (self.output_stride / self.input_scale)
+
+        return {
+            "peaks": peaks,
+            "peak_vals": values,
+        }

sleap_nn/export/wrappers/topdown.py

@@ -0,0 +1,180 @@
+"""Top-down ONNX wrapper."""
+
+from __future__ import annotations
+
+from typing import Dict, Tuple
+
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+from sleap_nn.export.wrappers.base import BaseExportWrapper
+
+
+class TopDownONNXWrapper(BaseExportWrapper):
+    """ONNX-exportable wrapper for top-down (centroid + centered-instance) inference.
+
+    Expects input images as uint8 tensors in [0, 255].
+    """
+
+    def __init__(
+        self,
+        centroid_model: nn.Module,
+        instance_model: nn.Module,
+        max_instances: int = 20,
+        crop_size: Tuple[int, int] = (192, 192),
+        centroid_output_stride: int = 2,
+        instance_output_stride: int = 4,
+        centroid_input_scale: float = 1.0,
+        instance_input_scale: float = 1.0,
+        n_nodes: int = 1,
+    ) -> None:
+        """Initialize top-down ONNX wrapper.
+
+        Args:
+            centroid_model: Centroid detection model.
+            instance_model: Instance pose estimation model.
+            max_instances: Maximum number of instances to detect.
+            crop_size: Size of instance crops (height, width).
+            centroid_output_stride: Centroid model output stride.
+            instance_output_stride: Instance model output stride.
+            centroid_input_scale: Centroid input scaling factor.
+            instance_input_scale: Instance input scaling factor.
+            n_nodes: Number of skeleton nodes.
+        """
+        super().__init__(centroid_model)
+        self.centroid_model = centroid_model
+        self.instance_model = instance_model
+        self.max_instances = max_instances
+        self.crop_size = crop_size
+        self.centroid_output_stride = centroid_output_stride
+        self.instance_output_stride = instance_output_stride
+        self.centroid_input_scale = centroid_input_scale
+        self.instance_input_scale = instance_input_scale
+        self.n_nodes = n_nodes
+
+        crop_h, crop_w = crop_size
+        y_crop = torch.linspace(-1, 1, crop_h, dtype=torch.float32)
+        x_crop = torch.linspace(-1, 1, crop_w, dtype=torch.float32)
+        grid_y, grid_x = torch.meshgrid(y_crop, x_crop, indexing="ij")
+        base_grid = torch.stack([grid_x, grid_y], dim=-1)
+        self.register_buffer("base_grid", base_grid, persistent=False)
+
+    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Run top-down inference and return fixed-size outputs."""
+        image = self._normalize_uint8(image)
+        batch_size, channels, height, width = image.shape
+
+        scaled_image = image
+        if self.centroid_input_scale != 1.0:
+            scaled_h = int(height * self.centroid_input_scale)
+            scaled_w = int(width * self.centroid_input_scale)
+            scaled_image = F.interpolate(
+                scaled_image,
+                size=(scaled_h, scaled_w),
+                mode="bilinear",
+                align_corners=False,
+            )
+
+        centroid_out = self.centroid_model(scaled_image)
+        centroid_cms = self._extract_tensor(centroid_out, ["centroid", "confmap"])
+
+        centroids, centroid_vals, instance_valid = self._find_topk_peaks(
+            centroid_cms, self.max_instances
+        )
+        centroids = centroids * (
+            self.centroid_output_stride / self.centroid_input_scale
+        )
+
+        crops = self._extract_crops(image, centroids)
+        crops_flat = crops.reshape(
+            batch_size * self.max_instances,
+            channels,
+            self.crop_size[0],
+            self.crop_size[1],
+        )
+
+        if self.instance_input_scale != 1.0:
+            scaled_h = int(self.crop_size[0] * self.instance_input_scale)
+            scaled_w = int(self.crop_size[1] * self.instance_input_scale)
+            crops_flat = F.interpolate(
+                crops_flat,
+                size=(scaled_h, scaled_w),
+                mode="bilinear",
+                align_corners=False,
+            )
+
+        instance_out = self.instance_model(crops_flat)
+        instance_cms = self._extract_tensor(
+            instance_out, ["centered", "instance", "confmap"]
+        )
+
+        crop_peaks, crop_peak_vals = self._find_global_peaks(instance_cms)
+        crop_peaks = crop_peaks * (
+            self.instance_output_stride / self.instance_input_scale
+        )
+
+        crop_peaks = crop_peaks.reshape(batch_size, self.max_instances, self.n_nodes, 2)
+        peak_vals = crop_peak_vals.reshape(batch_size, self.max_instances, self.n_nodes)
+
+        crop_offset = centroids.unsqueeze(2) - image.new_tensor(
+            [self.crop_size[1] / 2.0, self.crop_size[0] / 2.0]
+        )
+        peaks = crop_peaks + crop_offset
+
+        invalid_mask = ~instance_valid
+        centroids = centroids.masked_fill(invalid_mask.unsqueeze(-1), 0.0)
+        centroid_vals = centroid_vals.masked_fill(invalid_mask, 0.0)
+        peaks = peaks.masked_fill(invalid_mask.unsqueeze(-1).unsqueeze(-1), 0.0)
+        peak_vals = peak_vals.masked_fill(invalid_mask.unsqueeze(-1), 0.0)
+
+        return {
+            "centroids": centroids,
+            "centroid_vals": centroid_vals,
+            "peaks": peaks,
+            "peak_vals": peak_vals,
+            "instance_valid": instance_valid,
+        }
+
+    def _extract_crops(
+        self,
+        image: torch.Tensor,
+        centroids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Extract crops around centroids using grid_sample."""
+        batch_size, channels, height, width = image.shape
+        crop_h, crop_w = self.crop_size
+        n_instances = centroids.shape[1]
+
+        scale_x = crop_w / width
+        scale_y = crop_h / height
+        scale = image.new_tensor([scale_x, scale_y])
+        base_grid = self.base_grid.to(device=image.device, dtype=image.dtype)
+        scaled_grid = base_grid * scale
+
+        scaled_grid = scaled_grid.unsqueeze(0).unsqueeze(0)
+        scaled_grid = scaled_grid.expand(batch_size, n_instances, -1, -1, -1)
+
+        norm_centroids = torch.zeros_like(centroids)
+        norm_centroids[..., 0] = (centroids[..., 0] / (width - 1)) * 2 - 1
+        norm_centroids[..., 1] = (centroids[..., 1] / (height - 1)) * 2 - 1
+        offset = norm_centroids.unsqueeze(2).unsqueeze(2)
+
+        sample_grid = scaled_grid + offset
+
+        image_expanded = image.unsqueeze(1).expand(-1, n_instances, -1, -1, -1)
+        image_flat = image_expanded.reshape(
+            batch_size * n_instances, channels, height, width
+        )
+        grid_flat = sample_grid.reshape(batch_size * n_instances, crop_h, crop_w, 2)
+
+        crops_flat = F.grid_sample(
+            image_flat,
+            grid_flat,
+            mode="bilinear",
+            padding_mode="zeros",
+            align_corners=True,
+        )
+
+        crops = crops_flat.reshape(batch_size, n_instances, channels, crop_h, crop_w)
+        return crops