Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
29450139f9 | ||
|
|
ea400feb6a | ||
|
|
05495aca75 | ||
|
|
a3bea9ae37 | ||
|
|
651690f838 | ||
|
|
d87761f75d | ||
|
|
b20e4a6c19 | ||
|
|
0ef2ab7273 | ||
|
|
d1bb169593 | ||
|
|
c8f3f59a45 | ||
|
|
12523afe9a | ||
|
|
68412500fa | ||
|
|
49d4441419 | ||
|
|
0394a3f545 | ||
|
|
2bcd1e1dfd | ||
|
|
e62937c8ce | ||
|
|
b37aca52e2 | ||
|
|
790d2f6c5b | ||
|
|
0cb722e91e | ||
|
|
55cd04d0aa | ||
|
|
58181e6a7a | ||
|
|
ab00d2db0d | ||
|
|
6a76c96df5 | ||
|
|
d4a6a1a85b | ||
|
|
7c78feb4b8 | ||
|
|
73a7983719 | ||
|
|
1d6c2bc080 | ||
|
|
889461574f | ||
|
|
f95426540b | ||
|
|
c16599269c | ||
|
|
888edfbc67 | ||
|
|
7497f8d96a | ||
|
|
b34ecd53cf | ||
|
|
e9649acd90 | ||
|
|
2cf7a9b620 |
@@ -1,3 +0,0 @@
|
||||
## detections_to_voc_xml
|
||||
|
||||
:::supervision.annotation.voc.detections_to_voc_xml
|
||||
@@ -1,3 +1,15 @@
|
||||
### 0.6.0 <small>April 19, 2023</small>
|
||||
|
||||
- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: initial `Dataset` support and ability to save `Detections` in Pascal VOC XML format.
|
||||
- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
|
||||
- Added [[#72](https://github.com/roboflow/supervision/pull/72)]: ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
|
||||
- Changed [[#70](https://github.com/roboflow/supervision/pull/70)]: order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
|
||||
- Changed [[#71](https://github.com/roboflow/supervision/pull/71)]: `generate_2d_mask` to `polygon_to_mask`.
|
||||
|
||||
### 0.5.2 <small>April 13, 2023</small>
|
||||
|
||||
- Fixed [[#63](https://github.com/roboflow/supervision/pull/63)]: `LineZone.trigger` function expects 4 values instead of 5.
|
||||
|
||||
### 0.5.1 <small>April 12, 2023</small>
|
||||
|
||||
- Fixed `Detections.__getitem__` method did not return mask for selected item.
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
!!! warning
|
||||
|
||||
`Dataset` API is still fluid and may change. If you use Dataset in your project until further notice, freeze the
|
||||
`supervision` version in your `requirements.txt`.
|
||||
|
||||
## Dataset
|
||||
|
||||
:::supervision.dataset.core.Dataset
|
||||
+17
-5
@@ -1,7 +1,3 @@
|
||||
## generate_2d_mask
|
||||
|
||||
:::supervision.detection.utils.generate_2d_mask
|
||||
|
||||
## box_iou_batch
|
||||
|
||||
:::supervision.detection.utils.box_iou_batch
|
||||
@@ -10,6 +6,22 @@
|
||||
|
||||
:::supervision.detection.utils.non_max_suppression
|
||||
|
||||
## polygon_to_mask
|
||||
|
||||
:::supervision.detection.utils.polygon_to_mask
|
||||
|
||||
## mask_to_xyxy
|
||||
|
||||
:::supervision.detection.utils.mask_to_xyxy
|
||||
:::supervision.detection.utils.mask_to_xyxy
|
||||
|
||||
## mask_to_polygons
|
||||
|
||||
:::supervision.detection.utils.mask_to_polygons
|
||||
|
||||
## polygon_to_xyxy
|
||||
|
||||
:::supervision.detection.utils.polygon_to_xyxy
|
||||
|
||||
## filter_polygons_by_area
|
||||
|
||||
:::supervision.detection.utils.filter_polygons_by_area
|
||||
+3
-3
@@ -24,19 +24,19 @@ extra:
|
||||
nav:
|
||||
- Home: index.md
|
||||
- API reference:
|
||||
- Video: video.md
|
||||
- Detection:
|
||||
- Core: detection/core.md
|
||||
- Annotate: detection/annotate.md
|
||||
- Utils: detection/utils.md
|
||||
- Tools:
|
||||
- Polygon Zone: detection/tools/polygon_zone.md
|
||||
- Dataset:
|
||||
- Core: dataset/core.md
|
||||
- Draw:
|
||||
- Utils: draw/utils.md
|
||||
- Annotations:
|
||||
- Pascal VOC XML: annotation/voc.md
|
||||
- Notebook:
|
||||
- Utils: notebook/utils.md
|
||||
- Video: video.md
|
||||
- Changelog: changelog.md
|
||||
|
||||
theme:
|
||||
|
||||
+12
-3
@@ -1,13 +1,22 @@
|
||||
__version__ = "0.5.1"
|
||||
__version__ = "0.6.0"
|
||||
|
||||
from supervision.annotation.voc import detections_to_voc_xml
|
||||
from supervision.dataset.core import Dataset
|
||||
from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
|
||||
from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
|
||||
from supervision.detection.utils import generate_2d_mask, mask_to_xyxy
|
||||
from supervision.detection.utils import (
|
||||
box_iou_batch,
|
||||
filter_polygons_by_area,
|
||||
mask_to_polygons,
|
||||
mask_to_xyxy,
|
||||
non_max_suppression,
|
||||
polygon_to_mask,
|
||||
polygon_to_xyxy,
|
||||
)
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
|
||||
from supervision.file import list_files_with_extensions
|
||||
from supervision.geometry.core import Point, Position, Rect
|
||||
from supervision.geometry.utils import get_polygon_center
|
||||
from supervision.notebook.utils import plot_image, plot_images_grid
|
||||
|
||||
@@ -1,104 +0,0 @@
|
||||
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,128 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.formats.pascal_voc import (
|
||||
detections_to_pascal_voc,
|
||||
load_pascal_voc_annotations,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.file import list_files_with_extensions
|
||||
|
||||
|
||||
@dataclass
|
||||
class Dataset:
|
||||
"""
|
||||
Dataclass containing information about the dataset.
|
||||
|
||||
Attributes:
|
||||
classes (List[str]): List containing dataset class names.
|
||||
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
images: Dict[str, np.ndarray]
|
||||
annotations: Dict[str, Detections]
|
||||
|
||||
def as_pascal_voc(
|
||||
self,
|
||||
images_directory_path: Optional[str] = None,
|
||||
annotations_directory_path: Optional[str] = None,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
|
||||
PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area
|
||||
percentage.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_directory_path (Optional[str]): The path to the directory where the annotations in
|
||||
PASCAL VOC format should be saved. If not provided, annotations will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
|
||||
"""
|
||||
if images_directory_path:
|
||||
images_path = Path(images_directory_path)
|
||||
images_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if annotations_directory_path:
|
||||
annotations_path = Path(annotations_directory_path)
|
||||
annotations_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for image_name, image in self.images.items():
|
||||
detections = self.annotations[image_name]
|
||||
|
||||
if images_directory_path:
|
||||
cv2.imwrite(str(images_path / image_name), image)
|
||||
|
||||
if annotations_directory_path:
|
||||
annotation_name = Path(image_name).stem
|
||||
pascal_voc_xml = detections_to_pascal_voc(
|
||||
detections=detections,
|
||||
classes=self.classes,
|
||||
filename=image_name,
|
||||
image_shape=image.shape,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
with open(annotations_path / f"{annotation_name}.xml", "w") as f:
|
||||
f.write(pascal_voc_xml)
|
||||
|
||||
@classmethod
|
||||
def from_pascal_voc(
|
||||
cls, images_directory_path: str, annotations_directory_path: str
|
||||
) -> Dataset:
|
||||
"""
|
||||
Creates a Dataset instance from PASCAL VOC formatted data.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.
|
||||
|
||||
Returns:
|
||||
Dataset: A Dataset instance containing the loaded images and annotations.
|
||||
"""
|
||||
image_paths = list_files_with_extensions(
|
||||
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
|
||||
)
|
||||
annotation_paths = list_files_with_extensions(
|
||||
directory=annotations_directory_path, extensions=["xml"]
|
||||
)
|
||||
|
||||
raw_annotations: List[Tuple[str, Detections, List[str]]] = [
|
||||
load_pascal_voc_annotations(annotation_path=str(annotation_path))
|
||||
for annotation_path in annotation_paths
|
||||
]
|
||||
|
||||
classes = []
|
||||
for annotation in raw_annotations:
|
||||
classes.extend(annotation[2])
|
||||
classes = list(set(classes))
|
||||
|
||||
for annotation in raw_annotations:
|
||||
class_id = [classes.index(class_name) for class_name in annotation[2]]
|
||||
annotation[1].class_id = np.array(class_id)
|
||||
|
||||
images = {
|
||||
image_path.name: cv2.imread(str(image_path)) for image_path in image_paths
|
||||
}
|
||||
|
||||
annotations = {
|
||||
image_name: detections for image_name, detections, _ in raw_annotations
|
||||
}
|
||||
return Dataset(classes=classes, images=images, annotations=annotations)
|
||||
@@ -0,0 +1,171 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from xml.dom.minidom import parseString
|
||||
from xml.etree.ElementTree import Element, SubElement, parse, tostring
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import (
|
||||
approximate_polygon,
|
||||
filter_polygons_by_area,
|
||||
mask_to_polygons,
|
||||
polygon_to_xyxy,
|
||||
)
|
||||
|
||||
|
||||
def object_to_pascal_voc(
|
||||
xyxy: np.ndarray, name: str, polygon: Optional[np.ndarray] = None
|
||||
) -> Element:
|
||||
root = Element("object")
|
||||
|
||||
object_name = SubElement(root, "name")
|
||||
object_name.text = name
|
||||
|
||||
bndbox = SubElement(root, "bndbox")
|
||||
xmin = SubElement(bndbox, "xmin")
|
||||
xmin.text = str(int(xyxy[0]))
|
||||
ymin = SubElement(bndbox, "ymin")
|
||||
ymin.text = str(int(xyxy[1]))
|
||||
xmax = SubElement(bndbox, "xmax")
|
||||
xmax.text = str(int(xyxy[2]))
|
||||
ymax = SubElement(bndbox, "ymax")
|
||||
ymax.text = str(int(xyxy[3]))
|
||||
|
||||
if polygon is not None:
|
||||
object_polygon = SubElement(root, "polygon")
|
||||
for index, point in enumerate(polygon, start=1):
|
||||
x_coordinate, y_coordinate = point
|
||||
x = SubElement(object_polygon, f"x{index}")
|
||||
x.text = str(x_coordinate)
|
||||
y = SubElement(object_polygon, f"y{index}")
|
||||
y.text = str(y_coordinate)
|
||||
|
||||
return root
|
||||
|
||||
|
||||
def detections_to_pascal_voc(
|
||||
detections: Detections,
|
||||
classes: List[str],
|
||||
filename: str,
|
||||
image_shape: Tuple[int, int, int],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> 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.
|
||||
image_shape (Tuple[int, int, int]): The shape of the image file associated with the detections.
|
||||
min_image_area_percentage (float): Minimum detection area relative to area of image associated with it.
|
||||
max_image_area_percentage (float): Maximum detection area relative to area of image associated with it.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
|
||||
Returns:
|
||||
str: An XML string in Pascal VOC format representing the detections.
|
||||
"""
|
||||
height, width, depth = image_shape
|
||||
image_area = height * width
|
||||
minimum_detection_area = min_image_area_percentage * image_area
|
||||
maximum_detection_area = max_image_area_percentage * image_area
|
||||
|
||||
# Create root element
|
||||
annotation = Element("annotation")
|
||||
|
||||
# Add folder element
|
||||
folder = SubElement(annotation, "folder")
|
||||
folder.text = "VOC"
|
||||
|
||||
# Add filename element
|
||||
file_name = SubElement(annotation, "filename")
|
||||
file_name.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 xyxy, mask, _, class_id, _ in detections:
|
||||
name = classes[class_id]
|
||||
if mask is not None:
|
||||
polygons = mask_to_polygons(mask=mask)
|
||||
if len(polygons) == 1:
|
||||
polygons = filter_polygons_by_area(
|
||||
polygons=polygons, min_area=None, max_area=maximum_detection_area
|
||||
)
|
||||
else:
|
||||
polygons = filter_polygons_by_area(
|
||||
polygons=polygons,
|
||||
min_area=minimum_detection_area,
|
||||
max_area=maximum_detection_area,
|
||||
)
|
||||
for polygon in polygons:
|
||||
approx_polygon = approximate_polygon(
|
||||
polygon=polygon, percentage=approximation_percentage
|
||||
)
|
||||
xyxy = polygon_to_xyxy(polygon=approx_polygon)
|
||||
next_object = object_to_pascal_voc(
|
||||
xyxy=xyxy, name=name, polygon=approx_polygon
|
||||
)
|
||||
annotation.append(next_object)
|
||||
else:
|
||||
next_object = object_to_pascal_voc(xyxy=xyxy, name=name)
|
||||
annotation.append(next_object)
|
||||
|
||||
# Generate XML string
|
||||
xml_string = parseString(tostring(annotation)).toprettyxml(indent=" ")
|
||||
|
||||
return xml_string
|
||||
|
||||
|
||||
def load_pascal_voc_annotations(
|
||||
annotation_path: str,
|
||||
) -> Tuple[str, Detections, List[str]]:
|
||||
"""
|
||||
Loads PASCAL VOC XML annotations and returns the image name, a Detections instance, and a list of class names.
|
||||
|
||||
Args:
|
||||
annotation_path (str): The path to the PASCAL VOC XML annotations file.
|
||||
|
||||
Returns:
|
||||
Tuple[str, Detections, List[str]]: A tuple containing the image name, a Detections instance, and a list of class names of objects in the detections.
|
||||
"""
|
||||
tree = parse(annotation_path)
|
||||
root = tree.getroot()
|
||||
|
||||
image_name = root.find("filename").text
|
||||
|
||||
xyxy = []
|
||||
class_names = []
|
||||
for obj in root.findall("object"):
|
||||
class_name = obj.find("name").text
|
||||
class_names.append(class_name)
|
||||
|
||||
bbox = obj.find("bndbox")
|
||||
x1 = int(bbox.find("xmin").text)
|
||||
y1 = int(bbox.find("ymin").text)
|
||||
x2 = int(bbox.find("xmax").text)
|
||||
y2 = int(bbox.find("ymax").text)
|
||||
|
||||
xyxy.append([x1, y1, x2, y2])
|
||||
|
||||
xyxy = np.array(xyxy)
|
||||
detections = Detections(xyxy=xyxy)
|
||||
|
||||
return image_name, detections, class_names
|
||||
@@ -147,10 +147,10 @@ class MaskAnnotator:
|
||||
Returns:
|
||||
np.ndarray: The image with the masks overlaid
|
||||
"""
|
||||
for i in range(len(detections.xyxy)):
|
||||
if detections.mask is None:
|
||||
continue
|
||||
if detections.mask is None:
|
||||
return scene
|
||||
|
||||
for i in np.flip(np.argsort(detections.area)):
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Iterator, List, Optional, Tuple, Union
|
||||
from typing import Any, Iterator, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -51,19 +51,18 @@ def _validate_tracker_id(tracker_id: Any, n: int) -> None:
|
||||
class Detections:
|
||||
"""
|
||||
Data class containing information about the detections in a video frame.
|
||||
|
||||
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.
|
||||
class_id (Optional[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
|
||||
mask: np.Optional[np.ndarray] = None
|
||||
class_id: Optional[np.ndarray] = None
|
||||
confidence: Optional[np.ndarray] = None
|
||||
class_id: Optional[np.ndarray] = None
|
||||
tracker_id: Optional[np.ndarray] = None
|
||||
|
||||
def __post_init__(self):
|
||||
@@ -319,9 +318,7 @@ class Detections:
|
||||
):
|
||||
return Detections(
|
||||
xyxy=self.xyxy[index],
|
||||
mask=self.mask[index]
|
||||
if self.mask is not None
|
||||
else None,
|
||||
mask=self.mask[index] if self.mask is not None else None,
|
||||
confidence=self.confidence[index]
|
||||
if self.confidence is not None
|
||||
else None,
|
||||
|
||||
@@ -35,7 +35,7 @@ class LineZone:
|
||||
detections (Detections): The detections for which to update the counts.
|
||||
|
||||
"""
|
||||
for xyxy, confidence, class_id, tracker_id in detections:
|
||||
for xyxy, _, confidence, class_id, tracker_id in detections:
|
||||
# handle detections with no tracker_id
|
||||
if tracker_id is None:
|
||||
continue
|
||||
|
||||
@@ -5,7 +5,7 @@ import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision import Detections
|
||||
from supervision.detection.utils import clip_boxes, generate_2d_mask
|
||||
from supervision.detection.utils import clip_boxes, polygon_to_mask
|
||||
from supervision.draw.color import Color
|
||||
from supervision.draw.utils import draw_polygon, draw_text
|
||||
from supervision.geometry.core import Position
|
||||
@@ -36,7 +36,7 @@ class PolygonZone:
|
||||
self.current_count = 0
|
||||
|
||||
width, height = frame_resolution_wh
|
||||
self.mask = generate_2d_mask(
|
||||
self.mask = polygon_to_mask(
|
||||
polygon=polygon, resolution_wh=(width + 1, height + 1)
|
||||
)
|
||||
|
||||
|
||||
@@ -1,11 +1,13 @@
|
||||
from typing import Tuple
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
MIN_POLYGON_POINT_COUNT = 3
|
||||
|
||||
def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
|
||||
"""Generate a 2D mask from a polygon.
|
||||
|
||||
def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
|
||||
"""Generate a mask from a polygon.
|
||||
|
||||
Args:
|
||||
polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices.
|
||||
@@ -146,3 +148,106 @@ def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
|
||||
bboxes[i, :] = [x_min, y_min, x_max, y_max]
|
||||
|
||||
return bboxes
|
||||
|
||||
|
||||
def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
|
||||
"""
|
||||
Converts a binary mask to a list of polygons.
|
||||
|
||||
Parameters:
|
||||
mask (np.ndarray): A binary mask represented as a 2D NumPy array of shape `(H, W)`,
|
||||
where H and W are the height and width of the mask, respectively.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3`
|
||||
are excluded from the output.
|
||||
"""
|
||||
|
||||
contours, _ = cv2.findContours(
|
||||
mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
|
||||
)
|
||||
return [
|
||||
np.squeeze(contour, axis=1)
|
||||
for contour in contours
|
||||
if contour.shape[0] >= MIN_POLYGON_POINT_COUNT
|
||||
]
|
||||
|
||||
|
||||
def filter_polygons_by_area(
|
||||
polygons: List[np.ndarray],
|
||||
min_area: Optional[float] = None,
|
||||
max_area: Optional[float] = None,
|
||||
) -> List[np.ndarray]:
|
||||
"""
|
||||
Filters a list of polygons based on their area.
|
||||
|
||||
Parameters:
|
||||
polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points.
|
||||
min_area (Optional[float]): The minimum area threshold. Only polygons with an area greater than or equal to this value
|
||||
will be included in the output. If set to None, no minimum area constraint will be applied.
|
||||
max_area (Optional[float]): The maximum area threshold. Only polygons with an area less than or equal to this value
|
||||
will be included in the output. If set to None, no maximum area constraint will be applied.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A new list of polygons containing only those with areas within the specified thresholds.
|
||||
"""
|
||||
if min_area is None and max_area is None:
|
||||
return polygons
|
||||
ares = [cv2.contourArea(polygon) for polygon in polygons]
|
||||
return [
|
||||
polygon
|
||||
for polygon, area in zip(polygons, ares)
|
||||
if (min_area is None or area >= min_area)
|
||||
and (max_area is None or area <= max_area)
|
||||
]
|
||||
|
||||
|
||||
def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Converts a polygon represented by a NumPy array into a bounding box.
|
||||
|
||||
Parameters:
|
||||
polygon (np.ndarray): A polygon represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A 1D NumPy array containing the bounding box `(x_min, y_min, x_max, y_max)` of the input polygon.
|
||||
"""
|
||||
x_min, y_min = np.min(polygon, axis=0)
|
||||
x_max, y_max = np.max(polygon, axis=0)
|
||||
return np.array([x_min, y_min, x_max, y_max])
|
||||
|
||||
|
||||
def approximate_polygon(polygon: np.ndarray, percentage: float) -> np.ndarray:
|
||||
"""
|
||||
Approximates a given polygon by reducing a certain percentage of points.
|
||||
|
||||
This function uses the Ramer-Douglas-Peucker algorithm to simplify the input polygon by reducing the number of points
|
||||
while preserving the general shape.
|
||||
|
||||
Parameters:
|
||||
polygon (np.ndarray): A 2D NumPy array of shape (N, 2) containing the x, y coordinates of the input polygon's points.
|
||||
percentage (float): The percentage of points to be removed from the input polygon, in the range [0, 1).
|
||||
|
||||
Returns:
|
||||
np.ndarray: A new 2D NumPy array of shape (M, 2), where M <= N * (1 - percentage), containing the x, y coordinates of the
|
||||
approximated polygon's points.
|
||||
"""
|
||||
|
||||
if percentage < 0 or percentage >= 1:
|
||||
raise ValueError("Percentage must be in the range [0, 1).")
|
||||
|
||||
target_points = max(int(len(polygon) * (1 - percentage)), 3)
|
||||
|
||||
if len(polygon) <= target_points:
|
||||
return polygon
|
||||
|
||||
epsilon = 0
|
||||
approximated_points = polygon
|
||||
while len(approximated_points) > target_points:
|
||||
epsilon += 0.1
|
||||
approximated_points = cv2.approxPolyDP(polygon, epsilon, closed=True)
|
||||
|
||||
return np.squeeze(approximated_points, axis=1)
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
|
||||
def write_text_to_file(content: str, output_path: str) -> None:
|
||||
with open(output_path, "w") as file:
|
||||
file.write(content)
|
||||
|
||||
|
||||
def list_files_with_extensions(
|
||||
directory: Union[str, Path], extensions: Optional[List[str]] = None
|
||||
) -> List[Path]:
|
||||
directory = Path(directory)
|
||||
files_with_extensions = []
|
||||
|
||||
if extensions is not None:
|
||||
for ext in extensions:
|
||||
files_with_extensions.extend(directory.glob(f"*.{ext}"))
|
||||
else:
|
||||
files_with_extensions.extend(directory.glob("*"))
|
||||
|
||||
return files_with_extensions
|
||||
@@ -1,11 +1,11 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Optional, Tuple
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
import pytest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.utils import non_max_suppression, clip_boxes
|
||||
from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -186,3 +186,93 @@ def test_non_max_suppression(
|
||||
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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"polygons, min_area, max_area, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
None,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # single polygon without area constraints
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
50,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # single polygon with min_area constraint
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
None,
|
||||
50,
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # single polygon with max_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
200,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
None,
|
||||
200,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with max_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
200,
|
||||
200,
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # two polygons with both area constraints
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
100,
|
||||
100,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area and max_area equal to the area of the first polygon
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
400,
|
||||
400,
|
||||
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area and max_area equal to the area of the second polygon
|
||||
]
|
||||
)
|
||||
def test_filter_polygons_by_area(
|
||||
polygons: List[np.ndarray],
|
||||
min_area: Optional[float],
|
||||
max_area: Optional[float],
|
||||
expected_result: List[np.ndarray],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = filter_polygons_by_area(polygons=polygons, min_area=min_area, max_area=max_area)
|
||||
assert len(result) == len(expected_result)
|
||||
for result_polygon, expected_result_polygon in zip(result, expected_result):
|
||||
assert np.array_equal(result_polygon, expected_result_polygon)
|
||||
|
||||
Reference in New Issue
Block a user