From 7f9153e826ffbd317f14d3ded13868272295a696 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 30 Jan 2024 18:34:41 +0100 Subject: [PATCH 1/3] initial commit --- supervision/annotators/core.py | 6 +++--- supervision/config.py | 1 + supervision/detection/core.py | 18 +++++++++++++++--- 3 files changed, 19 insertions(+), 6 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 870b26f2..e9cd28ae 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -6,7 +6,7 @@ import numpy as np from supervision.annotators.base import BaseAnnotator from supervision.annotators.utils import ColorLookup, Trace, resolve_color -from supervision.config import CLASS_NAME_DATA_FIELD +from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from supervision.detection.core import Detections from supervision.detection.utils import clip_boxes, mask_to_polygons from supervision.draw.color import Color, ColorPalette @@ -153,11 +153,11 @@ class OrientedBoxAnnotator(BaseAnnotator): ``` """ # noqa E501 // docs - if detections.data is None or "xyxyxyxy" not in detections.data: + if detections.data is None or ORIENTED_BOX_COORDINATES not in detections.data: return scene for detection_idx in range(len(detections)): - bbox = np.int0(detections.data.get("xyxyxyxy")[detection_idx]) + bbox = np.int0(detections.data.get(ORIENTED_BOX_COORDINATES)[detection_idx]) color = resolve_color( color=self.color, detections=detections, diff --git a/supervision/config.py b/supervision/config.py index d1a83b55..b18d2e20 100644 --- a/supervision/config.py +++ b/supervision/config.py @@ -1 +1,2 @@ CLASS_NAME_DATA_FIELD = "class_name" +ORIENTED_BOX_COORDINATES = "xyxyxyxy" diff --git a/supervision/detection/core.py b/supervision/detection/core.py index fd2409d8..62ed8869 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -5,6 +5,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, Iterator, List, Optional, Tuple, Union import numpy as np +from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from supervision.detection.utils import ( calculate_masks_centroids, @@ -179,24 +180,35 @@ class Detections: """ # noqa: E501 // docs if ultralytics_results.obb is not None: + class_id = ultralytics_results.obb.cls.cpu().numpy().astype(int) + class_names = np.array([ultralytics_results.names[i] for i in class_id]) + oriented_box_coordinates = ultralytics_results.obb.xyxyxyxy.cpu().numpy() return cls( xyxy=ultralytics_results.obb.xyxy.cpu().numpy(), - data={"xyxyxyxy": ultralytics_results.obb.xyxyxyxy.cpu().numpy()}, confidence=ultralytics_results.obb.conf.cpu().numpy(), - class_id=ultralytics_results.obb.cls.cpu().numpy().astype(int), + class_id=class_id, tracker_id=ultralytics_results.obb.id.int().cpu().numpy() if ultralytics_results.obb.id is not None else None, + data={ + ORIENTED_BOX_COORDINATES: oriented_box_coordinates, + CLASS_NAME_DATA_FIELD: class_names + }, ) + class_id = ultralytics_results.boxes.cls.cpu().numpy().astype(int) + class_names = np.array([ultralytics_results.names[i] for i in class_id]) return cls( xyxy=ultralytics_results.boxes.xyxy.cpu().numpy(), confidence=ultralytics_results.boxes.conf.cpu().numpy(), - class_id=ultralytics_results.boxes.cls.cpu().numpy().astype(int), + class_id=class_id, mask=extract_ultralytics_masks(ultralytics_results), tracker_id=ultralytics_results.boxes.id.int().cpu().numpy() if ultralytics_results.boxes.id is not None else None, + data={ + CLASS_NAME_DATA_FIELD: class_names + } ) @classmethod From 33b7861c7080e216c409066969718e50337b1dad Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 30 Jan 2024 17:36:03 +0000 Subject: [PATCH 2/3] =?UTF-8?q?fix(pre=5Fcommit):=20=F0=9F=8E=A8=20auto=20?= =?UTF-8?q?format=20pre-commit=20hooks?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 62ed8869..dc730989 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -5,8 +5,8 @@ from dataclasses import dataclass, field from typing import Any, Dict, Iterator, List, Optional, Tuple, Union import numpy as np -from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES +from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES from supervision.detection.utils import ( calculate_masks_centroids, extract_ultralytics_masks, @@ -192,7 +192,7 @@ class Detections: else None, data={ ORIENTED_BOX_COORDINATES: oriented_box_coordinates, - CLASS_NAME_DATA_FIELD: class_names + CLASS_NAME_DATA_FIELD: class_names, }, ) @@ -206,9 +206,7 @@ class Detections: tracker_id=ultralytics_results.boxes.id.int().cpu().numpy() if ultralytics_results.boxes.id is not None else None, - data={ - CLASS_NAME_DATA_FIELD: class_names - } + data={CLASS_NAME_DATA_FIELD: class_names}, ) @classmethod From c6718d374ca616fa716a11caf5a18884ee1dfa88 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 31 Jan 2024 08:28:34 +0100 Subject: [PATCH 3/3] update docs --- docs/index.md | 4 ++-- supervision/detection/core.py | 29 ++++++++++++++++++++++++++++- 2 files changed, 30 insertions(+), 3 deletions(-) diff --git a/docs/index.md b/docs/index.md index 7f2e25ad..5f838eb8 100644 --- a/docs/index.md +++ b/docs/index.md @@ -34,7 +34,7 @@ We write your reusable computer vision tools. Whether you need to load your data Annotate predictions from a range of object detection and segmentation models - [:octicons-arrow-right-24: Tutorial](how_to/detect_and_annotate) + [:octicons-arrow-right-24: Tutorial](how_to/detect_and_annotate.md) - __Track Objects__ @@ -42,7 +42,7 @@ We write your reusable computer vision tools. Whether you need to load your data Discover how to enhance video analysis by implementing seamless object tracking - [:octicons-arrow-right-24: Tutorial](how_to/track_objects) + [:octicons-arrow-right-24: Tutorial](how_to/track_objects.md) - > __Count Objects Crossing Line__ diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 62ed8869..a1432182 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -25,7 +25,34 @@ from supervision.utils.internal import deprecated @dataclass class Detections: """ - A dataclass representing detection results. + 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. + + ```python + import cv2 + import supervision as sv + from ultralytics import YOLO + + image = cv2.imread() + model = YOLO('yolov8s.pt') + annotator = sv.BoundingBoxAnnotator() + + result = model(image)[0] + detections = sv.Detections.from_ultralytics(result) + + annotated_image = annotator.annotate(image, detections) + ``` + + !!! tip + + 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. Attributes: xyxy (np.ndarray): An array of shape `(n, 4)` containing