diff --git a/docs/how_to/track_objects.md b/docs/how_to/track_objects.md index 784cb7bf..29fa8398 100644 --- a/docs/how_to/track_objects.md +++ b/docs/how_to/track_objects.md @@ -1,5 +1,6 @@ --- comments: true +status: new --- # Track Objects @@ -317,7 +318,7 @@ movement patterns and interactions between objects in the video. ## Tracking Key Points -Keypoint tracking is currently supported via the conversion of `KeyPoints` to `Detections`. This is achieved with the [`keypoints_to_detections`](/latest/utils/datatypes/#supervision.utils.datatypes.keypoints_to_detections) function. We'll use a different video as well as [`DetectionsSmoother`](/latest/detection/tools/smoother/) to stabilize the boxes. +Keypoint tracking is currently supported via the conversion of `KeyPoints` to `Detections`. This is achieved with the [`KeyPoints.as_detections()`](/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.as_detections) function. We'll use a different video as well as [`DetectionsSmoother`](/latest/detection/tools/smoother/) to stabilize the boxes. !!! tip @@ -340,7 +341,7 @@ Keypoint tracking is currently supported via the conversion of `KeyPoints` to `D def callback(frame: np.ndarray, _: int) -> np.ndarray: results = model(frame)[0] keypoints = sv.KeyPoints.from_ultralytics(results) - detections = sv.keypoints_to_detections(keypoints) + detections = keypoints.as_detections() detections = tracker.update_with_detections(detections) detections = smoother.update_with_detections(detections) @@ -382,7 +383,7 @@ Keypoint tracking is currently supported via the conversion of `KeyPoints` to `D def callback(frame: np.ndarray, _: int) -> np.ndarray: results = model.infer(frame)[0] keypoints = sv.KeyPoints.from_inference(results) - detections = sv.keypoints_to_detections(keypoints) + detections = keypoints.as_detections() detections = tracker.update_with_detections(detections) detections = smoother.update_with_detections(detections) diff --git a/docs/keypoint/core.md b/docs/keypoint/core.md index 6f42c254..7354baba 100644 --- a/docs/keypoint/core.md +++ b/docs/keypoint/core.md @@ -1,5 +1,6 @@ --- comments: true +status: new --- # Keypoint Detection diff --git a/docs/utils/datatypes.md b/docs/utils/datatypes.md deleted file mode 100644 index 5e0560bd..00000000 --- a/docs/utils/datatypes.md +++ /dev/null @@ -1,12 +0,0 @@ ---- -comments: true -status: new ---- - -# Data Types Utils - -
-

keypoints_to_detections

-
- -:::supervision.utils.datatypes.keypoints_to_detections diff --git a/mkdocs.yml b/mkdocs.yml index 8c939cf0..b30dbcfc 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -79,7 +79,6 @@ nav: - File: utils/file.md - Draw: utils/draw.md - Geometry: utils/geometry.md - - Datatypes: utils/datatypes.md - Assets: assets.md - Cookbooks: cookbooks.md - Cheatsheet: https://roboflow.github.io/cheatsheet-supervision/ diff --git a/supervision/__init__.py b/supervision/__init__.py index 04813ef2..746b2f67 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -100,7 +100,6 @@ from supervision.keypoint.core import KeyPoints from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision from supervision.tracker.byte_tracker.core import ByteTrack from supervision.utils.conversion import cv2_to_pillow, pillow_to_cv2 -from supervision.utils.datatypes import keypoints_to_detections from supervision.utils.file import list_files_with_extensions from supervision.utils.image import ( ImageSink, diff --git a/supervision/keypoint/core.py b/supervision/keypoint/core.py index 252fb63f..4b8e9d55 100644 --- a/supervision/keypoint/core.py +++ b/supervision/keypoint/core.py @@ -2,12 +2,13 @@ from __future__ import annotations from contextlib import suppress from dataclasses import dataclass, field -from typing import Any, Dict, Iterator, List, Optional, Tuple, Union +from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union import numpy as np import numpy.typing as npt from supervision.config import CLASS_NAME_DATA_FIELD +from supervision.detection.core import Detections from supervision.detection.utils import get_data_item, is_data_equal from supervision.validators import validate_keypoints_fields @@ -620,3 +621,67 @@ class KeyPoints: empty_keypoints = KeyPoints.empty() empty_keypoints.data = self.data return self == empty_keypoints + + def as_detections( + self, selected_keypoint_indices: Optional[Iterable[int]] = None + ) -> Detections: + """ + Convert a KeyPoints object to a Detections object. This + approximates the bounding box of the detected object by + taking the bounding box that fits all keypoints. + + Arguments: + selected_keypoint_indices (Optional[Iterable[int]]): The + indices of the keypoints to include in the bounding box + calculation. This helps focus on a subset of keypoints, + e.g. when some are occluded. Captures all keypoints by default. + + Returns: + detections (Detections): The converted detections object. + + Example: + ```python + keypoints = sv.KeyPoints.from_inference(...) + detections = keypoints.as_detections() + ``` + """ + if self.is_empty(): + return Detections.empty() + + detections_list = [] + for i, xy in enumerate(self.xy): + if selected_keypoint_indices: + xy = xy[selected_keypoint_indices] + + # [0, 0] used by some frameworks to indicate missing keypoints + xy = xy[~np.all(xy == 0, axis=1)] + if len(xy) == 0: + xyxy = np.array([[0, 0, 0, 0]], dtype=np.float32) + else: + x_min = xy[:, 0].min() + x_max = xy[:, 0].max() + y_min = xy[:, 1].min() + y_max = xy[:, 1].max() + xyxy = np.array([[x_min, y_min, x_max, y_max]], dtype=np.float32) + + if self.confidence is None: + confidence = None + else: + confidence = self.confidence[i] + if selected_keypoint_indices: + confidence = confidence[selected_keypoint_indices] + confidence = np.array([confidence.mean()], dtype=np.float32) + + detections_list.append( + Detections( + xyxy=xyxy, + confidence=confidence, + ) + ) + + detections = Detections.merge(detections_list) + detections.class_id = self.class_id + detections.data = self.data + detections = detections[detections.area > 0] + + return detections diff --git a/supervision/utils/datatypes.py b/supervision/utils/datatypes.py deleted file mode 100644 index 6f2e8879..00000000 --- a/supervision/utils/datatypes.py +++ /dev/null @@ -1,72 +0,0 @@ -from typing import Iterable, Optional - -import numpy as np - -from supervision.detection.core import Detections -from supervision.keypoint.core import KeyPoints - - -def keypoints_to_detections( - keypoints: KeyPoints, selected_keypoint_indices: Optional[Iterable[int]] = None -) -> Detections: - """ - Convert a KeyPoints object to a Detections object. This - approximates the bounding box of the detected object by - taking the bounding box that fits all keypoints. - - Arguments: - keypoints (KeyPoints): The keypoints to convert to detections. - selected_keypoint_indices (Optional[Iterable[int]]): The - indices of the keypoints to include in the bounding box - calculation. This helps focus on a subset of keypoints, - e.g. when some are occluded. Captures all keypoints by default. - - Returns: - detections (Detections): The converted detections object. - - Example: - ```python - keypoints = sv.KeyPoints.from_inference(...) - detections = keypoints_to_detections(keypoints) - ``` - """ - if keypoints.is_empty(): - return Detections.empty() - - detections_list = [] - for i, xy in enumerate(keypoints.xy): - if selected_keypoint_indices: - xy = xy[selected_keypoint_indices] - - # [0, 0] used by some frameworks to indicate missing keypoints - xy = xy[~np.all(xy == 0, axis=1)] - if len(xy) == 0: - xyxy = np.array([[0, 0, 0, 0]], dtype=np.float32) - else: - x_min = xy[:, 0].min() - x_max = xy[:, 0].max() - y_min = xy[:, 1].min() - y_max = xy[:, 1].max() - xyxy = np.array([[x_min, y_min, x_max, y_max]], dtype=np.float32) - - if keypoints.confidence is None: - confidence = None - else: - confidence = keypoints.confidence[i] - if selected_keypoint_indices: - confidence = confidence[selected_keypoint_indices] - confidence = np.array([confidence.mean()], dtype=np.float32) - - detections_list.append( - Detections( - xyxy=xyxy, - confidence=confidence, - ) - ) - - detections = Detections.merge(detections_list) - detections.class_id = keypoints.class_id - detections.data = keypoints.data - detections = detections[detections.area > 0] - - return detections