xiaoxue_admin
/
supervision
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

fix: 🐛 200+ E501 doc error lines are fixed. 

Signed-off-by: Onuralp SEZER <thunderbirdtr@gmail.com>
This commit is contained in:
Onuralp SEZER 2023-07-25 22:28:26 +03:00
parent 57e17c34a0
commit af82b4e2d2
No known key found for this signature in database
GPG Key ID: CF0835DFDF14CA38
21 changed files with 506 additions and 259 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
.pre-commit-config.yaml
View File

@ -51,22 +51,22 @@ repos:
- mdformat-black
exclude: "README.md|CITATION|CONTRIBUTING.md"
- repo: https://github.com/PyCQA/flake8
rev: 6.0.0
hooks:
- id: flake8
name: PEP8 Max Line Length Check
args:
- --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics
# - repo: https://github.com/PyCQA/flake8
# rev: 6.0.0
# hooks:
# - id: flake8
# name: PEP8 Max Line Length Check
# args:
# - --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics
- repo: https://github.com/PyCQA/flake8
rev: 6.0.0
hooks:
- id: flake8
name: PEP8
args:
- --count --max-line-length=88 --select=E9,F63,F7,F82 --show-source --statistics
# - repo: https://github.com/PyCQA/flake8
# rev: 6.0.0
# hooks:
# - id: flake8
# name: PEP8
# args:
# - --count --max-line-length=88 --select=E9,F63,F7,F82 --show-source --statistics
- repo: https://github.com/PyCQA/bandit
rev: '1.7.5'
@ -74,10 +74,11 @@ repos:
- id: bandit
args: ["-c", "pyproject.toml"]
additional_dependencies: ["bandit[toml]"]
- repo: https://github.com/PyCQA/autoflake
rev: v2.2.0
hooks:
- id: autoflake
# - repo: https://github.com/PyCQA/autoflake
# rev: v2.2.0
# hooks:
# - id: autoflake
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.0.280

12
supervision/classification/core.py
View File

@ -42,10 +42,12 @@ class Classifications:
@classmethod
def from_yolov8(cls, yolov8_results) -> Classifications:
"""
Creates a Classifications instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
Creates a Classifications instance from a
[YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
Args:
yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
yolov8_results (ultralytics.yolo.engine.results.Results):
The output Results instance from YOLOv8
Returns:
Detections: A new Classifications object.
@ -67,13 +69,15 @@ class Classifications:
def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
"""
Retrieve the top k class IDs and confidences, ordered in descending order by confidence.
Retrieve the top k class IDs and confidences,
ordered in descending order by confidence.
Args:
k (int): The number of top class IDs and confidences to retrieve.
Returns:
Tuple[np.ndarray, np.ndarray]: A tuple containing the top k class IDs and confidences.
Tuple[np.ndarray, np.ndarray]: A tuple containing
the top k class IDs and confidences.
Example:
```python

181
supervision/dataset/core.py
View File

@ -55,7 +55,8 @@ class DetectionDataset(BaseDataset):
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.
annotations (Dict[str, Detections]): Dictionary mapping
image name to annotations.
"""
classes: List[str]
@ -76,8 +77,9 @@ class DetectionDataset(BaseDataset):
Iterate over the images and annotations in the dataset.
Yields:
Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name,
the image data, and its corresponding annotation.
Iterator[Tuple[str, np.ndarray, Detections]]:
An iterator that yields tuples containing the image name,
the image data, and its corresponding annotation.
"""
for image_name, image in self.images.items():
yield image_name, image, self.annotations.get(image_name, None)
@ -101,22 +103,27 @@ class DetectionDataset(BaseDataset):
self, split_ratio=0.8, random_state=None, shuffle: bool = True
) -> Tuple[DetectionDataset, DetectionDataset]:
"""
Splits the dataset into two parts (training and testing) using the provided split_ratio.
Splits the dataset into two parts (training and testing)
using the provided split_ratio.
Args:
split_ratio (float, optional): The ratio of the training set to the entire dataset.
random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
split_ratio (float, optional): The ratio of the training
set to the entire dataset.
random_state (int, optional): The seed for the random number generator.
This is used for reproducibility.
shuffle (bool, optional): Whether to shuffle the data before splitting.
Returns:
Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets.
Tuple[DetectionDataset, DetectionDataset]: A tuple containing
the training and testing datasets.
Example:
```python
>>> import supervision as sv
>>> ds = sv.DetectionDataset(...)
>>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
>>> train_ds, test_ds = ds.split(split_ratio=0.7,
... random_state=42, shuffle=True)
>>> len(train_ds), len(test_ds)
(700, 300)
```
@ -151,19 +158,27 @@ class DetectionDataset(BaseDataset):
approximation_percentage: float = 0.0,
) -> None:
"""
Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
PASCAL VOC format.
Exports the dataset to PASCAL VOC format. This method saves the images
and their corresponding annotations in PASCAL VOC format.
Args:
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
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. Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage of detection area relative to
the image area for a detection to be included. Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
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.
Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage
of detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of
polygon points to be removed from the input polygon,
in the range [0, 1). Argument is used only for segmentation datasets.
"""
if images_directory_path:
@ -203,11 +218,14 @@ class DetectionDataset(BaseDataset):
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.
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:
DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
DetectionDataset: A DetectionDataset instance containing
the loaded images and annotations.
Example:
```python
@ -273,13 +291,19 @@ class DetectionDataset(BaseDataset):
Creates a Dataset instance from YOLO 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 YOLO annotation files.
data_yaml_path (str): The path to the data YAML file containing class information.
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
images_directory_path (str): The path to the
directory containing the images.
annotations_directory_path (str): The path to the directory
containing the YOLO annotation files.
data_yaml_path (str): The path to the data
YAML file containing class information.
force_masks (bool, optional): If True, forces
masks to be loaded for all annotations,
regardless of whether they are present.
Returns:
DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
DetectionDataset: A DetectionDataset instance
containing the loaded images and annotations.
Example:
```python
@ -322,23 +346,32 @@ class DetectionDataset(BaseDataset):
approximation_percentage: float = 0.0,
) -> None:
"""
Exports the dataset to YOLO format. This method saves the images and their corresponding
annotations in YOLO format.
Exports the dataset to YOLO format. This method saves the
images and their corresponding annotations in YOLO format.
Args:
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
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
YOLO format should be saved. If not provided, annotations will not be saved.
data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
annotations_directory_path (Optional[str]): The path to the
directory where the annotations in
YOLO format should be saved. If not provided,
annotations will not be saved.
data_yaml_path (Optional[str]): The path where the data.yaml
file should be saved.
If not provided, the file 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. Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage of detection area relative to
the image area for a detection to be included. Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
segmentation datasets.
min_image_area_percentage (float): The minimum percentage of
detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage
of detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of polygon points to
be removed from the input polygon, in the range [0, 1).
This is useful for simplifying the annotations.
Argument is used only for segmentation datasets.
"""
if images_directory_path is not None:
save_dataset_images(
@ -367,12 +400,16 @@ class DetectionDataset(BaseDataset):
Creates a Dataset instance from COCO formatted data.
Args:
images_directory_path (str): The path to the directory containing the images.
images_directory_path (str): The path to the
directory containing the images.
annotations_path (str): The path to the json annotation files.
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
force_masks (bool, optional): If True,
forces masks to be loaded for all annotations,
regardless of whether they are present.
Returns:
DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
DetectionDataset: A DetectionDataset instance containing
the loaded images and annotations.
Example:
```python
@ -412,20 +449,26 @@ class DetectionDataset(BaseDataset):
approximation_percentage: float = 0.0,
) -> None:
"""
Exports the dataset to COCO format. This method saves the images and their corresponding
annotations in COCO format.
Exports the dataset to COCO format. This method saves the
images and their corresponding annotations in COCO format.
Args:
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
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_path (Optional[str]): The path to COCO annotation file.
min_image_area_percentage (float): The minimum percentage of detection area relative to
the image area for a detection to be included. Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage of detection area relative to
the image area for a detection to be included. Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
segmentation datasets.
min_image_area_percentage (float): The minimum percentage of
detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
max_image_area_percentage (float): The maximum percentage of
detection area relative to
the image area for a detection to be included.
Argument is used only for segmentation datasets.
approximation_percentage (float): The percentage of polygon points
to be removed from the input polygon,
in the range [0, 1). This is useful for simplifying the annotations.
Argument is used only for segmentation datasets.
"""
if images_directory_path is not None:
save_dataset_images(
@ -445,16 +488,20 @@ class DetectionDataset(BaseDataset):
@classmethod
def merge(cls, dataset_list: List[DetectionDataset]) -> DetectionDataset:
"""
Merge a list of `DetectionDataset` objects into a single `DetectionDataset` object.
Merge a list of `DetectionDataset` objects into a single
`DetectionDataset` object.
This method takes a list of `DetectionDataset` objects and combines their respective fields (`classes`, `images`,
This method takes a list of `DetectionDataset` objects and combines
their respective fields (`classes`, `images`,
`annotations`) into a single `DetectionDataset` object.
Args:
dataset_list (List[DetectionDataset]): A list of `DetectionDataset` objects to merge.
dataset_list (List[DetectionDataset]): A list of `DetectionDataset`
objects to merge.
Returns:
(DetectionDataset): A single `DetectionDataset` object containing the merged data from the input list.
(DetectionDataset): A single `DetectionDataset` object containing
the merged data from the input list.
Example:
```python
@ -512,7 +559,8 @@ class ClassificationDataset(BaseDataset):
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.
annotations (Dict[str, Detections]): Dictionary mapping
image name to annotations.
"""
classes: List[str]
@ -526,22 +574,28 @@ class ClassificationDataset(BaseDataset):
self, split_ratio=0.8, random_state=None, shuffle: bool = True
) -> Tuple[ClassificationDataset, ClassificationDataset]:
"""
Splits the dataset into two parts (training and testing) using the provided split_ratio.
Splits the dataset into two parts (training and testing)
using the provided split_ratio.
Args:
split_ratio (float, optional): The ratio of the training set to the entire dataset.
random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
split_ratio (float, optional): The ratio of the training
set to the entire dataset.
random_state (int, optional): The seed for the
random number generator.
This is used for reproducibility.
shuffle (bool, optional): Whether to shuffle the data before splitting.
Returns:
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing the training and testing datasets.
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing
the training and testing datasets.
Example:
```python
>>> import supervision as sv
>>> cd = sv.ClassificationDataset(...)
>>> train_cd, test_cd = cd.split(split_ratio=0.7, random_state=42, shuffle=True)
>>> train_cd,test_cd = cd.split(split_ratio=0.7,
... random_state=42,shuffle=True)
>>> len(train_cd), len(test_cd)
(700, 300)
```
@ -571,7 +625,8 @@ class ClassificationDataset(BaseDataset):
Saves the dataset as a multi-class folder structure.
Args:
root_directory_path (str): The path to the directory where the dataset will be saved.
root_directory_path (str): The path to the directory
where the dataset will be saved.
"""
os.makedirs(root_directory_path, exist_ok=True)

25
supervision/dataset/formats/pascal_voc.py
View File

@ -52,13 +52,19 @@ def detections_to_pascal_voc(
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.
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).
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.
"""
@ -123,13 +129,16 @@ 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.
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.
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()

17
supervision/dataset/formats/yolo.py
View File

@ -112,16 +112,23 @@ def load_yolo_annotations(
force_masks: bool = False,
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
"""
Loads YOLO annotations and returns class names, images, and their corresponding detections.
Loads YOLO annotations and returns class names, images,
and their corresponding detections.
Args:
images_directory_path (str): The path to the directory containing the images.
annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
data_yaml_path (str): The path to the data YAML file containing class information.
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
annotations_directory_path (str): The path to the directory
containing the YOLO annotation files.
data_yaml_path (str): The path to the data
YAML file containing class information.
force_masks (bool, optional): If True, forces masks to be loaded
for all annotations, regardless of whether they are present.
Returns:
Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values.
Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
A tuple containing a list of class names, a dictionary with
image names as keys and images as values, and a dictionary
with image names as keys and corresponding Detections instances as values.
"""
image_paths = list_files_with_extensions(
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]

25
supervision/detection/annotate.py
View File

@ -12,12 +12,15 @@ 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
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
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
"""
@ -49,8 +52,11 @@ class BoxAnnotator:
Args:
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` are not provided, corresponding `class_id` will be used as label.
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` are not provided,
corresponding `class_id` will be used as 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
@ -145,7 +151,8 @@ 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
color (Union[Color, ColorPalette]): The color to fill the mask,
can be a single color or a color palette
"""
def __init__(
@ -158,11 +165,13 @@ class MaskAnnotator:
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.
Overlays the masks on the given image based on the provided detections,
with a specified opacity.
Args:
scene (np.ndarray): The image on which the masks will be overlaid
detections (Detections): The detections for 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:

147
supervision/detection/core.py
View File

@ -57,11 +57,16 @@ 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.
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): 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.
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
@ -96,7 +101,8 @@ class Detections:
]
]:
"""
Iterates over the Detections object and yield a tuple of `(xyxy, mask, 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 (
@ -141,10 +147,12 @@ class Detections:
@classmethod
def from_yolov5(cls, yolov5_results) -> Detections:
"""
Creates a Detections instance from a [YOLOv5](https://github.com/ultralytics/yolov5) inference result.
Creates a Detections instance from a
[YOLOv5](https://github.com/ultralytics/yolov5) inference result.
Args:
yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5
yolov5_results (yolov5.models.common.Detections):
The output Detections instance from YOLOv5
Returns:
Detections: A new Detections object.
@ -171,10 +179,12 @@ class Detections:
@classmethod
def from_yolov8(cls, yolov8_results) -> Detections:
"""
Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
Creates a Detections instance from a
[YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
Args:
yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
yolov8_results (ultralytics.yolo.engine.results.Results):
The output Results instance from YOLOv8
Returns:
Detections: A new Detections object.
@ -201,10 +211,15 @@ class Detections:
@classmethod
def from_yolo_nas(cls, yolo_nas_results) -> Detections:
"""
Creates a Detections instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) inference result.
Creates a Detections instance from a
[YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md)
inference result.
Args:
yolo_nas_results (super_gradients.training.models.prediction_results.ImageDetectionPrediction): The output Results instance from YOLO-NAS
yolo_nas_results (ImageDetectionPrediction):
The output Results instance from YOLO-NAS
ImageDetectionPrediction is coming from
'super_gradients.training.models.prediction_results'
Returns:
Detections: A new Detections object.
@ -230,11 +245,13 @@ class Detections:
@classmethod
def from_mmdetection(cls, mmdet_results) -> Detections:
"""
Creates a Detections instance from a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result.
Creates a Detections instance from
a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result.
Also supported for [mmyolo](https://github.com/open-mmlab/mmyolo)
Args:
mmdet_results (mmdet.structures.DetDataSample): The output Results instance from MMDetection
mmdet_results (mmdet.structures.DetDataSample):
The output Results instance from MMDetection
Returns:
Detections: A new Detections object.
@ -246,7 +263,8 @@ class Detections:
>>> from mmdet.apis import DetInferencer
>>> inferencer = DetInferencer(model_name, checkpoint, device)
>>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output', return_datasample=True)["predictions"][0]
>>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output',
... return_datasample=True)["predictions"][0]
>>> detections = sv.Detections.from_mmdet(mmdet_result)
```
"""
@ -259,7 +277,8 @@ class Detections:
@classmethod
def from_transformers(cls, transformers_results: dict) -> Detections:
"""
Creates a Detections instance from object detection [transformer](https://github.com/huggingface/transformers) inference result.
Creates a Detections instance from object detection
[transformer](https://github.com/huggingface/transformers) inference result.
Returns:
Detections: A new Detections object.
@ -273,13 +292,16 @@ class Detections:
@classmethod
def from_detectron2(cls, detectron2_results) -> Detections:
"""
Create a Detections object from the [Detectron2](https://github.com/facebookresearch/detectron2) inference result.
Create a Detections object from the
[Detectron2](https://github.com/facebookresearch/detectron2) inference result.
Args:
detectron2_results: The output of a Detectron2 model containing instances with prediction data.
detectron2_results: The output of a
Detectron2 model containing instances with prediction data.
Returns:
(Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
(Detections): A Detections object containing the bounding boxes,
class IDs, and confidences of the predictions.
Example:
```python
@ -294,7 +316,6 @@ class Detections:
>>> cfg.MODEL.WEIGHTS = "path/to/model_weights.pth"
>>> predictor = DefaultPredictor(cfg)
>>> result = predictor(image)
>>> detections = sv.Detections.from_detectron2(result)
```
"""
@ -310,14 +331,18 @@ class Detections:
@classmethod
def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections:
"""
Create a Detections object from the [Roboflow](https://roboflow.com/) API inference result.
Create a Detections object from the [Roboflow](https://roboflow.com/)
API inference result.
Args:
roboflow_result (dict): The result from the Roboflow API containing predictions.
class_list (List[str]): A list of class names corresponding to the class IDs in the API result.
roboflow_result (dict): The result from the
Roboflow API containing predictions.
class_list (List[str]): A list of class names
corresponding to the class IDs in the API result.
Returns:
(Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
(Detections): A Detections object containing the bounding boxes, class IDs,
and confidences of the predictions.
Example:
```python
@ -354,7 +379,9 @@ class Detections:
@classmethod
def from_sam(cls, sam_result: List[dict]) -> Detections:
"""
Creates a Detections instance from [Segment Anything Model](https://github.com/facebookresearch/segment-anything) inference result.
Creates a Detections instance from
[Segment Anything Model](https://github.com/facebookresearch/segment-anything)
inference result.
Args:
sam_result (List[dict]): The output Results instance from SAM
@ -365,14 +392,19 @@ class Detections:
Example:
```python
>>> import supervision as sv
>>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
>>> from segment_anything import (
... sam_model_registry,
... SamAutomaticMaskGenerator
... )
>>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
>>> sam_model_reg = sam_model_registry[MODEL_TYPE]
>>> sam = sam_model_reg(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
)
@ -385,7 +417,8 @@ class Detections:
@classmethod
def empty(cls) -> Detections:
"""
Create an empty Detections object with no bounding boxes, confidences, or class IDs.
Create an empty Detections object with no bounding boxes,
confidences, or class IDs.
Returns:
(Detections): An empty Detections object.
@ -408,15 +441,18 @@ class Detections:
"""
Merge a list of Detections objects into a single Detections object.
This method takes a list of Detections objects and combines their respective fields (`xyxy`, `mask`,
`confidence`, `class_id`, and `tracker_id`) into a single Detections object. If all elements in a field are not
`None`, the corresponding field will be stacked. Otherwise, the field will be set to `None`.
This method takes a list of Detections objects and combines their
respective fields (`xyxy`, `mask`, `confidence`, `class_id`, and `tracker_id`)
into a single Detections object. If all elements in a field are not
`None`, the corresponding field will be stacked.
Otherwise, the field will be set to `None`.
Args:
detections_list (List[Detections]): A list of Detections objects to merge.
Returns:
(Detections): A single Detections object containing the merged data from the input list.
(Detections): A single Detections object containing
the merged data from the input list.
Example:
```python
@ -436,14 +472,14 @@ class Detections:
list(field) for field in zip(*detections_tuples_list)
]
def all_not_none(l):
return all(x is not None for x in l)
def __all_not_none(item_list: List[Any]):
return all(x is not None for x in item_list)
xyxy = np.vstack(xyxy)
mask = np.vstack(mask) if all_not_none(mask) else None
confidence = np.hstack(confidence) if all_not_none(confidence) else None
class_id = np.hstack(class_id) if all_not_none(class_id) else None
tracker_id = np.hstack(tracker_id) if all_not_none(tracker_id) else None
mask = np.vstack(mask) if __all_not_none(mask) else None
confidence = np.hstack(confidence) if __all_not_none(confidence) else None
class_id = np.hstack(class_id) if __all_not_none(class_id) else None
tracker_id = np.hstack(tracker_id) if __all_not_none(tracker_id) else None
return cls(
xyxy=xyxy,
@ -458,10 +494,12 @@ class Detections:
Returns the bounding box coordinates for a specific anchor.
Args:
anchor (Position): Position of bounding box anchor for which to return the coordinates.
anchor (Position): Position of bounding box anchor
for which to return the coordinates.
Returns:
np.ndarray: An array of shape `(n, 2)` containing the bounding box anchor coordinates in format `[x, y]`.
np.ndarray: An array of shape `(n, 2)` containing the bounding
box anchor coordinates in format `[x, y]`.
"""
if anchor == Position.CENTER:
return np.array(
@ -484,7 +522,8 @@ class Detections:
Get a subset of the Detections object.
Args:
index (Union[int, slice, List[int], np.ndarray]): The index or indices of the subset of the Detections
index (Union[int, slice, List[int], np.ndarray]):
The index or indices of the subset of the Detections
Returns:
(Detections): A subset of the Detections object.
@ -519,11 +558,14 @@ 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.
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.
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])
@ -536,7 +578,9 @@ class Detections:
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.
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])
@ -547,14 +591,19 @@ class 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.
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.
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.
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

3
supervision/detection/line_counter.py
View File

@ -113,7 +113,8 @@ class LineZoneAnnotator:
Attributes:
frame (np.ndarray): The image on which the line will be drawn.
line_counter (LineCounter): The line counter that will be used to draw the line.
line_counter (LineCounter): The line counter
that will be used to draw the line.
Returns:
np.ndarray: The image with the line drawn on it.

21
supervision/detection/tools/polygon_zone.py
View File

@ -17,9 +17,11 @@ class PolygonZone:
A class for defining a polygon-shaped zone within a frame for detecting objects.
Attributes:
polygon (np.ndarray): A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points.
polygon (np.ndarray): A polygon represented by a numpy array of shape
`(N, 2)`, containing the `x`, `y` coordinates of the points.
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)
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
"""
@ -45,10 +47,12 @@ class PolygonZone:
Determines if the detections are within the polygon zone.
Parameters:
detections (Detections): The detections to be checked against the polygon zone
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
np.ndarray: A boolean numpy array indicating
if each detection is within the polygon zone
"""
clipped_xyxy = clip_boxes(
@ -65,7 +69,8 @@ class PolygonZone:
class PolygonZoneAnnotator:
"""
A class for annotating a polygon-shaped zone within a frame with a count of detected objects.
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
@ -75,7 +80,8 @@ class PolygonZoneAnnotator:
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
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
"""
@ -105,7 +111,8 @@ class PolygonZoneAnnotator:
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)
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

105
supervision/detection/utils.py
View File

@ -10,11 +10,13 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n
"""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.
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.
Returns:
np.ndarray: The generated 2D mask, where the polygon is marked with `1`'s and the rest is filled with `0`'s.
np.ndarray: The generated 2D mask, where the polygon is marked with
`1`'s and the rest is filled with `0`'s.
"""
width, height = resolution_wh
mask = np.zeros((height, width), dtype=np.uint8)
@ -24,15 +26,20 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n
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.
Compute Intersection over Union (IoU) of two sets of bounding boxes -
`boxes_true` and `boxes_detection`. Both sets
of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
Args:
boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. `shape = (N, 4)` where `N` is number of true objects.
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. `shape = (M, 4)` where `M` is number of detected objects.
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.
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):
@ -55,14 +62,19 @@ def non_max_suppression(
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.
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.
np.ndarray: A boolean array indicating which predictions to keep after n
on-maximum suppression.
Raises:
AssertionError: If `iou_threshold` is not within the closed range from `0` to `1`.
AssertionError: If `iou_threshold` is not within the
closed range from `0` to `1`.
"""
assert 0 <= iou_threshold <= 1, (
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
@ -89,7 +101,8 @@ def non_max_suppression(
if not keep[index]:
continue
# drop detections with iou > iou_threshold and same category as current detections
# drop detections with iou > iou_threshold and
# same category as current detections
condition = (iou > iou_threshold) & (categories == category)
keep = keep & ~condition
@ -103,14 +116,16 @@ def clip_boxes(
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
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.
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.
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
@ -131,10 +146,12 @@ 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
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
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)
@ -155,12 +172,14 @@ 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.
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`
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.
"""
@ -183,15 +202,21 @@ def filter_polygons_by_area(
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)`,
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.
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.
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
@ -213,7 +238,8 @@ def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray:
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.
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)
@ -226,16 +252,23 @@ def approximate_polygon(
"""
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.
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)`.
epsilon_step (float): Approximation accuracy step. Epsilon is the maximum distance between the original curve and its approximation.
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)`.
epsilon_step (float): Approximation accuracy step.
Epsilon is the maximum distance between the original curve
and its approximation.
Returns:
np.ndarray: A new 2D NumPy array of shape `(M, 2)`, where `M <= N * (1 - percentage)`, containing the `x`, `y` coordinates of the
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.
"""

6
supervision/draw/color.py
View File

@ -42,7 +42,8 @@ class Color:
"""
Creates a Color instance from a color hex string
:param color_hex: str : The color hex string in the format of "fff", "ffffff", "#fff", or "#ffffff"
:param color_hex: str : The color hex string in the format
of "fff", "ffffff", "#fff", or "#ffffff"
:return: Color : A Color instance representing the color
Example:
@ -105,7 +106,8 @@ class ColorPalette:
"""
Creates a ColorPalette instance from a list of color hex strings
:param color_hex_list: List[str] : A list of color hex strings in the format of "fff", "ffffff", "#fff", or "#ffffff"
:param color_hex_list: List[str] : A list of color hex strings in the
format of "fff", "ffffff", "#fff", or "#ffffff"
:return: ColorPalette : A ColorPalette instance representing the color palette
Example:

16
supervision/draw/utils.py
View File

@ -115,15 +115,19 @@ def draw_text(
Draw text with background on a scene.
Parameters:
scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene.
scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene
text (str): The text to be drawn.
text_anchor (Point): The anchor point for the text, represented as a Point object with x and y attributes.
text_anchor (Point): The anchor point for the text, represented as a
Point object with x and y attributes.
text_color (Color, optional): The color of the text. Defaults to black.
text_scale (float, optional): The scale of the text. Defaults to 0.5.
text_thickness (int, optional): The thickness of the text. Defaults to 1.
text_padding (int, optional): The amount of padding to add around the text when drawing a rectangle in the background. Defaults to 10.
text_font (int, optional): The font to use for the text. Defaults to cv2.FONT_HERSHEY_SIMPLEX.
background_color (Color, optional): The color of the background rectangle, if one is to be drawn. Defaults to None.
text_padding (int, optional): The amount of padding to add around the text
when drawing a rectangle in the background. Defaults to 10.
text_font (int, optional): The font to use for the text.
Defaults to cv2.FONT_HERSHEY_SIMPLEX.
background_color (Color, optional): The color of the background rectangle,
if one is to be drawn. Defaults to None.
Returns:
np.ndarray: The input scene with the text drawn on it.
@ -132,7 +136,7 @@ def draw_text(
```python
>>> scene = np.zeros((100, 100, 3), dtype=np.uint8)
>>> text_anchor = Point(x=50, y=50)
>>> scene = draw_text(scene=scene, text="Hello, world!", text_anchor=text_anchor)
>>> scene = draw_text(scene=scene, text="Hello, world!",text_anchor=text_anchor)
```
"""
text_width, text_height = cv2.getTextSize(

11
supervision/geometry/utils.py
View File

@ -7,13 +7,18 @@ def get_polygon_center(polygon: np.ndarray) -> Point:
"""
Calculate the center of a polygon.
This function takes in a polygon as a 2-dimensional numpy ndarray and returns the center of the polygon as a Point object. The center is calculated as the mean of the polygon's vertices along each axis, and is rounded down to the nearest integer.
This function takes in a polygon as a 2-dimensional numpy ndarray and
returns the center of the polygon as a Point object.
The center is calculated as the mean of the polygon's vertices along each axis,
and is rounded down to the nearest integer.
Parameters:
polygon (np.ndarray): A 2-dimensional numpy ndarray representing the vertices of the polygon.
polygon (np.ndarray): A 2-dimensional numpy ndarray representing the
vertices of the polygon.
Returns:
Point: The center of the polygon, represented as a Point object with x and y attributes.
Point: The center of the polygon, represented as a
Point object with x and y attributes.
Examples:
```python

67
supervision/metrics/detection.py
View File

@ -18,10 +18,14 @@ class ConfusionMatrix:
Confusion matrix for object detection tasks.
Attributes:
matrix (np.ndarray): An 2D `np.ndarray` of shape `(len(classes) + 1, len(classes) + 1)` containing the number of `TP`, `FP`, `FN` and `TN` for each class.
matrix (np.ndarray): An 2D `np.ndarray` of shape
`(len(classes) + 1, len(classes) + 1)`
containing the number of `TP`, `FP`, `FN` and `TN` for each class.
classes (List[str]): Model class names.
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded from the matrix.
iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`.
conf_threshold (float): Detection confidence threshold between `0` and `1`.
Detections with lower confidence will be excluded from the matrix.
iou_threshold (float): Detection IoU threshold between `0` and `1`.
Detections with lower IoU will be classified as `FP`.
"""
matrix: np.ndarray
@ -45,8 +49,10 @@ class ConfusionMatrix:
targets (List[Detections]): Detections objects from ground-truth.
predictions (List[Detections]): Detections objects predicted by the model.
classes (List[str]): Model class names.
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`.
conf_threshold (float): Detection confidence threshold between `0` and `1`.
Detections with lower confidence will be excluded.
iou_threshold (float): Detection IoU threshold between `0` and `1`.
Detections with lower IoU will be classified as `FP`.
Returns:
ConfusionMatrix: New instance of ConfusionMatrix.
@ -132,11 +138,19 @@ class ConfusionMatrix:
Calculate confusion matrix based on predicted and ground-truth detections.
Args:
predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format.
targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format.
predictions (List[np.ndarray]): Each element of the list describes a single
image and has `shape = (M, 6)` where `M` is the number of detected
objects. Each row is expected to be in
`(x_min, y_min, x_max, y_max, class, conf)` format.
targets (List[np.ndarray]): Each element of the list describes a single
image and has `shape = (N, 5)` where `N` is the number of
ground-truth objects. Each row is expected to be in
`(x_min, y_min, x_max, y_max, class)` format.
classes (List[str]): Model class names.
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
iou_threshold (float): Detection iou threshold between `0` and `1`. Detections with lower iou will be classified as `FP`.
conf_threshold (float): Detection confidence threshold between `0` and `1`.
Detections with lower confidence will be excluded.
iou_threshold (float): Detection iou threshold between `0` and `1`.
Detections with lower iou will be classified as `FP`.
Returns:
ConfusionMatrix: New instance of ConfusionMatrix.
@ -246,11 +260,19 @@ class ConfusionMatrix:
Calculate confusion matrix for a batch of detections for a single image.
Args:
predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format.
targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format.
predictions (List[np.ndarray]): Each element of the list describes a single
image and has `shape = (M, 6)` where `M` is the number of detected
objects. Each row is expected to be in
`(x_min, y_min, x_max, y_max, class, conf)` format.
targets (List[np.ndarray]): Each element of the list describes a single
image and has `shape = (N, 5)` where `N` is the number of ground-truth
objects. Each row is expected to be in
`(x_min, y_min, x_max, y_max, class)` format.
num_classes (int): Number of classes.
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
iou_threshold (float): Detection iou threshold between `0` and `1`. Detections with lower iou will be classified as `FP`.
conf_threshold (float): Detection confidence threshold between `0` and `1`.
Detections with lower confidence will be excluded.
iou_threshold (float): Detection iou threshold between `0` and `1`.
Detections with lower iou will be classified as `FP`.
Returns:
np.ndarray: Confusion matrix based on a single image.
@ -304,8 +326,8 @@ class ConfusionMatrix:
@staticmethod
def _drop_extra_matches(matches: np.ndarray) -> np.ndarray:
"""
Deduplicate matches. If there are multiple matches for the same true or predicted box,
only the one with the highest IoU is kept.
Deduplicate matches. If there are multiple matches for the same true or
predicted box, only the one with the highest IoU is kept.
"""
if matches.shape[0] > 0:
matches = matches[matches[:, 2].argsort()[::-1]]
@ -327,9 +349,12 @@ class ConfusionMatrix:
Args:
dataset (DetectionDataset): Object detection dataset used for evaluation.
callback (Callable[[np.ndarray], Detections]): Function that takes an image as input and returns Detections object.
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`.
callback (Callable[[np.ndarray], Detections]): Function that takes an image
as input and returns Detections object.
conf_threshold (float): Detection confidence threshold between `0` and `1`.
Detections with lower confidence will be excluded.
iou_threshold (float): Detection IoU threshold between `0` and `1`.
Detections with lower IoU will be classified as `FP`.
Returns:
ConfusionMatrix: New instance of ConfusionMatrix.
@ -386,9 +411,11 @@ class ConfusionMatrix:
Create confusion matrix plot and save it at selected location.
Args:
save_path (Optional[str]): Path to save the plot. If not provided, plot will be displayed.
save_path (Optional[str]): Path to save the plot. If not provided,
plot will be displayed.
title (Optional[str]): Title of the plot.
classes (Optional[List[str]]): List of classes to be displayed on the plot. If not provided, all classes will be displayed.
classes (Optional[List[str]]): List of classes to be displayed on the plot.
If not provided, all classes will be displayed.
normalize (bool): If True, normalize the confusion matrix.
fig_size (Tuple[int, int]): Size of the plot.

10
supervision/utils/file.py
View File

@ -21,11 +21,13 @@ def list_files_with_extensions(
directory: Union[str, Path], extensions: Optional[List[str]] = None
) -> List[Path]:
"""
List files in a directory with specified extensions or all files if no extensions are provided.
List files in a directory with specified extensions or
all files if no extensions are provided.
Args:
directory (Union[str, Path]): The directory path as a string or Path object.
extensions (Optional[List[str]]): A list of file extensions to filter. Default is None, which lists all files.
extensions (Optional[List[str]]): A list of file extensions to filter.
Default is None, which lists all files.
Returns:
(List[Path]): A list of Path objects for the matching files.
@ -38,9 +40,11 @@ def list_files_with_extensions(
>>> files = sv.list_files_with_extensions(directory='my_directory')
>>> # List only files with '.txt' and '.md' extensions
>>> files = sv.list_files_with_extensions(directory='my_directory', extensions=['txt', 'md'])
>>> files = sv.list_files_with_extensions(
... directory='my_directory', extensions=['txt', 'md'])
```
"""
directory = Path(directory)
files_with_extensions = []

20
supervision/utils/image.py
View File

@ -12,7 +12,8 @@ def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
Args:
image (np.ndarray): The image to be cropped, represented as a numpy array.
xyxy (np.ndarray): A numpy array containing the bounding box coordinates in the format (x1, y1, x2, y2).
xyxy (np.ndarray): A numpy array containing the bounding box coordinates
in the format (x1, y1, x2, y2).
Returns:
(np.ndarray): The cropped image as a numpy array.
@ -47,18 +48,23 @@ class ImageSink:
Args:
target_dir_path (str): The target directory where images will be saved.
overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".
overwrite (bool, optional): Whether to overwrite the existing directory.
Defaults to False.
image_name_pattern (str, optional): The image file name pattern.
Defaults to "image_{:05d}.png".
Examples:
```python
>>> import supervision as sv
>>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
... for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
>>> with sv.ImageSink(target_dir_path='target/directory/path',
... overwrite=True) as sink:
... for image in sv.get_video_frames_generator(
... source_path='source_video.mp4', stride=2):
... sink.save_image(image=image)
```
"""
self.target_dir_path = target_dir_path
self.overwrite = overwrite
self.image_name_pattern = image_name_pattern
@ -80,7 +86,9 @@ class ImageSink:
Args:
image (np.ndarray): The image to be saved.
image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
image_name (str, optional): The name to use for the saved image.
If not provided, a name will be
generated using the `image_name_pattern`.
"""
if image_name is None:
image_name = self.image_name_pattern.format(self.image_count)

9
supervision/utils/notebook.py
View File

@ -50,9 +50,12 @@ def plot_images_grid(
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.
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:

35
supervision/utils/video.py
View File

@ -10,13 +10,15 @@ import numpy as np
@dataclass
class VideoInfo:
"""
A class to store video information, including width, height, fps and total number of frames.
A class to store video information, including width, height, fps and
total number of frames.
Attributes:
width (int): width of the video in pixels
height (int): height of the video in pixels
fps (int): frames per second of the video
total_frames (int, optional): total number of frames in the video, default is None
total_frames (int, optional): total number of frames in the video,
default is None
Examples:
```python
@ -61,7 +63,8 @@ class VideoSink:
Attributes:
target_path (str): The path to the output file where the video will be saved.
video_info (VideoInfo): Information about the video resolution, fps, and total frame count.
video_info (VideoInfo): Information about the video resolution, fps,
and total frame count.
Example:
```python
@ -69,8 +72,10 @@ class VideoSink:
>>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4')
>>> with sv.VideoSink(target_path='target_video.mp4', video_info=video_info) as sink:
... for frame in get_video_frames_generator(source_path='source_video.mp4', stride=2):
>>> with sv.VideoSink(target_path='target_video.mp4',
... video_info=video_info) as sink:
... for frame in get_video_frames_generator(source_path='source_video.mp4',
... stride=2):
... sink.write_frame(frame=frame)
```
"""
@ -118,12 +123,16 @@ def get_video_frames_generator(
Args:
source_path (str): The path of the video file.
stride (int): Indicates the interval at which frames are returned, skipping stride - 1 frames between each.
start (int): Indicates the starting position from which video should generate frames
end (Optional[int]): Indicates the ending position at which video should stop generating frames. If None, video will be read to the end.
stride (int): Indicates the interval at which frames are returned,
skipping stride - 1 frames between each.
start (int): Indicates the starting position from which
video should generate frames
end (Optional[int]): Indicates the ending position at which video
should stop generating frames. If None, video will be read to the end.
Returns:
(Generator[np.ndarray, None, None]): A generator that yields the frames of the video.
(Generator[np.ndarray, None, None]): A generator that yields the
frames of the video.
Examples:
```python
@ -154,12 +163,16 @@ def process_video(
callback: Callable[[np.ndarray, int], np.ndarray],
) -> None:
"""
Process a video file by applying a callback function on each frame and saving the result to a target video file.
Process a video file by applying a callback function on each frame
and saving the result to a target video file.
Args:
source_path (str): The path to the source video file.
target_path (str): The path to the target video file.
callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in a numpy ndarray representation of a video frame and an int index of the frame and returns a processed numpy ndarray representation of the frame.
callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in
a numpy ndarray representation of a video frame and an
int index of the frame and returns a processed numpy ndarray
representation of the frame.
Examples:
```python

9
test/dataset/formats/test_coco.py
View File

@ -47,7 +47,8 @@ def mock_cock_coco_annotation(
],
["fashion-assistant", "baseball cap"],
DoesNotRaise(),
), # two coco categories; one with supercategory == "none" and one with supercategory != "none"
), # two coco categories; one with supercategory == "none" and
# one with supercategory != "none"
(
[
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
@ -56,7 +57,8 @@ def mock_cock_coco_annotation(
],
["fashion-assistant", "baseball cap", "hoodie"],
DoesNotRaise(),
), # three coco categories; one with supercategory == "none" and two with supercategory != "none"
), # three coco categories; one with supercategory == "none" and
# two with supercategory != "none"
(
[
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
@ -65,7 +67,8 @@ def mock_cock_coco_annotation(
],
["fashion-assistant", "baseball cap", "hoodie"],
DoesNotRaise(),
), # three coco categories; one with supercategory == "none" and two with supercategory != "none" (different order)
), # three coco categories; one with supercategory == "none" and
# two with supercategory != "none" (different order)
],
)
def test_coco_categories_to_classes(

3
test/dataset/formats/test_yolo.py
View File

@ -154,7 +154,8 @@ def test_with_mask(
),
),
DoesNotRaise(),
), # yolo annotation file with two lines - one box and one polygon in with_masks mode
), # yolo annotation file with two lines -
# one box and one polygon in with_masks mode
(
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6", "1 0.11 0.47 0.22 0.30"],
(1000, 1000),

6
test/detection/test_utils.py
View File

@ -224,7 +224,8 @@ def test_clip_boxes(
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
), # 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]]),
@ -234,7 +235,8 @@ def test_clip_boxes(
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
), # two polygons with min_area and
# max_area equal to the area of the second polygon
],
)
def test_filter_polygons_by_area(