From 64302f450c3f4949650cd8f5fd2ba62dccfcdb5a Mon Sep 17 00:00:00 2001 From: Onuralp SEZER Date: Tue, 23 Jan 2024 19:54:52 +0300 Subject: [PATCH 1/5] =?UTF-8?q?feat:=20=E2=9C=A8=20yolov8-obb=20support=20?= =?UTF-8?q?and=20OrientedBoxAnnotator=20added=20for=20obb?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Onuralp SEZER --- supervision/__init__.py | 1 + supervision/annotators/core.py | 78 ++++++++++++++++++++++++++++++++++ supervision/detection/core.py | 11 +++++ 3 files changed, 90 insertions(+) diff --git a/supervision/__init__.py b/supervision/__init__.py index da28b124..3370e13d 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -18,6 +18,7 @@ from supervision.annotators.core import ( HeatMapAnnotator, LabelAnnotator, MaskAnnotator, + OrientedBoxAnnotator, PercentageBarAnnotator, PixelateAnnotator, PolygonAnnotator, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 1c4a2527..2b3bebd3 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -92,6 +92,84 @@ class BoundingBoxAnnotator(BaseAnnotator): return scene +class OrientedBoxAnnotator(BaseAnnotator): + """ + A class for drawing oriented bounding boxes on an image using provided detections. + """ + + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + thickness: int = 2, + color_lookup: ColorLookup = ColorLookup.CLASS, + ): + """ + Args: + color (Union[Color, ColorPalette]): The color or color palette to use for + annotating detections. + thickness (int): Thickness of the bounding box lines. + color_lookup (str): Strategy for mapping colors to annotations. + Options are `INDEX`, `CLASS`, `TRACK`. + """ + self.color: Union[Color, ColorPalette] = color + self.thickness: int = thickness + self.color_lookup: ColorLookup = color_lookup + + def annotate( + self, + scene: np.ndarray, + detections: Detections, + custom_color_lookup: Optional[np.ndarray] = None, + ) -> np.ndarray: + """ + Annotates the given scene with oriented bounding boxes based on the provided detections. + + Args: + scene (np.ndarray): The image where bounding boxes will be drawn. + detections (Detections): Object detections to annotate. + custom_color_lookup (Optional[np.ndarray]): Custom color lookup array. + Allows to override the default color mapping strategy. + + Returns: + The annotated image. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> oriented_box_annotator = sv.OrientedBoxAnnotator() + >>> annotated_frame = oriented_box_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + + ![oriented-box-annotator-example](https://media.roboflow.com/ + supervision-annotator-examples/oriented-box-annotator-example-purple.png) + """ # noqa E501 // docs + + if detections.data is None or "xyxyxyxy" not in detections.data: + return scene + + for detection_idx in range(len(detections)): + bbox = np.int0(detections.data.get("xyxyxyxy")[detection_idx]) + color = resolve_color( + color=self.color, + detections=detections, + detection_idx=detection_idx, + color_lookup=self.color_lookup + if custom_color_lookup is None + else custom_color_lookup, + ) + + cv2.drawContours(scene, [bbox], 0, color.as_bgr(), self.thickness) + + return scene + + class MaskAnnotator(BaseAnnotator): """ A class for drawing masks on an image using provided detections. diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 39e129e1..53c1e46c 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -170,6 +170,17 @@ class Detections: ``` """ + if ultralytics_results.obb is not None: + 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), + tracker_id=ultralytics_results.obb.id.int().cpu().numpy() + if ultralytics_results.obb.id is not None + else None, + ) + return cls( xyxy=ultralytics_results.boxes.xyxy.cpu().numpy(), confidence=ultralytics_results.boxes.conf.cpu().numpy(), From 10f3b73e79849fb5055f6f13117bffd888e28087 Mon Sep 17 00:00:00 2001 From: Onuralp SEZER Date: Tue, 23 Jan 2024 21:07:09 +0300 Subject: [PATCH 2/5] =?UTF-8?q?docs:=20=F0=9F=93=9D=20OrientedBox=20Annota?= =?UTF-8?q?tor=20docs=20added?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Onuralp SEZER --- docs/annotators.md | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/docs/annotators.md b/docs/annotators.md index f7dc8ece..0c521f1e 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -65,6 +65,27 @@ comments: true +=== "OrientedBox" + + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> oriented_box_annotator = sv.OrientedBoxAnnotator() + >>> annotated_frame = oriented_box_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + +
+ + ![oriented-box-annotator-example](https://media.roboflow.com/supervision-annotator-examples/oriented-box-annotator-example-purple.png){ align=center width="800" } + +
+ === "Color" ```python @@ -389,6 +410,10 @@ comments: true :::supervision.annotators.core.BoxCornerAnnotator +## OrientedBoxAnnotator + +:::supervision.annotators.core.OrientedBoxAnnotator + ## ColorAnnotator :::supervision.annotators.core.ColorAnnotator From a5bef3b15e245f4ae4f62bc467c55ab27c58b3d5 Mon Sep 17 00:00:00 2001 From: Onuralp SEZER Date: Wed, 24 Jan 2024 19:31:46 +0300 Subject: [PATCH 3/5] =?UTF-8?q?docs:=20=F0=9F=93=9D=20>>>/...=20removed=20?= =?UTF-8?q?from=20obb=20docs=20and=20code=20example?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Onuralp SEZER --- docs/annotators.md | 16 ++++++++-------- supervision/annotators/core.py | 16 ++++++++-------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/docs/annotators.md b/docs/annotators.md index 796f1b24..4b4df152 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -69,16 +69,16 @@ status: new === "OrientedBox" ```python - >>> import supervision as sv + import supervision as sv - >>> image = ... - >>> detections = sv.Detections(...) + image = ... + detections = sv.Detections(...) - >>> oriented_box_annotator = sv.OrientedBoxAnnotator() - >>> annotated_frame = oriented_box_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) + oriented_box_annotator = sv.OrientedBoxAnnotator() + annotated_frame = oriented_box_annotator.annotate( + scene=image.copy(), + detections=detections + ) ```
diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 5b206e35..7b11c7a6 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -135,16 +135,16 @@ class OrientedBoxAnnotator(BaseAnnotator): Example: ```python - >>> import supervision as sv + import supervision as sv - >>> image = ... - >>> detections = sv.Detections(...) + image = ... + detections = sv.Detections(...) - >>> oriented_box_annotator = sv.OrientedBoxAnnotator() - >>> annotated_frame = oriented_box_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) + oriented_box_annotator = sv.OrientedBoxAnnotator() + annotated_frame = oriented_box_annotator.annotate( + scene=image.copy(), + detections=detections + ) ``` ![oriented-box-annotator-example](https://media.roboflow.com/ From 6565fae1c2aac40a9c1745cb42d017e8db3e945d Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Wed, 24 Jan 2024 17:47:11 +0100 Subject: [PATCH 4/5] final docs updates --- docs/annotators.md | 21 --------- docs/index.md | 11 ++++- supervision/annotators/core.py | 3 -- supervision/detection/core.py | 60 +++++++++++++++--------- supervision/tracker/byte_tracker/core.py | 43 ++++++++++------- 5 files changed, 72 insertions(+), 66 deletions(-) diff --git a/docs/annotators.md b/docs/annotators.md index 796f1b24..f10944f1 100644 --- a/docs/annotators.md +++ b/docs/annotators.md @@ -66,27 +66,6 @@ status: new
-=== "OrientedBox" - - ```python - >>> import supervision as sv - - >>> image = ... - >>> detections = sv.Detections(...) - - >>> oriented_box_annotator = sv.OrientedBoxAnnotator() - >>> annotated_frame = oriented_box_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - -
- - ![oriented-box-annotator-example](https://media.roboflow.com/supervision-annotator-examples/oriented-box-annotator-example-purple.png){ align=center width="800" } - -
- === "Color" ```python diff --git a/docs/index.md b/docs/index.md index 2e04331d..7f2e25ad 100644 --- a/docs/index.md +++ b/docs/index.md @@ -7,7 +7,7 @@ comments: true

@@ -17,6 +17,15 @@ comments: true We write your reusable computer vision tools. Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us! + + +## 🚀 Quickstart +
- __Detect and Annotate__ diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 5b206e35..2b2ca6ba 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -146,9 +146,6 @@ class OrientedBoxAnnotator(BaseAnnotator): ... detections=detections ... ) ``` - - ![oriented-box-annotator-example](https://media.roboflow.com/ - supervision-annotator-examples/oriented-box-annotator-example-purple.png) """ # noqa E501 // docs if detections.data is None or "xyxyxyxy" not in detections.data: diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 44575253..b9914b99 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -130,7 +130,7 @@ class Detections: import torch import supervision as sv - image = cv2.imread(SOURCE_IMAGE_PATH) + image = cv2.imread() model = torch.hub.load('ultralytics/yolov5', 'yolov5s') result = model(image) detections = sv.Detections.from_yolov5(result) @@ -150,6 +150,13 @@ class Detections: Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result. + !!! Note + + `from_ultralytics` is compatible with + [detection](https://docs.ultralytics.com/tasks/detect/), + [segmentation](https://docs.ultralytics.com/tasks/segment/), and + [OBB](https://docs.ultralytics.com/tasks/obb/) models. + Args: ultralytics_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8 @@ -163,12 +170,13 @@ class Detections: import supervision as sv from ultralytics import YOLO - image = cv2.imread() + image = cv2.imread() model = YOLO('yolov8s.pt') + result = model(image)[0] detections = sv.Detections.from_ultralytics(result) ``` - """ + """ # noqa: E501 // docs if ultralytics_results.obb is not None: return cls( @@ -213,8 +221,9 @@ class Detections: from super_gradients.training import models import supervision as sv - image = cv2.imread(SOURCE_IMAGE_PATH) + image = cv2.imread() model = models.get('yolo_nas_l', pretrained_weights="coco") + result = list(model.predict(image, conf=0.35))[0] detections = sv.Detections.from_yolo_nas(result) ``` @@ -309,9 +318,9 @@ class Detections: @classmethod def from_mmdetection(cls, mmdet_results) -> Detections: """ - Creates a Detections instance from - a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result. - Also supported for [mmyolo](https://github.com/open-mmlab/mmyolo) + Creates a Detections instance from a + [mmdetection](https://github.com/open-mmlab/mmdetection) and + [mmyolo](https://github.com/open-mmlab/mmyolo) inference result. Args: mmdet_results (mmdet.structures.DetDataSample): @@ -324,14 +333,15 @@ class Detections: ```python import cv2 import supervision as sv - from mmdet.apis import DetInferencer + from mmdet.apis import init_detector, inference_detector + + image = cv2.imread() + model = init_detector(, , device=) - inferencer = DetInferencer(model_name, checkpoint, device) - mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output', - return_datasamples=True)["predictions"][0] - detections = sv.Detections.from_mmdetection(mmdet_result) + result = inference_detector(model, image) + detections = sv.Detections.from_mmdetection(result) ``` - """ + """ # noqa: E501 // docs return cls( xyxy=mmdet_results.pred_instances.bboxes.cpu().numpy(), @@ -372,15 +382,17 @@ class Detections: Example: ```python import cv2 + import supervision as sv from detectron2.engine import DefaultPredictor from detectron2.config import get_cfg - import supervision as sv - image = cv2.imread(SOURCE_IMAGE_PATH) + + image = cv2.imread() cfg = get_cfg() - cfg.merge_from_file("path/to/config.yaml") - cfg.MODEL.WEIGHTS = "path/to/model_weights.pth" + cfg.merge_from_file() + cfg.MODEL.WEIGHTS = predictor = DefaultPredictor(cfg) + result = predictor(image) detections = sv.Detections.from_detectron2(result) ``` @@ -423,8 +435,9 @@ class Detections: import supervision as sv from inference.models.utils import get_roboflow_model - image = cv2.imread() + image = cv2.imread() model = get_roboflow_model(model_id="yolov8s-640") + result = model.infer(image)[0] detections = sv.Detections.from_inference(result) ``` @@ -472,8 +485,9 @@ class Detections: import supervision as sv from inference.models.utils import get_roboflow_model - image = cv2.imread() + image = cv2.imread() model = get_roboflow_model(model_id="yolov8s-640") + result = model.infer(image)[0] detections = sv.Detections.from_roboflow(result) ``` @@ -888,13 +902,12 @@ class Detections: Example: ```python import cv2 - from ultralytics import YOLO import supervision as sv + from ultralytics import YOLO + image = cv2.imread() model = YOLO('yolov8s.pt') - image = cv2.imread(SOURCE_IMAGE_PATH) - result = model(image)[0] detections = sv.Detections.from_ultralytics(result) @@ -950,7 +963,8 @@ class Detections: Args: threshold (float, optional): The intersection-over-union threshold - to use for non-maximum suppression. Defaults to 0.5. + to use for non-maximum suppression. I'm the lower the value the more + restrictive the NMS becomes. Defaults to 0.5. class_agnostic (bool, optional): Whether to perform class-agnostic non-maximum suppression. If True, the class_id of each detection will be ignored. Defaults to False. diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index bf831daa..8b74aa02 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -166,6 +166,10 @@ def detections2boxes(detections: Detections) -> np.ndarray: class ByteTrack: """ Initialize the ByteTrack object. + + Parameters: track_thresh (float, optional): Detection confidence threshold @@ -173,7 +177,7 @@ class ByteTrack: track_buffer (int, optional): Number of frames to buffer when a track is lost. match_thresh (float, optional): Threshold for matching tracks with detections. frame_rate (int, optional): The frame rate of the video. - """ + """ # noqa: E501 // docs def __init__( self, @@ -196,36 +200,39 @@ class ByteTrack: def update_with_detections(self, detections: Detections) -> Detections: """ - Updates the tracker with the provided detections and - returns the updated detection results. + Updates the tracker with the provided detections and returns the updated + detection results. + + Args: + detections (Detections): The detections to pass through the tracker. - Parameters: - detections: The new detections to update with. - Returns: - Detection: The updated detection results that now include tracking IDs. Example: ```python import supervision as sv from ultralytics import YOLO - model = YOLO(...) - byte_tracker = sv.ByteTrack() - annotator = sv.BoxAnnotator() + model = YOLO() + tracker = sv.ByteTrack() + + bounding_box_annotator = sv.BoundingBoxAnnotator() + label_annotator = sv.LabelAnnotator() def callback(frame: np.ndarray, index: int) -> np.ndarray: results = model(frame)[0] detections = sv.Detections.from_ultralytics(results) detections = byte_tracker.update_with_detections(detections) - labels = [ - f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}" - for _, _, confidence, class_id, tracker_id in detections - ] - return annotator.annotate(scene=frame.copy(), - detections=detections, labels=labels) + + labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id] + + annotated_frame = bounding_box_annotator.annotate( + scene=frame.copy(), detections=detections) + annotated_frame = label_annotator.annotate( + scene=annotated_frame, detections=detections, labels=labels) + return annotated_frame sv.process_video( - source_path='...', - target_path='...', + source_path=, + target_path=, callback=callback ) ``` From 9d56d8a9fe695a687d07ccc9300a9a3764d043bd Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Wed, 24 Jan 2024 16:47:54 +0000 Subject: [PATCH 5/5] =?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 | 12 ++++++------ supervision/tracker/byte_tracker/core.py | 2 +- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index b9914b99..fd2409d8 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -152,9 +152,9 @@ class Detections: !!! Note - `from_ultralytics` is compatible with - [detection](https://docs.ultralytics.com/tasks/detect/), - [segmentation](https://docs.ultralytics.com/tasks/segment/), and + `from_ultralytics` is compatible with + [detection](https://docs.ultralytics.com/tasks/detect/), + [segmentation](https://docs.ultralytics.com/tasks/segment/), and [OBB](https://docs.ultralytics.com/tasks/obb/) models. Args: @@ -318,8 +318,8 @@ class Detections: @classmethod def from_mmdetection(cls, mmdet_results) -> Detections: """ - Creates a Detections instance from a - [mmdetection](https://github.com/open-mmlab/mmdetection) and + Creates a Detections instance from a + [mmdetection](https://github.com/open-mmlab/mmdetection) and [mmyolo](https://github.com/open-mmlab/mmyolo) inference result. Args: @@ -334,7 +334,7 @@ class Detections: import cv2 import supervision as sv from mmdet.apis import init_detector, inference_detector - + image = cv2.imread() model = init_detector(, , device=) diff --git a/supervision/tracker/byte_tracker/core.py b/supervision/tracker/byte_tracker/core.py index 8b74aa02..9e3c0167 100644 --- a/supervision/tracker/byte_tracker/core.py +++ b/supervision/tracker/byte_tracker/core.py @@ -166,7 +166,7 @@ def detections2boxes(detections: Detections) -> np.ndarray: class ByteTrack: """ Initialize the ByteTrack object. - +