diff --git a/docs/annotators.md b/docs/annotators.md
index 19379770..f10944f1 100644
--- a/docs/annotators.md
+++ b/docs/annotators.md
@@ -390,6 +390,10 @@ status: new
:::supervision.annotators.core.BoxCornerAnnotator
+## OrientedBoxAnnotator
+
+:::supervision.annotators.core.OrientedBoxAnnotator
+
## ColorAnnotator
:::supervision.annotators.core.ColorAnnotator
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/__init__.py b/supervision/__init__.py
index 4a1776e1..d44f6723 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 93f92410..e021df53 100644
--- a/supervision/annotators/core.py
+++ b/supervision/annotators/core.py
@@ -92,6 +92,81 @@ 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
+ )
+ ```
+ """ # 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 6ade4bbf..fd2409d8 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,24 @@ 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(
+ 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(),
@@ -202,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)
```
@@ -298,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):
@@ -313,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
- 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)
+ image = cv2.imread()
+ model = init_detector(, , device=)
+
+ result = inference_detector(model, image)
+ detections = sv.Detections.from_mmdetection(result)
```
- """
+ """ # noqa: E501 // docs
return cls(
xyxy=mmdet_results.pred_instances.bboxes.cpu().numpy(),
@@ -361,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)
```
@@ -412,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)
```
@@ -461,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)
```
@@ -877,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)
@@ -939,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..9e3c0167 100644
--- a/supervision/tracker/byte_tracker/core.py
+++ b/supervision/tracker/byte_tracker/core.py
@@ -167,13 +167,17 @@ class ByteTrack:
"""
Initialize the ByteTrack object.
+
+
Parameters:
track_thresh (float, optional): Detection confidence threshold
for track activation.
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
)
```
