DeepLabCut · MMathisLab · May 23, 2025 · May 15, 2025 · May 15, 2025 · May 15, 2025
diff --git a/deeplabcut/pose_estimation_pytorch/data/preprocessor.py b/deeplabcut/pose_estimation_pytorch/data/preprocessor.py
@@ -13,16 +13,13 @@
 
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Any, TypeVar
+from typing import Any, TypeVar, Callable
 
 import albumentations as A
 import numpy as np
 import torch
 
-from deeplabcut.pose_estimation_pytorch.data.image import (
-    load_image,
-    top_down_crop
-)
+from deeplabcut.pose_estimation_pytorch.data.image import load_image, top_down_crop
 from deeplabcut.pose_estimation_pytorch.data.utils import bbox_from_keypoints
 
 
@@ -151,7 +148,9 @@ def build_conditional_top_down_preprocessor(
     return ComposePreprocessor(
         components=[
             LoadImage(color_mode),
+            FilterLowConfidencePoses(),
             ComputeBoundingBoxesFromCondKeypoints(bbox_margin=bbox_margin),
+            FilterInvalidBoundingBoxes(),
             TopDownCrop(
                 output_size=top_down_crop_size,
                 margin=top_down_crop_margin,
@@ -351,12 +350,71 @@ def __call__(self, image: Image, context: Context) -> tuple[np.ndarray, Context]
 
 
 class ToBatch(Preprocessor):
-    """TODO"""
+    """Adds a batch dimension to the image tensor.
+
+    This preprocessor is used to convert a single image tensor into a batched format
+    by unsqueezing along the 0th dimension. This is typically required before passing
+    the image to models that expect batched input (i.e., shape `[B, C, H, W]`).
+    """
 
     def __call__(self, image: Image, context: Context) -> tuple[np.ndarray, Context]:
         return image.unsqueeze(0), context
 
 
+class FilterLowConfidencePoses(Preprocessor):
+    """
+    Filters out poses with low confidence scores.
+    By default, the confidence associated to the pose is the max confidence value.
+    """
+
+    def __init__(
+        self,
+        confidence_threshold: float = 0.05,
+        aggregate_func: Callable[[np.ndarray], float] = lambda arr: np.max(arr, axis=1),
+    ) -> None:
+        self.confidence_threshold = confidence_threshold
+        self.aggregate_func = aggregate_func
+
+    def __call__(
+        self, image: np.ndarray, context: Context
+    ) -> tuple[np.ndarray, Context]:
+        if "cond_kpts" not in context:
+            raise ValueError(f"Must include cond_kpts, found {context}")
+
+        keypoints = context["cond_kpts"]
+        mask = self.aggregate_func(keypoints[:, :, 2]) >= self.confidence_threshold
+        context["cond_kpts"] = keypoints[mask]
+
+        return image, context
+
+
+class FilterInvalidBoundingBoxes(Preprocessor):
+    """Filters out poses and bounding boxes that are invalid (e.g., area too small)."""
+
+    def __init__(self, min_area: int = 1) -> None:
+        self.min_area = min_area
+
+    def __call__(
+        self, image: np.ndarray, context: Context
+    ) -> tuple[np.ndarray, Context]:
+        bboxes = context.get("bboxes", [])
+        keypoints = context.get("cond_kpts", [])
+
+        valid_bboxes = []
+        valid_indices = []
+
+        for i, bbox in enumerate(bboxes):
+            _, _, w, h = bbox
+            if w * h >= self.min_area:
+                valid_bboxes.append(bbox)
+                valid_indices.append(i)
+
+        context["bboxes"] = valid_bboxes
+        context["cond_kpts"] = keypoints[valid_indices]
+
+        return image, context
+
+
 class TopDownCrop(Preprocessor):
     """Crops bounding boxes out of images for top-down pose estimation
 

diff --git a/deeplabcut/pose_estimation_pytorch/data/utils.py b/deeplabcut/pose_estimation_pytorch/data/utils.py
@@ -10,6 +10,7 @@
 #
 from __future__ import annotations
 
+import warnings
 from collections import defaultdict
 from functools import reduce, lru_cache
 from pathlib import Path
@@ -56,8 +57,10 @@ def bbox_from_keypoints(
         keypoints = np.expand_dims(keypoints, axis=0)
 
     bboxes = np.full((keypoints.shape[0], 4), np.nan)
-    bboxes[:, :2] = np.nanmin(keypoints[..., :2], axis=1) - margin  # X1, Y1
-    bboxes[:, 2:4] = np.nanmax(keypoints[..., :2], axis=1) + margin  # X2, Y2
+    with warnings.catch_warnings():  # silence warnings when all pose confidence levels are <= 0
+        warnings.simplefilter("ignore", category=RuntimeWarning)
+        bboxes[:, :2] = np.nanmin(keypoints[..., :2], axis=1) - margin  # X1, Y1
+        bboxes[:, 2:4] = np.nanmax(keypoints[..., :2], axis=1) + margin  # X2, Y2
 
     # can have NaNs if some individuals have no visible keypoints
     bboxes = np.nan_to_num(bboxes, nan=0)
@@ -409,7 +412,7 @@ def _annotation_to_keypoints(annotation: dict, h: int, w: int) -> np.array:
 
     Returns:
         keypoints: np.array where the first two columns are x and y coordinates of the
-    
+
     """
     # we don't mess up visibility flags here
     return annotation["keypoints"].reshape(-1, 3)

diff --git a/tests/pose_estimation_pytorch/data/test_preprocessor.py b/tests/pose_estimation_pytorch/data/test_preprocessor.py
@@ -12,9 +12,13 @@
 import albumentations as A
 import numpy as np
 import pytest
+from albumentations import BaseCompose
 
 from deeplabcut.pose_estimation_pytorch.data.transforms import build_resize_transforms
-from deeplabcut.pose_estimation_pytorch.data.preprocessor import AugmentImage
+from deeplabcut.pose_estimation_pytorch.data.preprocessor import (
+    AugmentImage,
+    build_conditional_top_down_preprocessor,
+)
 
 
 @pytest.mark.parametrize(
@@ -25,21 +29,21 @@
             "resize_transform": {"height": 5, "width": 4, "keep_ratio": True},
             "output_shape": (2, 4, 4),
             "padded_shape": (5, 4, 4),  # single offset as not a batch
-            "output_context": {"offsets": (0, 0), "scales": (1, 1)}
+            "output_context": {"offsets": (0, 0), "scales": (1, 1)},
         },
         {
             "image_shape": (1, 2, 4, 4),  # as batch
             "resize_transform": {"height": 10, "width": 4, "keep_ratio": True},
             "output_shape": (1, 2, 4, 4),
             "padded_shape": (1, 10, 4, 4),
-            "output_context": {"offsets": [(0, 0)], "scales": [(1, 1)]}
+            "output_context": {"offsets": [(0, 0)], "scales": [(1, 1)]},
         },
         {
             "image_shape": (2, 4, 3),
             "resize_transform": {"height": 10, "width": 8, "keep_ratio": True},
             "output_shape": (4, 8, 3),
             "padded_shape": (10, 8, 3),
-            "output_context": {"offsets": (0, 0), "scales": (0.5, 0.5)}
+            "output_context": {"offsets": (0, 0), "scales": (0.5, 0.5)},
         },
     ],
 )
@@ -59,3 +63,116 @@ def test_augment_image_rescaling(data):
     assert np.sum(transformed_image) == np.sum(np.ones(data["output_shape"]))
     assert context == data["output_context"]
     assert transformed_image.shape == data["padded_shape"]
+
+
+ctd_preprocessor = build_conditional_top_down_preprocessor(
+    color_mode="RGB",
+    transform=A.Compose(
+        build_resize_transforms({"height": 100, "width": 100, "keep_ratio": True}),
+        keypoint_params=A.KeypointParams("xy", remove_invisible=False),
+        bbox_params=A.BboxParams(format="coco", label_fields=["bbox_labels"]),
+    ),
+    bbox_margin=0,
+    top_down_crop_size=(256, 256),
+)
+
+
+@pytest.mark.parametrize(
+    "data",
+    [
+        # two well-defined individuals
+        {
+            "image_shape": (100, 100, 3),
+            "context": {
+                "cond_kpts": np.array(
+                    [[[10, 10, 0.8], [20, 20, 0.8]], [[60, 60, 0.8], [70, 70, 0.8]]]
+                )
+            },
+            "output_context": {
+                "cond_kpts": np.array(
+                    [[[10, 10, 0.8], [20, 20, 0.8]], [[60, 60, 0.8], [70, 70, 0.8]]]
+                ),
+                "bboxes": [np.array([10, 10, 10, 10]), np.array([60, 60, 10, 10])],
+                "offsets": [(10, 10), (60, 60)],
+                "scales": [(0.1, 0.1), (0.1, 0.1)],
+            },
+        },
+        # one individual has 0 keypoints
+        {
+            "image_shape": (100, 100, 3),
+            "context": {
+                "cond_kpts": np.array(
+                    [[[10, 10, 0.8], [20, 20, 0.8]], [[60, 60, 0.0], [70, 70, 0.0]]]
+                )
+            },
+            "output_context": {
+                "cond_kpts": np.array(
+                    [
+                        [[10, 10, 0.8], [20, 20, 0.8]],
+                    ]
+                ),
+                "bboxes": [np.array([10, 10, 10, 10])],
+                "offsets": [(10, 10)],
+                "scales": [(0.1, 0.1)],
+            },
+        },
+        # one individual has only 1 keypoints
+        {
+            "image_shape": (100, 100, 3),
+            "context": {
+                "cond_kpts": np.array(
+                    [[[10, 10, 0.8], [20, 20, 0.8]], [[60, 60, 0.0], [70, 70, 0.9]]]
+                )
+            },
+            "output_context": {
+                "cond_kpts": np.array(
+                    [
+                        [[10, 10, 0.8], [20, 20, 0.8]],
+                    ]
+                ),
+                "bboxes": [np.array([10, 10, 10, 10])],
+                "offsets": [(10, 10)],
+                "scales": [(0.1, 0.1)],
+            },
+        },
+        # two individuals but one is low confidence
+        {
+            "image_shape": (100, 100, 3),
+            "context": {
+                "cond_kpts": np.array(
+                    [[[10, 10, 0.8], [20, 20, 0.8]], [[60, 60, 0.01], [70, 70, 0.01]]]
+                )
+            },
+            "output_context": {
+                "cond_kpts": np.array(
+                    [
+                        [[10, 10, 0.8], [20, 20, 0.8]],
+                    ]
+                ),
+                "bboxes": [np.array([10, 10, 10, 10])],
+                "offsets": [(10, 10)],
+                "scales": [(0.1, 0.1)],
+            },
+        },
+    ],
+)
+def test_conditional_top_down_preprocessor(data):
+    input_img = np.ones(data["image_shape"])
+
+    output_img, output_context = ctd_preprocessor(input_img, context=data["context"])
+
+    for context_key in ["cond_kpts", "bboxes", "offsets", "scales"]:
+        assert deep_equal(
+            output_context[context_key], data["output_context"][context_key]
+        )
+
+
+def deep_equal(a, b):
+    if isinstance(a, np.ndarray) and isinstance(b, np.ndarray):
+        return np.array_equal(a, b)
+    elif isinstance(a, list) and isinstance(b, list):
+        if len(a) != len(b):
+            return False
+        return all(deep_equal(x, y) for x, y in zip(a, b))
+    else:
+        return a == b