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

sleap_nn/export/wrappers/topdown_multiclass.py

@@ -0,0 +1,304 @@
+"""ONNX wrapper for top-down 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 TopDownMultiClassONNXWrapper(BaseExportWrapper):
+    """ONNX-exportable wrapper for top-down multiclass (supervised ID) models.
+
+    This wrapper handles models that output both confidence maps for keypoint
+    detection and class logits for identity classification. It runs on instance
+    crops (centered around detected centroids).
+
+    Expects input images as uint8 tensors in [0, 255].
+
+    Attributes:
+        model: The underlying PyTorch model (centered instance + class vectors heads).
+        output_stride: Output stride of the confmap head.
+        input_scale: Scale factor applied to input images before inference.
+        n_classes: Number of identity classes.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        output_stride: int = 2,
+        input_scale: float = 1.0,
+        n_classes: int = 2,
+    ):
+        """Initialize the wrapper.
+
+        Args:
+            model: The underlying PyTorch model.
+            output_stride: Output stride of the confidence maps.
+            input_scale: Scale factor for input images.
+            n_classes: Number of identity classes (e.g., 2 for male/female).
+        """
+        super().__init__(model)
+        self.output_stride = output_stride
+        self.input_scale = input_scale
+        self.n_classes = n_classes
+
+    def forward(self, image: torch.Tensor) -> Dict[str, torch.Tensor]:
+        """Run top-down multiclass inference on crops.
+
+        Args:
+            image: Input image tensor of shape (batch, channels, height, width).
+                   Expected to be uint8 in [0, 255].
+
+        Returns:
+            Dictionary with keys:
+                - "peaks": Predicted peak coordinates (batch, n_nodes, 2) in (x, y).
+                - "peak_vals": Peak confidence values (batch, n_nodes).
+                - "class_logits": Raw class logits (batch, n_classes).
+
+            The class assignment is done on CPU using Hungarian matching
+            via `get_class_inds_from_vectors()`.
+        """
+        # 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
+            )
+
+        # Forward pass
+        out = self.model(image)
+
+        # Extract outputs
+        confmaps = self._extract_tensor(out, ["centered", "instance", "confmap"])
+        class_logits = self._extract_tensor(out, ["class", "vector"])
+
+        # Find global peaks (one per node)
+        peaks, peak_vals = self._find_global_peaks(confmaps)
+
+        # Scale peaks back to input coordinates
+        peaks = peaks * (self.output_stride / self.input_scale)
+
+        return {
+            "peaks": peaks,
+            "peak_vals": peak_vals,
+            "class_logits": class_logits,
+        }
+
+
+class TopDownMultiClassCombinedONNXWrapper(BaseExportWrapper):
+    """ONNX-exportable wrapper for combined centroid + multiclass instance models.
+
+    This wrapper combines a centroid detection model with a centered instance
+    multiclass model. It performs:
+    1. Centroid detection on full images
+    2. Cropping around each centroid using vectorized grid_sample
+    3. Instance keypoint detection + identity classification on each crop
+
+    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 = (192, 192),
+        centroid_output_stride: int = 4,
+        instance_output_stride: int = 2,
+        centroid_input_scale: float = 1.0,
+        instance_input_scale: float = 1.0,
+        n_nodes: int = 13,
+        n_classes: int = 2,
+    ):
+        """Initialize the combined wrapper.
+
+        Args:
+            centroid_model: Model for centroid detection.
+            instance_model: Model for instance keypoints + class prediction.
+            max_instances: Maximum number of instances to detect.
+            crop_size: Size of crops around centroids (height, width).
+            centroid_output_stride: Output stride of centroid model.
+            instance_output_stride: Output stride of instance model.
+            centroid_input_scale: Input scale for centroid model.
+            instance_input_scale: Input scale for instance model.
+            n_nodes: Number of keypoint nodes per instance.
+            n_classes: Number of identity classes.
+        """
+        super().__init__(centroid_model)  # Primary model is centroid
+        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
+        self.n_classes = n_classes
+
+        # Pre-compute base grid for crop extraction (same as TopDownONNXWrapper)
+        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 combined top-down multiclass inference.
+
+        Args:
+            image: Input image tensor of shape (batch, channels, height, width).
+                   Expected to be uint8 in [0, 255].
+
+        Returns:
+            Dictionary with keys:
+                - "centroids": Detected centroids (batch, max_instances, 2).
+                - "centroid_vals": Centroid confidence values (batch, max_instances).
+                - "peaks": Instance peaks (batch, max_instances, n_nodes, 2).
+                - "peak_vals": Peak values (batch, max_instances, n_nodes).
+                - "class_logits": Class logits per instance (batch, max_instances, n_classes).
+                - "instance_valid": Validity mask (batch, max_instances).
+        """
+        # Normalize input
+        image = self._normalize_uint8(image)
+        batch_size, channels, height, width = image.shape
+
+        # Apply centroid input scaling
+        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 detection
+        centroid_out = self.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
+        )
+
+        # Extract crops using vectorized grid_sample (same as TopDownONNXWrapper)
+        crops = self._extract_crops(image, centroids)
+        crops_flat = crops.reshape(
+            batch_size * self.max_instances,
+            channels,
+            self.crop_size[0],
+            self.crop_size[1],
+        )
+
+        # Apply instance input scaling if needed
+        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 model forward (batch all crops)
+        instance_out = self.instance_model(crops_flat)
+        instance_cms = self._extract_tensor(
+            instance_out, ["centered", "instance", "confmap"]
+        )
+        instance_class = self._extract_tensor(instance_out, ["class", "vector"])
+
+        # Find peaks in all crops
+        crop_peaks, crop_peak_vals = self._find_global_peaks(instance_cms)
+        crop_peaks = crop_peaks * (
+            self.instance_output_stride / self.instance_input_scale
+        )
+
+        # Reshape to batch x instances x nodes x 2
+        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)
+
+        # Reshape class logits
+        class_logits = instance_class.reshape(
+            batch_size, self.max_instances, self.n_classes
+        )
+
+        # Transform peaks from crop coordinates to full image coordinates
+        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
+
+        # Zero out invalid instances
+        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)
+        class_logits = class_logits.masked_fill(invalid_mask.unsqueeze(-1), 0.0)
+
+        return {
+            "centroids": centroids,
+            "centroid_vals": centroid_vals,
+            "peaks": peaks,
+            "peak_vals": peak_vals,
+            "class_logits": class_logits,
+            "instance_valid": instance_valid,
+        }
+
+    def _extract_crops(
+        self,
+        image: torch.Tensor,
+        centroids: torch.Tensor,
+    ) -> torch.Tensor:
+        """Extract crops around centroids using grid_sample.
+
+        This is the same vectorized implementation as TopDownONNXWrapper.
+        """
+        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

sleap_nn/inference/__init__.py

@@ -1 +1,7 @@
 """Inference-related modules."""
+
+from sleap_nn.inference.provenance import (
+    build_inference_provenance,
+    build_tracking_only_provenance,
+    merge_provenance,
+)

sleap_nn/inference/bottomup.py

@@ -1,5 +1,6 @@
 """Inference modules for BottomUp models."""
 
+import logging
 from typing import Dict, Optional
 import torch
 import lightning as L
@@ -7,6 +8,8 @@ from sleap_nn.inference.peak_finding import find_local_peaks
 from sleap_nn.inference.paf_grouping import PAFScorer
 from sleap_nn.inference.identity import classify_peaks_from_maps
 
+logger = logging.getLogger(__name__)
+
 
 class BottomUpInferenceModel(L.LightningModule):
     """BottomUp Inference model.
@@ -63,8 +66,28 @@ class BottomUpInferenceModel(L.LightningModule):
         return_pafs: Optional[bool] = False,
         return_paf_graph: Optional[bool] = False,
         input_scale: float = 1.0,
+        max_peaks_per_node: Optional[int] = None,
     ):
-        """Initialise the model attributes."""
+        """Initialise the model attributes.
+
+        Args:
+            torch_model: A `nn.Module` that accepts images and predicts confidence maps.
+            paf_scorer: A `PAFScorer` instance for grouping instances.
+            cms_output_stride: Output stride of confidence maps relative to images.
+            pafs_output_stride: Output stride of PAFs relative to images.
+            peak_threshold: Minimum confidence map value for valid peaks.
+            refinement: Peak refinement method: None, "integral", or "local".
+            integral_patch_size: Size of patches for integral refinement.
+            return_confmaps: If True, return confidence maps in output.
+            return_pafs: If True, return PAFs in output.
+            return_paf_graph: If True, return intermediate PAF graph in output.
+            input_scale: Scale factor applied to input images.
+            max_peaks_per_node: Maximum number of peaks allowed per node before
+                skipping PAF scoring. If any node has more peaks than this limit,
+                empty predictions are returned. This prevents combinatorial explosion
+                during early training when confidence maps are noisy. Set to None to
+                disable this check (default). Recommended value: 100.
+        """
         super().__init__()
         self.torch_model = torch_model
         self.paf_scorer = paf_scorer
@@ -77,6 +100,7 @@ class BottomUpInferenceModel(L.LightningModule):
         self.return_pafs = return_pafs
         self.return_paf_graph = return_paf_graph
         self.input_scale = input_scale
+        self.max_peaks_per_node = max_peaks_per_node
 
     def _generate_cms_peaks(self, cms):
         # TODO: append nans to batch them -> tensor (vectorize the initial paf grouping steps)
@@ -124,26 +148,68 @@ class BottomUpInferenceModel(L.LightningModule):
         )  # (batch, h, w, 2*edges)
         cms_peaks, cms_peak_vals, cms_peak_channel_inds = self._generate_cms_peaks(cms)
 
-        (
-            predicted_instances,
-            predicted_peak_scores,
-            predicted_instance_scores,
-            edge_inds,
-            edge_peak_inds,
-            line_scores,
-        ) = self.paf_scorer.predict(
-            pafs=pafs,
-            peaks=cms_peaks,
-            peak_vals=cms_peak_vals,
-            peak_channel_inds=cms_peak_channel_inds,
-        )
-
-        predicted_instances = [p / self.input_scale for p in predicted_instances]
-        predicted_instances_adjusted = []
-        for idx, p in enumerate(predicted_instances):
-            predicted_instances_adjusted.append(
-                p / inputs["eff_scale"][idx].to(p.device)
+        # Check if too many peaks per node (prevents combinatorial explosion)
+        skip_paf_scoring = False
+        if self.max_peaks_per_node is not None:
+            n_nodes = cms.shape[1]
+            for b in range(self.batch_size):
+                for node_idx in range(n_nodes):
+                    n_peaks = int((cms_peak_channel_inds[b] == node_idx).sum().item())
+                    if n_peaks > self.max_peaks_per_node:
+                        logger.warning(
+                            f"Skipping PAF scoring: node {node_idx} has {n_peaks} peaks "
+                            f"(max_peaks_per_node={self.max_peaks_per_node}). "
+                            f"Model may need more training."
+                        )
+                        skip_paf_scoring = True
+                        break
+                if skip_paf_scoring:
+                    break
+
+        if skip_paf_scoring:
+            # Return empty predictions for each sample
+            device = cms.device
+            n_nodes = cms.shape[1]
+            predicted_instances_adjusted = []
+            predicted_peak_scores = []
+            predicted_instance_scores = []
+            for _ in range(self.batch_size):
+                predicted_instances_adjusted.append(
+                    torch.full((0, n_nodes, 2), float("nan"), device=device)
+                )
+                predicted_peak_scores.append(
+                    torch.full((0, n_nodes), float("nan"), device=device)
+                )
+                predicted_instance_scores.append(torch.tensor([], device=device))
+            edge_inds = [
+                torch.tensor([], dtype=torch.int32, device=device)
+            ] * self.batch_size
+            edge_peak_inds = [
+                torch.tensor([], dtype=torch.int32, device=device).reshape(0, 2)
+            ] * self.batch_size
+            line_scores = [torch.tensor([], device=device)] * self.batch_size
+        else:
+            (
+                predicted_instances,
+                predicted_peak_scores,
+                predicted_instance_scores,
+                edge_inds,
+                edge_peak_inds,
+                line_scores,
+            ) = self.paf_scorer.predict(
+                pafs=pafs,
+                peaks=cms_peaks,
+                peak_vals=cms_peak_vals,
+                peak_channel_inds=cms_peak_channel_inds,
             )
+
+            predicted_instances = [p / self.input_scale for p in predicted_instances]
+            predicted_instances_adjusted = []
+            for idx, p in enumerate(predicted_instances):
+                predicted_instances_adjusted.append(
+                    p / inputs["eff_scale"][idx].to(p.device)
+                )
+
         out = {
             "pred_instance_peaks": predicted_instances_adjusted,
             "pred_peak_values": predicted_peak_scores,

sleap_nn/inference/peak_finding.py

@@ -2,18 +2,60 @@
 
 from typing import Optional, Tuple
 
-import kornia as K
-import numpy as np
 import torch
-from kornia.geometry.transform import crop_and_resize
+import torch.nn.functional as F
 
 from sleap_nn.data.instance_cropping import make_centered_bboxes
 
 
+def morphological_dilation(image: torch.Tensor, kernel: torch.Tensor) -> torch.Tensor:
+    """Apply morphological dilation using max pooling.
+
+    This is a pure PyTorch replacement for kornia.morphology.dilation.
+    For non-maximum suppression, it computes the maximum of 8 neighbors
+    (excluding the center pixel).
+
+    Args:
+        image: Input tensor of shape (B, 1, H, W).
+        kernel: Dilation kernel (3x3 expected for NMS).
+
+    Returns:
+        Dilated tensor of same shape as input.
+    """
+    # Pad the image to handle border pixels
+    padded = F.pad(image, (1, 1, 1, 1), mode="constant", value=float("-inf"))
+
+    # Extract 3x3 patches using unfold
+    # Shape: (B, 1, H, W, 3, 3)
+    patches = padded.unfold(2, 3, 1).unfold(3, 3, 1)
+
+    # Reshape to (B, 1, H, W, 9)
+    b, c, h, w, kh, kw = patches.shape
+    patches = patches.reshape(b, c, h, w, -1)
+
+    # Apply kernel mask (kernel has 0 at center, 1 elsewhere for NMS)
+    # Reshape kernel to (1, 1, 1, 1, 9)
+    kernel_flat = kernel.reshape(-1).to(patches.device)
+    kernel_mask = kernel_flat > 0
+
+    # Set non-kernel positions to -inf so they don't affect max
+    patches_masked = patches.clone()
+    patches_masked[..., ~kernel_mask] = float("-inf")
+
+    # Take max over the kernel neighborhood
+    max_vals = patches_masked.max(dim=-1)[0]
+
+    return max_vals
+
+
 def crop_bboxes(
     images: torch.Tensor, bboxes: torch.Tensor, sample_inds: torch.Tensor
 ) -> torch.Tensor:
-    """Crop bounding boxes from a batch of images.
+    """Crop bounding boxes from a batch of images using fast tensor indexing.
+
+    This uses tensor unfold operations to extract patches, which is significantly
+    faster than kornia's crop_and_resize (17-51x speedup) as it avoids perspective
+    transform computations.
 
     Args:
         images: Tensor of shape (samples, channels, height, width) of a batch of images.
@@ -27,7 +69,7 @@ def crop_bboxes(
             box should be cropped from.
 
     Returns:
-        A tensor of shape (n_bboxes, crop_height, crop_width, channels) of the same
+        A tensor of shape (n_bboxes, channels, crop_height, crop_width) of the same
         dtype as the input image. The crop size is inferred from the bounding box
         coordinates.
 
@@ -42,20 +84,55 @@ def crop_bboxes(
 
     See also: `make_centered_bboxes`
     """
+    n_crops = bboxes.shape[0]
+    if n_crops == 0:
+        # Return empty tensor; use default crop size since we can't infer from bboxes
+        return torch.empty(
+            0, images.shape[1], 0, 0, device=images.device, dtype=images.dtype
+        )
+
     # Compute bounding box size to use for crops.
-    height = abs(bboxes[0, 3, 1] - bboxes[0, 0, 1])
-    width = abs(bboxes[0, 1, 0] - bboxes[0, 0, 0])
-    box_size = tuple(torch.round(torch.Tensor((height + 1, width + 1))).to(torch.int32))
-
-    # Crop.
-    crops = crop_and_resize(
-        images[sample_inds],  # (n_boxes, channels, height, width)
-        boxes=bboxes,
-        size=box_size,
+    height = int(abs(bboxes[0, 3, 1] - bboxes[0, 0, 1]).item()) + 1
+    width = int(abs(bboxes[0, 1, 0] - bboxes[0, 0, 0]).item()) + 1
+
+    # Store original dtype for conversion back after cropping.
+    original_dtype = images.dtype
+    device = images.device
+    n_samples, channels, img_h, img_w = images.shape
+    half_h, half_w = height // 2, width // 2
+
+    # Pad images for edge handling.
+    images_padded = F.pad(
+        images.float(), (half_w, half_w, half_h, half_h), mode="constant", value=0
     )
 
+    # Extract all possible patches using unfold (creates a view, no copy).
+    # Shape after unfold: (n_samples, channels, img_h, img_w, height, width)
+    patches = images_padded.unfold(2, height, 1).unfold(3, width, 1)
+
+    # Get crop centers from bboxes.
+    # The bbox top-left is at index 0, with (x, y) coordinates.
+    # We need the center of the crop (peak location), which is top-left + half_size.
+    # Ensure bboxes are on the same device as images for index computation.
+    bboxes_on_device = bboxes.to(device)
+    crop_x = (bboxes_on_device[:, 0, 0] + half_w).to(torch.long)
+    crop_y = (bboxes_on_device[:, 0, 1] + half_h).to(torch.long)
+
+    # Clamp indices to valid bounds to handle edge cases where centroids
+    # might be at or beyond image boundaries.
+    crop_x = torch.clamp(crop_x, 0, patches.shape[3] - 1)
+    crop_y = torch.clamp(crop_y, 0, patches.shape[2] - 1)
+
+    # Select crops using advanced indexing.
+    # Convert sample_inds to tensor if it's a list.
+    if not isinstance(sample_inds, torch.Tensor):
+        sample_inds = torch.tensor(sample_inds, device=device)
+    sample_inds_long = sample_inds.to(device=device, dtype=torch.long)
+    crops = patches[sample_inds_long, :, crop_y, crop_x]
+    # Shape: (n_crops, channels, height, width)
+
     # Cast back to original dtype and return.
-    crops = crops.to(images.dtype)
+    crops = crops.to(original_dtype)
     return crops
 
 
@@ -236,7 +313,7 @@ def find_local_peaks_rough(
     flat_img = cms.reshape(-1, 1, height, width)
 
     # Perform dilation filtering to find local maxima per channel and reshape back.
-    max_img = K.morphology.dilation(flat_img, kernel.to(flat_img.device))
+    max_img = morphological_dilation(flat_img, kernel.to(flat_img.device))
     max_img = max_img.reshape(-1, channels, height, width)
 
     # Filter for maxima and threshold.