okooooooo/supervision
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #770 from roboflow/yolov8-obb

Browse Source 
feat: ✨ yolov8-obb support and OrientedBoxAnnotator added for obb
This commit is contained in:
Piotr Skalski authored and GitHub committed 2024-01-24 17:50:41 +01:00
2 parents aabe60be6a 9d56d8a9fe
commit 35d5738622
6 files changed
+163 -42
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/annotators.md
+4
View File
@@ -390,6 +390,10 @@ status: new
:::supervision.annotators.core.BoxCornerAnnotator
## OrientedBoxAnnotator
:::supervision.annotators.core.OrientedBoxAnnotator
## ColorAnnotator
:::supervision.annotators.core.ColorAnnotator
docs/index.md
+10 -1
View File
@@ -7,7 +7,7 @@ comments: true
<a align="center" href="" target="_blank">
<img
width="850"
src="https://media.roboflow.com/open-source/supervision/roboflow-supervision-banner.png?ik-sdk-version=javascript-1.4.3&updatedAt=1674062891088"
src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529"
>
</a>
</p>
@@ -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!
<video controls>
<source
src="https://media.roboflow.com/traffic_analysis_result.mp4"
type="video/mp4"
>
</video>
## 🚀 Quickstart
<div class="grid cards" markdown>
- __Detect and Annotate__
supervision/__init__.py
+1
View File
@@ -18,6 +18,7 @@ from supervision.annotators.core import (
HeatMapAnnotator,
LabelAnnotator,
MaskAnnotator,
OrientedBoxAnnotator,
PercentageBarAnnotator,
PixelateAnnotator,
PolygonAnnotator,
supervision/annotators/core.py
+75
View File
@@ -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.
supervision/detection/core.py
+48 -23
View File
@@ -130,7 +130,7 @@ class Detections:
import torch
import supervision as sv
image = cv2.imread(SOURCE_IMAGE_PATH)
image = cv2.imread(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
model = init_detector(<CONFIG_PATH>, <WEIGHTS_PATH>, device=<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(<SOURCE_IMAGE_PATH>)
cfg = get_cfg()
cfg.merge_from_file("path/to/config.yaml")
cfg.MODEL.WEIGHTS = "path/to/model_weights.pth"
cfg.merge_from_file(<CONFIG_PATH>)
cfg.MODEL.WEIGHTS = <WEIGHTS_PATH>
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(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
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(<SOURCE_IMAGE_PATH>)
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.
supervision/tracker/byte_tracker/core.py
+25 -18
View File
@@ -167,13 +167,17 @@ class ByteTrack:
"""
Initialize the ByteTrack object.
<video controls>
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-traces.mp4" type="video/mp4">
</video>
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(<MODEL_PATH>)
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=<SOURCE_VIDEO_PATH>,
target_path=<TARGET_VIDEO_PATH>,
callback=callback
)
```