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

sleap_nn/export/wrappers/bottomup_multiclass.py

@@ -1,195 +0,0 @@
-"""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

@@ -1,56 +0,0 @@
-"""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

@@ -1,58 +0,0 @@
-"""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

@@ -1,83 +0,0 @@
-"""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

@@ -1,180 +0,0 @@
-"""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