From 03e18c3da26dca5b62bfd42d23be09e691bb8302 Mon Sep 17 00:00:00 2001
From: Dripto Saha <sahadripto21@gmail.com>
Date: Sun, 20 Oct 2024 10:57:39 +0530
Subject: [PATCH] Updated req files for 3D annotation

---
 supervision/annotators/core.py | 84 +++++++++++++++++++++++++++++++++-
 supervision/draw/utils.py      | 13 ++----
 2 files changed, 88 insertions(+), 9 deletions(-)

diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py
index 1910ac9f4..ea529fd2c 100644
--- a/supervision/annotators/core.py
+++ b/supervision/annotators/core.py
@@ -18,7 +18,7 @@
 from supervision.detection.core import Detections
 from supervision.detection.utils import clip_boxes, mask_to_polygons
 from supervision.draw.color import Color, ColorPalette
-from supervision.draw.utils import draw_polygon
+from supervision.draw.utils import draw_polygon, draw_rectangle_3D
 from supervision.geometry.core import Position
 from supervision.utils.conversion import (
     ensure_cv2_image_for_annotation,
@@ -115,6 +115,88 @@ def annotate(
             )
         return scene
 
+class BoxAnnotator3D(BaseAnnotator):
+    """
+    A class for drawing 3D 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 (ColorLookup): 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
+
+    @ensure_cv2_image_for_annotation
+    def annotate(
+        self,
+        scene: ImageType,
+        detections: Detections,
+        custom_color_lookup: Optional[np.ndarray] = None,
+    ) -> ImageType:
+        """
+        Annotates the given scene with bounding boxes based on the provided detections.
+
+        Args:
+            scene (ImageType): The image where bounding boxes will be drawn. `ImageType`
+                is a flexible type, accepting either `numpy.ndarray` or
+                `PIL.Image.Image`.
+            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, matching the type of `scene` (`numpy.ndarray`
+                or `PIL.Image.Image`)
+
+        Example:
+            ```python
+            import supervision as sv
+
+            image = ...
+            detections = sv.Detections(...)
+
+            box_annotator = sv.BoxAnnotator()
+            annotated_frame = box_annotator.annotate(
+                scene=image.copy(),
+                detections=detections
+            )
+            ```
+
+        ![bounding-box-annotator-example](https://media.roboflow.com/
+        supervision-annotator-examples/bounding-box-annotator-example-purple.png)
+        """
+        assert isinstance(scene, np.ndarray)
+        for detection_idx in range(len(detections)):
+            x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
+            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.rectangle(
+            #     img=scene,
+            #     pt1=(x1, y1),
+            #     pt2=(x2, y2),
+            #     color=color.as_bgr(),
+            #     thickness=self.thickness,
+            # )
+            draw_rectangle_3D(scene, (x1,y1), (x2,y2), color=color, thickness=self.thickness)
+        return scene
 
 @deprecated(
     "`BoundingBoxAnnotator` is deprecated and has been renamed to `BoxAnnotator`."
diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py
index 83d989c0d..2f829871f 100644
--- a/supervision/draw/utils.py
+++ b/supervision/draw/utils.py
@@ -64,7 +64,7 @@ def draw_rectangle(
 
 
 def draw_rectangle_3D(
-    scene: np.ndarray, rect: Rect, color: Color, thickness: int = 2
+    scene: np.ndarray, pt1, pt2, color: Color, thickness : int
 ) -> np.ndarray:
     """
     Draws a cuboid by connecting two rectangles.
@@ -81,18 +81,15 @@ def draw_rectangle_3D(
     """
     rec1 = cv2.rectangle(
         scene,
-        rect.top_left.as_xy_int_tuple(),
-        rect.bottom_right.as_xy_int_tuple(),
+        pt1,
+        pt2,
         color.as_bgr(),
         -1,
     )
-    tlx,tly = rect.top_left.as_xy_int_tuple()
-    brx, bry = rect.bottom_right.as_xy_int_tuple()
-
     rec2 = cv2.rectangle(
         scene, 
-        (tlx+10, tly+10),
-        (brx+10, bry+10),
+        (pt1[0]+10, pt1[1]+10),
+        (pt2[0]+10, pt2[1]+10),
         color.as_bgr(),
         -1,
     )