feat: add RF-DETR keypoint uncertainty visualization

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

Author: Borda

src/supervision/__init__.py

@@ -121,6 +121,7 @@
 from supervision.key_points.annotators import (
     EdgeAnnotator,
     VertexAnnotator,
+    VertexEllipseAnnotator,
     VertexLabelAnnotator,
 )
 from supervision.key_points.core import KeyPoints
@@ -204,6 +205,7 @@
     "TraceAnnotator",
     "TriangleAnnotator",
     "VertexAnnotator",
+    "VertexEllipseAnnotator",
     "VertexLabelAnnotator",
     "VideoInfo",
     "VideoSink",

src/supervision/key_points/annotators.py

@@ -2,7 +2,7 @@
 
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from typing import Any
+from typing import Any, Literal
 
 import cv2
 import numpy as np
@@ -192,6 +192,193 @@ def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
         return scene
 
 
+class VertexEllipseAnnotator(BaseKeyPointAnnotator):
+    """
+    A class that draws covariance ellipses around skeleton vertices.
+
+    The annotator expects per-keypoint covariance matrices stored in
+    ``key_points.data[covariance_data_key]`` with shape ``(N, K, 2, 2)`` in pixel
+    coordinates, where ``N`` is the number of keypoint sets and ``K`` is the
+    number of vertices per set.
+    """
+
+    def __init__(
+        self,
+        color: Color = Color.ROBOFLOW,
+        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,
+    ) -> None:
+        """
+        Args:
+            color: The color to use for covariance ellipses.
+            thickness: The line thickness used to draw the ellipses.
+            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.
+            line_style: Ellipse line style. Use ``"dashed"`` for less visually
+                dominant uncertainty overlays.
+            dash_length: Arc length in degrees for each dashed segment. Only used
+                when ``line_style="dashed"``.
+        """
+        if sigma <= 0:
+            raise ValueError("sigma must be positive")
+        if thickness <= 0:
+            raise ValueError("thickness must be positive")
+        if max_axis_length is not None and max_axis_length <= 0:
+            raise ValueError("max_axis_length must be positive when provided")
+        if line_style not in {"solid", "dashed"}:
+            raise ValueError("line_style must be 'solid' or 'dashed'")
+        if dash_length <= 0:
+            raise ValueError("dash_length must be positive")
+
+        self.color = color
+        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
+
+    @ensure_cv2_image_for_class_method
+    def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType:
+        """
+        Annotates the given scene with covariance ellipses around keypoints.
+
+        Args:
+            scene: The image where covariance ellipses will be drawn. ``ImageType``
+                accepts either ``numpy.ndarray`` or ``PIL.Image.Image``.
+            key_points: A collection of key points. Covariance matrices must be
+                stored in ``key_points.data[covariance_data_key]``.
+
+        Returns:
+            The annotated image, matching the type of ``scene``.
+
+        Example:
+            ```pycon
+            >>> import numpy as np
+            >>> import supervision as sv
+            >>> image = np.zeros((100, 100, 3), dtype=np.uint8)
+            >>> key_points = sv.KeyPoints(
+            ...     xy=np.array([[[50, 50], [60, 60]]], dtype=np.float32),
+            ...     data={"covariance": np.array([[[[25, 0], [0, 9]], [[9, 0], [0, 4]]]], dtype=np.float32)}
+            ... )
+            >>> annotator = sv.VertexEllipseAnnotator(color=sv.Color.GREEN)
+            >>> annotated_frame = annotator.annotate(image.copy(), key_points)
+            >>> annotated_frame.shape
+            (100, 100, 3)
+
+            ```
+        """  # noqa: E501 // docs
+        assert isinstance(scene, np.ndarray)
+        if len(key_points) == 0:
+            return scene
+
+        covariances = self._get_covariances(key_points=key_points)
+        for detection_index, xy in enumerate(key_points.xy):
+            for point_index, (x, y) in enumerate(xy):
+                if np.allclose((x, y), 0):
+                    continue
+                if key_points.confidence is not None:
+                    confidence = key_points.confidence[detection_index, point_index]
+                    if confidence < self.confidence_threshold:
+                        continue
+                ellipse = self._covariance_to_ellipse(
+                    covariance=covariances[detection_index, point_index]
+                )
+                if ellipse is None:
+                    continue
+                axis_lengths, angle = ellipse
+                self._draw_ellipse(
+                    scene=scene,
+                    center=(round(x), round(y)),
+                    axes=axis_lengths,
+                    angle=angle,
+                )
+
+        return scene
+
+    def _get_covariances(self, key_points: KeyPoints) -> npt.NDArray[np.float32]:
+        covariances = key_points.data.get(self.covariance_data_key)
+        if covariances is None:
+            raise ValueError(
+                f"key_points.data must contain {self.covariance_data_key!r} "
+                "with shape (N, K, 2, 2)."
+            )
+        covariances = np.asarray(covariances, dtype=np.float32)
+        expected_shape = (*key_points.xy.shape[:2], 2, 2)
+        if covariances.shape != expected_shape:
+            raise ValueError(
+                f"Expected covariance shape {expected_shape}, got {covariances.shape}."
+            )
+        return covariances
+
+    def _covariance_to_ellipse(
+        self, covariance: npt.NDArray[np.float32]
+    ) -> tuple[tuple[int, int], float] | None:
+        if not np.isfinite(covariance).all():
+            return None
+        try:
+            eigenvalues, eigenvectors = np.linalg.eigh(covariance.astype(np.float64))
+        except np.linalg.LinAlgError:
+            return None
+        if not np.isfinite(eigenvalues).all() or np.any(eigenvalues <= 0):
+            return None
+
+        order = np.argsort(eigenvalues)[::-1]
+        eigenvalues = eigenvalues[order]
+        eigenvectors = eigenvectors[:, order]
+        axes = self.sigma * np.sqrt(eigenvalues)
+        if self.max_axis_length is not None:
+            axes = np.minimum(axes, self.max_axis_length)
+        axis_lengths = tuple(max(1, round(axis)) for axis in axes)
+        angle = float(np.degrees(np.arctan2(eigenvectors[1, 0], eigenvectors[0, 0])))
+        return axis_lengths, angle
+
+    def _draw_ellipse(
+        self,
+        scene: npt.NDArray[np.uint8],
+        center: tuple[int, int],
+        axes: tuple[int, int],
+        angle: float,
+    ) -> None:
+        if self.line_style == "solid":
+            cv2.ellipse(
+                img=scene,
+                center=center,
+                axes=axes,
+                angle=angle,
+                startAngle=0,
+                endAngle=360,
+                color=self.color.as_bgr(),
+                thickness=self.thickness,
+                lineType=cv2.LINE_AA,
+            )
+            return
+
+        step = self.dash_length * 2
+        for start_angle in range(0, 360, step):
+            cv2.ellipse(
+                img=scene,
+                center=center,
+                axes=axes,
+                angle=angle,
+                startAngle=start_angle,
+                endAngle=min(start_angle + self.dash_length, 360),
+                color=self.color.as_bgr(),
+                thickness=self.thickness,
+                lineType=cv2.LINE_AA,
+            )
+
+
 class VertexLabelAnnotator:
     """
     A class that draws labels of skeleton vertices on images. It uses specified key

src/supervision/key_points/core.py

@@ -23,6 +23,72 @@
 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)."
+        )
+
+    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]}."
+        )
+
+    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():
+                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:
+                continue
+
+            pixel_covariance = scale @ covariance @ scale
+            if np.isfinite(pixel_covariance).all():
+                covariances[detection_index, keypoint_index] = pixel_covariance
+    return covariances
+
+
 @dataclass
 class KeyPoints:
     """
@@ -231,6 +297,58 @@ 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`.
+
+        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.
+        """
+        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}."
+            )
+
+        data: dict[str, npt.NDArray[np.generic] | list[Any]] = {}
+        precision_cholesky = rfdetr_detections.data.get("keypoint_precision_cholesky")
+        if precision_cholesky is not None:
+            source_shape = _rfdetr_source_shape(
+                rfdetr_detections, detections_count=keypoints.shape[0]
+            )
+            data["covariance"] = _rfdetr_precision_cholesky_to_pixel_covariance(
+                precision_cholesky=np.asarray(precision_cholesky, dtype=np.float32),
+                source_shape=source_shape,
+            )
+
+        return cls(
+            xy=keypoints[:, :, :2].astype(np.float32),
+            confidence=keypoints[:, :, 2].astype(np.float32),
+            class_id=rfdetr_detections.class_id,
+            data=data,
+        )
+
     @classmethod
     def from_inference(cls, inference_result: Any) -> KeyPoints:
         """