Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
bc12a8ea0b | ||
|
|
2be4981c76 | ||
|
|
fe68327b75 | ||
|
|
141809ae92 | ||
|
|
1fa4a1c72e | ||
|
|
2df4261355 | ||
|
|
a88eff868e | ||
|
|
8bf686e23b | ||
|
|
b2f6520b87 | ||
|
|
420dae9e70 | ||
|
|
550c913dbb | ||
|
|
4dcca5309c | ||
|
|
d30cc180d0 | ||
|
|
ace8301b12 | ||
|
|
b91bc91897 | ||
|
|
e04c531c05 | ||
|
|
63fd5a3f8a | ||
|
|
38f52a04ae | ||
|
|
65090f612c | ||
|
|
9401210aa5 | ||
|
|
59d8655014 | ||
|
|
4dde484141 | ||
|
|
e6f250dbfb | ||
|
|
e6090be9eb | ||
|
|
25c19b4422 |
+14
-1
@@ -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.
|
||||
@@ -2,8 +2,8 @@
|
||||
<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>
|
||||
@@ -62,9 +62,9 @@
|
||||
|
||||
## 👋 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
|
||||
|
||||
@@ -97,32 +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/l_kf9CfZ_8M">
|
||||
<img src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg" alt="count-objects-in-zone">
|
||||
<img
|
||||
width="90%"
|
||||
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="count-objects-crossing-line">
|
||||
<img
|
||||
width="90%"
|
||||
src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg"
|
||||
alt="count-objects-crossing-line"
|
||||
>
|
||||
</a>
|
||||
</p>
|
||||
|
||||
## 🧹 code quality
|
||||
## 🏆 contribution
|
||||
|
||||
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.
|
||||
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!
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
## detections_to_voc_xml
|
||||
|
||||
:::supervision.annotation.voc.detections_to_voc_xml
|
||||
@@ -0,0 +1,33 @@
|
||||
### 0.4.0 <small>April 5, 2023</small>
|
||||
|
||||
- Added [[#46](https://github.com/roboflow/supervision/discussions/48)]: `Detections.empty` to allow easy creation of empty `Detections` objects.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `plot_images_grid` to allow easy plotting of multiple images on single plot.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
|
||||
- Changed [[#56](https://github.com/roboflow/supervision/pull/56)]: `show_frame_in_notebook` refactored and renamed to `plot_image`.
|
||||
|
||||
### 0.3.2 <small>March 23, 2023</small>
|
||||
|
||||
- Changed [[#50](https://github.com/roboflow/supervision/issues/50)]: Allow `Detections.class_id` to be `None`.
|
||||
|
||||
### 0.3.1 <small>March 6, 2023</small>
|
||||
|
||||
- Fixed [[#41](https://github.com/roboflow/supervision/issues/41)]: `PolygonZone` throws an exception when the object touches the bottom edge of the image.
|
||||
- Fixed [[#42](https://github.com/roboflow/supervision/issues/42)]: `Detections.wth_nms` method throws an exception when `Detections` is empty.
|
||||
- Changed [[#36](https://github.com/roboflow/supervision/pull/36)]: `Detections.wth_nms` support class agnostic and non-class agnostic case.
|
||||
|
||||
### 0.3.0 <small>March 6, 2023</small>
|
||||
|
||||
- Changed: Allow `Detections.confidence` to be `None`.
|
||||
- Added: `Detections.from_transformers` and `Detections.from_detectron2` to enable seamless integration with Transformers and Detectron2 models.
|
||||
- Added: `Detections.area` to dynamically calculate bounding box area.
|
||||
- Added: `Detections.wth_nms` to filter out double detections with NMS. Initial - only class agnostic - implementation.
|
||||
|
||||
### 0.2.0 <small>February 2, 2023</small>
|
||||
|
||||
- Added: Advanced `Detections` filtering with pandas-like API.
|
||||
- Added: `Detections.from_yolov5` and `Detections.from_yolov8` to enable seamless integration with YOLOv5 and YOLOv8 models.
|
||||
|
||||
### 0.1.0 <small>January 19, 2023</small>
|
||||
|
||||
Say hello to Supervision 👋
|
||||
@@ -0,0 +1,3 @@
|
||||
## BoxAnnotator
|
||||
|
||||
:::supervision.detection.annotate.BoxAnnotator
|
||||
@@ -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
|
||||
@@ -1,3 +0,0 @@
|
||||
## generate_2d_mask
|
||||
|
||||
:::supervision.detection.utils.generate_2d_mask
|
||||
@@ -1,7 +0,0 @@
|
||||
## draw_line
|
||||
|
||||
:::supervision.draw.utils.draw_line
|
||||
|
||||
## draw_rectangle
|
||||
|
||||
:::supervision.draw.utils.draw_rectangle
|
||||
@@ -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
|
||||
@@ -1,3 +0,0 @@
|
||||
## show_frame_in_notebook
|
||||
|
||||
:::supervision.notebook.utils.show_frame_in_notebook
|
||||
@@ -0,0 +1,7 @@
|
||||
## plot_image
|
||||
|
||||
:::supervision.notebook.utils.plot_image
|
||||
|
||||
## plot_images_grid
|
||||
|
||||
:::supervision.notebook.utils.plot_images_grid
|
||||
+15
-8
@@ -26,16 +26,21 @@ nav:
|
||||
- API reference:
|
||||
- Video: video.md
|
||||
- Detection:
|
||||
- Core: detection_core.md
|
||||
- Utils: detection_utils.md
|
||||
- Draw: draw.md
|
||||
- Geometry: geometry.md
|
||||
- Notebook: notebook.md
|
||||
- Core: detection/core.md
|
||||
- Annotate: detection/annotate.md
|
||||
- Utils: detection/utils.md
|
||||
- Draw:
|
||||
- Utils: draw/utils.md
|
||||
- Annotations:
|
||||
- Pascal VOC XML: annotation/voc.md
|
||||
- Notebook:
|
||||
- Utils: notebook/utils.md
|
||||
- Changelog: changelog.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
|
||||
favicon: https://media.roboflow.com/open-source/supervision/supervision-lenny.png
|
||||
palette:
|
||||
primary: 'deep purple'
|
||||
accent: 'teal'
|
||||
@@ -50,4 +55,6 @@ plugins:
|
||||
markdown_extensions:
|
||||
- admonition
|
||||
- pymdownx.details
|
||||
- pymdownx.superfences
|
||||
- pymdownx.superfences
|
||||
- toc:
|
||||
permalink: true
|
||||
@@ -1,6 +1,8 @@
|
||||
__version__ = "0.3.0"
|
||||
__version__ = "0.4.0"
|
||||
|
||||
from supervision.detection.core import BoxAnnotator, Detections
|
||||
from supervision.annotation.voc import detections_to_voc_xml
|
||||
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
|
||||
@@ -8,7 +10,7 @@ from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
|
||||
from supervision.geometry.core import Point, Position, Rect
|
||||
from supervision.geometry.utils import get_polygon_center
|
||||
from supervision.notebook.utils import show_frame_in_notebook
|
||||
from supervision.notebook.utils import plot_image, plot_images_grid
|
||||
from supervision.video import (
|
||||
VideoInfo,
|
||||
VideoSink,
|
||||
|
||||
@@ -0,0 +1,104 @@
|
||||
from typing import List
|
||||
from xml.dom.minidom import parseString
|
||||
from xml.etree.ElementTree import Element, SubElement, tostring
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
|
||||
def detections_to_voc_xml(
|
||||
detections: Detections,
|
||||
classes: List[str],
|
||||
filename: str,
|
||||
width: int,
|
||||
height: int,
|
||||
depth: int = 3,
|
||||
) -> str:
|
||||
"""
|
||||
Converts Detections object to Pascal VOC XML format.
|
||||
|
||||
Args:
|
||||
detections (Detections): A Detections object containing bounding boxes, class ids, and other relevant information.
|
||||
classes (List[str]): A list of class names corresponding to the class ids in the Detections object.
|
||||
filename (str): The name of the image file associated with the detections.
|
||||
width (int): The width of the image in pixels.
|
||||
height (int): The height of the image in pixels.
|
||||
depth (int, optional): The number of color channels in the image. Defaults to 3 for RGB images.
|
||||
|
||||
Returns:
|
||||
str: An XML string in Pascal VOC format representing the detections.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import numpy as np
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> xyxy = np.array([
|
||||
... [50, 30, 200, 180],
|
||||
... [20, 40, 150, 190]
|
||||
... ])
|
||||
>>> class_id = np.array([1, 0])
|
||||
>>> detections = Detections(xyxy=xyxy, class_id=class_id)
|
||||
|
||||
>>> classes = ["dog", "cat"]
|
||||
|
||||
>>> voc_xml = detections_to_voc_xml(
|
||||
... detections=detections,
|
||||
... classes=classes,
|
||||
... filename="image1.jpg",
|
||||
... width=500,
|
||||
... height=400
|
||||
... )
|
||||
```
|
||||
"""
|
||||
|
||||
# Create root element
|
||||
annotation = Element("annotation")
|
||||
|
||||
# Add folder element
|
||||
folder = SubElement(annotation, "folder")
|
||||
folder.text = "VOC"
|
||||
|
||||
# Add filename element
|
||||
fname = SubElement(annotation, "filename")
|
||||
fname.text = filename
|
||||
|
||||
# Add source element
|
||||
source = SubElement(annotation, "source")
|
||||
database = SubElement(source, "database")
|
||||
database.text = "roboflow.ai"
|
||||
|
||||
# Add size element
|
||||
size = SubElement(annotation, "size")
|
||||
w = SubElement(size, "width")
|
||||
w.text = str(width)
|
||||
h = SubElement(size, "height")
|
||||
h.text = str(height)
|
||||
d = SubElement(size, "depth")
|
||||
d.text = str(depth)
|
||||
|
||||
# Add segmented element
|
||||
segmented = SubElement(annotation, "segmented")
|
||||
segmented.text = "0"
|
||||
|
||||
# Add object elements
|
||||
for i in range(detections.xyxy.shape[0]):
|
||||
obj = SubElement(annotation, "object")
|
||||
|
||||
class_id = detections.class_id[i] if detections.class_id is not None else None
|
||||
name = SubElement(obj, "name")
|
||||
name.text = classes[class_id] if class_id is not None else "unknown"
|
||||
|
||||
bndbox = SubElement(obj, "bndbox")
|
||||
xmin = SubElement(bndbox, "xmin")
|
||||
xmin.text = str(int(detections.xyxy[i, 0]))
|
||||
ymin = SubElement(bndbox, "ymin")
|
||||
ymin.text = str(int(detections.xyxy[i, 1]))
|
||||
xmax = SubElement(bndbox, "xmax")
|
||||
xmax.text = str(int(detections.xyxy[i, 2]))
|
||||
ymax = SubElement(bndbox, "ymax")
|
||||
ymax.text = str(int(detections.xyxy[i, 3]))
|
||||
|
||||
# Generate XML string
|
||||
xml_string = parseString(tostring(annotation)).toprettyxml(indent=" ")
|
||||
|
||||
return xml_string
|
||||
@@ -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
|
||||
+108
-157
@@ -3,11 +3,9 @@ from __future__ import annotations
|
||||
from dataclasses import dataclass
|
||||
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
|
||||
|
||||
|
||||
@@ -18,13 +16,13 @@ class Detections:
|
||||
|
||||
Attributes:
|
||||
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`
|
||||
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.
|
||||
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
|
||||
class_id: np.ndarray
|
||||
class_id: Optional[np.ndarray] = None
|
||||
confidence: Optional[np.ndarray] = None
|
||||
tracker_id: Optional[np.ndarray] = None
|
||||
|
||||
@@ -32,7 +30,8 @@ class Detections:
|
||||
n = len(self.xyxy)
|
||||
validators = [
|
||||
(isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)),
|
||||
(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)
|
||||
@@ -47,8 +46,8 @@ class Detections:
|
||||
if not all(validators):
|
||||
raise ValueError(
|
||||
"xyxy must be 2d np.ndarray with (n, 4) shape, "
|
||||
"class_id must be None or 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"
|
||||
)
|
||||
|
||||
@@ -68,7 +67,7 @@ class Detections:
|
||||
yield (
|
||||
self.xyxy[i],
|
||||
self.confidence[i] if self.confidence is not None else None,
|
||||
self.class_id[i],
|
||||
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,
|
||||
)
|
||||
|
||||
@@ -76,13 +75,18 @@ class 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),
|
||||
]
|
||||
),
|
||||
np.array_equal(self.class_id, other.class_id),
|
||||
any(
|
||||
[
|
||||
self.tracker_id is None and other.tracker_id is None,
|
||||
@@ -93,14 +97,15 @@ class Detections:
|
||||
)
|
||||
|
||||
@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
|
||||
@@ -112,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],
|
||||
@@ -120,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
|
||||
@@ -146,7 +152,13 @@ class Detections:
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_transformers(cls, transformers_results: dict):
|
||||
def from_transformers(cls, transformers_results: dict) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from Object Detection Transformer output Results
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
"""
|
||||
return cls(
|
||||
xyxy=transformers_results["boxes"].cpu().numpy(),
|
||||
confidence=transformers_results["scores"].cpu().numpy(),
|
||||
@@ -154,7 +166,7 @@ class Detections:
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_detectron2(cls, detectron2_results):
|
||||
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(),
|
||||
@@ -165,7 +177,32 @@ class Detections:
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_coco_annotations(cls, coco_annotation: dict):
|
||||
def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections:
|
||||
xyxy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
|
||||
for prediction in roboflow_result["predictions"]:
|
||||
x = prediction["x"]
|
||||
y = prediction["y"]
|
||||
width = prediction["width"]
|
||||
height = prediction["height"]
|
||||
x_min = x - width / 2
|
||||
y_min = y - height / 2
|
||||
x_max = x_min + width
|
||||
y_max = y_min + height
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(class_list.index(prediction["class"]))
|
||||
confidence.append(prediction["confidence"])
|
||||
|
||||
return Detections(
|
||||
xyxy=np.array(xyxy),
|
||||
confidence=np.array(confidence),
|
||||
class_id=np.array(class_id).astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
|
||||
xyxy, class_id = [], []
|
||||
|
||||
for annotation in coco_annotation:
|
||||
@@ -175,40 +212,19 @@ class Detections:
|
||||
|
||||
return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
|
||||
|
||||
def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[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.
|
||||
|
||||
Returns:
|
||||
Optional[np.ndarray]: A new instance of Detections with the filtered detections, if inplace is set to `False`. `None` otherwise.
|
||||
"""
|
||||
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,
|
||||
)
|
||||
@classmethod
|
||||
def empty(cls) -> Detections:
|
||||
return cls(
|
||||
xyxy=np.empty((0, 4), dtype=np.float32),
|
||||
confidence=np.array([], dtype=np.float32),
|
||||
class_id=np.array([], dtype=int),
|
||||
)
|
||||
|
||||
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:
|
||||
@@ -234,8 +250,10 @@ class Detections:
|
||||
):
|
||||
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,
|
||||
@@ -246,118 +264,51 @@ class Detections:
|
||||
|
||||
@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) -> Detections:
|
||||
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."
|
||||
indices = non_max_suppression(self.xyxy, self.confidence, threshold=threshold)
|
||||
|
||||
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__(
|
||||
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,
|
||||
):
|
||||
"""
|
||||
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
|
||||
|
||||
"""
|
||||
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)
|
||||
color = (
|
||||
self.color.by_idx(class_id)
|
||||
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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -20,38 +20,98 @@ def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.
|
||||
return mask
|
||||
|
||||
|
||||
def non_max_suppression(boxes: np.ndarray, scores: np.ndarray, threshold: float):
|
||||
assert boxes.shape[0] == scores.shape[0]
|
||||
ys1 = boxes[:, 0]
|
||||
xs1 = boxes[:, 1]
|
||||
ys2 = boxes[:, 2]
|
||||
xs2 = boxes[:, 3]
|
||||
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.
|
||||
|
||||
areas = (ys2 - ys1) * (xs2 - xs1)
|
||||
scores_indexes = scores.argsort().tolist()
|
||||
boxes_keep_index = []
|
||||
while len(scores_indexes):
|
||||
index = scores_indexes.pop()
|
||||
boxes_keep_index.append(index)
|
||||
if not len(scores_indexes):
|
||||
break
|
||||
iou = compute_iou(
|
||||
boxes[index], boxes[scores_indexes], areas[index], areas[scores_indexes]
|
||||
)
|
||||
filtered_indexes = set((iou > threshold).nonzero()[0])
|
||||
scores_indexes = [
|
||||
v for (i, v) in enumerate(scores_indexes) if i not in filtered_indexes
|
||||
]
|
||||
return np.array(boxes_keep_index)
|
||||
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 compute_iou(box, boxes, box_area, boxes_area):
|
||||
assert boxes.shape[0] == boxes_area.shape[0]
|
||||
ys1 = np.maximum(box[0], boxes[:, 0])
|
||||
xs1 = np.maximum(box[1], boxes[:, 1])
|
||||
ys2 = np.minimum(box[2], boxes[:, 2])
|
||||
xs2 = np.minimum(box[3], boxes[:, 3])
|
||||
intersections = np.maximum(ys2 - ys1, 0) * np.maximum(xs2 - xs1, 0)
|
||||
unions = box_area + boxes_area - intersections
|
||||
iou = intersections / unions
|
||||
return iou
|
||||
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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -1,33 +1,92 @@
|
||||
from typing import Tuple
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
def show_frame_in_notebook(
|
||||
frame: np.ndarray, size: Tuple[int, int] = (10, 10), cmap: str = "gray"
|
||||
):
|
||||
def plot_image(
|
||||
image: np.ndarray, size: Tuple[int, int] = (10, 10), cmap: Optional[str] = "gray"
|
||||
) -> None:
|
||||
"""
|
||||
Display a frame in Jupyter Notebook using Matplotlib
|
||||
Plots image using matplotlib.
|
||||
|
||||
Attributes:
|
||||
frame (np.ndarray): The frame to be displayed.
|
||||
size (Tuple[int, int]): The size of the plot. default:(10,10)
|
||||
cmap (str): the colormap to use for single channel images. default:gray
|
||||
Args:
|
||||
image (np.ndarray): The frame to be displayed.
|
||||
size (Tuple[int, int]): The size of the plot.
|
||||
cmap (str): the colormap to use for single channel images.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> from supervision.notebook.utils import show_frame_in_notebook
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread("path/to/image.jpg")
|
||||
|
||||
%matplotlib inline
|
||||
show_frame_in_notebook(frame, (16, 16))
|
||||
>>> sv.plot_image(image, (16, 16))
|
||||
```
|
||||
"""
|
||||
if frame.ndim == 2:
|
||||
if image.ndim == 2:
|
||||
plt.figure(figsize=size)
|
||||
plt.imshow(frame, cmap=cmap)
|
||||
plt.imshow(image, cmap=cmap)
|
||||
else:
|
||||
plt.figure(figsize=size)
|
||||
plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
plt.show()
|
||||
|
||||
|
||||
def plot_images_grid(
|
||||
images: List[np.ndarray],
|
||||
grid_size: Tuple[int, int],
|
||||
titles: Optional[List[str]] = None,
|
||||
size: Tuple[int, int] = (12, 12),
|
||||
) -> None:
|
||||
"""
|
||||
Plots images in a grid using matplotlib.
|
||||
|
||||
Args:
|
||||
images (List[np.ndarray]): A list of images as numpy arrays.
|
||||
grid_size (Tuple[int, int]): A tuple specifying the number of rows and columns for the grid.
|
||||
titles (Optional[List[str]]): A list of titles for each image. Defaults to None.
|
||||
size (Tuple[int, int]): A tuple specifying the width and height of the entire plot in inches.
|
||||
|
||||
Raises:
|
||||
ValueError: If the number of images exceeds the grid size.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image1 = cv2.imread("path/to/image1.jpg")
|
||||
>>> image2 = cv2.imread("path/to/image2.jpg")
|
||||
>>> image3 = cv2.imread("path/to/image3.jpg")
|
||||
|
||||
>>> images = [image1, image2, image3]
|
||||
>>> titles = ["Image 1", "Image 2", "Image 3"]
|
||||
|
||||
%matplotlib inline
|
||||
>>> plot_images_grid(images, grid_size=(2, 2), titles=titles, figsize=(16, 16))
|
||||
```
|
||||
"""
|
||||
|
||||
nrows, ncols = grid_size
|
||||
|
||||
if len(images) > nrows * ncols:
|
||||
raise ValueError(
|
||||
"The number of images exceeds the grid size. Please increase the grid size or reduce the number of images."
|
||||
)
|
||||
|
||||
fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=size)
|
||||
|
||||
for idx, ax in enumerate(axes.flat):
|
||||
if idx < len(images):
|
||||
ax.imshow(cv2.cvtColor(images[idx], cv2.COLOR_BGR2RGB))
|
||||
if titles is not None and idx < len(titles):
|
||||
ax.set_title(titles[idx])
|
||||
ax.axis("off")
|
||||
else:
|
||||
ax.axis("off")
|
||||
|
||||
plt.show()
|
||||
|
||||
@@ -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)
|
||||
Reference in New Issue
Block a user