From 64302f450c3f4949650cd8f5fd2ba62dccfcdb5a Mon Sep 17 00:00:00 2001
From: Onuralp SEZER <thunderbirdtr@gmail.com>
Date: Tue, 23 Jan 2024 19:54:52 +0300
Subject: [PATCH] =?UTF-8?q?feat:=20=E2=9C=A8=20yolov8-obb=20support=20and?=
 =?UTF-8?q?=20OrientedBoxAnnotator=20added=20for=20obb?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Onuralp SEZER <thunderbirdtr@gmail.com>
---
 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(),