From 77793571b094f646eaec2bebd99369e203b87159 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 24 Apr 2024 15:39:08 +0200 Subject: [PATCH] major docs refactor --- README.md | 7 +- docs/deprecated.md | 2 +- docs/{ => detection}/annotators.md | 0 docs/detection/metrics.md | 21 +++ docs/draw/color.md | 13 -- docs/how_to/detect_and_annotate.md | 28 +-- docs/how_to/detect_small_objects.md | 16 +- docs/keypoint/annotate.md | 10 -- docs/keypoint/annotators.md | 10 ++ docs/metrics/detection.md | 22 --- docs/{draw/utils.md => utils/draw.md} | 13 ++ docs/{geometry/core.md => utils/geometry.md} | 0 mkdocs.yml | 22 +-- supervision/__init__.py | 2 +- supervision/assets/list.py | 5 + supervision/detection/core.py | 127 ++++++++------ supervision/keypoint/annotate.py | 104 ------------ supervision/keypoint/annotators.py | 170 +++++++++++++++++++ supervision/keypoint/core.py | 43 ++--- 19 files changed, 348 insertions(+), 267 deletions(-) rename docs/{ => detection}/annotators.md (100%) create mode 100644 docs/detection/metrics.md delete mode 100644 docs/draw/color.md delete mode 100644 docs/keypoint/annotate.md create mode 100644 docs/keypoint/annotators.md delete mode 100644 docs/metrics/detection.md rename docs/{draw/utils.md => utils/draw.md} (83%) rename docs/{geometry/core.md => utils/geometry.md} (100%) delete mode 100644 supervision/keypoint/annotate.py create mode 100644 supervision/keypoint/annotators.py diff --git a/README.md b/README.md index beed655c..0894a6a4 100644 --- a/README.md +++ b/README.md @@ -71,16 +71,15 @@ len(detections) ```python import cv2 import supervision as sv - from inference.models.utils import get_roboflow_model + from inference import get_model image = cv2.imread(...) - model = get_roboflow_model(model_id="yolov8s-640", api_key=) + model = get_model(model_id="yolov8s-640", api_key=) result = model.infer(image)[0] detections = sv.Detections.from_inference(result) len(detections) - # 5 - + # 5 ``` diff --git a/docs/deprecated.md b/docs/deprecated.md index 35c332e4..cc1e9649 100644 --- a/docs/deprecated.md +++ b/docs/deprecated.md @@ -12,7 +12,7 @@ These features are phased out due to better alternatives or potential issues in - The method `Color.green()` is deprecated and will be removed in `supervision-0.22.0`. Use the constant `Color.GREEN` instead. - The method `Color.blue()` is deprecated and will be removed in `supervision-0.22.0`. Use the constant `Color.BLUE` instead. - The method [`ColorPalette.default()`](draw/color.md/#supervision.draw.color.ColorPalette.default) is deprecated and will be removed in `supervision-0.22.0`. Use the constant [`ColorPalette.DEFAULT`](draw/color.md/#supervision.draw.color.ColorPalette.DEFAULT) instead. -- `BoxAnnotator` is deprecated and will be removed in `supervision-0.22.0`. Use [`BoundingBoxAnnotator`](annotators.md/#supervision.annotators.core.BoundingBoxAnnotator) and [`LabelAnnotator`](annotators.md/#supervision.annotators.core.LabelAnnotator) instead. +- `BoxAnnotator` is deprecated and will be removed in `supervision-0.22.0`. Use [`BoundingBoxAnnotator`](detection/annotators.md/#supervision.annotators.core.BoundingBoxAnnotator) and [`LabelAnnotator`](detection/annotators.md/#supervision.annotators.core.LabelAnnotator) instead. - The method [`FPSMonitor.__call__`](utils/video.md/#supervision.utils.video.FPSMonitor.__call__) is deprecated and will be removed in `supervision-0.22.0`. Use the attribute [`FPSMonitor.fps`](utils/video.md/#supervision.utils.video.FPSMonitor.fps) instead. - The `track_buffer`, `track_thresh`, and `match_thresh` parameters in [`ByterTrack`](trackers.md/#supervision.tracker.byte_tracker.core.ByteTrack) are deprecated and will be removed in `supervision-0.23.0`. Use `lost_track_buffer,` `track_activation_threshold`, and `minimum_matching_threshold` instead. - The `triggering_position ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) is deprecated and will be removed in `supervision-0.23.0`. Use `triggering_anchors ` instead. diff --git a/docs/annotators.md b/docs/detection/annotators.md similarity index 100% rename from docs/annotators.md rename to docs/detection/annotators.md diff --git a/docs/detection/metrics.md b/docs/detection/metrics.md new file mode 100644 index 00000000..981f399c --- /dev/null +++ b/docs/detection/metrics.md @@ -0,0 +1,21 @@ +--- +comments: true +--- + +# Metrics + +!!! warning + + Evaluation API is still fluid and may change. If you use Evaluation API in your project until further notice, freeze the `supervision` version in your `requirements.txt` or `setup.py`. + +
+

ConfusionMatrix

+
+ +:::supervision.metrics.detection.ConfusionMatrix + +
+

MeanAveragePrecision

+
+ +:::supervision.metrics.detection.MeanAveragePrecision diff --git a/docs/draw/color.md b/docs/draw/color.md deleted file mode 100644 index c702814a..00000000 --- a/docs/draw/color.md +++ /dev/null @@ -1,13 +0,0 @@ ---- -comments: true ---- - -# Color - -:::supervision.draw.color.Color - -
-

ColorPalette

-
- -:::supervision.draw.color.ColorPalette diff --git a/docs/how_to/detect_and_annotate.md b/docs/how_to/detect_and_annotate.md index adea95cf..a9a4405e 100644 --- a/docs/how_to/detect_and_annotate.md +++ b/docs/how_to/detect_and_annotate.md @@ -27,7 +27,7 @@ model. from inference import get_model model = get_model(model_id="yolov8n-640") - image = cv2.imread() + image = cv2.imread() results = model.infer(image)[0] ``` @@ -38,7 +38,7 @@ model. from ultralytics import YOLO model = YOLO("yolov8n.pt") - image = cv2.imread() + image = cv2.imread() results = model(image)[0] ``` @@ -52,7 +52,7 @@ model. processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") - image = Image.open() + image = Image.open() inputs = processor(images=image, return_tensors="pt") with torch.no_grad(): @@ -78,7 +78,7 @@ Now that we have predictions from a model, we can load them into Supervision. from inference import get_model model = get_model(model_id="yolov8n-640") - image = cv2.imread() + image = cv2.imread() results = model.infer(image)[0] detections = sv.Detections.from_inference(results) ``` @@ -93,7 +93,7 @@ Now that we have predictions from a model, we can load them into Supervision. from ultralytics import YOLO model = YOLO("yolov8n.pt") - image = cv2.imread() + image = cv2.imread() results = model(image)[0] detections = sv.Detections.from_ultralytics(results) ``` @@ -111,7 +111,7 @@ Now that we have predictions from a model, we can load them into Supervision. processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") - image = Image.open() + image = Image.open() inputs = processor(images=image, return_tensors="pt") with torch.no_grad(): @@ -146,7 +146,7 @@ Finally, we can annotate the image with the predictions. Since we are working wi from inference import get_model model = get_model(model_id="yolov8n-640") - image = cv2.imread() + image = cv2.imread() results = model.infer(image)[0] detections = sv.Detections.from_inference(results) @@ -167,7 +167,7 @@ Finally, we can annotate the image with the predictions. Since we are working wi from ultralytics import YOLO model = YOLO("yolov8n.pt") - image = cv2.imread() + image = cv2.imread() results = model(image)[0] detections = sv.Detections.from_ultralytics(results) @@ -191,7 +191,7 @@ Finally, we can annotate the image with the predictions. Since we are working wi processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") - image = Image.open() + image = Image.open() inputs = processor(images=image, return_tensors="pt") with torch.no_grad(): @@ -230,7 +230,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me from inference import get_model model = get_model(model_id="yolov8n-640") - image = cv2.imread() + image = cv2.imread() results = model.infer(image)[0] detections = sv.Detections.from_inference(results) @@ -257,7 +257,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me from ultralytics import YOLO model = YOLO("yolov8n.pt") - image = cv2.imread() + image = cv2.imread() results = model(image)[0] detections = sv.Detections.from_ultralytics(results) @@ -287,7 +287,7 @@ override this behavior by passing a list of custom `labels` to the `annotate` me processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") - image = Image.open() + image = Image.open() inputs = processor(images=image, return_tensors="pt") with torch.no_grad(): @@ -334,7 +334,7 @@ that will allow you to draw masks instead of boxes. from inference import get_model model = get_model(model_id="yolov8n-seg-640") - image = cv2.imread() + image = cv2.imread() results = model.infer(image)[0] detections = sv.Detections.from_inference(results) @@ -355,7 +355,7 @@ that will allow you to draw masks instead of boxes. from ultralytics import YOLO model = YOLO("yolov8n-seg.pt") - image = cv2.imread() + image = cv2.imread() results = model(image)[0] detections = sv.Detections.from_ultralytics(results) diff --git a/docs/how_to/detect_small_objects.md b/docs/how_to/detect_small_objects.md index 9f35a96c..e2d02328 100644 --- a/docs/how_to/detect_small_objects.md +++ b/docs/how_to/detect_small_objects.md @@ -28,7 +28,7 @@ size relative to the image resolution. from inference import get_model model = get_model(model_id="yolov8x-640") - image = cv2.imread() + image = cv2.imread() results = model.infer(image)[0] detections = sv.Detections.from_inference(results) @@ -49,7 +49,7 @@ size relative to the image resolution. from ultralytics import YOLO model = YOLO("yolov8x.pt") - image = cv2.imread() + image = cv2.imread() results = model(image)[0] detections = sv.Detections.from_ultralytics(results) @@ -73,7 +73,7 @@ size relative to the image resolution. processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") - image = Image.open() + image = Image.open() inputs = processor(images=image, return_tensors="pt") with torch.no_grad(): @@ -116,7 +116,7 @@ is less effective for ultra-high-resolution images (4K and above). from inference import get_model model = get_model(model_id="yolov8x-1280") - image = cv2.imread() + image = cv2.imread() results = model.infer(image)[0] detections = sv.Detections.from_inference(results) @@ -137,7 +137,7 @@ is less effective for ultra-high-resolution images (4K and above). from ultralytics import YOLO model = YOLO("yolov8x.pt") - image = cv2.imread() + image = cv2.imread() results = model(image, imgsz=1280)[0] detections = sv.Detections.from_ultralytics(results) @@ -171,7 +171,7 @@ objects within each, and aggregating the results. from inference import get_model model = get_model(model_id="yolov8x-640") - image = cv2.imread() + image = cv2.imread() def callback(image_slice: np.ndarray) -> sv.Detections: results = model.infer(image_slice)[0] @@ -198,7 +198,7 @@ objects within each, and aggregating the results. from ultralytics import YOLO model = YOLO("yolov8x.pt") - image = cv2.imread() + image = cv2.imread() def callback(image_slice: np.ndarray) -> sv.Detections: result = model(image_slice)[0] @@ -229,7 +229,7 @@ objects within each, and aggregating the results. processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") - image = cv2.imread() + image = cv2.imread() def callback(image_slice: np.ndarray) -> sv.Detections: image_slice = cv2.cvtColor(image_slice, cv2.COLOR_BGR2RGB) diff --git a/docs/keypoint/annotate.md b/docs/keypoint/annotate.md deleted file mode 100644 index 4f9cd861..00000000 --- a/docs/keypoint/annotate.md +++ /dev/null @@ -1,10 +0,0 @@ ---- -comments: true -status: new ---- - -# Annotate - -:::supervision.keypoint.annotate.VertexAnnotator - -:::supervision.keypoint.annotate.EdgeAnnotator diff --git a/docs/keypoint/annotators.md b/docs/keypoint/annotators.md new file mode 100644 index 00000000..3efc5b03 --- /dev/null +++ b/docs/keypoint/annotators.md @@ -0,0 +1,10 @@ +--- +comments: true +status: new +--- + +# Annotators + +:::supervision.keypoint.annotators.VertexAnnotator + +:::supervision.keypoint.annotators.EdgeAnnotator diff --git a/docs/metrics/detection.md b/docs/metrics/detection.md deleted file mode 100644 index a953ba9b..00000000 --- a/docs/metrics/detection.md +++ /dev/null @@ -1,22 +0,0 @@ ---- -comments: true ---- - -#  Detection Metrics - -!!! warning - - Evaluation API is still fluid and may change. If you use Evaluation API in your project until further notice, freeze the - `supervision` version in your `requirements.txt` or `setup.py`. - -
-

ConfusionMatrix

-
- -:::supervision.metrics.detection.ConfusionMatrix - -
-

MeanAveragePrecision

-
- -:::supervision.metrics.detection.MeanAveragePrecision diff --git a/docs/draw/utils.md b/docs/utils/draw.md similarity index 83% rename from docs/draw/utils.md rename to docs/utils/draw.md index 7975952c..7b47a6a7 100644 --- a/docs/draw/utils.md +++ b/docs/utils/draw.md @@ -51,3 +51,16 @@ comments: true :::supervision.draw.utils.calculate_optimal_line_thickness + +
+

Color

+
+ +:::supervision.draw.color.Color + +
+

ColorPalette

+
+ +:::supervision.draw.color.ColorPalette + diff --git a/docs/geometry/core.md b/docs/utils/geometry.md similarity index 100% rename from docs/geometry/core.md rename to docs/utils/geometry.md diff --git a/mkdocs.yml b/mkdocs.yml index 256c6bef..cf206a82 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -44,15 +44,16 @@ nav: - Track Objects: how_to/track_objects.md - API: - - Annotators: annotators.md - - Classifications: - - Core: classification/core.md - - Detections: + - Detection and Segmentation: - Core: detection/core.md + - Annotators: detection/annotators.md + - Metrics: detection/metrics.md - Utils: detection/utils.md - - KeyPoints: + - Keypoint Detection: - Core: keypoint/core.md - - Annotators: keypoint/annotate.md + - Annotators: keypoint/annotators.md + - Classification: + - Core: classification/core.md - Tools: - Line Zone: detection/tools/line_zone.md - Polygon Zone: detection/tools/polygon_zone.md @@ -61,19 +62,14 @@ nav: - Save Detections: detection/tools/save_detections.md - Trackers: trackers.md - Datasets: datasets.md - - Metrics: - - Object Detection: metrics/detection.md - - Draw: - - Color: draw/color.md - - Utils: draw/utils.md - - Geometry: - - Position: geometry/core.md - Utils: - Video: utils/video.md - Image: utils/image.md - Iterables: utils/iterables.md - Notebook: utils/notebook.md - File: utils/file.md + - Draw: utils/draw.md + - Geometry: utils/geometry.md - Assets: assets.md - Cookbooks: cookbooks.md - Contribute: diff --git a/supervision/__init__.py b/supervision/__init__.py index 053e3718..bb526514 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -69,7 +69,7 @@ from supervision.draw.utils import ( ) from supervision.geometry.core import Point, Position, Rect from supervision.geometry.utils import get_polygon_center -from supervision.keypoint.annotate import EdgeAnnotator, VertexAnnotator +from supervision.keypoint.annotators import EdgeAnnotator, VertexAnnotator from supervision.keypoint.core import KeyPoints from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision from supervision.tracker.byte_tracker.core import ByteTrack diff --git a/supervision/assets/list.py b/supervision/assets/list.py index 77b1e461..0866431a 100644 --- a/supervision/assets/list.py +++ b/supervision/assets/list.py @@ -29,6 +29,7 @@ class VideoAssets(Enum): MARKET_SQUARE = "market-square.mp4" PEOPLE_WALKING = "people-walking.mp4" BEACH = "beach-1.mp4" + BASKETBALL = "basketball-1.mp4" @classmethod def list(cls): @@ -68,4 +69,8 @@ VIDEO_ASSETS: Dict[str, Tuple[str, str]] = { f"{BASE_VIDEO_URL}{VideoAssets.BEACH.value}", "4175d42fec4d450ed081523fd39e0cf8", ), + VideoAssets.BASKETBALL.value: ( + f"{BASE_VIDEO_URL}{VideoAssets.BASKETBALL.value}", + "60d94a3c7c47d16f09d342b088012ecc", + ), } diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 53286375..3e3d0114 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -27,34 +27,68 @@ from supervision.validators import validate_detections_fields @dataclass class Detections: """ - The `sv.Detections` allows you to convert results from a variety of object detection - and segmentation models into a single, unified format. The `sv.Detections` class - enables easy data manipulation and filtering, and provides a consistent API for - Supervision's tools like trackers, annotators, and zones. + The `sv.Detections` class in the Supervision library standardizes results from + various object detection and segmentation models into a consistent format. This + class simplifies data manipulation and filtering, providing a uniform API for + integration with Supervision [trackers](/trackers/), [annotators](/detection/annotators/), and [tools](/detection/tools/line_zone/). - ```python - import cv2 - import supervision as sv - from ultralytics import YOLO + === "Inference" - image = cv2.imread() - model = YOLO('yolov8s.pt') - annotator = sv.BoundingBoxAnnotator() + Use [`sv.Detections.from_inference`](/detection/core/#supervision.detection.core.Detections.from_inference) method, which accepts model results from both detection and segmentation models. - result = model(image)[0] - detections = sv.Detections.from_ultralytics(result) + ```python + import cv2 + import supervision as sv + from inference import get_model - annotated_image = annotator.annotate(image, detections) - ``` + model = get_model(model_id="yolov8n-640") + image = cv2.imread() + results = model.infer(image)[0] + detections = sv.Detections.from_inference(results) + ``` - !!! tip + === "Ultralytics" - In `sv.Detections`, detection data is categorized into two main field types: - fixed and custom. The fixed fields include `xyxy`, `mask`, `confidence`, - `class_id`, and `tracker_id`. For any additional data requirements, custom - fields come into play, stored in the data field. These custom fields are easily - accessible using the `detections[]` syntax, providing flexibility - for diverse data handling needs. + Use [`sv.Detections.from_ultralytics`](/detection/core/#supervision.detection.core.Detections.from_ultralytics) method, which accepts model results from both detection and segmentation models. + + ```python + import cv2 + import supervision as sv + from ultralytics import YOLO + + model = YOLO("yolov8n.pt") + image = cv2.imread() + results = model(image)[0] + detections = sv.Detections.from_ultralytics(results) + ``` + + === "Transformers" + + Use [`sv.Detections.from_transformers`](/detection/core/#supervision.detection.core.Detections.from_transformers) method, which accepts model results from both detection and segmentation models. + + ```python + import torch + import supervision as sv + from PIL import Image + from transformers import DetrImageProcessor, DetrForObjectDetection + + processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50") + model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50") + + image = Image.open() + inputs = processor(images=image, return_tensors="pt") + + with torch.no_grad(): + outputs = model(**inputs) + + width, height = image.size + target_size = torch.tensor([[height, width]]) + results = processor.post_process_object_detection( + outputs=outputs, target_sizes=target_size)[0] + detections = sv.Detections.from_transformers( + transformers_results=results, + id2label=model.config.id2label) + ``` Attributes: xyxy (np.ndarray): An array of shape `(n, 4)` containing @@ -70,14 +104,6 @@ class Detections: data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional data where each key is a string representing the data type, and the value is either a NumPy array or a list of corresponding data. - - !!! warning - - The `data` field in the `sv.Detections` class is currently in an experimental - phase. Please be aware that its API and functionality are subject to change in - future updates as we continue to refine and improve its capabilities. - We encourage users to experiment with this feature and provide feedback, but - also to be prepared for potential modifications in upcoming releases. """ xyxy: np.ndarray @@ -177,8 +203,8 @@ class Detections: @classmethod def from_ultralytics(cls, ultralytics_results) -> Detections: """ - Creates a Detections instance from a - [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. + Creates a `sv.Detections` instance from a + [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. !!! Note @@ -202,10 +228,13 @@ class Detections: image = cv2.imread() model = YOLO('yolov8s.pt') - - result = model(image)[0] - detections = sv.Detections.from_ultralytics(result) + results = model(image)[0] + detections = sv.Detections.from_ultralytics(results) ``` + + !!! tip + + Class names values can be accessed using `detections["class_name"]`. """ # noqa: E501 // docs if "obb" in ultralytics_results and ultralytics_results.obb is not None: @@ -396,11 +425,6 @@ class Detections: Creates a Detections instance from object detection or segmentation [Transformer](https://github.com/huggingface/transformers) inference result. - !!! note - - Class names can be accessed using the key `class_name` in the returned - object's data attribute. - Args: transformers_results (dict): The output of Transformers model inference. A dictionary containing the `scores`, `labels`, `boxes` and `masks` keys. @@ -437,6 +461,10 @@ class Detections: id2label=model.config.id2label ) ``` + + !!! tip + + Class names values can be accessed using `detections["class_name"]`. """ # noqa: E501 // docs class_ids = transformers_results["labels"].cpu().detach().numpy().astype(int) @@ -510,17 +538,12 @@ class Detections: @classmethod def from_inference(cls, roboflow_result: Union[dict, Any]) -> Detections: """ - Create a Detections object from the [Roboflow](https://roboflow.com/) + Create a `sv.Detections` object from the [Roboflow](https://roboflow.com/) API inference result or the [Inference](https://inference.roboflow.com/) package results. This method extracts bounding boxes, class IDs, confidences, and class names from the Roboflow API result and encapsulates them into a Detections object. - !!! note - - Class names can be accessed using the key `class_name` in the returned - object's data attribute. - Args: roboflow_result (dict, any): The result from the Roboflow API or Inference package containing predictions. @@ -533,14 +556,18 @@ class Detections: ```python import cv2 import supervision as sv - from inference.models.utils import get_roboflow_model + from inference import get_model image = cv2.imread() - model = get_roboflow_model(model_id="yolov8s-640") + model = get_model(model_id="yolov8s-640") result = model.infer(image)[0] detections = sv.Detections.from_inference(result) ``` + + !!! tip + + Class names values can be accessed using `detections["class_name"]`. """ with suppress(AttributeError): roboflow_result = roboflow_result.dict(exclude_none=True, by_alias=True) @@ -590,10 +617,10 @@ class Detections: ```python import cv2 import supervision as sv - from inference.models.utils import get_roboflow_model + from inference import get_model image = cv2.imread() - model = get_roboflow_model(model_id="yolov8s-640") + model = get_model(model_id="yolov8s-640") result = model.infer(image)[0] detections = sv.Detections.from_roboflow(result) diff --git a/supervision/keypoint/annotate.py b/supervision/keypoint/annotate.py deleted file mode 100644 index b1f3530f..00000000 --- a/supervision/keypoint/annotate.py +++ /dev/null @@ -1,104 +0,0 @@ -from abc import ABC, abstractmethod -from logging import warn -from typing import List, Optional, Tuple - -import cv2 -import numpy as np - -from supervision.annotators.base import ImageType -from supervision.draw.color import Color -from supervision.keypoint.core import KeyPoints -from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT -from supervision.utils.conversion import convert_for_annotation_method - - -class BaseKeyPointAnnotator(ABC): - @abstractmethod - def annotate(self, scene: ImageType, detections: KeyPoints) -> ImageType: - pass - - -class VertexAnnotator(BaseKeyPointAnnotator): - def __init__( - self, - color: Color = Color.ROBOFLOW, - radius: int = 4, - ) -> None: - """ - Most basic keypoint annotator. - - Args: - color (Color, optional): The color of the keypoint. - radius (int, optional): The radius of the keypoint. - """ - self.color = color - self.radius = radius - - @convert_for_annotation_method - def annotate(self, scene: ImageType, keypoints: KeyPoints) -> ImageType: - if len(keypoints) == 0: - return scene - - for xy in keypoints.xy: - for x, y in xy: - cv2.circle( - img=scene, - center=(int(x), int(y)), - radius=self.radius, - color=self.color.as_bgr(), - thickness=-1, - ) - - return scene - - -class EdgeAnnotator(BaseKeyPointAnnotator): - def __init__( - self, - color: Color = Color.ROBOFLOW, - thickness: int = 2, - edges: Optional[List[Tuple[int, int]]] = None, - ) -> None: - """ - Draw the lines between points of the image. - - Args: - color (Color, optional): The color of the lines. - thickness (int, optional): The thickness of the lines. - edge (Optional[List[Tuple[int, int]]]): The edges to draw. - If set to `None`, will attempt to select automatically. - """ - self.color = color - self.thickness = thickness - self.edges = edges - - @convert_for_annotation_method - def annotate(self, scene: ImageType, keypoints: KeyPoints) -> ImageType: - if len(keypoints) == 0: - return scene - - for xy in keypoints.xy: - edges = self.edges - if not edges: - edges = SKELETONS_BY_VERTEX_COUNT.get(len(xy)) - if not edges: - warn(f"No skeleton found with {len(xy)} vertices") - return scene - - for class_a, class_b in edges: - xy_a = xy[class_a - 1] - xy_b = xy[class_b - 1] - missing_a = np.allclose(xy_a, 0) - missing_b = np.allclose(xy_b, 0) - if missing_a or missing_b: - continue - - cv2.line( - img=scene, - pt1=(int(xy_a[0]), int(xy_a[1])), - pt2=(int(xy_b[0]), int(xy_b[1])), - color=self.color.as_bgr(), - thickness=self.thickness, - ) - - return scene diff --git a/supervision/keypoint/annotators.py b/supervision/keypoint/annotators.py new file mode 100644 index 00000000..5cadc377 --- /dev/null +++ b/supervision/keypoint/annotators.py @@ -0,0 +1,170 @@ +from abc import ABC, abstractmethod +from logging import warn +from typing import List, Optional, Tuple + +import cv2 +import numpy as np + +from supervision.annotators.base import ImageType +from supervision.draw.color import Color +from supervision.keypoint.core import KeyPoints +from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT +from supervision.utils.conversion import convert_for_annotation_method + + +class BaseKeyPointAnnotator(ABC): + @abstractmethod + def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType: + pass + + +class VertexAnnotator(BaseKeyPointAnnotator): + """ + A class that specializes in drawing skeleton vertices on images. It uses + specified key points to determine the locations where the vertices should be + drawn. + """ + + def __init__( + self, + color: Color = Color.ROBOFLOW, + radius: int = 4, + ) -> None: + """ + Args: + color (Color, optional): The color to use for annotating key points. + radius (int, optional): The radius of the circles used to represent the key + points. + """ + self.color = color + self.radius = radius + + @convert_for_annotation_method + def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType: + """ + Annotates the given scene with skeleton vertices based on the provided key + points. It draws circles at each key point location. + + Args: + scene (ImageType): The image where bounding boxes will be drawn. `ImageType` + is a flexible type, accepting either `numpy.ndarray` or + `PIL.Image.Image`. + key_points (KeyPoints): A collection of key points where each key point + consists of x and y coordinates. + + Returns: + The annotated image, matching the type of `scene` (`numpy.ndarray` + or `PIL.Image.Image`) + + Example: + ```python + import supervision as sv + + image = ... + key_points = sv.KeyPoints(...) + + vertex_annotator = sv.VertexAnnotator() + annotated_frame = vertex_annotator.annotate( + scene=image.copy(), + key_points=key_points + ) + ``` + """ + if len(key_points) == 0: + return scene + + for xy in key_points.xy: + for x, y in xy: + cv2.circle( + img=scene, + center=(int(x), int(y)), + radius=self.radius, + color=self.color.as_bgr(), + thickness=-1, + ) + + return scene + + +class EdgeAnnotator(BaseKeyPointAnnotator): + """ + A class that specializes in drawing skeleton edges on images using specified key + points. It connects key points with lines to form the skeleton structure. + """ + def __init__( + self, + color: Color = Color.ROBOFLOW, + thickness: int = 2, + edges: Optional[List[Tuple[int, int]]] = None, + ) -> None: + """ + Args: + color (Color, optional): The color to use for the edges. + thickness (int, optional): The thickness of the edges. + edges (Optional[List[Tuple[int, int]]]): The edges to draw. + If set to `None`, will attempt to select automatically. + """ + self.color = color + self.thickness = thickness + self.edges = edges + + @convert_for_annotation_method + def annotate(self, scene: ImageType, key_points: KeyPoints) -> ImageType: + """ + Annotates the given scene by drawing lines between specified key points to form + edges. + + Args: + scene (ImageType): The image where bounding boxes will be drawn. `ImageType` + is a flexible type, accepting either `numpy.ndarray` or + `PIL.Image.Image`. + key_points (KeyPoints): A collection of key points where each key point + consists of x and y coordinates. + + Returns: + Returns: + The annotated image, matching the type of `scene` (`numpy.ndarray` + or `PIL.Image.Image`) + + Example: + ```python + import supervision as sv + + image = ... + key_points = sv.KeyPoints(...) + + edge_annotator = sv.EdgeAnnotator() + annotated_frame = edge_annotator.annotate( + scene=image.copy(), + key_points=key_points + ) + ``` + """ + if len(key_points) == 0: + return scene + + for xy in key_points.xy: + edges = self.edges + if not edges: + edges = SKELETONS_BY_VERTEX_COUNT.get(len(xy)) + if not edges: + warn(f"No skeleton found with {len(xy)} vertices") + return scene + + for class_a, class_b in edges: + xy_a = xy[class_a - 1] + xy_b = xy[class_b - 1] + missing_a = np.allclose(xy_a, 0) + missing_b = np.allclose(xy_b, 0) + if missing_a or missing_b: + continue + + cv2.line( + img=scene, + pt1=(int(xy_a[0]), int(xy_a[1])), + pt2=(int(xy_b[0]), int(xy_b[1])), + color=self.color.as_bgr(), + thickness=self.thickness, + ) + + return scene diff --git a/supervision/keypoint/core.py b/supervision/keypoint/core.py index 6749f845..b44d3b7c 100644 --- a/supervision/keypoint/core.py +++ b/supervision/keypoint/core.py @@ -14,28 +14,25 @@ from supervision.validators import validate_keypoints_fields @dataclass class KeyPoints: """ - The `sv.KeyPoints` allows you to convert results from a variety of keypoint - keypoints models into a single, unified format. + 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. - ```python - import cv2 - import supervision as sv - from ultralytics import YOLO + === "Ultralytics" - image = cv2.imread() - model = YOLO('yolov8s-pose.pt') - result = model(image)[0] - keypoints = sv.KeyPoints.from_ultralytics(result) - ``` + Use [`sv.KeyPoints.from_ultralytics`](/keypoint/core/#supervision.keypoint.core.KeyPoints.from_ultralytics) method, which accepts model results. - !!! tip + ```python + import cv2 + import supervision as sv + from ultralytics import YOLO - In `sv.KeyPoints`, detection data is categorized into two main field types: - fixed and custom. The fixed fields include `xy`, `confidence`, - `class_id`. For any additional data requirements, custom - fields come into play, stored in the data field. These custom fields are easily - accessible using the `keypoints[]` syntax, providing flexibility - for diverse data handling needs. + image = cv2.imread() + model = YOLO('yolov8s-pose.pt') + result = model(image)[0] + key_points = sv.KeyPoints.from_ultralytics(result) + ``` Attributes: xy (np.ndarray): An array of shape `(n, 2)` containing @@ -47,14 +44,6 @@ class KeyPoints: data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional data where each key is a string representing the data type, and the value is either a NumPy array or a list of corresponding data. - - !!! warning - - The `data` field in the `sv.KeyPoints` class is currently in an experimental - phase. Please be aware that its API and functionality are subject to change in - future updates as we continue to refine and improve its capabilities. - We encourage users to experiment with this feature and provide feedback, but - also to be prepared for potential modifications in upcoming releases. """ xy: npt.NDArray[np.float32] @@ -232,7 +221,7 @@ class KeyPoints: Create an empty Keypoints object with no keypoints. Returns: - (Keypoints): An empty Keypoints object. + (KeyPoints): An empty Keypoints object. Example: ```python