xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

docs improvements

This commit is contained in:
SkalskiP 2024-06-13 13:50:27 -04:00
parent 93b21a8d2d
commit 3725beb2db
1 changed files with 123 additions and 66 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

189
supervision/keypoint/core.py
View File

@ -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(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
model = get_model(model_id=<POSE_MODEL_ID>, api_key=<ROBOFLOW_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(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
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))