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

sleap_nn/inference/postprocessing.py

@@ -0,0 +1,284 @@
+"""Inference-level postprocessing filters for pose predictions.
+
+This module provides filters that run after model inference but before tracking.
+These filters are independent of tracking configuration and can be used standalone.
+"""
+
+from typing import List, Literal
+
+import numpy as np
+import sleap_io as sio
+
+
+def filter_overlapping_instances(
+    labels: sio.Labels,
+    threshold: float = 0.8,
+    method: Literal["iou", "oks"] = "iou",
+) -> sio.Labels:
+    """Filter overlapping instances using greedy non-maximum suppression.
+
+    Removes duplicate/overlapping instances by applying greedy NMS based on
+    either bounding box IOU or Object Keypoint Similarity (OKS). When two
+    instances overlap above the threshold, the lower-scoring one is removed.
+
+    This filter runs independently of tracking and can be used to clean up
+    model outputs before saving or further processing.
+
+    Args:
+        labels: Labels object with predicted instances to filter.
+        threshold: Similarity threshold for considering instances as overlapping.
+            Instances with similarity > threshold are candidates for removal.
+            Lower values are more aggressive (remove more).
+            Typical values: 0.3 (aggressive) to 0.8 (permissive).
+        method: Similarity metric to use for comparing instances.
+            "iou": Bounding box intersection-over-union.
+            "oks": Object Keypoint Similarity (pose-based).
+
+    Returns:
+        The input Labels object with overlapping instances removed.
+        Modification is done in place, but the object is also returned
+        for convenience.
+
+    Example:
+        >>> # Filter instances with >80% bounding box overlap
+        >>> labels = filter_overlapping_instances(labels, threshold=0.8, method="iou")
+        >>> # Filter using OKS similarity
+        >>> labels = filter_overlapping_instances(labels, threshold=0.5, method="oks")
+
+    Note:
+        - Only affects frames with 2+ predicted instances
+        - Uses instance.score for ranking; higher scores are preferred
+        - For IOU: bounding boxes computed from non-NaN keypoints
+        - For OKS: uses standard COCO OKS formula with bbox-derived scale
+    """
+    for lf in labels.labeled_frames:
+        if len(lf.instances) <= 1:
+            continue
+
+        # Separate predicted instances (have scores) from other instances
+        predicted = []
+        other = []
+        for inst in lf.instances:
+            if isinstance(inst, sio.PredictedInstance):
+                predicted.append(inst)
+            else:
+                other.append(inst)
+
+        # Only filter predicted instances
+        if len(predicted) <= 1:
+            continue
+
+        # Get scores
+        scores = np.array([_instance_score(inst) for inst in predicted])
+
+        # Apply greedy NMS with selected method
+        if method == "iou":
+            bboxes = np.array([_instance_bbox(inst) for inst in predicted])
+            keep_indices = _nms_greedy_iou(bboxes, scores, threshold)
+        elif method == "oks":
+            points = [inst.numpy() for inst in predicted]
+            keep_indices = _nms_greedy_oks(points, scores, threshold)
+        else:
+            raise ValueError(f"Unknown method: {method}. Use 'iou' or 'oks'.")
+
+        # Reconstruct instance list: kept predicted + other instances
+        kept_predicted = [predicted[i] for i in keep_indices]
+        lf.instances = kept_predicted + other
+
+    return labels
+
+
+def _instance_bbox(instance: sio.PredictedInstance) -> np.ndarray:
+    """Compute axis-aligned bounding box from instance keypoints.
+
+    Args:
+        instance: Instance with keypoints.
+
+    Returns:
+        Bounding box as [xmin, ymin, xmax, ymax].
+        Returns [0, 0, 0, 0] if no valid keypoints.
+    """
+    pts = instance.numpy()  # (n_nodes, 2)
+    valid = ~np.isnan(pts).any(axis=1)
+
+    if not valid.any():
+        return np.array([0.0, 0.0, 0.0, 0.0])
+
+    pts = pts[valid]
+    return np.array(
+        [pts[:, 0].min(), pts[:, 1].min(), pts[:, 0].max(), pts[:, 1].max()]
+    )
+
+
+def _instance_score(instance: sio.PredictedInstance) -> float:
+    """Get instance confidence score.
+
+    Args:
+        instance: Predicted instance.
+
+    Returns:
+        Instance score, or 1.0 if not available.
+    """
+    return getattr(instance, "score", 1.0)
+
+
+def _nms_greedy_iou(
+    bboxes: np.ndarray,
+    scores: np.ndarray,
+    threshold: float,
+) -> List[int]:
+    """Apply greedy NMS using bounding box IOU.
+
+    Args:
+        bboxes: Bounding boxes of shape (N, 4) as [xmin, ymin, xmax, ymax].
+        scores: Confidence scores of shape (N,).
+        threshold: IOU threshold for suppression.
+
+    Returns:
+        List of indices to keep, in order of decreasing score.
+    """
+    if len(bboxes) == 0:
+        return []
+
+    # Sort by score descending
+    order = scores.argsort()[::-1].tolist()
+
+    keep = []
+    while order:
+        # Take highest scoring remaining instance
+        i = order.pop(0)
+        keep.append(i)
+
+        if not order:
+            break
+
+        # Compute IOU with all remaining instances
+        remaining_indices = np.array(order)
+        similarities = _compute_iou_one_to_many(bboxes[i], bboxes[remaining_indices])
+
+        # Keep only instances with similarity <= threshold
+        mask = similarities <= threshold
+        order = [order[j] for j in range(len(order)) if mask[j]]
+
+    return keep
+
+
+def _nms_greedy_oks(
+    points_list: List[np.ndarray],
+    scores: np.ndarray,
+    threshold: float,
+) -> List[int]:
+    """Apply greedy NMS using Object Keypoint Similarity (OKS).
+
+    Args:
+        points_list: List of keypoint arrays, each of shape (n_nodes, 2).
+        scores: Confidence scores of shape (N,).
+        threshold: OKS threshold for suppression.
+
+    Returns:
+        List of indices to keep, in order of decreasing score.
+    """
+    if len(points_list) == 0:
+        return []
+
+    # Sort by score descending
+    order = scores.argsort()[::-1].tolist()
+
+    keep = []
+    while order:
+        # Take highest scoring remaining instance
+        i = order.pop(0)
+        keep.append(i)
+
+        if not order:
+            break
+
+        # Compute OKS with all remaining instances
+        similarities = np.array(
+            [_compute_oks(points_list[i], points_list[j]) for j in order]
+        )
+
+        # Keep only instances with similarity <= threshold
+        mask = similarities <= threshold
+        order = [order[j] for j in range(len(order)) if mask[j]]
+
+    return keep
+
+
+def _compute_iou_one_to_many(box: np.ndarray, boxes: np.ndarray) -> np.ndarray:
+    """Compute IOU between one box and multiple boxes.
+
+    Args:
+        box: Single box of shape (4,) as [xmin, ymin, xmax, ymax].
+        boxes: Multiple boxes of shape (N, 4).
+
+    Returns:
+        IOU values of shape (N,).
+    """
+    # Intersection coordinates
+    inter_xmin = np.maximum(box[0], boxes[:, 0])
+    inter_ymin = np.maximum(box[1], boxes[:, 1])
+    inter_xmax = np.minimum(box[2], boxes[:, 2])
+    inter_ymax = np.minimum(box[3], boxes[:, 3])
+
+    # Intersection area (0 if no overlap)
+    inter_w = np.maximum(0.0, inter_xmax - inter_xmin)
+    inter_h = np.maximum(0.0, inter_ymax - inter_ymin)
+    inter_area = inter_w * inter_h
+
+    # Individual areas
+    area_a = (box[2] - box[0]) * (box[3] - box[1])
+    area_b = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+
+    # Union area
+    union_area = area_a + area_b - inter_area
+
+    # IOU (avoid division by zero)
+    return np.where(union_area > 0, inter_area / union_area, 0.0)
+
+
+def _compute_oks(
+    points_a: np.ndarray,
+    points_b: np.ndarray,
+    kappa: float = 0.1,
+) -> float:
+    """Compute Object Keypoint Similarity (OKS) between two instances.
+
+    Uses a simplified OKS formula where all keypoints have equal weight
+    and scale is derived from the bounding box of the reference instance.
+
+    Args:
+        points_a: Keypoints of first instance, shape (n_nodes, 2).
+        points_b: Keypoints of second instance, shape (n_nodes, 2).
+        kappa: Per-keypoint constant controlling falloff. Default 0.1.
+
+    Returns:
+        OKS value in [0, 1]. Higher means more similar.
+    """
+    # Find valid keypoints (present in both instances)
+    valid_a = ~np.isnan(points_a).any(axis=1)
+    valid_b = ~np.isnan(points_b).any(axis=1)
+    valid = valid_a & valid_b
+
+    if not valid.any():
+        return 0.0
+
+    # Compute scale from bounding box area of instance A
+    pts_a_valid = points_a[valid_a]
+    if len(pts_a_valid) < 2:
+        return 0.0
+
+    bbox_w = pts_a_valid[:, 0].max() - pts_a_valid[:, 0].min()
+    bbox_h = pts_a_valid[:, 1].max() - pts_a_valid[:, 1].min()
+    scale_sq = bbox_w * bbox_h
+
+    if scale_sq <= 0:
+        return 0.0
+
+    # Compute squared distances for valid keypoints
+    d_sq = np.sum((points_a[valid] - points_b[valid]) ** 2, axis=1)
+
+    # OKS formula: mean of exp(-d^2 / (2 * s^2 * k^2))
+    oks_per_kpt = np.exp(-d_sq / (2 * scale_sq * kappa**2))
+
+    return float(np.mean(oks_per_kpt))

sleap_nn/predict.py

@@ -74,6 +74,9 @@ def run_inference(
     frames: Optional[list] = None,
     crop_size: Optional[int] = None,
     peak_threshold: Union[float, List[float]] = 0.2,
+    filter_overlapping: bool = False,
+    filter_overlapping_method: str = "iou",
+    filter_overlapping_threshold: float = 0.8,
     integral_refinement: Optional[str] = "integral",
     integral_patch_size: int = 5,
     return_confmaps: bool = False,
@@ -160,6 +163,15 @@ def run_inference(
                 centroid and centered-instance model, where the first element corresponds
                 to centroid model peak finding threshold and the second element is for
                 centered-instance model peak finding.
+        filter_overlapping: (bool) If True, removes overlapping instances after
+                inference using greedy NMS. Applied independently of tracking.
+                Default: False.
+        filter_overlapping_method: (str) Similarity metric for filtering overlapping
+                instances. One of "iou" (bounding box) or "oks" (keypoint similarity).
+                Default: "iou".
+        filter_overlapping_threshold: (float) Similarity threshold for filtering.
+                Instances with similarity > threshold are removed (keeping higher-scoring).
+                Typical values: 0.3 (aggressive) to 0.8 (permissive). Default: 0.8.
         integral_refinement: (str) If `None`, returns the grid-aligned peaks with no refinement.
                 If `"integral"`, peaks will be refined with integral regression.
                 Default: `"integral"`.
@@ -553,6 +565,20 @@ def run_inference(
             make_labels=make_labels,
         )
 
+        # Filter overlapping instances (independent of tracking)
+        if filter_overlapping and make_labels:
+            from sleap_nn.inference.postprocessing import filter_overlapping_instances
+
+            output = filter_overlapping_instances(
+                output,
+                threshold=filter_overlapping_threshold,
+                method=filter_overlapping_method,
+            )
+            logger.info(
+                f"Filtered overlapping instances with {filter_overlapping_method.upper()} "
+                f"threshold: {filter_overlapping_threshold}"
+            )
+
         if tracking:
             lfs = [x for x in output]
             if tracking_clean_instance_count > 0:
@@ -607,6 +633,9 @@ def run_inference(
         # Build inference parameters for provenance
         inference_params = {
             "peak_threshold": peak_threshold,
+            "filter_overlapping": filter_overlapping,
+            "filter_overlapping_method": filter_overlapping_method,
+            "filter_overlapping_threshold": filter_overlapping_threshold,
             "integral_refinement": integral_refinement,
             "integral_patch_size": integral_patch_size,
             "batch_size": batch_size,

sleap_nn/train.py

@@ -118,6 +118,70 @@ def run_training(
                 logger.info(f"p90 dist: {metrics['distance_metrics']['p90']}")
                 logger.info(f"p50 dist: {metrics['distance_metrics']['p50']}")
 
+                # Log test metrics to wandb summary
+                if (
+                    d_name.startswith("test")
+                    and trainer.config.trainer_config.use_wandb
+                ):
+                    import wandb
+
+                    if wandb.run is not None:
+                        summary_metrics = {
+                            f"eval/{d_name}/mOKS": metrics["mOKS"]["mOKS"],
+                            f"eval/{d_name}/oks_voc_mAP": metrics["voc_metrics"][
+                                "oks_voc.mAP"
+                            ],
+                            f"eval/{d_name}/oks_voc_mAR": metrics["voc_metrics"][
+                                "oks_voc.mAR"
+                            ],
+                            f"eval/{d_name}/mPCK": metrics["pck_metrics"]["mPCK"],
+                            f"eval/{d_name}/PCK_5": metrics["pck_metrics"]["PCK@5"],
+                            f"eval/{d_name}/PCK_10": metrics["pck_metrics"]["PCK@10"],
+                            f"eval/{d_name}/distance_avg": metrics["distance_metrics"][
+                                "avg"
+                            ],
+                            f"eval/{d_name}/distance_p50": metrics["distance_metrics"][
+                                "p50"
+                            ],
+                            f"eval/{d_name}/distance_p95": metrics["distance_metrics"][
+                                "p95"
+                            ],
+                            f"eval/{d_name}/distance_p99": metrics["distance_metrics"][
+                                "p99"
+                            ],
+                            f"eval/{d_name}/visibility_precision": metrics[
+                                "visibility_metrics"
+                            ]["precision"],
+                            f"eval/{d_name}/visibility_recall": metrics[
+                                "visibility_metrics"
+                            ]["recall"],
+                        }
+                        for key, value in summary_metrics.items():
+                            wandb.run.summary[key] = value
+
+            # Finish wandb run and cleanup after all evaluation is complete
+            if trainer.config.trainer_config.use_wandb:
+                import wandb
+                import shutil
+
+                if wandb.run is not None:
+                    wandb.finish()
+
+                # Delete local wandb logs if configured
+                wandb_config = trainer.config.trainer_config.wandb
+                should_delete_wandb_logs = wandb_config.delete_local_logs is True or (
+                    wandb_config.delete_local_logs is None
+                    and wandb_config.wandb_mode != "offline"
+                )
+                if should_delete_wandb_logs:
+                    wandb_dir = run_path / "wandb"
+                    if wandb_dir.exists():
+                        logger.info(
+                            f"Deleting local wandb logs at {wandb_dir}... "
+                            "(set trainer_config.wandb.delete_local_logs=false to disable)"
+                        )
+                        shutil.rmtree(wandb_dir, ignore_errors=True)
+
 
 def train(
     train_labels_path: Optional[List[str]] = None,