feat: store keypoints on detections

Add validated keypoints support to Detections, including slicing, equality, merge, and inner-object merge preservation.

Remove the RF-DETR-specific KeyPoints adapter and make VertexEllipseAnnotator rely on caller-selected keypoints instead of an annotator confidence threshold.

Co-authored-by: Codex <codex@openai.com>

Author: Borda

src/supervision/detection/core.py

@@ -135,6 +135,9 @@ class simplifies data manipulation and filtering, providing a uniform API for
         mask: An array of shape `(n, H, W)` containing the segmentation masks
             (`bool` data type), or `None` when masks are not available, or as
             :class:`~supervision.detection.compact_mask.CompactMask`.
+        keypoints: An array of shape `(n, K, 2)` or `(n, K, 3)` containing
+            keypoint coordinates for each detection, or `None` when keypoints
+            are not available.
         confidence: An array of shape `(n,)` containing the confidence scores
             of the detections, or `None` when confidence values are not available.
         class_id: An array of shape `(n,)` containing the class ids of the
@@ -156,6 +159,7 @@ class simplifies data manipulation and filtering, providing a uniform API for
     tracker_id: npt.NDArray[np.generic] | None = None
     data: dict[str, npt.NDArray[np.generic] | list[Any]] = field(default_factory=dict)
     metadata: dict[str, Any] = field(default_factory=dict)
+    keypoints: npt.NDArray[np.generic] | None = None
 
     def __post_init__(self) -> None:
         validate_detections_fields(
@@ -165,6 +169,7 @@ def __post_init__(self) -> None:
             class_id=self.class_id,
             tracker_id=self.tracker_id,
             data=self.data,
+            keypoints=self.keypoints,
         )
 
     def __len__(self) -> int:
@@ -206,6 +211,7 @@ def __eq__(self, other: object) -> bool:
             [
                 np.array_equal(self.xyxy, other.xyxy),
                 np.array_equal(self.mask, other.mask),
+                np.array_equal(self.keypoints, other.keypoints),
                 np.array_equal(self.class_id, other.class_id),
                 np.array_equal(self.confidence, other.confidence),
                 np.array_equal(self.tracker_id, other.tracker_id),
@@ -2109,8 +2115,8 @@ def merge(cls, detections_list: list[Detections]) -> Detections:
         Merge a list of Detections objects into a single Detections object.
 
         This method takes a list of Detections objects and combines their
-        respective fields (`xyxy`, `mask`, `confidence`, `class_id`, and `tracker_id`)
-        into a single Detections object.
+        respective fields (`xyxy`, `mask`, `keypoints`, `confidence`, `class_id`, and
+        `tracker_id`) into a single Detections object.
 
         For example, if merging Detections with 3 and 4 detected objects, this method
         will return a Detections with 7 objects (7 entries in `xyxy`, `mask`, etc).
@@ -2171,6 +2177,7 @@ def merge(cls, detections_list: list[Detections]) -> Detections:
                 class_id=detections.class_id,
                 tracker_id=detections.tracker_id,
                 data=detections.data,
+                keypoints=detections.keypoints,
             )
 
         xyxy = np.vstack([d.xyxy for d in detections_list])
@@ -2188,9 +2195,12 @@ def stack_or_none(
                     return CompactMask.merge(masks)
                 # Mixed or all-ndarray: __array__ auto-converts any CompactMask.
                 return np.vstack([np.asarray(m) for m in masks])
+            if name == "keypoints":
+                return np.vstack([d.__getattribute__(name) for d in detections_list])
             return np.hstack([d.__getattribute__(name) for d in detections_list])
 
         mask = stack_or_none("mask")
+        keypoints = stack_or_none("keypoints")
         confidence = stack_or_none("confidence")
         class_id = stack_or_none("class_id")
         tracker_id = stack_or_none("tracker_id")
@@ -2208,6 +2218,7 @@ def stack_or_none(
             tracker_id=tracker_id,
             data=data,
             metadata=metadata,
+            keypoints=keypoints,
         )
 
     def get_anchors_coordinates(self, anchor: Position) -> npt.NDArray[np.generic]:
@@ -2322,6 +2333,7 @@ def __getitem__(
             tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
             data=get_data_item(self.data, index),
             metadata=self.metadata,
+            keypoints=self.keypoints[index] if self.keypoints is not None else None,
         )
 
     def __setitem__(self, key: str, value: npt.NDArray[np.generic] | list[Any]) -> None:
@@ -2582,7 +2594,8 @@ def merge_inner_detection_object_pair(
     The resulting `confidence` of the merged object is calculated by the weighted
     contribution of each detection to the merged object.
     The bounding boxes and masks of the two input detections are merged into a
-    single bounding box and mask, respectively.
+    single bounding box and mask, respectively. If keypoints are present, keypoints
+    from the winning detection are preserved.
 
     Args:
         detections_1: The first Detections object.
@@ -2657,6 +2670,7 @@ def merge_inner_detection_object_pair(
         tracker_id=winning_detection.tracker_id,
         data=winning_detection.data,
         metadata=metadata,
+        keypoints=winning_detection.keypoints,
     )
 
 

src/supervision/key_points/annotators.py

@@ -208,7 +208,6 @@ def __init__(
         thickness: int = 2,
         sigma: float = 2.0,
         covariance_data_key: str = "covariance",
-        confidence_threshold: float = 0.0,
         max_axis_length: float | None = None,
         line_style: Literal["solid", "dashed"] = "solid",
         dash_length: int = 16,
@@ -220,8 +219,6 @@ def __init__(
             sigma: Number of standard deviations represented by the ellipse axes.
             covariance_data_key: Key in ``key_points.data`` containing covariance
                 matrices with shape ``(N, K, 2, 2)``.
-            confidence_threshold: Minimum keypoint confidence required for drawing.
-                Ignored when ``key_points.confidence`` is ``None``.
             max_axis_length: Optional cap for ellipse semi-axis lengths in pixels.
                 When ``None`` (default), near-singular precision matrices can produce
                 extremely large eigenvalues and frame-spanning ellipses. Set this to
@@ -247,7 +244,6 @@ def __init__(
         self.thickness = thickness
         self.sigma = sigma
         self.covariance_data_key = covariance_data_key
-        self.confidence_threshold = confidence_threshold
         self.max_axis_length = max_axis_length
         self.line_style = line_style
         self.dash_length = dash_length
@@ -300,8 +296,6 @@ def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
                     confidence = key_points.confidence[detection_index, point_index]
                     if not np.isfinite(confidence):
                         continue
-                    if confidence < self.confidence_threshold:
-                        continue
                 ellipse = self._covariance_to_ellipse(
                     covariance=covariances[detection_index, point_index]
                 )

src/supervision/key_points/core.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import logging
 from collections.abc import Iterable, Iterator
 from dataclasses import dataclass, field
 from typing import Any, Union, cast
@@ -13,8 +12,6 @@
 from supervision.detection.utils.internal import get_data_item, is_data_equal
 from supervision.validators import validate_key_points_fields
 
-logger = logging.getLogger(__name__)
-
 Index1D = Union[
     int,
     slice,
@@ -26,94 +23,6 @@
 Index2D = tuple[Index1D, Index1D]
 
 
-def _rfdetr_source_shape(
-    rfdetr_detections: Detections,
-    detections_count: int,
-) -> npt.NDArray[np.float32]:
-    source_shape = rfdetr_detections.data.get("source_shape")
-    if source_shape is None:
-        raise ValueError(
-            "RF-DETR detections with keypoint precision data must contain "
-            "data['source_shape'] with shape (N, 2) where each row is "
-            "(height, width) in pixels."
-        )
-
-    source_shape_array = np.asarray(source_shape, dtype=np.float32)
-    expected_shape = (detections_count, 2)
-    if source_shape_array.shape != expected_shape:
-        raise ValueError(
-            "Expected RF-DETR source_shape shape "
-            f"{expected_shape}, got {source_shape_array.shape}."
-        )
-    return source_shape_array
-
-
-def _rfdetr_precision_cholesky_to_pixel_covariance(
-    precision_cholesky: npt.NDArray[np.float32],
-    source_shape: npt.NDArray[np.float32],
-) -> npt.NDArray[np.float32]:
-    if precision_cholesky.ndim != 3 or precision_cholesky.shape[2] != 3:
-        raise ValueError(
-            "Expected RF-DETR keypoint precision shape (N, K, 3), "
-            f"got {precision_cholesky.shape}."
-        )
-    if precision_cholesky.shape[0] != source_shape.shape[0]:
-        raise ValueError(
-            "RF-DETR keypoint precision and source_shape must contain the same "
-            "number of detections, got "
-            f"{precision_cholesky.shape[0]} and {source_shape.shape[0]}."
-        )
-
-    n_total = precision_cholesky.shape[0] * precision_cholesky.shape[1]
-    n_non_finite = 0
-    n_singular = 0
-    n_overflow = 0
-
-    covariances = np.full(
-        (*precision_cholesky.shape[:2], 2, 2), np.nan, dtype=np.float32
-    )
-    for detection_index, detection_precision in enumerate(precision_cholesky):
-        height, width = source_shape[detection_index]
-        scale = np.diag([width, height]).astype(np.float64)
-        for keypoint_index, params in enumerate(detection_precision):
-            if not np.isfinite(params).all():
-                n_non_finite += 1
-                continue
-            log_l11 = float(np.clip(params[0], -20.0, 20.0))
-            l21 = float(np.clip(params[1], -1.0e4, 1.0e4))
-            log_l22 = float(np.clip(params[2], -20.0, 20.0))
-            l11 = float(np.exp(log_l11))
-            l22 = float(np.exp(log_l22))
-            precision = np.array(
-                [[l11 * l11, l11 * l21], [l11 * l21, l21 * l21 + l22 * l22]],
-                dtype=np.float64,
-            )
-            try:
-                covariance = np.linalg.inv(precision)
-            except np.linalg.LinAlgError:
-                n_singular += 1
-                continue
-
-            pixel_covariance = scale @ covariance @ scale
-            if np.isfinite(pixel_covariance).all():
-                covariances[detection_index, keypoint_index] = pixel_covariance
-            else:
-                n_overflow += 1
-
-    n_failed = n_non_finite + n_singular + n_overflow
-    if n_failed > 0:
-        logger.warning(
-            "%d of %d precision matrices failed: "
-            "non_finite=%d, singular=%d, overflow=%d",
-            n_failed,
-            n_total,
-            n_non_finite,
-            n_singular,
-            n_overflow,
-        )
-    return covariances
-
-
 def _optional_array_equal(
     first: npt.NDArray[np.generic] | None,
     second: npt.NDArray[np.generic] | None,
@@ -250,13 +159,6 @@ class simplifies data manipulation and filtering, providing a uniform API for
         key_point = sv.KeyPoints.from_transformers(results[0])
         ```
 
-    Note:
-        [`sv.KeyPoints.from_rfdetr`][supervision.key_points.core.KeyPoints.from_rfdetr]
-        accepts ``sv.Detections`` (not native RF-DETR output) because RF-DETR keypoints
-        are attached as extra fields inside a ``sv.Detections`` object returned by
-        ``model.predict()``. Run that conversion first, then pass the result to
-        ``from_rfdetr``.
-
     Attributes:
         xy: An array of shape `(n, m, 2)` containing
             `n` detected objects, each composed of `m` equally-sized
@@ -338,111 +240,6 @@ def __eq__(self, other: object) -> bool:
             ]
         )
 
-    @classmethod
-    def from_rfdetr(cls, rfdetr_detections: Detections) -> KeyPoints:
-        """
-        Create a `sv.KeyPoints` object from RF-DETR `sv.Detections` output.
-
-        RF-DETR attaches keypoint coordinates to ``detections.data["keypoints"]``
-        with shape ``(N, K, 3)`` where the last dimension stores ``[x, y,
-        confidence]`` in pixel coordinates. When RF-DETR also provides
-        ``detections.data["keypoint_precision_cholesky"]``, this method converts
-        those per-keypoint precision parameters into pixel-space covariance matrices
-        and stores them in ``key_points.data["covariance"]`` for use with
-        `sv.VertexEllipseAnnotator`.
-
-        Note:
-            ``detections.data["source_shape"]`` must have shape ``(N, 2)`` where each
-            row is ``(height, width)`` in pixels — note this is HW order, not the WH
-            order used by ``resolution_wh`` elsewhere in supervision.
-
-            Keypoint confidence values are stored as-is from RF-DETR output and are
-            expected to be probabilities in the range ``[0, 1]``. If RF-DETR returns
-            logits instead, user-supplied ``confidence_threshold`` values in
-            `sv.VertexEllipseAnnotator` should be adjusted accordingly.
-
-        Args:
-            rfdetr_detections: RF-DETR prediction returned by ``model.predict()``.
-
-        Returns:
-            A `sv.KeyPoints` object containing RF-DETR keypoints and optional
-                covariance matrices.
-
-        Raises:
-            ValueError: If the RF-DETR detections do not contain valid keypoints,
-                or if precision parameters are present without source shape data.
-
-        Examples:
-            Basic usage — keypoints only:
-
-            >>> import numpy as np
-            >>> import supervision as sv
-            >>> kp_arr = np.array([[[50, 80, 0.9], [60, 90, 0.8]]], dtype=np.float32)
-            >>> detections = sv.Detections(
-            ...     xyxy=np.array([[10, 20, 100, 200]], dtype=np.float32),
-            ...     data={"keypoints": kp_arr},
-            ... )
-            >>> key_points = sv.KeyPoints.from_rfdetr(detections)
-            >>> key_points.xy.shape
-            (1, 2, 2)
-
-            With precision Cholesky parameters (produces covariance data):
-
-            >>> kp_arr2 = np.array([[[50, 80, 0.9], [60, 90, 0.8]]], dtype=np.float32)
-            >>> chol = np.zeros((1, 2, 3), dtype=np.float32)
-            >>> src = np.array([[480, 640]], dtype=np.float32)
-            >>> detections_with_cov = sv.Detections(
-            ...     xyxy=np.array([[10, 20, 100, 200]], dtype=np.float32),
-            ...     data={
-            ...         "keypoints": kp_arr2,
-            ...         "keypoint_precision_cholesky": chol,
-            ...         "source_shape": src,
-            ...     },
-            ... )
-            >>> kp = sv.KeyPoints.from_rfdetr(detections_with_cov)
-            >>> "covariance" in kp.data
-            True
-        """
-        rfdetr_keypoints = rfdetr_detections.data.get("keypoints")
-        if rfdetr_keypoints is None:
-            raise ValueError("RF-DETR detections must contain data['keypoints'].")
-
-        keypoints = np.asarray(rfdetr_keypoints, dtype=np.float32)
-        if keypoints.ndim != 3 or keypoints.shape[2] != 3:
-            raise ValueError(
-                f"Expected RF-DETR keypoints shape (N, K, 3), got {keypoints.shape}."
-            )
-        if keypoints.shape[0] == 0:
-            return cls.empty()
-
-        data: dict[str, npt.NDArray[np.generic] | list[Any]] = {}
-        precision_cholesky = rfdetr_detections.data.get("keypoint_precision_cholesky")
-        if precision_cholesky is not None:
-            precision_cholesky_array = np.asarray(precision_cholesky, dtype=np.float32)
-            if precision_cholesky_array.shape[:2] != keypoints.shape[:2]:
-                raise ValueError(
-                    "keypoint_precision_cholesky shape "
-                    f"{precision_cholesky_array.shape[:2]} does not match "
-                    f"keypoints shape {keypoints.shape[:2]}."
-                )
-            source_shape = _rfdetr_source_shape(
-                rfdetr_detections, detections_count=keypoints.shape[0]
-            )
-            data["covariance"] = _rfdetr_precision_cholesky_to_pixel_covariance(
-                precision_cholesky=precision_cholesky_array,
-                source_shape=source_shape,
-            )
-        class_id: npt.NDArray[np.int_] | None = None
-        if rfdetr_detections.class_id is not None:
-            class_id = rfdetr_detections.class_id.astype(np.int_)
-
-        return cls(
-            xy=keypoints[:, :, :2].astype(np.float32),
-            confidence=keypoints[:, :, 2].astype(np.float32),
-            class_id=class_id,
-            data=data,
-        )
-
     @classmethod
     def from_inference(cls, inference_result: Any) -> KeyPoints:
         """