Compare commits

...
26 Commits
Author SHA1 Message Date
SkalskiP e091f9c959 📚 Docs changelog.md update 2023-04-12 17:54:06 +02:00
SkalskiP f12504b0dd 🛠️ fix Detections.__getitem__ and Detections.area bugs 2023-04-11 13:50:42 +02:00
Piotr Skalski dba4d9f8ad Merge pull request #58 from roboflow/feature/update_to_support_sam
feature/update_to_support_sam
2023-04-10 23:52:35 +02:00
SkalskiP ddc8a8ccce 🖤 Make black happy 2023-04-10 23:16:04 +02:00
SkalskiP c1ccff474f Documentation for PolygonZone and PolygonZoneAnnotator 2023-04-10 23:15:35 +02:00
SkalskiP d02bac75c4 Document changes in docs 2023-04-10 23:07:53 +02:00
SkalskiP f3c4a966c4 changelog added 2023-04-10 22:57:08 +02:00
SkalskiP 5bc436b1c5 add area and box_area to Detections object 2023-04-10 19:58:44 +02:00
SkalskiP c37590ef3a add area and box_area to Detections object 2023-04-10 19:17:53 +02:00
SkalskiP a14bb94b46 😷 to 📦 initial version of mask_to_xyxy util 2023-04-10 18:58:07 +02:00
SkalskiP aab7e238e4 😷 to 📦 initial version of mask_to_xyxy util 2023-04-10 18:46:13 +02:00
SkalskiP 34dba17641 🖤 Make Black happy 2023-04-10 12:00:00 +02:00
SkalskiP ef5ee71f0b Update sv.Detections.from_sam API 2023-04-10 11:58:09 +02:00
SkalskiP bed0e41cb7 👊 version bump 2023-04-09 23:24:10 +02:00
SkalskiP e439b29f07 Initial from_sam connector for Detections 2023-04-09 23:21:04 +02:00
SkalskiP 511375f24b 😷 Top level import of MaskAnnotator 2023-04-09 23:05:56 +02:00
SkalskiP ff2788945d 😷 Initial implementation of MaskAnnotator 2023-04-09 22:57:29 +02:00
SkalskiP 5b9eef363d small fix :) 2023-04-09 22:16:12 +02:00
SkalskiP bcc4b25eb9 🖤bake black happy 2023-04-08 19:05:07 +02:00
SkalskiP d5d2035048 😷add mask support to Detection object 2023-04-08 19:01:48 +02:00
SkalskiP 14795a6626 🏞 updated plot_image and plot_images_grid 2023-04-08 17:31:52 +02:00
Piotr Skalski bc12a8ea0b Merge pull request #56 from roboflow/feature/rafactor_of_functions_in_notebooks_package_and_intitial_version_of_voc_xml_support
feature/rafactor_of_functions_in_notebooks_package_and_intitial_version_of_voc_xml_support
2023-04-05 17:18:38 +02:00
SkalskiP 2be4981c76 🖤 Make black happy 2023-04-05 16:27:22 +02:00
SkalskiP fe68327b75 Changelog updated 2023-04-05 16:25:29 +02:00
SkalskiP 141809ae92 - 📓 Refactor of notebooks utilities
- Initial voc xml sing support
- Documentation improvements
2023-04-05 15:48:21 +02:00
SkalskiP 1fa4a1c72e 📜 Supervision docs changelog added 2023-03-29 15:59:02 +02:00
18 changed files with 560 additions and 62 deletions
+3
View File
@@ -0,0 +1,3 @@
## detections_to_voc_xml
:::supervision.annotation.voc.detections_to_voc_xml
+45
View File
@@ -0,0 +1,45 @@
### 0.5.1 <small>April 12, 2023</small>
- Fixed `Detections.__getitem__` method did not return mask for selected item.
- Fixed `Detections.area` crashed for mask detections.
### 0.5.0 <small>April 10, 2023</small>
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.mask` to enable segmentation support.
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `MaskAnnotator` to allow easy `Detections.mask` annotation.
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
- Changed [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.area` behaviour to work not only with boxes but also with masks.
### 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 👋
+5 -1
View File
@@ -1,3 +1,7 @@
## BoxAnnotator
:::supervision.detection.annotate.BoxAnnotator
:::supervision.detection.annotate.BoxAnnotator
## MaskAnnotator
:::supervision.detection.annotate.MaskAnnotator
+7
View File
@@ -0,0 +1,7 @@
## PolygonZone
:::supervision.detection.tools.polygon_zone.PolygonZone
## PolygonZoneAnnotator
:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator
+5 -1
View File
@@ -8,4 +8,8 @@
## non_max_suppression
:::supervision.detection.utils.non_max_suppression
:::supervision.detection.utils.non_max_suppression
## mask_to_xyxy
:::supervision.detection.utils.mask_to_xyxy
-3
View File
@@ -1,3 +0,0 @@
## show_frame_in_notebook
:::supervision.notebook.utils.show_frame_in_notebook
+7
View File
@@ -0,0 +1,7 @@
## plot_image
:::supervision.notebook.utils.plot_image
## plot_images_grid
:::supervision.notebook.utils.plot_images_grid
+13 -5
View File
@@ -29,14 +29,20 @@ nav:
- Core: detection/core.md
- Annotate: detection/annotate.md
- Utils: detection/utils.md
- Tools:
- Polygon Zone: detection/tools/polygon_zone.md
- Draw:
- Utils: draw_utils.md
- Notebook: notebook.md
- Utils: draw/utils.md
- Annotations:
- Pascal VOC XML: annotation/voc.md
- Notebook:
- Utils: notebook/utils.md
- Changelog: changelog.md
theme:
name: 'material'
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
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'
@@ -51,4 +57,6 @@ plugins:
markdown_extensions:
- admonition
- pymdownx.details
- pymdownx.superfences
- pymdownx.superfences
- toc:
permalink: true
+6 -5
View File
@@ -1,15 +1,16 @@
__version__ = "0.3.2"
__version__ = "0.5.1"
from supervision.detection.annotate import BoxAnnotator
from supervision.annotation.voc import detections_to_voc_xml
from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
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.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
from supervision.detection.utils import generate_2d_mask, mask_to_xyxy
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,
View File
+104
View File
@@ -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
+60 -2
View File
@@ -56,8 +56,11 @@ class BoxAnnotator:
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)
for i in range(len(detections)):
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
idx = class_id if class_id is not None else i
color = (
self.color.by_idx(idx)
@@ -114,3 +117,58 @@ class BoxAnnotator:
lineType=cv2.LINE_AA,
)
return scene
class MaskAnnotator:
"""
A class for overlaying masks on an image using detections provided.
Attributes:
color (Union[Color, ColorPalette]): The color to fill the mask, can be a single color or a color palette
"""
def __init__(
self,
color: Union[Color, ColorPalette] = ColorPalette.default(),
):
self.color: Union[Color, ColorPalette] = color
def annotate(
self, scene: np.ndarray, detections: Detections, opacity: float = 0.5
) -> np.ndarray:
"""
Overlays the masks on the given image based on the provided detections, with a specified opacity.
Parameters:
scene (np.ndarray): The image on which the masks will be overlaid
detections (Detections): The detections for which the masks will be overlaid
opacity (float): The opacity of the masks, between 0 and 1, default is 0.5
Returns:
np.ndarray: The image with the masks overlaid
"""
for i in range(len(detections.xyxy)):
if detections.mask is None:
continue
class_id = (
detections.class_id[i] if detections.class_id is not None else None
)
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
)
mask = detections.mask[i]
colored_mask = np.zeros_like(scene, dtype=np.uint8)
colored_mask[:] = color.as_bgr()
scene = np.where(
np.expand_dims(mask, axis=-1),
np.uint8(opacity * colored_mask + (1 - opacity) * scene),
scene,
)
return scene
+148 -29
View File
@@ -1,14 +1,52 @@
from __future__ import annotations
from dataclasses import dataclass
from typing import Iterator, Optional, Tuple, Union
from typing import Any, Iterator, List, Optional, Tuple, Union
import numpy as np
from supervision.detection.utils import non_max_suppression
from supervision.detection.utils import non_max_suppression, xywh_to_xyxy
from supervision.geometry.core import Position
def _validate_xyxy(xyxy: Any, n: int) -> None:
is_valid = isinstance(xyxy, np.ndarray) and xyxy.shape == (n, 4)
if not is_valid:
raise ValueError("xyxy must be 2d np.ndarray with (n, 4) shape")
def _validate_mask(mask: Any, n: int) -> None:
is_valid = mask is None or (
isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
)
if not is_valid:
raise ValueError("mask must be 3d np.ndarray with (n, W, H) shape")
def _validate_class_id(class_id: Any, n: int) -> None:
is_valid = class_id is None or (
isinstance(class_id, np.ndarray) and class_id.shape == (n,)
)
if not is_valid:
raise ValueError("class_id must be None or 1d np.ndarray with (n,) shape")
def _validate_confidence(confidence: Any, n: int) -> None:
is_valid = confidence is None or (
isinstance(confidence, np.ndarray) and confidence.shape == (n,)
)
if not is_valid:
raise ValueError("confidence must be None or 1d np.ndarray with (n,) shape")
def _validate_tracker_id(tracker_id: Any, n: int) -> None:
is_valid = tracker_id is None or (
isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
)
if not is_valid:
raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape")
@dataclass
class Detections:
"""
@@ -16,40 +54,25 @@ class Detections:
Attributes:
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`
mask: (Optional[np.ndarray]): An array of shape `(n, W, H)` containing the segmentation masks.
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
mask: np.Optional[np.ndarray] = None
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)),
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)
and self.tracker_id.shape == (n,)
),
]
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, "
"tracker_id must be None or 1d np.ndarray with (n,) shape"
)
_validate_xyxy(xyxy=self.xyxy, n=n)
_validate_mask(mask=self.mask, n=n)
_validate_class_id(class_id=self.class_id, n=n)
_validate_confidence(confidence=self.confidence, n=n)
_validate_tracker_id(tracker_id=self.tracker_id, n=n)
def __len__(self):
"""
@@ -59,13 +82,22 @@ class Detections:
def __iter__(
self,
) -> Iterator[Tuple[np.ndarray, Optional[float], int, Optional[Union[str, int]]]]:
) -> Iterator[
Tuple[
np.ndarray,
Optional[np.ndarray],
Optional[float],
Optional[int],
Optional[int],
]
]:
"""
Iterates over the Detections object and yield a tuple of `(xyxy, confidence, class_id, tracker_id)` for each detection.
Iterates over the Detections object and yield a tuple of `(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
"""
for i in range(len(self.xyxy)):
yield (
self.xyxy[i],
self.mask[i] if self.mask is not None else None,
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,
@@ -75,6 +107,12 @@ class Detections:
return all(
[
np.array_equal(self.xyxy, other.xyxy),
any(
[
self.mask is None and other.mask is None,
np.array_equal(self.mask, other.mask),
]
),
any(
[
self.class_id is None and other.class_id is None,
@@ -113,7 +151,7 @@ class Detections:
>>> from supervision import Detections
>>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
>>> results = model(frame)
>>> results = model(IMAGE)
>>> detections = Detections.from_yolov5(results)
```
"""
@@ -141,8 +179,8 @@ class Detections:
>>> from supervision import Detections
>>> model = YOLO('yolov8s.pt')
>>> results = model(frame)[0]
>>> detections = Detections.from_yolov8(results)
>>> yolov8_results = model(IMAGE)[0]
>>> detections = Detections.from_yolov8(yolov8_results)
```
"""
return cls(
@@ -176,6 +214,62 @@ class Detections:
.astype(int),
)
@classmethod
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_sam(cls, sam_result: List[dict]) -> Detections:
"""
Creates a Detections instance from Segment Anything Model (SAM) by Meta AI.
Args:
sam_result (List[dict]): The output Results instance from SAM
Returns:
Detections: A new Detections object.
Example:
```python
>>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
>>> import supervision as sv
>>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
>>> mask_generator = SamAutomaticMaskGenerator(sam)
>>> sam_result = mask_generator.generate(IMAGE)
>>> detections = sv.Detections.from_sam(sam_result=sam_result)
```
"""
sorted_generated_masks = sorted(
sam_result, key=lambda x: x["area"], reverse=True
)
xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
mask = np.array([mask["segmentation"] for mask in sorted_generated_masks])
return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
@classmethod
def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
xyxy, class_id = [], []
@@ -187,6 +281,14 @@ class Detections:
return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
@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.
@@ -217,6 +319,9 @@ class Detections:
):
return Detections(
xyxy=self.xyxy[index],
mask=self.mask[index]
if self.mask is not None
else None,
confidence=self.confidence[index]
if self.confidence is not None
else None,
@@ -231,6 +336,20 @@ class Detections:
@property
def area(self) -> np.ndarray:
"""
Calculate the area of each detection in the set of object detections. If masks field is defined property
returns are of each mask. If only box is given property return area of each box.
Returns:
np.ndarray: An array of floats containing the area of each detection in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections.
"""
if self.mask is not None:
return np.array([np.sum(mask) for mask in self.mask])
else:
return self.box_area
@property
def box_area(self) -> np.ndarray:
"""
Calculate the area of each bounding box in the set of object detections.
@@ -13,6 +13,17 @@ from supervision.geometry.utils import get_polygon_center
class PolygonZone:
"""
A class for defining a polygon-shaped zone within a frame for detecting objects.
Attributes:
polygon (np.ndarray): A numpy array defining the polygon vertices
frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height)
triggering_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER)
current_count (int): The current count of detected objects within the zone
mask (np.ndarray): The 2D bool mask for the polygon zone
"""
def __init__(
self,
polygon: np.ndarray,
@@ -30,6 +41,16 @@ class PolygonZone:
)
def trigger(self, detections: Detections) -> np.ndarray:
"""
Determines if the detections are within the polygon zone.
Parameters:
detections (Detections): The detections to be checked against the polygon zone
Returns:
np.ndarray: A boolean numpy array indicating if each detection is within the polygon zone
"""
clipped_xyxy = clip_boxes(
boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh
)
@@ -43,6 +64,21 @@ class PolygonZone:
class PolygonZoneAnnotator:
"""
A class for annotating a polygon-shaped zone within a frame with a count of detected objects.
Attributes:
zone (PolygonZone): The polygon zone to be annotated
color (Color): The color to draw the polygon lines
thickness (int): The thickness of the polygon lines, default is 2
text_color (Color): The color of the text on the polygon, default is black
text_scale (float): The scale of the text on the polygon, default is 0.5
text_thickness (int): The thickness of the text on the polygon, default is 1
text_padding (int): The padding around the text on the polygon, default is 10
font (int): The font type for the text on the polygon, default is cv2.FONT_HERSHEY_SIMPLEX
center (Tuple[int, int]): The center of the polygon for text placement
"""
def __init__(
self,
zone: PolygonZone,
@@ -64,6 +100,16 @@ class PolygonZoneAnnotator:
self.center = get_polygon_center(polygon=zone.polygon)
def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray:
"""
Annotates the polygon zone within a frame with a count of detected objects.
Parameters:
scene (np.ndarray): The image on which the polygon zone will be annotated
label (Optional[str]): An optional label for the count of detected objects within the polygon zone (default: None)
Returns:
np.ndarray: The image with the polygon zone and count of detected objects
"""
annotated_frame = draw_polygon(
scene=scene,
polygon=self.zone.polygon,
+31
View File
@@ -115,3 +115,34 @@ def clip_boxes(
result[:, [0, 2]] = result[:, [0, 2]].clip(0, width)
result[:, [1, 3]] = result[:, [1, 3]].clip(0, height)
return result
def xywh_to_xyxy(boxes_xywh: np.ndarray) -> np.ndarray:
xyxy = boxes_xywh.copy()
xyxy[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2]
xyxy[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3]
return xyxy
def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
"""
Converts a 3D `np.array` of 2D bool masks into a 2D `np.array` of bounding boxes.
Parameters:
masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)` containing 2D bool masks
Returns:
np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes `(x_min, y_min, x_max, y_max)` for each mask
"""
n = masks.shape[0]
bboxes = np.zeros((n, 4), dtype=int)
for i, mask in enumerate(masks):
rows, cols = np.where(mask)
if len(rows) > 0 and len(cols) > 0:
x_min, x_max = np.min(cols), np.max(cols)
y_min, y_max = np.min(rows), np.max(rows)
bboxes[i, :] = [x_min, y_min, x_max, y_max]
return bboxes
+80 -16
View File
@@ -1,33 +1,97 @@
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] = (12, 12), 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:
plt.figure(figsize=size)
plt.imshow(frame, cmap=cmap)
plt.figure(figsize=size)
if image.ndim == 2:
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.axis("off")
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),
cmap: Optional[str] = "gray",
) -> 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.
cmap (str): the colormap to use for single channel images.
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):
if images[idx].ndim == 2:
ax.imshow(images[idx], cmap=cmap)
else:
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")
plt.show()