Compare commits

...
29 Commits
Author SHA1 Message Date
SkalskiP d30cc180d0 👊 bump version before release 2023-03-14 13:34:46 +01:00
Piotr Skalski ace8301b12 Merge pull request #46 from roboflow/fix/041_index_error_on_polygon_zones
fix/041_index_error_on_polygon_zones
2023-03-14 13:09:22 +01:00
SkalskiP b91bc91897 🖤 make black happy 2023-03-14 12:41:12 +01:00
SkalskiP e04c531c05 Merge remote-tracking branch 'origin/fix/041_index_error_on_polygon_zones' into fix/041_index_error_on_polygon_zones 2023-03-14 12:36:33 +01:00
SkalskiP 63fd5a3f8a 🧪 ready for tests 2023-03-14 12:36:06 +01:00
SkalskiP 38f52a04ae 🧪 ready for tests 2023-03-14 11:33:31 +01:00
Piotr Skalski 65090f612c Merge pull request #45 from roboflow/fix/042_fix_with_nms_when_detections_is_empty
🧪 ready for tests
2023-03-13 22:08:15 +01:00
SkalskiP 9401210aa5 🧪 ready for tests 2023-03-13 21:45:12 +01:00
Piotr Skalski 59d8655014 Merge pull request #36 from roboflow/feature/add_class_agnostic_nms
feature/add_class_agnostic_nms
2023-03-10 13:44:11 +01:00
SkalskiP 4dde484141 🖤 make black happy 2023-03-08 11:31:18 +01:00
SkalskiP e6f250dbfb 🧹 docs refinement and clean up 2023-03-08 11:28:54 +01:00
SkalskiP e6090be9eb 📄 small type changes in docs 2023-03-08 11:08:07 +01:00
SkalskiP 25c19b4422 👋 initial commit 2023-03-07 22:35:09 +01:00
Piotr Skalski ac1658268c Merge pull request #32 from roboflow/feature/detections_api_refinement
feature/detections_api_refinement
2023-03-01 12:38:19 +01:00
SkalskiP 8e7228313b make black happy 🖤 2023-03-01 12:36:15 +01:00
SkalskiP a27377aaf3 add area property to Detections 2023-02-28 15:27:08 +01:00
SkalskiP eefc5ddf7d fix redundant assert 2023-02-27 22:09:41 +01:00
SkalskiP 96d8badf1a 🥥 add from_coco_annotations support 2023-02-27 17:35:07 +01:00
SkalskiP c0ae62a460 making score/confidence optional to allow for using Detections to process labels 2023-02-27 15:12:33 +01:00
SkalskiP 1c2d946527 nms support 2023-02-27 12:35:07 +01:00
SkalskiP 95b44db2bd from_transformers added 2023-02-27 11:50:17 +01:00
SkalskiP ba732330a5 🫵 allow integer array indexing 2023-02-27 11:43:01 +01:00
Piotr Skalski 3ea95b3966 Merge pull request #31 from dceejay/main
Update __init__.py to include linezone
2023-02-24 17:00:38 +01:00
Dave Conway-Jones c2a5161ff9 Update __init__.py
Add export of linezone
2023-02-24 14:32:51 +00:00
SkalskiP c3eec02b6c 👊 bump version from 0.2.0 to 0.2.1 2023-02-12 23:08:04 +01:00
Piotr Skalski 8c68999324 Merge pull request #23 from roboflow/fix/22-uses-deprecated-alias-bool
🛠️ fix `core.py uses deprecated alias np.bool` #22 + tests
2023-02-12 23:05:37 +01:00
SkalskiP 4a43f4ec69 🖤 make black happy 2023-02-12 23:01:40 +01:00
SkalskiP 78d267e169 🛠️ fix core.py uses deprecated alias np.bool #22 + tests 2023-02-12 22:54:54 +01:00
Piotr Skalski 7780308533 Update README.md 2023-02-09 23:08:41 +01:00
15 changed files with 586 additions and 106 deletions
+12 -1
View File
@@ -7,6 +7,7 @@
>
</a>
</p>
<br>
<div align="center">
@@ -53,6 +54,10 @@
</a>
</div>
<br>
[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
</div>
## 👋 hello
@@ -97,9 +102,15 @@ Curious how Supervision can help you solve problems on your project? Visit our [
Learn how to use YOLOv8, ByteTrack and **Supervision** to detect, track and count objects. 🔥
[Subscribe](https://www.youtube.com/@Roboflow), and stay up to date with our latest YouTube videos!
<p align="center">
<a href="https://youtu.be/l_kf9CfZ_8M">
<img src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg" alt="count-objects-in-zone">
</a>
</p>
<p align="center">
<a href="https://youtu.be/OS5qI9YBkfk">
<img src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" alt="latest-roboflow-tutorial">
<img src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg" alt="count-objects-crossing-line">
</a>
</p>
+9 -1
View File
@@ -1,3 +1,11 @@
## generate_2d_mask
:::supervision.detection.utils.generate_2d_mask
:::supervision.detection.utils.generate_2d_mask
## box_iou_batch
:::supervision.detection.utils.box_iou_batch
## non_max_suppression
:::supervision.detection.utils.non_max_suppression
-7
View File
@@ -1,7 +0,0 @@
## draw_line
:::supervision.draw.utils.draw_line
## draw_rectangle
:::supervision.draw.utils.draw_rectangle
+19
View File
@@ -0,0 +1,19 @@
## draw_line
:::supervision.draw.utils.draw_line
## draw_rectangle
:::supervision.draw.utils.draw_rectangle
## draw_filled_rectangle
:::supervision.draw.utils.draw_filled_rectangle
## draw_polygon
:::supervision.draw.utils.draw_polygon
## draw_text
:::supervision.draw.utils.draw_text
View File
+2 -2
View File
@@ -28,8 +28,8 @@ nav:
- Detection:
- Core: detection_core.md
- Utils: detection_utils.md
- Draw: draw.md
- Geometry: geometry.md
- Draw:
- Utils: draw_utils.md
- Notebook: notebook.md
theme:
-30
View File
@@ -1,30 +0,0 @@
certifi==2022.12.7
charset-normalizer==3.0.1
click==8.1.3
colorama==0.4.6
ghp-import==2.1.0
griffe==0.25.4
idna==3.4
importlib-metadata==6.0.0
Jinja2==3.1.2
Markdown==3.3.7
MarkupSafe==2.1.2
mergedeep==1.3.4
mkdocs==1.4.2
mkdocs-autorefs==0.4.1
mkdocs-material==9.0.9
mkdocs-material-extensions==1.1.1
mkdocstrings==0.20.0
mkdocstrings-python==0.8.3
packaging==23.0
Pygments==2.14.0
pymdown-extensions==9.9.2
python-dateutil==2.8.2
PyYAML==6.0
pyyaml_env_tag==0.1
regex==2022.10.31
requests==2.28.2
six==1.16.0
urllib3==1.26.14
watchdog==2.2.1
zipp==3.12.0
+1 -1
View File
@@ -24,7 +24,7 @@ setuptools.setup(
long_description_content_type='text/markdown',
url='https://github.com/roboflow/supervision',
install_requires=[
'numpy',
'numpy>=1.20.0',
'opencv-python',
'matplotlib'
],
+2 -1
View File
@@ -1,6 +1,7 @@
__version__ = "0.2.0"
__version__ = "0.3.1"
from supervision.detection.core import BoxAnnotator, Detections
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator
from supervision.detection.utils import generate_2d_mask
from supervision.draw.color import Color, ColorPalette
+130 -41
View File
@@ -1,11 +1,12 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import List, Optional, Union
from typing import Iterator, List, Optional, Tuple, Union
import cv2
import numpy as np
from supervision.detection.utils import non_max_suppression
from supervision.draw.color import Color, ColorPalette
from supervision.geometry.core import Position
@@ -17,22 +18,26 @@ class Detections:
Attributes:
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`
confidence (np.ndarray): An array of shape `(n,)` containing the confidence scores of the detections.
confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores of the detections.
class_id (np.ndarray): An array of shape `(n,)` containing the class ids of the detections.
tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections.
"""
xyxy: np.ndarray
confidence: np.ndarray
class_id: np.ndarray
confidence: Optional[np.ndarray] = None
tracker_id: Optional[np.ndarray] = None
def __post_init__(self):
n = len(self.xyxy)
validators = [
(isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)),
(isinstance(self.confidence, np.ndarray) and self.confidence.shape == (n,)),
(isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)),
self.confidence is None
or (
isinstance(self.confidence, np.ndarray)
and self.confidence.shape == (n,)
),
self.tracker_id is None
or (
isinstance(self.tracker_id, np.ndarray)
@@ -42,7 +47,7 @@ class Detections:
if not all(validators):
raise ValueError(
"xyxy must be 2d np.ndarray with (n, 4) shape, "
"confidence must be 1d np.ndarray with (n,) shape, "
"confidence must be None or 1d np.ndarray with (n,) shape, "
"class_id must be 1d np.ndarray with (n,) shape, "
"tracker_id must be None or 1d np.ndarray with (n,) shape"
)
@@ -53,27 +58,50 @@ class Detections:
"""
return len(self.xyxy)
def __iter__(self):
def __iter__(
self,
) -> Iterator[Tuple[np.ndarray, Optional[float], int, Optional[Union[str, int]]]]:
"""
Iterates over the Detections object and yield a tuple of `(xyxy, confidence, class_id, tracker_id)` for each detection.
"""
for i in range(len(self.xyxy)):
yield (
self.xyxy[i],
self.confidence[i],
self.confidence[i] if self.confidence is not None else None,
self.class_id[i],
self.tracker_id[i] if self.tracker_id is not None else None,
)
def __eq__(self, other: Detections):
return all(
[
np.array_equal(self.xyxy, other.xyxy),
any(
[
self.confidence is None and other.confidence is None,
np.array_equal(self.confidence, other.confidence),
]
),
np.array_equal(self.class_id, other.class_id),
any(
[
self.tracker_id is None and other.tracker_id is None,
np.array_equal(self.tracker_id, other.tracker_id),
]
),
]
)
@classmethod
def from_yolov5(cls, yolov5_detections):
def from_yolov5(cls, yolov5_results) -> Detections:
"""
Creates a Detections instance from a YOLOv5 output Detections
Attributes:
yolov5_detections (yolov5.models.common.Detections): The output Detections instance from YOLOv5
Args:
yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5
Returns:
Detections: A new Detections object.
Example:
```python
@@ -85,7 +113,7 @@ class Detections:
>>> detections = Detections.from_yolov5(results)
```
"""
yolov5_detections_predictions = yolov5_detections.pred[0].cpu().cpu().numpy()
yolov5_detections_predictions = yolov5_results.pred[0].cpu().cpu().numpy()
return cls(
xyxy=yolov5_detections_predictions[:, :4],
confidence=yolov5_detections_predictions[:, 4],
@@ -93,14 +121,15 @@ class Detections:
)
@classmethod
def from_yolov8(cls, yolov8_results):
def from_yolov8(cls, yolov8_results) -> Detections:
"""
Creates a Detections instance from a YOLOv8 output Results
Attributes:
Args:
yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
Returns:
Detections: A new Detections object.
Example:
```python
@@ -108,7 +137,7 @@ class Detections:
>>> from supervision import Detections
>>> model = YOLO('yolov8s.pt')
>>> results = model(frame)
>>> results = model(frame)[0]
>>> detections = Detections.from_yolov8(results)
```
"""
@@ -118,40 +147,47 @@ class Detections:
class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
)
def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]:
@classmethod
def from_transformers(cls, transformers_results: dict) -> Detections:
"""
Filter the detections by applying a mask.
Attributes:
mask (np.ndarray): A mask of shape `(n,)` containing a boolean value for each detection indicating if it should be included in the filtered detections
inplace (bool): If True, the original data will be modified and self will be returned.
Creates a Detections instance from Object Detection Transformer output Results
Returns:
Optional[np.ndarray]: A new instance of Detections with the filtered detections, if inplace is set to `False`. `None` otherwise.
Detections: A new Detections object.
"""
if inplace:
self.xyxy = self.xyxy[mask]
self.confidence = self.confidence[mask]
self.class_id = self.class_id[mask]
self.tracker_id = (
self.tracker_id[mask] if self.tracker_id is not None else None
)
return self
else:
return Detections(
xyxy=self.xyxy[mask],
confidence=self.confidence[mask],
class_id=self.class_id[mask],
tracker_id=self.tracker_id[mask]
if self.tracker_id is not None
else None,
)
return cls(
xyxy=transformers_results["boxes"].cpu().numpy(),
confidence=transformers_results["scores"].cpu().numpy(),
class_id=transformers_results["labels"].cpu().numpy().astype(int),
)
@classmethod
def from_detectron2(cls, detectron2_results) -> Detections:
return cls(
xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
confidence=detectron2_results["instances"].scores.cpu().numpy(),
class_id=detectron2_results["instances"]
.pred_classes.cpu()
.numpy()
.astype(int),
)
@classmethod
def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
xyxy, class_id = [], []
for annotation in coco_annotation:
x_min, y_min, width, height = annotation["bbox"]
xyxy.append([x_min, y_min, x_min + width, y_min + height])
class_id.append(annotation["category_id"])
return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
def get_anchor_coordinates(self, anchor: Position) -> np.ndarray:
"""
Returns the bounding box coordinates for a specific anchor.
Properties:
Args:
anchor (Position): Position of bounding box anchor for which to return the coordinates.
Returns:
@@ -172,7 +208,9 @@ class Detections:
raise ValueError(f"{anchor} is not supported.")
def __getitem__(self, index: np.ndarray) -> Detections:
if isinstance(index, np.ndarray) and index.dtype == np.bool:
if isinstance(index, np.ndarray) and (
index.dtype == bool or index.dtype == int
):
return Detections(
xyxy=self.xyxy[index],
confidence=self.confidence[index],
@@ -185,6 +223,57 @@ class Detections:
f"Detections.__getitem__ not supported for index of type {type(index)}."
)
@property
def area(self) -> np.ndarray:
"""
Calculate the area of each bounding box in the set of object detections.
Returns:
np.ndarray: An array of floats containing the area of each bounding box in the format of (area_1, area_2, ..., area_n), where n is the number of detections.
"""
return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0])
def with_nms(
self, threshold: float = 0.5, class_agnostic: bool = False
) -> Detections:
"""
Perform non-maximum suppression on the current set of object detections.
Args:
threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression. 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.
Returns:
Detections: A new Detections object containing the subset of detections after non-maximum suppression.
Raises:
AssertionError: If `confidence` is None and class_agnostic is False. If `class_id` is None and class_agnostic is False.
"""
if len(self) == 0:
return self
assert (
self.confidence is not None
), f"Detections confidence must be given for NMS to be executed."
if class_agnostic:
predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
indices = non_max_suppression(
predictions=predictions, iou_threshold=threshold
)
return self[indices]
assert self.class_id is not None, (
f"Detections class_id must be given for NMS to be executed. If you intended to perform class agnostic "
f"NMS set class_agnostic=True."
)
predictions = np.hstack(
(self.xyxy, self.confidence.reshape(-1, 1), self.class_id.reshape(-1, 1))
)
indices = non_max_suppression(predictions=predictions, iou_threshold=threshold)
return self[indices]
class BoxAnnotator:
def __init__(
@@ -252,7 +341,7 @@ class BoxAnnotator:
continue
text = (
f"{confidence:0.2f}"
f"{class_id}"
if (labels is None or len(detections) != len(labels))
else labels[i]
)
+15 -9
View File
@@ -1,10 +1,11 @@
from dataclasses import replace
from typing import Optional, Tuple
import cv2
import numpy as np
from supervision import Detections
from supervision.detection.utils import generate_2d_mask
from supervision.detection.utils import clip_boxes, generate_2d_mask
from supervision.draw.color import Color
from supervision.draw.utils import draw_polygon, draw_text
from supervision.geometry.core import Position
@@ -21,17 +22,22 @@ class PolygonZone:
self.polygon = polygon
self.frame_resolution_wh = frame_resolution_wh
self.triggering_position = triggering_position
self.mask = generate_2d_mask(polygon=polygon, resolution_wh=frame_resolution_wh)
self.current_count = 0
def trigger(self, detections: Detections) -> np.ndarray:
anchors = (
np.ceil(
detections.get_anchor_coordinates(anchor=self.triggering_position)
).astype(int)
- 1
width, height = frame_resolution_wh
self.mask = generate_2d_mask(
polygon=polygon, resolution_wh=(width + 1, height + 1)
)
is_in_zone = self.mask[anchors[:, 1], anchors[:, 0]]
def trigger(self, detections: Detections) -> np.ndarray:
clipped_xyxy = clip_boxes(
boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh
)
clipped_detections = replace(detections, xyxy=clipped_xyxy)
clipped_anchors = np.ceil(
clipped_detections.get_anchor_coordinates(anchor=self.triggering_position)
).astype(int)
is_in_zone = self.mask[clipped_anchors[:, 1], clipped_anchors[:, 0]]
self.current_count = np.sum(is_in_zone)
return is_in_zone.astype(bool)
+98 -1
View File
@@ -7,7 +7,7 @@ import numpy as np
def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
"""Generate a 2D mask from a polygon.
Properties:
Args:
polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices.
resolution_wh (Tuple[int, int]): The width and height of the desired resolution.
@@ -18,3 +18,100 @@ def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.
mask = np.zeros((height, width), dtype=np.uint8)
cv2.fillPoly(mask, [polygon], color=1)
return mask
def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray:
"""
Compute Intersection over Union (IoU) of two sets of bounding boxes - `boxes_true` and `boxes_detection`. Both sets
of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
Args:
boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. `shape = (N, 4)` where `N` is number of true objects.
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. `shape = (M, 4)` where `M` is number of detected objects.
Returns:
np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`. `shape = (N, M)` where `N` is number of true objects and `M` is number of detected objects.
"""
def box_area(box):
return (box[2] - box[0]) * (box[3] - box[1])
area_true = box_area(boxes_true.T)
area_detection = box_area(boxes_detection.T)
top_left = np.maximum(boxes_true[:, None, :2], boxes_detection[:, :2])
bottom_right = np.minimum(boxes_true[:, None, 2:], boxes_detection[:, 2:])
area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)
return area_inter / (area_true[:, None] + area_detection - area_inter)
def non_max_suppression(
predictions: np.ndarray, iou_threshold: float = 0.5
) -> np.ndarray:
"""
Perform Non-Maximum Suppression (NMS) on object detection predictions.
Args:
predictions (np.ndarray): An array of object detection predictions in the format of `(x_min, y_min, x_max, y_max, score)` or `(x_min, y_min, x_max, y_max, score, class)`.
iou_threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression.
Returns:
np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.
Raises:
AssertionError: If `iou_threshold` is not within the closed range from `0` to `1`.
"""
assert 0 <= iou_threshold <= 1, (
f"Value of `iou_threshold` must be in the closed range from 0 to 1, "
f"{iou_threshold} given."
)
rows, columns = predictions.shape
# add column #5 - category filled with zeros for agnostic nms
if columns == 5:
predictions = np.c_[predictions, np.zeros(rows)]
# sort predictions column #4 - score
sort_index = np.flip(predictions[:, 4].argsort())
predictions = predictions[sort_index]
boxes = predictions[:, :4]
categories = predictions[:, 5]
ious = box_iou_batch(boxes, boxes)
ious = ious - np.eye(rows)
keep = np.ones(rows, dtype=bool)
for index, (iou, category) in enumerate(zip(ious, categories)):
if not keep[index]:
continue
# drop detections with iou > iou_threshold and same category as current detections
condition = (iou > iou_threshold) & (categories == category)
keep = keep & ~condition
return keep[sort_index.argsort()]
def clip_boxes(
boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int]
) -> np.ndarray:
"""
Clips bounding boxes coordinates to fit within the frame resolution.
Args:
boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each row corresponds to a bounding box in
the format `(x_min, y_min, x_max, y_max)`.
frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the
frame.
Returns:
np.ndarray: A numpy array of shape `(N, 4)` where each row corresponds to a bounding box with coordinates
clipped to fit within the frame resolution.
"""
result = np.copy(boxes_xyxy)
width, height = frame_resolution_wh
result[:, [0, 2]] = result[:, [0, 2]].clip(0, width)
result[:, [1, 3]] = result[:, [1, 3]].clip(0, height)
return result
+5 -12
View File
@@ -39,7 +39,7 @@ def draw_rectangle(
"""
Draws a rectangle on an image.
Attributes:
Parameters:
scene (np.ndarray): The scene on which the rectangle will be drawn
rect (Rect): The rectangle to be drawn
color (Color): The color of the rectangle
@@ -60,14 +60,9 @@ def draw_rectangle(
def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.ndarray:
"""
Draws a filled rectangle on the given scene.
Draws a filled rectangle on an image.
:param scene: np.ndarray : The scene on which to draw the rectangle.
:param rect: Rect : The rectangle to be drawn.
:param color: Color : The color of the rectangle.
:return: np.ndarray : The updated scene with the filled rectangle drawn on it.
Attributes:
Parameters:
scene (np.ndarray): The scene on which the rectangle will be drawn
rect (Rect): The rectangle to be drawn
color (Color): The color of the rectangle
@@ -90,7 +85,7 @@ def draw_polygon(
) -> np.ndarray:
"""Draw a polygon on a scene.
Attributes:
Parameters:
scene (np.ndarray): The scene to draw the polygon on.
polygon (np.ndarray): The polygon to be drawn, given as a list of vertices.
color (Color): The color of the polygon.
@@ -117,9 +112,7 @@ def draw_text(
background_color: Optional[Color] = None,
) -> np.ndarray:
"""
Draw text on a scene.
This function takes in a 2-dimensional numpy ndarray representing an image or scene, and draws text on it using OpenCV's putText function. The text is anchored at a specified Point, and its appearance can be customized using arguments such as color, scale, and font. An optional background color and padding can be specified to draw a rectangle behind the text.
Draw text with background on a scene.
Parameters:
scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene.
+105
View File
@@ -0,0 +1,105 @@
from contextlib import ExitStack as DoesNotRaise
import pytest
from supervision import Detections
from typing import Optional, Union
import numpy as np
PREDICTIONS = np.array([
[ 2254, 906, 2447, 1353, 0.90538, 0],
[ 2049, 1133, 2226, 1371, 0.59002, 56],
[ 727, 1224, 838, 1601, 0.51119, 39],
[ 808, 1214, 910, 1564, 0.45287, 39],
[ 6, 52, 1131, 2133, 0.45057, 72],
[ 299, 1225, 512, 1663, 0.45029, 39],
[ 529, 874, 645, 945, 0.31101, 39],
[ 8, 47, 1935, 2135, 0.28192, 72],
[ 2265, 813, 2328, 901, 0.2714, 62]
], dtype=np.float32)
DETECTIONS = Detections(
xyxy=PREDICTIONS[:, :4],
confidence=PREDICTIONS[:, 4],
class_id=PREDICTIONS[:, 5].astype(int)
)
PREDICTIONS = np.array([
[ 2254, 906, 2447, 1353, 0.90538, 0],
[ 2049, 1133, 2226, 1371, 0.59002, 56],
[ 727, 1224, 838, 1601, 0.51119, 39],
], dtype=np.float32)
@pytest.mark.parametrize(
'detections, index, expected_result, exception',
[
(
DETECTIONS,
DETECTIONS.class_id == 0,
Detections(
xyxy=np.array([
[ 2254, 906, 2447, 1353]
], dtype=np.float32),
confidence=np.array([
0.90538
], dtype=np.float32),
class_id=np.array([
0
], dtype=int)
),
DoesNotRaise()
), # take only detections with class_id = 0
(
DETECTIONS,
DETECTIONS.confidence > 0.5,
Detections(
xyxy=np.array([
[ 2254, 906, 2447, 1353],
[ 2049, 1133, 2226, 1371],
[ 727, 1224, 838, 1601]
], dtype=np.float32),
confidence=np.array([
0.90538,
0.59002,
0.51119
], dtype=np.float32),
class_id=np.array([
0,
56,
39
], dtype=int)
),
DoesNotRaise()
), # take only detections with confidence > 0.5
(
DETECTIONS,
np.array([True, True, True, True, True, True, True, True, True], dtype=bool),
DETECTIONS,
DoesNotRaise()
), # take all detections
(
DETECTIONS,
np.array([False, False, False, False, False, False, False, False, False], dtype=bool),
Detections(
xyxy=np.empty((0, 4), dtype=np.float32),
confidence=np.array([], dtype=np.float32),
class_id=np.array([], dtype=int)
),
DoesNotRaise()
), # take no detections
]
)
def test_getitem(
detections: Detections,
index: Union[int, slice, np.ndarray],
expected_result: Optional[Detections],
exception: Exception
) -> None:
with exception:
result = detections[index]
assert result == expected_result
+188
View File
@@ -0,0 +1,188 @@
from contextlib import ExitStack as DoesNotRaise
from typing import Optional, Tuple
import pytest
import numpy as np
from supervision.detection.utils import non_max_suppression, clip_boxes
@pytest.mark.parametrize(
"predictions, iou_threshold, expected_result, exception",
[
(
np.empty(shape=(0, 5)),
0.5,
np.array([]),
DoesNotRaise()
), # single box with no category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8]
]),
0.5,
np.array([
True
]),
DoesNotRaise()
), # single box with no category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8, 0]
]),
0.5,
np.array([
True
]),
DoesNotRaise()
), # single box with category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8],
[15.0, 15.0, 40.0, 40.0, 0.9],
]),
0.5,
np.array([
False,
True
]),
DoesNotRaise()
), # two boxes with no category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
]),
0.5,
np.array([
True,
True
]),
DoesNotRaise()
), # two boxes with different category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
]),
0.5,
np.array([
False,
True
]),
DoesNotRaise()
), # two boxes with same category
(
np.array([
[0.0, 0.0, 30.0, 40.0, 0.8],
[5.0, 5.0, 35.0, 45.0, 0.9],
[10.0, 10.0, 40.0, 50.0, 0.85],
]),
0.5,
np.array([
False,
True,
False
]),
DoesNotRaise()
), # three boxes with no category
(
np.array([
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
]),
0.5,
np.array([
True,
True,
True
]),
DoesNotRaise()
), # three boxes with same category
(
np.array([
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
]),
0.5,
np.array([
False,
True,
True
]),
DoesNotRaise()
), # three boxes with different category
]
)
def test_non_max_suppression(
predictions: np.ndarray,
iou_threshold: float,
expected_result: Optional[np.ndarray],
exception: Exception
) -> None:
with exception:
result = non_max_suppression(predictions=predictions, iou_threshold=iou_threshold)
assert np.array_equal(result, expected_result)
@pytest.mark.parametrize(
"boxes_xyxy, frame_resolution_wh, expected_result",
[
(
np.empty(shape=(0, 4)),
(1280, 720),
np.empty(shape=(0, 4)),
),
(
np.array([
[1.0, 1.0, 1279.0, 719.0]
]),
(1280, 720),
np.array([
[1.0, 1.0, 1279.0, 719.0]
]),
),
(
np.array([
[-1.0, 1.0, 1279.0, 719.0]
]),
(1280, 720),
np.array([
[0.0, 1.0, 1279.0, 719.0]
]),
),
(
np.array([
[1.0, -1.0, 1279.0, 719.0]
]),
(1280, 720),
np.array([
[1.0, 0.0, 1279.0, 719.0]
]),
),
(
np.array([
[1.0, 1.0, 1281.0, 719.0]
]),
(1280, 720),
np.array([
[1.0, 1.0, 1280.0, 719.0]
]),
),
(
np.array([
[1.0, 1.0, 1279.0, 721.0]
]),
(1280, 720),
np.array([
[1.0, 1.0, 1279.0, 720.0]
]),
),
]
)
def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None:
result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh)
assert np.array_equal(result, expected_result)