Skip to content

Commit

Permalink
Updated req files for 3D annotation
Browse files Browse the repository at this point in the history
  • Loading branch information
dsaha21 committed Oct 20, 2024
1 parent 8c5ca59 commit 03e18c3
Show file tree
Hide file tree
Showing 2 changed files with 88 additions and 9 deletions.
84 changes: 83 additions & 1 deletion supervision/annotators/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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`."
Expand Down
13 changes: 5 additions & 8 deletions supervision/draw/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
)
Expand Down

0 comments on commit 03e18c3

Please sign in to comment.