From 3725beb2dbe6c98db7cb41ecdbd806f259658c74 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 13 Jun 2024 13:50:27 -0400 Subject: [PATCH] docs improvements --- supervision/keypoint/core.py | 189 +++++++++++++++++++++++------------ 1 file changed, 123 insertions(+), 66 deletions(-) diff --git a/supervision/keypoint/core.py b/supervision/keypoint/core.py index b4243c18..79008683 100644 --- a/supervision/keypoint/core.py +++ b/supervision/keypoint/core.py @@ -18,12 +18,13 @@ class KeyPoints: The `sv.KeyPoints` class in the Supervision library standardizes results from various keypoint detection and pose estimation models into a consistent format. This class simplifies data manipulation and filtering, providing a uniform API for - integration with Supervision annotators. + integration with Supervision [keypoints annotators](/keypoint/annotators). === "Ultralytics" Use [`sv.KeyPoints.from_ultralytics`](/keypoint/core/#supervision.keypoint.core.KeyPoints.from_ultralytics) - method, which accepts model results. + method, which accepts [YOLOv8](https://github.com/ultralytics/ultralytics) + pose result. ```python import cv2 @@ -32,9 +33,59 @@ class KeyPoints: image = cv2.imread() model = YOLO('yolov8s-pose.pt') + result = model(image)[0] key_points = sv.KeyPoints.from_ultralytics(result) ``` + + === "Inference" + + Use [`sv.KeyPoints.from_inference`](/keypoint/core/#supervision.keypoint.core.KeyPoints.from_inference) + method, which accepts [Inference](https://inference.roboflow.com/) pose result. + + ```python + import cv2 + import supervision as sv + from inference import get_model + + image = cv2.imread() + model = get_model(model_id=, api_key=) + + result = model.infer(image)[0] + key_points = sv.KeyPoints.from_inference(result) + ``` + + === "MediaPipe" + + Use [`sv.KeyPoints.from_mediapipe`](/keypoint/core/#supervision.keypoint.core.KeyPoints.from_mediapipe) + method, which accepts [MediaPipe](https://github.com/google-ai-edge/mediapipe) + pose result. + + ```python + import cv2 + import mediapipe as mp + import supervision as sv + + image = cv2.imread() + image_height, image_width, _ = image.shape + mediapipe_image = mp.Image( + image_format=mp.ImageFormat.SRGB, + data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) + + options = mp.tasks.vision.PoseLandmarkerOptions( + base_options=mp.tasks.BaseOptions( + model_asset_path="pose_landmarker_heavy.task" + ), + running_mode=mp.tasks.vision.RunningMode.IMAGE, + num_poses=2) + + PoseLandmarker = mp.tasks.vision.PoseLandmarker + with PoseLandmarker.create_from_options(options) as landmarker: + pose_landmarker_result = landmarker.detect(mediapipe_image) + + key_points = sv.KeyPoints.from_mediapipe( + pose_landmarker_result, (image_width, image_height)) + ``` Attributes: xy (np.ndarray): An array of shape `(n, 2)` containing @@ -63,7 +114,7 @@ class KeyPoints: def __len__(self) -> int: """ - Returns the number of keypoints in the keypoints object. + Returns the number of keypoints in the `sv.KeyPoints` object. """ return len(self.xy) @@ -106,18 +157,17 @@ class KeyPoints: """ Create a `sv.KeyPoints` object from the [Roboflow](https://roboflow.com/) API inference result or the [Inference](https://inference.roboflow.com/) - package results. When a keypoint detection model is used, this method - extracts the keypoint coordinates, class IDs, confidences, and class names. + package results. Args: inference_result (dict, any): The result from the Roboflow API or Inference package containing predictions with keypoints. Returns: - (KeyPoints): A KeyPoints object containing the keypoint coordinates, - class IDs, and confidences of each keypoint. + A `sv.KeyPoints` object containing the keypoint coordinates, class IDs, + and class names, and confidences of each keypoint. - Example: + Examples: ```python import cv2 import supervision as sv @@ -151,8 +201,6 @@ class KeyPoints: "You can retrieve it like so: inference_result = model.infer(image)[0]" ) - # Unpack the result if received from inference.get_model, - # rather than inference_sdk.InferenceHTTPClient with suppress(AttributeError): inference_result = inference_result.dict(exclude_none=True, by_alias=True) @@ -190,25 +238,26 @@ class KeyPoints: cls, mediapipe_results, resolution_wh: Tuple[int, int] ) -> KeyPoints: """ - Creates a KeyPoints instance from a - [Pose landmark detection](https://ai.google.dev/edge/mediapipe/solutions/vision/pose_landmarker/python) - inference result. - + Creates a `sv.KeyPoints` instance from a + [MediaPipe](https://github.com/google-ai-edge/mediapipe) + pose landmark detection inference result. Args: - mediapipe_results - (Union[mediapipe.tasks.python.vision.pose_landmarker.PoseLandmarkerResult, - mediapipe.python.solution_base.SolutionOutputs]): + mediapipe_results (Union[PoseLandmarkerResult, SolutionOutputs]): The output results from Mediapipe. It supports both: the inference - result from mp.tasks.vision.pose_landmaker.PoseLandmarker and the legacy - one from mp.solutions.pose.Pose. + result `PoseLandmarker` and the legacy one from `Pose`. resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the frame. Returns: - KeyPoints: A new KeyPoints object. + A `sv.KeyPoints` object containing the keypoint coordinates and + confidences of each keypoint. + + !!! tip + Before you start, download model bundles from the + [MediaPipe website](https://ai.google.dev/edge/mediapipe/solutions/vision/pose_landmarker/index#models). - Example: + Examples: ```python import cv2 import mediapipe as mp @@ -218,27 +267,23 @@ class KeyPoints: image_height, image_width, _ = image.shape mediapipe_image = mp.Image( image_format=mp.ImageFormat.SRGB, - data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB), - ) + data=cv2.cvtColor(image, cv2.COLOR_BGR2RGB)) - # Download model task from: https://ai.google.dev/edge/mediapipe/solutions/vision/pose_landmarker/index#models options = mp.tasks.vision.PoseLandmarkerOptions( base_options=mp.tasks.BaseOptions( model_asset_path="pose_landmarker_heavy.task" ), running_mode=mp.tasks.vision.RunningMode.IMAGE, - num_poses=2, - ) + num_poses=2) PoseLandmarker = mp.tasks.vision.PoseLandmarker with PoseLandmarker.create_from_options(options) as landmarker: pose_landmarker_result = landmarker.detect(mediapipe_image) - keypoints = sv.KeyPoints.from_mediapipe( - pose_landmarker_result, (image_width, image_height) - ) + key_points = sv.KeyPoints.from_mediapipe( + pose_landmarker_result, (image_width, image_height)) ``` - """ + """ # noqa: E501 // docs results = mediapipe_results.pose_landmarks if not isinstance(mediapipe_results.pose_landmarks, list): if mediapipe_results.pose_landmarks is None: @@ -275,17 +320,18 @@ class KeyPoints: @classmethod def from_ultralytics(cls, ultralytics_results) -> KeyPoints: """ - Creates a KeyPoints instance from a - [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. + Creates a `sv.KeyPoints` instance from a + [YOLOv8](https://github.com/ultralytics/ultralytics) pose inference result. Args: ultralytics_results (ultralytics.engine.results.Keypoints): The output Results instance from YOLOv8 Returns: - KeyPoints: A new KeyPoints object. + A `sv.KeyPoints` object containing the keypoint coordinates, class IDs, + and class names, and confidences of each keypoint. - Example: + Examples: ```python import cv2 import supervision as sv @@ -293,8 +339,9 @@ class KeyPoints: image = cv2.imread() model = YOLO('yolov8s-pose.pt') + result = model(image)[0] - keypoints = sv.KeyPoints.from_ultralytics(result) + key_points = sv.KeyPoints.from_ultralytics(result) ``` """ if ultralytics_results.keypoints.xy.numel() == 0: @@ -311,16 +358,18 @@ class KeyPoints: @classmethod def from_yolo_nas(cls, yolo_nas_results) -> KeyPoints: """ - Create a KeyPoints instance from a YOLO NAS results. + Create a `sv.KeyPoints` instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS-POSE.md) + pose inference results. Args: - yolo_nas_results (ImagePoseEstimationPrediction): - The output object from YOLO NAS. + yolo_nas_results (ImagePoseEstimationPrediction): The output object from + YOLO NAS. Returns: - KeyPoints: A new KeyPoints object. + A `sv.KeyPoints` object containing the keypoint coordinates, class IDs, + and class names, and confidences of each keypoint. - Example: + Examples: ```python import cv2 import torch @@ -330,13 +379,13 @@ class KeyPoints: image = cv2.imread() device = "cuda" if torch.cuda.is_available() else "cpu" - yolo_nas = super_gradients.training.models.get( + model = super_gradients.training.models.get( "yolo_nas_pose_s", pretrained_weights="coco_pose").to(device) - results = yolo_nas.predict(image, conf=0.1) - keypoints = sv.KeyPoints.from_yolo_nas(results) + results = model.predict(image, conf=0.1) + key_points = sv.KeyPoints.from_yolo_nas(results) ``` - """ + """ # noqa: E501 // docs if len(yolo_nas_results.prediction.poses) == 0: return cls.empty() @@ -369,34 +418,41 @@ class KeyPoints: self, index: Union[int, slice, List[int], np.ndarray, str] ) -> Union[KeyPoints, List, np.ndarray, None]: """ - Get a subset of the KeyPoints object or access an item from its data field. + Get a subset of the `sv.KeyPoints` object or access an item from its data field. When provided with an integer, slice, list of integers, or a numpy array, this - method returns a new KeyPoints object that represents a subset of the original - keypoints. When provided with a string, it accesses the corresponding item in - the data dictionary. + method returns a new `sv.KeyPoints` object that represents a subset of the + original `sv.KeyPoints`. When provided with a string, it accesses the + corresponding item in the data dictionary. Args: index (Union[int, slice, List[int], np.ndarray, str]): The index, indices, - or key to access a subset of the KeyPoints or an item from the data. + or key to access a subset of the `sv.KeyPoints` or an item from the + data. Returns: - Union[KeyPoints, Any]: A subset of the KeyPoints object or an item from - the data field. + A subset of the `sv.KeyPoints` object or an item from the data field. - Example: + Examples: ```python import supervision as sv - keypoints = sv.KeyPoints() + key_points = sv.KeyPoints() - first_detection = keypoints[0] - first_10_keypoints = keypoints[0:10] - some_keypoints = keypoints[[0, 2, 4]] - class_0_keypoints = keypoints[keypoints.class_id == 0] - high_confidence_keypoints = keypoints[keypoints.confidence > 0.5] + # access the first keypoint using an integer index + key_points[0] - feature_vector = keypoints['feature_vector'] + # access the first 10 keypoints using index slice + key_points[0:10] + + # access selected keypoints using a list of indices + key_points[[0, 2, 4]] + + # access keypoints with selected class_id + key_points[key_points.class_id == 0] + + # access keypoints with confidence greater than 0.5 + key_points[key_points.confidence > 0.5] ``` """ if isinstance(index, str): @@ -412,13 +468,13 @@ class KeyPoints: def __setitem__(self, key: str, value: Union[np.ndarray, List]): """ - Set a value in the data dictionary of the KeyPoints object. + Set a value in the data dictionary of the `sv.KeyPoints` object. Args: key (str): The key in the data dictionary to set. value (Union[np.ndarray, List]): The value to set for the key. - Example: + Examples: ```python import cv2 import supervision as sv @@ -430,7 +486,7 @@ class KeyPoints: result = model(image)[0] keypoints = sv.KeyPoints.from_ultralytics(result) - keypoints['names'] = [ + keypoints['class_name'] = [ model.model.names[class_id] for class_id in keypoints.class_id @@ -451,12 +507,13 @@ class KeyPoints: Create an empty Keypoints object with no keypoints. Returns: - (KeyPoints): An empty Keypoints object. + An empty `sv.KeyPoints` object. - Example: + Examples: ```python - from supervision import Keypoints - empty_keypoints = Keypoints.empty() + import supervision as sv + + key_points = sv.KeyPoints.empty() ``` """ return cls(xy=np.empty((0, 0, 2), dtype=np.float32))