Compare commits

...
36 Commits
Author SHA1 Message Date
Piotr Skalski 2df4261355 Merge pull request #50 from roboflow/feature/drop_requirement_for_class_id_in_detections
feature/drop_requirement_for_class_id_in_detections
2023-03-23 14:42:52 +01:00
SkalskiP a88eff868e 📖 improve documentation 2023-03-23 14:41:03 +01:00
SkalskiP 8bf686e23b 🖤 make black happy 2023-03-23 14:18:46 +01:00
SkalskiP b2f6520b87 🛠️ fix idx error in annotate 2023-03-23 13:43:03 +01:00
SkalskiP 420dae9e70 🛠️ fix circular import 2023-03-23 13:02:52 +01:00
SkalskiP 550c913dbb 💧 drop requirement for class_id in Detections constructor to make it more flexible 2023-03-23 12:31:23 +01:00
SkalskiP 4dcca5309c 🔥 docs and README.md changes 2023-03-21 13:43:40 +01:00
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
20 changed files with 750 additions and 239 deletions
+14 -1
View File
@@ -35,4 +35,17 @@ When creating new functions, please ensure you have the following:
All pull requests will be reviewed by the maintainers of the project. We will provide feedback and ask for changes if necessary.
PRs must pass all tests and linting requirements before they can be merged.
PRs must pass all tests and linting requirements before they can be merged.
## 🧹 code quality
We provide two handy commands inside the `Makefile`, namely:
- `make style` to format the code
- `make check_code_quality` to check code quality (PEP8 basically)
So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
## 🧪 tests
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
+28 -19
View File
@@ -2,11 +2,12 @@
<p>
<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"
width="100%"
src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529"
>
</a>
</p>
<br>
<div align="center">
@@ -53,13 +54,17 @@
</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
A set of easy-to-use utils that will come in handy in any Computer Vision project. **Supervision** is still in
A set of easy-to-use utils that will come in handy in any computer vision project. **Supervision is still in
pre-release stage. 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid
and may change.
and may change.**
## 💻 install
@@ -92,26 +97,30 @@ pip install -e ".[dev]"
Curious how Supervision can help you solve problems on your project? Visit our [documentation](https://roboflow.github.io/supervision) page!
## 🎬 videos
## 🎬 tutorials
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!
🔥 [Subscribe](https://www.youtube.com/@Roboflow) and stay up to date with the latest changes to the Supervision library.
<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">
<a href="https://youtu.be/l_kf9CfZ_8M">
<img
width="90%"
src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg"
alt="count-objects-in-zone"
>
</a>
</p>
## 🧹 code quality
<p align="center">
<a href="https://youtu.be/OS5qI9YBkfk">
<img
width="90%"
src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg"
alt="count-objects-crossing-line"
>
</a>
</p>
We provide two handy commands inside the `Makefile`, namely:
## 🏆 contribution
- `make style` to format the code
- `make check_code_quality` to check code quality (PEP8 basically)
So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
## 🧪 tests
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
We love your input! Please see our [contributing guide](https://github.com/roboflow/supervision/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!
+3
View File
@@ -0,0 +1,3 @@
## BoxAnnotator
:::supervision.detection.annotate.BoxAnnotator
+11
View File
@@ -0,0 +1,11 @@
## 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
-3
View File
@@ -1,3 +0,0 @@
## generate_2d_mask
:::supervision.detection.utils.generate_2d_mask
-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
+7 -6
View File
@@ -26,16 +26,17 @@ nav:
- API reference:
- Video: video.md
- Detection:
- Core: detection_core.md
- Utils: detection_utils.md
- Draw: draw.md
- Geometry: geometry.md
- Core: detection/core.md
- Annotate: detection/annotate.md
- Utils: detection/utils.md
- Draw:
- Utils: draw_utils.md
- Notebook: notebook.md
theme:
name: 'material'
logo: https://raw.githubusercontent.com/roboflow/supervision/main/docs/assets/roboflow_logomark_white.svg
favicon: https://raw.githubusercontent.com/roboflow/supervision/main/docs/assets/roboflow_logomark_color.svg
logo: https://media.roboflow.com/open-source/supervision/supervision-lenny.png?updatedAt=1678995918671
favicon: https://media.roboflow.com/open-source/supervision/supervision-lenny.png?updatedAt=1678995918671
palette:
primary: 'deep purple'
accent: 'teal'
-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'
],
+4 -2
View File
@@ -1,6 +1,8 @@
__version__ = "0.2.0"
__version__ = "0.3.2"
from supervision.detection.core import BoxAnnotator, Detections
from supervision.detection.annotate import BoxAnnotator
from supervision.detection.core import 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
+116
View File
@@ -0,0 +1,116 @@
from typing import List, Optional, Union
import cv2
import numpy as np
from supervision.detection.core import Detections
from supervision.draw.color import Color, ColorPalette
class BoxAnnotator:
"""
A class for drawing bounding boxes on an image using detections provided.
Attributes:
color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette
thickness (int): The thickness of the bounding box lines, default is 2
text_color (Color): The color of the text on the bounding box, default is white
text_scale (float): The scale of the text on the bounding box, default is 0.5
text_thickness (int): The thickness of the text on the bounding box, default is 1
text_padding (int): The padding around the text on the bounding box, default is 5
"""
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
text_color: Color = Color.black(),
text_scale: float = 0.5,
text_thickness: int = 1,
text_padding: int = 10,
):
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.text_color: Color = text_color
self.text_scale: float = text_scale
self.text_thickness: int = text_thickness
self.text_padding: int = text_padding
def annotate(
self,
scene: np.ndarray,
detections: Detections,
labels: Optional[List[str]] = None,
skip_label: bool = False,
) -> np.ndarray:
"""
Draws bounding boxes on the frame using the detections provided.
Parameters:
scene (np.ndarray): The image on which the bounding boxes will be drawn
detections (Detections): The detections for which the bounding boxes will be drawn
labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.
skip_label (bool): Is set to True, skips bounding box label annotation.
Returns:
np.ndarray: The image with the bounding boxes drawn on it
"""
font = cv2.FONT_HERSHEY_SIMPLEX
for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections):
x1, y1, x2, y2 = xyxy.astype(int)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
if isinstance(self.color, ColorPalette)
else self.color
)
cv2.rectangle(
img=scene,
pt1=(x1, y1),
pt2=(x2, y2),
color=color.as_bgr(),
thickness=self.thickness,
)
if skip_label:
continue
text = (
f"{class_id}"
if (labels is None or len(detections) != len(labels))
else labels[i]
)
text_width, text_height = cv2.getTextSize(
text=text,
fontFace=font,
fontScale=self.text_scale,
thickness=self.text_thickness,
)[0]
text_x = x1 + self.text_padding
text_y = y1 - self.text_padding
text_background_x1 = x1
text_background_y1 = y1 - 2 * self.text_padding - text_height
text_background_x2 = x1 + 2 * self.text_padding + text_width
text_background_y2 = y1
cv2.rectangle(
img=scene,
pt1=(text_background_x1, text_background_y1),
pt2=(text_background_x2, text_background_y2),
color=color.as_bgr(),
thickness=cv2.FILLED,
)
cv2.putText(
img=scene,
text=text,
org=(text_x, text_y),
fontFace=font,
fontScale=self.text_scale,
color=self.text_color.as_rgb(),
thickness=self.text_thickness,
lineType=cv2.LINE_AA,
)
return scene
+136 -148
View File
@@ -1,12 +1,11 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import List, Optional, Union
from typing import Iterator, Optional, Tuple, Union
import cv2
import numpy as np
from supervision.draw.color import Color, ColorPalette
from supervision.detection.utils import non_max_suppression
from supervision.geometry.core import Position
@@ -17,22 +16,27 @@ 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.
class_id (np.ndarray): An array of shape `(n,)` containing the class ids of the detections.
class_id (Optional[np.ndarray]): An array of shape `(n,)` containing the class ids of the detections.
confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores 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
class_id: Optional[np.ndarray] = None
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.class_id is None
or (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,8 +46,8 @@ 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, "
"class_id must be 1d np.ndarray with (n,) shape, "
"class_id must be None or 1d np.ndarray with (n,) shape, "
"confidence must be None or 1d np.ndarray with (n,) shape, "
"tracker_id must be None or 1d np.ndarray with (n,) shape"
)
@@ -53,27 +57,55 @@ 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.class_id[i],
self.confidence[i] if self.confidence is not None else None,
self.class_id[i] if self.class_id is not None else None,
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.class_id is None and other.class_id is None,
np.array_equal(self.class_id, other.class_id),
]
),
any(
[
self.confidence is None and other.confidence is None,
np.array_equal(self.confidence, other.confidence),
]
),
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 +117,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 +125,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 +141,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 +151,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,11 +212,15 @@ 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],
class_id=self.class_id[index],
confidence=self.confidence[index]
if self.confidence is not None
else None,
class_id=self.class_id[index] if self.class_id is not None else None,
tracker_id=self.tracker_id[index]
if self.tracker_id is not None
else None,
@@ -185,109 +229,53 @@ class Detections:
f"Detections.__getitem__ not supported for index of type {type(index)}."
)
class BoxAnnotator:
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
thickness: int = 2,
text_color: Color = Color.black(),
text_scale: float = 0.5,
text_thickness: int = 1,
text_padding: int = 10,
):
@property
def area(self) -> np.ndarray:
"""
A class for drawing bounding boxes on an image using detections provided.
Calculate the area of each bounding box in the set of object detections.
Attributes:
color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette
thickness (int): The thickness of the bounding box lines, default is 2
text_color (Color): The color of the text on the bounding box, default is white
text_scale (float): The scale of the text on the bounding box, default is 0.5
text_thickness (int): The thickness of the text on the bounding box, default is 1
text_padding (int): The padding around the text on the bounding box, default is 5
"""
self.color: Union[Color, ColorPalette] = color
self.thickness: int = thickness
self.text_color: Color = text_color
self.text_scale: float = text_scale
self.text_thickness: int = text_thickness
self.text_padding: int = text_padding
def annotate(
self,
scene: np.ndarray,
detections: Detections,
labels: Optional[List[str]] = None,
skip_label: bool = False,
) -> np.ndarray:
"""
Draws bounding boxes on the frame using the detections provided.
Parameters:
scene (np.ndarray): The image on which the bounding boxes will be drawn
detections (Detections): The detections for which the bounding boxes will be drawn
labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.
skip_label (bool): Is set to True, skips bounding box label annotation.
Returns:
np.ndarray: The image with the bounding boxes drawn on it
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.
"""
font = cv2.FONT_HERSHEY_SIMPLEX
for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections):
x1, y1, x2, y2 = xyxy.astype(int)
color = (
self.color.by_idx(class_id)
if isinstance(self.color, ColorPalette)
else self.color
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
)
cv2.rectangle(
img=scene,
pt1=(x1, y1),
pt2=(x2, y2),
color=color.as_bgr(),
thickness=self.thickness,
)
if skip_label:
continue
return self[indices]
text = (
f"{confidence:0.2f}"
if (labels is None or len(detections) != len(labels))
else labels[i]
)
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."
)
text_width, text_height = cv2.getTextSize(
text=text,
fontFace=font,
fontScale=self.text_scale,
thickness=self.text_thickness,
)[0]
text_x = x1 + self.text_padding
text_y = y1 - self.text_padding
text_background_x1 = x1
text_background_y1 = y1 - 2 * self.text_padding - text_height
text_background_x2 = x1 + 2 * self.text_padding + text_width
text_background_y2 = y1
cv2.rectangle(
img=scene,
pt1=(text_background_x1, text_background_y1),
pt2=(text_background_x2, text_background_y2),
color=color.as_bgr(),
thickness=cv2.FILLED,
)
cv2.putText(
img=scene,
text=text,
org=(text_x, text_y),
fontFace=font,
fontScale=self.text_scale,
color=self.text_color.as_rgb(),
thickness=self.text_thickness,
lineType=cv2.LINE_AA,
)
return scene
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]
+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)