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

Merge remote-tracking branch 'origin/develop' into feature/fps

This commit is contained in:
Hardik Dava 2023-10-03 17:53:52 +02:00
parent 5d71ab2961 b47bc808f7
commit 3b00e04101
65 changed files with 3964 additions and 2722 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.gitattributes vendored
View File

@ -1 +1 @@
*.ipynb linguist-vendored
*.ipynb linguist-vendored

2
.github/ISSUE_TEMPLATE/bug-report.yml vendored
View File

@ -62,4 +62,4 @@ body:
description: >
(Optional) We encourage you to submit a [Pull Request](https://github.com/roboflow/supervision/pulls) (PR) to help improve Supervision for everyone, especially if you have a good understanding of how to implement a fix or feature.
options:
- label: Yes I'd like to help by submitting a PR!
- label: Yes I'd like to help by submitting a PR!

2
.github/ISSUE_TEMPLATE/feature-request.yml vendored
View File

@ -46,4 +46,4 @@ body:
description: >
(Optional) We encourage you to submit a [Pull Request](https://github.com/roboflow/supervision/pulls) (PR) to help improve Supervision for everyone, especially if you have a good understanding of how to implement a fix or feature.
options:
- label: Yes I'd like to help by submitting a PR!
- label: Yes I'd like to help by submitting a PR!

2
.github/ISSUE_TEMPLATE/question.yml vendored
View File

@ -30,4 +30,4 @@ body:
- type: textarea
attributes:
label: Additional
description: Anything else you would like to share?
description: Anything else you would like to share?

5
.github/workflows/test.yml vendored
View File

@ -25,8 +25,5 @@ jobs:
pip install isort
pip install flake8
pip install "black==22.3.0"
- name: 🧹 Lint with flake8
run: |
make check_code_quality
- name: 🧪 Test
run: "python -m pytest ./test"
run: "python -m pytest ./test"

2
.github/workflows/welcome.yml vendored
View File

@ -15,4 +15,4 @@ jobs:
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
issue-message: "Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap."
pr-message: "Hello there, thank you for opening an PR ! 🙏🏻 The team was notified and they will get back to you asap."
pr-message: "Hello there, thank you for opening an PR ! 🙏🏻 The team was notified and they will get back to you asap."

80
.pre-commit-config.yaml Normal file
View File

@ -0,0 +1,80 @@
ci:
autofix_prs: true
autoupdate_schedule: weekly
autofix_commit_msg: "fix(pre_commit): 🎨 auto format pre-commit hooks"
autoupdate_commit_msg: "chore(pre_commit): ⬆ pre_commit autoupdate"
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.4.0
hooks:
- id: end-of-file-fixer
- id: trailing-whitespace
- id: check-yaml
- id: check-docstring-first
- id: check-executables-have-shebangs
- id: check-toml
- id: check-case-conflict
- id: check-added-large-files
args: ['--maxkb=2048']
exclude: ^logo/
- id: detect-private-key
- id: forbid-new-submodules
- id: pretty-format-json
args: ['--autofix', '--no-sort-keys', '--indent=4']
- id: end-of-file-fixer
- id: mixed-line-ending
# - repo: https://github.com/asottile/pyupgrade
# rev: v3.9.0
# hooks:
# - id: pyupgrade
# name: Upgrade code
# args: [--py38-plus]
- repo: https://github.com/PyCQA/isort
rev: 5.12.0
hooks:
- id: isort
name: Sort imports
- repo: https://github.com/PyCQA/flake8
rev: 6.0.0
hooks:
- id: flake8
name: Flake8 Checks
entry: pflake8
additional_dependencies: [pyproject-flake8]
- repo: https://github.com/PyCQA/bandit
rev: '1.7.5'
hooks:
- id: bandit
args: ["-c", "pyproject.toml"]
additional_dependencies: ["bandit[toml]"]
- repo: https://github.com/pycqa/isort
rev: 5.12.0
hooks:
- id: isort
name: isort (python)
- id: isort
name: isort (cython)
types: [cython]
- id: isort
name: isort (pyi)
types: [pyi]
- repo: https://github.com/psf/black
rev: 23.7.0
hooks:
- id: black
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.0.280
hooks:
- id: ruff
args: [--fix, --exit-non-zero-on-fix]

27
CITATION.cff
View File

@ -1,12 +1,23 @@
# This CITATION.cff file was generated with cffinit.
# Visit https://bit.ly/cffinit to generate yours today!
cff-version: 1.2.0
author: Roboflow
message: If you use this software, please cite it as below.
title: supervision
version: 0.11.1
date-released: 2023-01-19
license: MIT
repository-code: https://github.com/roboflow/supervision
title: Supervision
message: >-
If you use this software, please cite it using the
metadata from this file.
type: software
authors:
- given-names: Roboflow
email: support@roboflow.com
repository-code: 'https://github.com/roboflow/supervision'
url: 'https://roboflow.github.io/supervision/'
abstract: >-
supervision features a range of utilities for use in
computer vision projects, from detections processing and
filtering to confusion matrix calcuation.
keywords:
- computer vision
- image processing
- video processing
- video processing
license: MIT

8
CONTRIBUTING.md
View File

@ -32,21 +32,21 @@ When creating new functions, please ensure you have the following:
2. Unit tests for the function.
3. Examples in the documentation for the function.
4. Created an entry in our docs to autogenerate the documentation for the function.
5. Please share google colab with minimal code to test new feature or reproduce PR whenever it is possible. Please ensure that google colab can be accessed without any issue.
5. Please share google colab with minimal code to test new feature or reproduce PR whenever it is possible. Please ensure that google colab can be accessed without any issue.
All pull requests will be reviewed by the maintainers of the project. We will provide feedback and ask for changes if necessary.
PRs must pass all tests and linting requirements before they can be merged.
## 🧹 code quality
## 🧹 code quality
We provide two handy commands inside the `Makefile`, namely:
- `make style` to format the code
- `make check_code_quality` to check code quality (PEP8 basically)
So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
## 🧪 tests
## 🧪 tests
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.

8
Makefile
View File

@ -12,11 +12,11 @@ check_code_quality:
isort --check-only --profile black $(check_dirs)
# stop the build if there are Python syntax errors or undefined names
flake8 $(check_dirs) --count --select=E9,F63,F7,F82 --show-source --statistics
# exit-zero treats all errors as warnings. E203 for black, E501 for docstring, W503 for line breaks before logical operators
# exit-zero treats all errors as warnings. E203 for black, E501 for docstring, W503 for line breaks before logical operators
flake8 $(check_dirs) --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics
publish:
poetry build
twine upload -r testpypi dist/* -u ${PYPI_USERNAME} -p ${PYPI_TEST_PASSWORD} --verbose
twine upload -r testpypi dist/* -u ${PYPI_USERNAME} -p ${PYPI_TEST_PASSWORD} --verbose
twine check dist/*
twine upload dist/* -u ${PYPI_USERNAME} -p ${PYPI_PASSWORD} --verbose
twine upload dist/* -u ${PYPI_USERNAME} -p ${PYPI_PASSWORD} --verbose

92
README.md
View File

@ -31,7 +31,7 @@ Pip install the supervision package in a
pip install supervision[desktop]
```
Read more about desktop, headless and local installation in our [guide](https://roboflow.github.io/supervision/).
Read more about desktop, headless, and local installation in our [guide](https://roboflow.github.io/supervision/).
## 🔥 quickstart
@ -43,7 +43,7 @@ Read more about desktop, headless and local installation in our [guide](https://
>>> model = YOLO('yolov8s.pt')
>>> result = model(IMAGE)[0]
>>> detections = sv.Detections.from_yolov8(result)
>>> detections = sv.Detections.from_ultralytics(result)
>>> len(detections)
5
@ -51,9 +51,9 @@ Read more about desktop, headless and local installation in our [guide](https://
<details close>
<summary>👉 more detections utils</summary>
- Easily switch inference pipeline between supported object detection / instance segmentation models
- Easily switch inference pipeline between supported object detection/instance segmentation models
```python
>>> import supervision as sv
>>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
@ -63,17 +63,17 @@ Read more about desktop, headless and local installation in our [guide](https://
>>> sam_result = mask_generator.generate(IMAGE)
>>> detections = sv.Detections.from_sam(sam_result=sam_result)
```
- [Advanced filtering](https://roboflow.github.io/supervision/quickstart/detections/)
```python
>>> detections = detections[detections.class_id == 0]
>>> detections = detections[detections.confidence > 0.5]
>>> detections = detections[detections.area > 1000]
```
- Image annotation
```python
>>> import supervision as sv
@ -83,7 +83,7 @@ Read more about desktop, headless and local installation in our [guide](https://
... detections=detections
... )
```
</details>
### [datasets processing](https://roboflow.github.io/supervision/dataset/core/)
@ -107,7 +107,7 @@ Read more about desktop, headless and local installation in our [guide](https://
<details close>
<summary>👉 more dataset utils</summary>
- Load object detection / instance segmentation datasets in one of supported formats
- Load object detection/instance segmentation datasets in one of the supported formats
```python
>>> dataset = sv.DetectionDataset.from_yolo(
@ -126,7 +126,7 @@ Read more about desktop, headless and local installation in our [guide](https://
... annotations_path='...'
... )
```
- Loop over dataset entries
```python
@ -137,18 +137,18 @@ Read more about desktop, headless and local installation in our [guide](https://
[155. , 497. , 404. , 833.5 ],
[ 20.154999, 347.825 , 416.125 , 915.895 ]], dtype=float32)
```
- Split dataset for training, testing and validation
- Split dataset for training, testing, and validation
```python
>>> train_dataset, test_dataset = dataset.split(split_ratio=0.7)
>>> test_dataset, valid_dataset = test_dataset.split(split_ratio=0.5)
>>> len(train_dataset), len(test_dataset), len(valid_dataset)
(700, 150, 150)
```
- Merge multiple datasets together
- Merge multiple datasets
```python
>>> ds_1 = sv.DetectionDataset(...)
@ -156,22 +156,22 @@ Read more about desktop, headless and local installation in our [guide](https://
100
>>> ds_1.classes
['dog', 'person']
>>> ds_2 = sv.DetectionDataset(...)
>>> len(ds_2)
200
>>> ds_2.classes
['cat']
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
>>> len(ds_merged)
300
>>> ds_merged.classes
['cat', 'dog', 'person']
```
- Save object detection / instance segmentation datasets in one of supported formats
- Save object detection/instance segmentation datasets in one of the supported formats
```python
>>> dataset.as_yolo(
... images_directory_path='...',
@ -189,9 +189,9 @@ Read more about desktop, headless and local installation in our [guide](https://
... annotations_path='...'
... )
```
- Convert labels between supported formats
```python
>>> sv.DetectionDataset.from_yolo(
... images_directory_path='...',
@ -202,8 +202,8 @@ Read more about desktop, headless and local installation in our [guide](https://
... annotations_directory_path='...'
... )
```
- Load classification datasets in one of supported formats
- Load classification datasets in one of the supported formats
```python
>>> cs = sv.ClassificationDataset.from_folder_structure(
@ -211,7 +211,7 @@ Read more about desktop, headless and local installation in our [guide](https://
... )
```
- Save classification datasets in one of supported formats
- Save classification datasets in one of the supported formats
```python
>>> cs.as_folder_structure(
@ -245,15 +245,45 @@ array([
])
```
<details close>
<summary>👉 more metrics</summary>
- Mean average precision (mAP) for object detection tasks.
```python
>>> import supervision as sv
>>> dataset = sv.DetectionDataset.from_yolo(...)
>>> def callback(image: np.ndarray) -> sv.Detections:
... ...
>>> mean_average_precision = sv.MeanAveragePrecision.benchmark(
... dataset = dataset,
... callback = callback
... )
>>> mean_average_precision.map50_95
0.433
```
</details>
## 🛠️ built with supervision
Did you build something cool using supervision? [Let us know!](https://github.com/roboflow/supervision/discussions/categories/built-with-supervision)
https://user-images.githubusercontent.com/26109316/207858600-ee862b22-0353-440b-ad85-caa0c4777904.mp4
## 🎬 tutorials
<p align="left">
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/ae1ca38e-40b7-4b35-8582-e8ea5de3806e" alt="Accelerate Image Annotation with SAM and Grounding DINO" width="300px" align="left" /></a>
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><strong>Accelerate Image Annotation with SAM and Grounding DINO</strong></a>
<div><strong>Created: 20 Apr 2023</strong> | <strong>Updated: 20 Apr 2023</strong></div>
<br/> Discover how to speed up your image annotation process using Grounding DINO and Segment Anything Model (SAM). Learn how to convert object detection datasets into instance segmentation datasets, and see the potential of using these models to automatically annotate your datasets for real-time detectors like YOLOv8... </p>
<br/> Discover how to speed up your image annotation process using Grounding DINO and Segment Anything Model (SAM). Learn how to convert object detection datasets into instance segmentation datasets, and see the potential of using these models to automatically annotate your datasets for real-time detectors like YOLOv8... </p>
<br/>
<br/>
<p align="left">
<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/6913ff11-53c6-4341-8d90-eaff3023c3fd" alt="SAM - Segment Anything Model by Meta AI: Complete Guide" width="300px" align="left" /></a>
@ -263,7 +293,7 @@ array([
## 📚 documentation
Visit our [documentation](https://roboflow.github.io/supervision) page to learn how supervision can help you build computer vision applications faster and more reliably.
Visit our [documentation](https://roboflow.github.io/supervision) page to learn how supervision can help you build computer vision applications faster and more reliably.
## 🏆 contribution

2576
demo.ipynb vendored
View File

File diff suppressed because one or more lines are too long

72
docs/changelog.md
View File

@ -1,3 +1,39 @@
### 0.13.0 <small>August 8, 2023</small>
- Added [#236](https://github.com/roboflow/supervision/pull/236): support for mean average precision (mAP) for object detection models with [`sv.MeanAveragePrecision`](https://roboflow.github.io/supervision/metrics/detection/#meanaverageprecision).
```python
>>> import supervision as sv
>>> from ultralytics import YOLO
>>> dataset = sv.DetectionDataset.from_yolo(...)
>>> model = YOLO(...)
>>> def callback(image: np.ndarray) -> sv.Detections:
... result = model(image)[0]
... return sv.Detections.from_yolov8(result)
>>> mean_average_precision = sv.MeanAveragePrecision.benchmark(
... dataset = dataset,
... callback = callback
... )
>>> mean_average_precision.map50_95
0.433
```
- Added [#256](https://github.com/roboflow/supervision/pull/256): support for ByteTrack for object tracking with [`sv.ByteTrack`](https://roboflow.github.io/supervision/tracker/core/#bytetrack).
- Added [#222](https://github.com/roboflow/supervision/pull/222): [`sv.Detections.from_ultralytics`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_ultralytics) to enable seamless integration with [Ultralytics](https://github.com/ultralytics/ultralytics) framework. This will enable you to use `supervision` with all [models](https://docs.ultralytics.com/models/) that Ultralytics supports.
!!! warning
[`sv.Detections.from_yolov8`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_yolov8) is now deprecated and will be removed with `supervision-0.15.0` release.
- Added [#191](https://github.com/roboflow/supervision/pull/191): [`sv.Detections.from_paddledet`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_paddledet) to enable seamless integration with [PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection) framework.
- Added [#245](https://github.com/roboflow/supervision/pull/245): support for loading PASCAL VOC segmentation datasets with [`sv.DetectionDataset.`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_pascal_voc).
### 0.12.0 <small>July 24, 2023</small>
!!! warning
@ -47,7 +83,7 @@ array([
### 0.11.0 <small>June 28, 2023</small>
- Added [#150](https://github.com/roboflow/supervision/pull/150): ability to load and save [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) in COCO format using [`as_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_coco) and [`from_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_coco) methods.
- Added [#150](https://github.com/roboflow/supervision/pull/150): ability to load and save [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) in COCO format using [`as_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_coco) and [`from_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_coco) methods.
```python
>>> import supervision as sv
@ -63,7 +99,7 @@ array([
... )
```
- Added [#158](https://github.com/roboflow/supervision/pull/158): ability to marge multiple [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) together using [`merge`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.merge) method.
- Added [#158](https://github.com/roboflow/supervision/pull/158): ability to marge multiple [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) together using [`merge`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.merge) method.
```python
>>> import supervision as sv
@ -109,11 +145,11 @@ array([
- Added [#125](https://github.com/roboflow/supervision/pull/125): support for [`sv.ClassificationDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.ClassificationDataset.split) allowing to divide `sv.ClassificationDataset` into two parts.
- Added [#110](https://github.com/roboflow/supervision/pull/110): ability to extract masks from Roboflow API results using [`sv.Detections.from_roboflow`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_roboflow).
- Added [#110](https://github.com/roboflow/supervision/pull/110): ability to extract masks from Roboflow API results using [`sv.Detections.from_roboflow`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_roboflow).
- Added [commit hash](https://github.com/roboflow/supervision/commit/d000292eb2f2342544e0947b65528082e60fb8d6): Supervision Quickstart [notebook](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb) where you can learn more about Detection, Dataset and Video APIs.
- Changed [#135](https://github.com/roboflow/supervision/pull/135): `sv.get_video_frames_generator` documentation to better describe actual behavior.
- Changed [#135](https://github.com/roboflow/supervision/pull/135): `sv.get_video_frames_generator` documentation to better describe actual behavior.
### 0.9.0 <small>June 7, 2023</small>
@ -132,7 +168,7 @@ array([
```
- Added [#101](https://github.com/roboflow/supervision/pull/101): ability to extract masks from YOLOv8 result using [`sv.Detections.from_yolov8`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_yolov8). Here is an example illustrating how to extract boolean masks from the result of the YOLOv8 model inference.
- Added [#122](https://github.com/roboflow/supervision/pull/122): ability to crop image using [`sv.crop`](https://roboflow.github.io/supervision/utils/image/#crop). Here is an example showing how to get a separate crop for each detection in `sv.Detections`.
- Added [#120](https://github.com/roboflow/supervision/pull/120): ability to conveniently save multiple images into directory using [`sv.ImageSink`](https://roboflow.github.io/supervision/utils/image/#imagesink). Here is an example showing how to save every tenth video frame as a separate image.
@ -150,7 +186,7 @@ array([
### 0.8.0 <small>May 17, 2023</small>
- Added [#100](https://github.com/roboflow/supervision/pull/100): support for dataset inheritance. The current `Dataset` got renamed to `DetectionDataset`. Now [`DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) inherits from `BaseDataset`. This change was made to enforce the future consistency of APIs of different types of computer vision datasets.
- Added [#100](https://github.com/roboflow/supervision/pull/100): ability to save datasets in YOLO format using [`DetectionDataset.as_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_yolo).
- Added [#100](https://github.com/roboflow/supervision/pull/100): ability to save datasets in YOLO format using [`DetectionDataset.as_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_yolo).
```python
>>> import roboflow
@ -174,7 +210,7 @@ array([
['dog', 'person']
```
- Added [#102](https://github.com/roboflow/supervision/pull/103): support for [`DetectionDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.split) allowing to divide `DetectionDataset` into two parts.
- Added [#102](https://github.com/roboflow/supervision/pull/103): support for [`DetectionDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.split) allowing to divide `DetectionDataset` into two parts.
```python
>>> import supervision as sv
@ -191,14 +227,14 @@ array([
### 0.7.0 <small>May 11, 2023</small>
- Added [#91](https://github.com/roboflow/supervision/pull/91): `Detections.from_yolo_nas` to enable seamless integration with [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) model.
- Added [#86](https://github.com/roboflow/supervision/pull/86): ability to load datasets in YOLO format using `Dataset.from_yolo`.
- Added [#86](https://github.com/roboflow/supervision/pull/86): ability to load datasets in YOLO format using `Dataset.from_yolo`.
- Added [#84](https://github.com/roboflow/supervision/pull/84): `Detections.merge` to merge multiple `Detections` objects together.
- Fixed [#81](https://github.com/roboflow/supervision/pull/81): `LineZoneAnnotator.annotate` does not return annotated frame.
- Changed [#44](https://github.com/roboflow/supervision/pull/44): `LineZoneAnnotator.annotate` to allow for custom text for the in and out tags.
### 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): 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.
@ -220,7 +256,7 @@ array([
- Added [#58](https://github.com/roboflow/supervision/pull/58): `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
- Changed [#58](https://github.com/roboflow/supervision/pull/58): `Detections.area` behaviour to work not only with boxes but also with masks.
### 0.4.0 <small>April 5, 2023</small>
### 0.4.0 <small>April 5, 2023</small>
- Added [#46](https://github.com/roboflow/supervision/discussions/48): `Detections.empty` to allow easy creation of empty `Detections` objects.
- Added [#56](https://github.com/roboflow/supervision/pull/56): `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
@ -228,28 +264,28 @@ array([
- Added [#56](https://github.com/roboflow/supervision/pull/56): initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
- Changed [#56](https://github.com/roboflow/supervision/pull/56): `show_frame_in_notebook` refactored and renamed to `plot_image`.
### 0.3.2 <small>March 23, 2023</small>
### 0.3.2 <small>March 23, 2023</small>
- Changed [#50](https://github.com/roboflow/supervision/issues/50): Allow `Detections.class_id` to be `None`.
- Changed [#50](https://github.com/roboflow/supervision/issues/50): Allow `Detections.class_id` to be `None`.
### 0.3.1 <small>March 6, 2023</small>
### 0.3.1 <small>March 6, 2023</small>
- Fixed [#41](https://github.com/roboflow/supervision/issues/41): `PolygonZone` throws an exception when the object touches the bottom edge of the image.
- Fixed [#42](https://github.com/roboflow/supervision/issues/42): `Detections.wth_nms` method throws an exception when `Detections` is empty.
- Changed [#36](https://github.com/roboflow/supervision/pull/36): `Detections.wth_nms` support class agnostic and non-class agnostic case.
### 0.3.0 <small>March 6, 2023</small>
### 0.3.0 <small>March 6, 2023</small>
- Changed: Allow `Detections.confidence` to be `None`.
- Added: `Detections.from_transformers` and `Detections.from_detectron2` to enable seamless integration with Transformers and Detectron2 models.
- Added: `Detections.from_transformers` and `Detections.from_detectron2` to enable seamless integration with Transformers and Detectron2 models.
- Added: `Detections.area` to dynamically calculate bounding box area.
- Added: `Detections.wth_nms` to filter out double detections with NMS. Initial - only class agnostic - implementation.
- Added: `Detections.wth_nms` to filter out double detections with NMS. Initial - only class agnostic - implementation.
### 0.2.0 <small>February 2, 2023</small>
### 0.2.0 <small>February 2, 2023</small>
- Added: Advanced `Detections` filtering with pandas-like API.
- Added: `Detections.from_yolov5` and `Detections.from_yolov8` to enable seamless integration with YOLOv5 and YOLOv8 models.
### 0.1.0 <small>January 19, 2023</small>
### 0.1.0 <small>January 19, 2023</small>
Say hello to Supervision 👋

4
docs/dataset/core.md
View File

@ -1,6 +1,6 @@
!!! warning
Dataset API is still fluid and may change. If you use Dataset API in your project until further notice, freeze the
Dataset API is still fluid and may change. If you use Dataset API in your project until further notice, freeze the
`supervision` version in your `requirements.txt` or `setup.py`.
## DetectionDataset
@ -9,4 +9,4 @@
## ClassificationDataset
:::supervision.dataset.core.ClassificationDataset
:::supervision.dataset.core.ClassificationDataset

2
docs/detection/annotate.md
View File

@ -4,4 +4,4 @@
## MaskAnnotator
:::supervision.detection.annotate.MaskAnnotator
:::supervision.detection.annotate.MaskAnnotator

2
docs/detection/tools/polygon_zone.md
View File

@ -4,4 +4,4 @@
## PolygonZoneAnnotator
:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator
:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator

2
docs/detection/utils.md
View File

@ -24,4 +24,4 @@
## filter_polygons_by_area
:::supervision.detection.utils.filter_polygons_by_area
:::supervision.detection.utils.filter_polygons_by_area

2
docs/draw/utils.md
View File

@ -16,4 +16,4 @@
## draw_text
:::supervision.draw.utils.draw_text
:::supervision.draw.utils.draw_text

22
docs/index.md
View File

@ -11,22 +11,22 @@
## 👋 Hello
We write your reusable computer vision tools. Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us!
We write your reusable computer vision tools. Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us!
## 💻 Install
You can install `supervision` with pip in a
You can install `supervision` with pip in a
[**3.11>=Python>=3.8**](https://www.python.org/) environment.
!!! example "pip install (recommended)"
=== "headless"
The headless installation of `supervision` is designed for environments where graphical user interfaces (GUI) are not needed, making it more lightweight and suitable for server-side applications.
The headless installation of `supervision` is designed for environments where graphical user interfaces (GUI) are not needed, making it more lightweight and suitable for server-side applications.
```bash
pip install supervision
```
=== "desktop"
If you require the full version of `supervision` with GUI support you can install the desktop version. This version includes the GUI components of OpenCV, allowing you to display images and videos on the screen.
@ -42,32 +42,32 @@ You can install `supervision` with pip in a
# clone repository and navigate to root directory
git clone https://github.com/roboflow/supervision.git
cd supervision
# setup python environment and activate it
python3 -m venv venv
source venv/bin/activate
# headless install
pip install -e "."
# desktop install
pip install -e ".[desktop]"
```
=== "poetry"
```bash
# clone repository and navigate to root directory
git clone https://github.com/roboflow/supervision.git
cd supervision
# setup python environment and activate it
poetry env use python 3.10
poetry shell
# headless install
poetry install
# desktop install
poetry install --extras "desktop"
```
```

2
docs/metrics/detection.md
View File

@ -1,6 +1,6 @@
!!! warning
Evaluation API is still fluid and may change. If you use Evaluation API in your project until further notice, freeze the
Evaluation API is still fluid and may change. If you use Evaluation API in your project until further notice, freeze the
`supervision` version in your `requirements.txt` or `setup.py`.
## ConfusionMatrix

36
docs/quickstart/detections.md
View File

@ -1,10 +1,10 @@
## advanced filtering
The advanced filtering capabilities of the `Detections` class offer users a versatile and efficient way to narrow down
and refine object detections. This section outlines various filtering methods, including filtering by specific class
or a set of classes, confidence, object area, bounding box area, relative area, box dimensions, and designated zones.
Each method is demonstrated with concise code examples to provide users with a clear understanding of how to implement
The advanced filtering capabilities of the `Detections` class offer users a versatile and efficient way to narrow down
and refine object detections. This section outlines various filtering methods, including filtering by specific class
or a set of classes, confidence, object area, bounding box area, relative area, box dimensions, and designated zones.
Each method is demonstrated with concise code examples to provide users with a clear understanding of how to implement
the filters in their applications.
### by specific class
@ -44,14 +44,14 @@ Allows you to select detections that belong only to one selected class.
### by set of classes
Allows you to select detections that belong only to selected set of classes.
Allows you to select detections that belong only to selected set of classes.
=== "After"
```python
import numpy as np
import supervision as sv
selected_classes = [0, 2, 3]
detections = sv.Detections(...)
detections = detections[np.isin(detections.class_id, selected_classes)]
@ -68,7 +68,7 @@ Allows you to select detections that belong only to selected set of classes.
```python
import numpy as np
import supervision as sv
class_id = [0, 2, 3]
detections = sv.Detections(...)
detections = detections[np.isin(detections.class_id, class_id)]
@ -116,8 +116,8 @@ Allows you to select detections with specific confidence value, for example high
### by area
Allows you to select detections based on their size. We define the area as the number of pixels occupied by the
detection in the image. In the example below, we have sifted out the detections that are too small.
Allows you to select detections based on their size. We define the area as the number of pixels occupied by the
detection in the image. In the example below, we have sifted out the detections that are too small.
=== "After"
@ -151,9 +151,9 @@ detection in the image. In the example below, we have sifted out the detections
### by relative area
Allows you to select detections based on their size in relation to the size of whole image. Sometimes the concept of
detection size changes depending on the image. Detection occupying 10000 square px can be large on a 1280x720 image
but small on a 3840x2160 image. In such cases, we can filter out detections based on the percentage of the image area
Allows you to select detections based on their size in relation to the size of whole image. Sometimes the concept of
detection size changes depending on the image. Detection occupying 10000 square px can be large on a 1280x720 image
but small on a 3840x2160 image. In such cases, we can filter out detections based on the percentage of the image area
occupied by them. In the example below, we remove too large detections.
=== "After"
@ -164,7 +164,7 @@ occupied by them. In the example below, we remove too large detections.
image = ...
height, width, channels = image.shape
image_area = height * width
detections = sv.Detections(...)
detections = detections[(detections.area / image_area) < 0.8]
```
@ -183,7 +183,7 @@ occupied by them. In the example below, we remove too large detections.
image = ...
height, width, channels = image.shape
image_area = height * width
detections = sv.Detections(...)
detections = detections[(detections.area / image_area) < 0.8]
```
@ -196,8 +196,8 @@ occupied by them. In the example below, we remove too large detections.
### by box dimensions
Allows you to select detections based on their dimensions. The size of the bounding box, as well as its coordinates,
can be criteria for rejecting detection. Implementing such filtering requires a bit of custom code but is relatively
Allows you to select detections based on their dimensions. The size of the bounding box, as well as its coordinates,
can be criteria for rejecting detection. Implementing such filtering requires a bit of custom code but is relatively
simple and fast.
=== "After"
@ -236,7 +236,7 @@ simple and fast.
### by `PolygonZone`
Allows you to use `Detections` in combination with `PolygonZone` to weed out bounding boxes that are in and out of the
Allows you to use `Detections` in combination with `PolygonZone` to weed out bounding boxes that are in and out of the
zone. In the example below you can see how to filter out all detections located in the lower part of the image.
=== "After"
@ -309,4 +309,4 @@ zone. In the example below you can see how to filter out all detections located
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
</div>

2
docs/stylesheets/extra.css
View File

@ -1,4 +1,4 @@
:root {
--md-primary-fg-color: #8315F9;
--md-accent-fg-color: #00FFCE;
}
}

3
docs/tracker/core.md Normal file
View File

@ -0,0 +1,3 @@
## ByteTrack
:::supervision.tracker.byte_tracker.core.ByteTrack

2
docs/utils/image.md
View File

@ -4,4 +4,4 @@
## crop
:::supervision.utils.image.crop
:::supervision.utils.image.crop

2
docs/utils/notebook.md
View File

@ -4,4 +4,4 @@
## plot_images_grid
:::supervision.utils.notebook.plot_images_grid
:::supervision.utils.notebook.plot_images_grid

2
docs/utils/video.md
View File

@ -12,4 +12,4 @@
## process_video
:::supervision.utils.video.process_video
:::supervision.utils.video.process_video

4
mkdocs.yml
View File

@ -37,6 +37,8 @@ nav:
- Utils: detection/utils.md
- Tools:
- Polygon Zone: detection/tools/polygon_zone.md
- Trackers:
- Core: tracker/core.md
- Dataset:
- Core: dataset/core.md
- Metrics:
@ -86,4 +88,4 @@ markdown_extensions:
- pymdownx.tabbed:
alternate_style: true
- toc:
permalink: true
permalink: true

40
poetry.lock generated
View File

@ -3078,6 +3078,44 @@ files = [
{file = "ruff-0.0.280.tar.gz", hash = "sha256:581c43e4ac5e5a7117ad7da2120d960a4a99e68ec4021ec3cd47fe1cf78f8380"},
]
[[package]]
name = "scipy"
version = "1.10.1"
description = "Fundamental algorithms for scientific computing in Python"
optional = false
python-versions = "<3.12,>=3.8"
files = [
{file = "scipy-1.10.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e7354fd7527a4b0377ce55f286805b34e8c54b91be865bac273f527e1b839019"},
{file = "scipy-1.10.1-cp310-cp310-macosx_12_0_arm64.whl", hash = "sha256:4b3f429188c66603a1a5c549fb414e4d3bdc2a24792e061ffbd607d3d75fd84e"},
{file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1553b5dcddd64ba9a0d95355e63fe6c3fc303a8fd77c7bc91e77d61363f7433f"},
{file = "scipy-1.10.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4c0ff64b06b10e35215abce517252b375e580a6125fd5fdf6421b98efbefb2d2"},
{file = "scipy-1.10.1-cp310-cp310-win_amd64.whl", hash = "sha256:fae8a7b898c42dffe3f7361c40d5952b6bf32d10c4569098d276b4c547905ee1"},
{file = "scipy-1.10.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0f1564ea217e82c1bbe75ddf7285ba0709ecd503f048cb1236ae9995f64217bd"},
{file = "scipy-1.10.1-cp311-cp311-macosx_12_0_arm64.whl", hash = "sha256:d925fa1c81b772882aa55bcc10bf88324dadb66ff85d548c71515f6689c6dac5"},
{file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:aaea0a6be54462ec027de54fca511540980d1e9eea68b2d5c1dbfe084797be35"},
{file = "scipy-1.10.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:15a35c4242ec5f292c3dd364a7c71a61be87a3d4ddcc693372813c0b73c9af1d"},
{file = "scipy-1.10.1-cp311-cp311-win_amd64.whl", hash = "sha256:43b8e0bcb877faf0abfb613d51026cd5cc78918e9530e375727bf0625c82788f"},
{file = "scipy-1.10.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:5678f88c68ea866ed9ebe3a989091088553ba12c6090244fdae3e467b1139c35"},
{file = "scipy-1.10.1-cp38-cp38-macosx_12_0_arm64.whl", hash = "sha256:39becb03541f9e58243f4197584286e339029e8908c46f7221abeea4b749fa88"},
{file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bce5869c8d68cf383ce240e44c1d9ae7c06078a9396df68ce88a1230f93a30c1"},
{file = "scipy-1.10.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:07c3457ce0b3ad5124f98a86533106b643dd811dd61b548e78cf4c8786652f6f"},
{file = "scipy-1.10.1-cp38-cp38-win_amd64.whl", hash = "sha256:049a8bbf0ad95277ffba9b3b7d23e5369cc39e66406d60422c8cfef40ccc8415"},
{file = "scipy-1.10.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:cd9f1027ff30d90618914a64ca9b1a77a431159df0e2a195d8a9e8a04c78abf9"},
{file = "scipy-1.10.1-cp39-cp39-macosx_12_0_arm64.whl", hash = "sha256:79c8e5a6c6ffaf3a2262ef1be1e108a035cf4f05c14df56057b64acc5bebffb6"},
{file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:51af417a000d2dbe1ec6c372dfe688e041a7084da4fdd350aeb139bd3fb55353"},
{file = "scipy-1.10.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1b4735d6c28aad3cdcf52117e0e91d6b39acd4272f3f5cd9907c24ee931ad601"},
{file = "scipy-1.10.1-cp39-cp39-win_amd64.whl", hash = "sha256:7ff7f37b1bf4417baca958d254e8e2875d0cc23aaadbe65b3d5b3077b0eb23ea"},
{file = "scipy-1.10.1.tar.gz", hash = "sha256:2cf9dfb80a7b4589ba4c40ce7588986d6d5cebc5457cad2c2880f6bc2d42f3a5"},
]
[package.dependencies]
numpy = ">=1.19.5,<1.27.0"
[package.extras]
dev = ["click", "doit (>=0.36.0)", "flake8", "mypy", "pycodestyle", "pydevtool", "rich-click", "typing_extensions"]
doc = ["matplotlib (>2)", "numpydoc", "pydata-sphinx-theme (==0.9.0)", "sphinx (!=4.1.0)", "sphinx-design (>=0.2.0)"]
test = ["asv", "gmpy2", "mpmath", "pooch", "pytest", "pytest-cov", "pytest-timeout", "pytest-xdist", "scikit-umfpack", "threadpoolctl"]
[[package]]
name = "secretstorage"
version = "3.3.3"
@ -3472,4 +3510,4 @@ desktop = ["opencv-python"]
[metadata]
lock-version = "2.0"
python-versions = ">=3.8,<3.12.0"
content-hash = "4917c08576fa8226c0593bac637f8442171d16cb50912a0cfe225959dcaa4e5e"
content-hash = "fcefb11e72c40fa7defc13f4a9cb1fe24b9d3d171b64bb54fc263d8add9ef40f"

108
pyproject.toml
View File

@ -1,6 +1,6 @@
[tool.poetry]
name = "supervision"
version = "0.12.0"
version = "0.13.0"
description = "A set of easy-to-use utils that will come in handy in any Computer Vision project"
authors = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
maintainers = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
@ -41,6 +41,7 @@ pyyaml = "^6.0"
pillow = "^9.4.0"
opencv-python = { version = "^4.8.0.74", optional = true }
opencv-python-headless = "^4.8.0.74"
scipy = "^1.9.0"
[tool.poetry.extras]
@ -65,6 +66,111 @@ flake8 = { version = "*", python = ">=3.8.1,<3.12.0" }
mkdocs-material = "^9.1.4"
mkdocstrings = {extras = ["python"], version = "^0.20.0"}
[tool.flake8]
exclude = ".venv"
max-complexity = 10
max-line-length = 88
extend-ignore = """
W503,
E203,
E701,
C901,
"""
per-file-ignores = """
__init__.py: F401
"""
[tool.isort]
line_length = 88
profile = "black"
[tool.bandit]
target = ["test", "supervision"]
tests = ["B201", "B301"]
[tool.autoflake]
check = true
imports = ["cv2", "supervision"]
[tool.black]
target-version = ["py38"]
line-length = 88
include = '\.pyi?$'
exclude = '''
/(
\.git
| \.hg
| \.mypy_cache
| \.tox
| \.venv
| _build
| buck-out
| build
| dist
| docs
)/
'''
[tool.ruff]
target-version = "py38"
# Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default.
select = ["E", "F"]
ignore = []
# Allow autofix for all enabled rules (when `--fix`) is provided.
fixable = ["A", "B", "C", "D", "E", "F", "G", "I", "N", "Q", "S", "T", "W", "ANN", "ARG", "BLE", "COM", "DJ", "DTZ", "EM", "ERA", "EXE", "FBT", "ICN", "INP", "ISC", "NPY", "PD", "PGH", "PIE", "PL", "PT", "PTH", "PYI", "RET", "RSE", "RUF", "SIM", "SLF", "TCH", "TID", "TRY", "UP", "YTT"]
unfixable = []
# Exclude a variety of commonly ignored directories.
exclude = [
".bzr",
".direnv",
".eggs",
".git",
".git-rewrite",
".hg",
".mypy_cache",
".nox",
".pants.d",
".pytype",
".ruff_cache",
".svn",
".tox",
".venv",
"__pypackages__",
"_build",
"buck-out",
"build",
"dist",
"node_modules",
"venv",
"yarn-error.log",
"yarn.lock",
"docs",
]
# Same as Black.
line-length = 88
# Allow unused variables when underscore-prefixed.
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
[tool.ruff.flake8-quotes]
inline-quotes = "double"
multiline-quotes = "double"
docstring-quotes = "double"
[tool.ruff.pydocstyle]
convention = "google"
[tool.ruff.per-file-ignores]
"__init__.py" = ["E402","F401"]
[tool.ruff.pylint]
max-args = 20
[tool.setuptools]
include-package-data = false

1
supervision/__init__.py
View File

@ -31,6 +31,7 @@ from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_tex
from supervision.geometry.core import Point, Position, Rect
from supervision.geometry.utils import get_polygon_center
from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision
from supervision.tracker.byte_tracker.core import ByteTrack
from supervision.utils.file import list_files_with_extensions
from supervision.utils.fps import FpsMonitor
from supervision.utils.image import ImageSink, crop

53
supervision/classification/core.py
View File

@ -5,6 +5,8 @@ from typing import Any, Optional, Tuple
import numpy as np
from supervision.utils.internal import deprecated
def _validate_class_ids(class_id: Any, n: int) -> None:
"""
@ -40,15 +42,24 @@ class Classifications:
_validate_confidence(self.confidence, n)
@classmethod
@deprecated(
"""
This method is deprecated and removed in 0.16.0 release.
Use sv.Classifications.from_ultralytics() instead as it is more generic and
can be used for detections from any ultralytics.engine.results.Results Object
"""
)
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.
Classifications: A new Classifications object.
Example:
```python
@ -65,15 +76,47 @@ class Classifications:
confidence = yolov8_results.probs.data.cpu().numpy()
return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
@classmethod
def from_ultralytics(cls, ultralytics_results) -> Classifications:
"""
Creates a Classifications instance from a
(https://github.com/ultralytics/ultralytics) inference result.
Args:
ultralytics_results (ultralytics.engine.results.Results):
The output Results instance from ultralytics model
Returns:
Classifications: A new Classifications object.
Example:
```python
>>> import cv2
>>> from ultralytics import YOLO
>>> import supervision as sv
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
>>> model = YOLO('yolov8n-cls.pt')
>>> model = YOLO('yolov8s-cls.pt')
>>> result = model(image)[0]
>>> classifications = sv.Classifications.from_ultralytics(result)
```
"""
confidence = ultralytics_results.probs.data.cpu().numpy()
return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
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

183
supervision/dataset/core.py
View File

@ -54,7 +54,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]
@ -75,8 +76,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)
@ -100,22 +102,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)
```
@ -150,19 +157,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:
@ -205,12 +220,15 @@ 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.
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
images_directory_path (str): Path to the directory containing the images.
annotations_directory_path (str): Path to the directory
containing the PASCAL VOC XML annotations.
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
@ -255,13 +273,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
@ -304,23 +328,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(
@ -349,12 +382,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
@ -394,20 +431,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(
@ -427,16 +470,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
@ -494,7 +541,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]
@ -508,22 +556,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)
```
@ -553,7 +607,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)

2
supervision/dataset/formats/coco.py
View File

@ -215,7 +215,7 @@ def save_coco_annotations(
coco_images.append(coco_image)
detections = annotations[image_name]
coco_annotation, label_id = detections_to_coco_annotations(
coco_annotation, annotation_id = detections_to_coco_annotations(
detections=detections,
image_id=image_id,
annotation_id=annotation_id,

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

@ -21,6 +21,9 @@ def object_to_pascal_voc(
object_name = SubElement(root, "name")
object_name.text = name
# https://github.com/roboflow/supervision/issues/144
xyxy += 1
bndbox = SubElement(root, "bndbox")
xmin = SubElement(bndbox, "xmin")
xmin.text = str(int(xyxy[0]))
@ -32,6 +35,8 @@ def object_to_pascal_voc(
ymax.text = str(int(xyxy[3]))
if polygon is not None:
# https://github.com/roboflow/supervision/issues/144
polygon += 1
object_polygon = SubElement(root, "polygon")
for index, point in enumerate(polygon, start=1):
x_coordinate, y_coordinate = point
@ -56,13 +61,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.
"""
@ -129,15 +140,22 @@ def load_pascal_voc_annotations(
force_masks: bool = False,
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
"""
Loads PASCAL VOC annotations and returns class names, images, and their corresponding detections.
Loads PASCAL VOC XML annotations and returns the image name,
a Detections instance, and a list of class names.
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 annotation files.
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
PASCAL VOC annotation files.
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(
@ -202,7 +220,9 @@ def detections_from_xml_obj(
</annotation>
Returns:
Tuple[Detections, List[str]]: A tuple containing a Detections object and an updated list of class names, extended with the class names from the XML object.
Tuple[Detections, List[str]]: A tuple containing a Detections object and an
updated list of class names, extended with the class names
from the XML object.
"""
xyxy = []
class_names = []
@ -225,15 +245,21 @@ def detections_from_xml_obj(
with_masks = force_masks if force_masks else with_masks
for polygon in obj.findall("polygon"):
polygon_points = parse_polygon_points(polygon)
polygon = parse_polygon_points(polygon)
# https://github.com/roboflow/supervision/issues/144
polygon -= 1
mask_from_polygon = polygon_to_mask(
polygon=np.array(polygon_points),
polygon=polygon,
resolution_wh=resolution_wh,
)
masks.append(mask_from_polygon)
xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
# https://github.com/roboflow/supervision/issues/144
xyxy -= 1
for k in set(class_names):
if k not in extended_classes:
extended_classes.append(k)
@ -250,11 +276,8 @@ def detections_from_xml_obj(
return annotation, extended_classes
def parse_polygon_points(polygon: Element) -> List[List[int]]:
polygon_points = []
coords = polygon.findall(".//*")
for i in range(0, len(coords), 2):
x = int(coords[i].text)
y = int(coords[i + 1].text)
polygon_points.append([x, y])
return polygon_points
def parse_polygon_points(polygon: Element) -> np.ndarray:
coordinates = [int(coord.text) for coord in polygon.findall(".//*")]
return np.array(
[(coordinates[i], coordinates[i + 1]) for i in range(0, len(coordinates), 2)]
)

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

@ -1,6 +1,6 @@
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Union
from typing import Dict, List, Optional, Tuple
import cv2
import numpy as np
@ -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"]

2
supervision/dataset/utils.py
View File

@ -64,7 +64,7 @@ def build_class_index_mapping(
if class_name not in target_classes:
raise ValueError(
f"Class {class_name} not found in target classes. "
f"source_classes must be a subset of target_classes."
"source_classes must be a subset of target_classes."
)
corresponding_index = target_classes.index(class_name)
index_mapping[i] = corresponding_index

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:

184
supervision/detection/core.py
View File

@ -3,7 +3,6 @@ from __future__ import annotations
from dataclasses import astuple, dataclass
from typing import Any, Iterator, List, Optional, Tuple, Union
import cv2
import numpy as np
from supervision.detection.utils import (
@ -59,11 +58,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, H, W)` 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, H, W)` 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
@ -98,7 +102,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 (
@ -143,10 +148,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.
@ -172,14 +179,20 @@ class Detections:
@classmethod
@deprecated(
"This method is deprecated and removed in 0.15.0 release. Use sv.Detections.from_ultralytics() instead."
"""
This method is deprecated and removed in 0.16.0 release.
Use sv.Classifications.from_ultralytics() instead as it is more generic and
can be used for detections from any ultralytics.engine.results.Results Object
"""
)
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.
@ -206,10 +219,12 @@ class Detections:
@classmethod
def from_ultralytics(cls, ultralytics_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
ultralytics_results (ultralytics.yolo.engine.results.Results):
The output Results instance from YOLOv8
Returns:
Detections: A new Detections object.
@ -226,8 +241,10 @@ class Detections:
>>> model = SAM('mobile_sam.pt')
>>> model = FastSAM('FastSAM-s.pt')
>>> model = RTDETR('rtdetr-l.pt')
>>> # model inferences
>>> result = model(image)[0]
>>> # if tracker is enabled
>>> result = model.track(image)[0]
>>> detections = sv.Detections.from_ultralytics(result)
```
"""
@ -236,15 +253,23 @@ class Detections:
confidence=ultralytics_results.boxes.conf.cpu().numpy(),
class_id=ultralytics_results.boxes.cls.cpu().numpy().astype(int),
mask=extract_ultralytics_masks(ultralytics_results),
tracker_id=ultralytics_results.boxes.id.int().cpu().numpy()
if ultralytics_results.boxes.id is not None
else None,
)
@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.
@ -270,14 +295,16 @@ 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.
Detections: A new Detections object.
Example:
```python
@ -286,7 +313,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)
```
"""
@ -299,7 +327,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.
@ -313,13 +342,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
@ -334,7 +366,6 @@ class Detections:
>>> cfg.MODEL.WEIGHTS = "path/to/model_weights.pth"
>>> predictor = DefaultPredictor(cfg)
>>> result = predictor(image)
>>> detections = sv.Detections.from_detectron2(result)
```
"""
@ -350,14 +381,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
@ -394,7 +429,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
@ -405,14 +442,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
)
@ -423,12 +465,14 @@ class Detections:
return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
@classmethod
def from_paddledet(cls, paddledet_result):
def from_paddledet(cls, paddledet_result) -> Detections:
"""
Creates a Detections instance from [PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection) inference result.
Creates a Detections instance from
[PaddleDetection](https://github.com/PaddlePaddle/PaddleDetection)
inference result.
Args:
paddledet_result (List[dict]): The output Results instance from SAM
paddledet_result (List[dict]): The output Results instance from PaddleDet
Returns:
Detections: A new Detections object.
@ -449,7 +493,7 @@ class Detections:
>>> paddledet_result = trainer.predict([images])[0]
>>> detections = sv.Detections.from_paddledet(paddledet_result=paddledet_result)
>>> detections = sv.Detections.from_paddledet(paddledet_result)
```
"""
return cls(
@ -461,7 +505,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.
@ -484,15 +529,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
@ -512,13 +560,14 @@ class Detections:
list(field) for field in zip(*detections_tuples_list)
]
all_not_none = lambda l: 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,
@ -533,10 +582,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(
@ -559,7 +610,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.
@ -594,11 +646,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])
@ -611,7 +666,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])
@ -622,21 +679,26 @@ 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
assert (
self.confidence is not None
), f"Detections confidence must be given for NMS to be executed."
), "Detections confidence must be given for NMS to be executed."
if class_agnostic:
predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
@ -646,8 +708,8 @@ class Detections:
return self[indices]
assert self.class_id is not None, (
f"Detections class_id must be given for NMS to be executed. If you intended to perform class agnostic "
f"NMS set class_agnostic=True."
"Detections class_id must be given for NMS to be executed. If you intended"
" to perform class agnostic NMS set class_agnostic=True."
)
predictions = np.hstack(

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

107
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))
@ -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,17 +62,22 @@ 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, (
f"Value of `iou_threshold` must be in the closed range from 0 to 1, "
"Value of `iou_threshold` must be in the closed range from 0 to 1, "
f"{iou_threshold} given."
)
rows, columns = predictions.shape
@ -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

339
supervision/metrics/detection.py
View File

@ -21,7 +21,8 @@ def detections_to_tensor(
detections (sv.Detections): Detections/Targets in the format of sv.Detections
with_confidence (bool): Whether to include confidence in the tensor
Returns:
(np.ndarray): Detections as numpy tensors as in (xyxy, class_id, confidence) order
(np.ndarray): Detections as numpy tensors as in (xyxy, class_id,
confidence) order
"""
if detections.class_id is None:
raise ValueError(
@ -46,18 +47,21 @@ def validate_input_tensors(predictions: List[np.ndarray], targets: List[np.ndarr
"""
if len(predictions) != len(targets):
raise ValueError(
f"Number of predictions ({len(predictions)}) and targets ({len(targets)}) must be equal."
f"Number of predictions ({len(predictions)}) and"
f"targets ({len(targets)}) must be equal."
)
if len(predictions) > 0:
if not isinstance(predictions[0], np.ndarray) or not isinstance(
targets[0], np.ndarray
):
raise ValueError(
f"Predictions and targets must be lists of numpy arrays. Got {type(predictions[0])} and {type(targets[0])} instead."
f"Predictions and targets must be lists of numpy arrays."
f"Got {type(predictions[0])} and {type(targets[0])} instead."
)
if predictions[0].shape[1] != 6:
raise ValueError(
f"Predictions must have shape (N, 6). Got {predictions[0].shape} instead."
f"Predictions must have shape (N, 6)."
f"Got {predictions[0].shape} instead."
)
if targets[0].shape[1] != 5:
raise ValueError(
@ -71,10 +75,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
@ -98,8 +106,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.
@ -162,11 +172,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 +264,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 (np.ndarray): Batch prediction. 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 (np.ndarray): Batch target labels. 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 +330,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 +353,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.
@ -344,7 +373,7 @@ class ConfusionMatrix:
>>> model = YOLO(...)
>>> def callback(image: np.ndarray) -> sv.Detections:
... result = model(image)[0]
... return sv.Detections.from_yolov8(result)
... return sv.Detections.from_ultralytics(result)
>>> confusion_matrix = sv.ConfusionMatrix.benchmark(
... dataset = dataset,
@ -386,9 +415,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.
@ -468,16 +499,21 @@ class MeanAveragePrecision:
Mean Average Precision for object detection tasks.
Attributes:
map (float): mAP value.
map50 (float): mAP value at IoU `threshold = 0.5`.
map75 (float): mAP value at IoU `threshold = 0.75`.
per_class_ap (np.ndarray): values for every classes.
map50_95 (float): Mean Average Precision (mAP) calculated over IoU thresholds
ranging from `0.50` to `0.95` with a step size of `0.05`.
map50 (float): Mean Average Precision (mAP) calculated specifically at
an IoU threshold of `0.50`.
map75 (float): Mean Average Precision (mAP) calculated specifically at
an IoU threshold of `0.75`.
per_class_ap50_95 (np.ndarray): Average Precision (AP) values calculated over
IoU thresholds ranging from `0.50` to `0.95` with a step size of `0.05`,
provided for each individual class.
"""
map: float
map50_95: float
map50: float
map75: float
per_class_ap: np.ndarray
per_class_ap50_95: np.ndarray
@classmethod
def from_detections(
@ -513,7 +549,7 @@ class MeanAveragePrecision:
... targets=target,
... )
>>> mean_average_precison.map
>>> mean_average_precison.map50_95
0.2899
```
"""
@ -540,7 +576,8 @@ class MeanAveragePrecision:
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.
callback (Callable[[np.ndarray], Detections]): Function that takes
an image as input and returns Detections object.
Returns:
MeanAveragePrecision: New instance of MeanAveragePrecision.
@ -554,14 +591,14 @@ class MeanAveragePrecision:
>>> model = YOLO(...)
>>> def callback(image: np.ndarray) -> sv.Detections:
... result = model(image)[0]
... return sv.Detections.from_yolov8(result)
... return sv.Detections.from_ultralytics(result)
>>> mean_average_precision = sv.MeanAveragePrecision.benchmark(
... dataset = dataset,
... callback = callback
... )
>>> mean_average_precision.map
>>> mean_average_precision.map50_95
0.433
```
"""
@ -583,11 +620,18 @@ class MeanAveragePrecision:
targets: List[np.ndarray],
) -> MeanAveragePrecision:
"""
Calculate Mean Average Precision based on predicted and ground-truth detections at different threshold.
Calculate Mean Average Precision based on predicted and ground-truth
detections at different threshold.
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.
Returns:
MeanAveragePrecision: New instance of MeanAveragePrecision.
@ -625,98 +669,63 @@ class MeanAveragePrecision:
... targets=targets,
... )
>>> mean_average_precison.map
>>> mean_average_precison.map50_95
0.2899
```
"""
validate_input_tensors(predictions, targets)
map, map50, map75 = 0, 0, 0
iou_thresholds = np.linspace(0.5, 0.95, 10)
stats = []
class_index = 4
conf_index = 5
stats, average_precisions = [], []
iou_levels = np.linspace(0.5, 0.95, 10)
num_ious = iou_levels.size
for true_batch, detection_batch in zip(targets, predictions):
nl, npr = (
true_batch.shape[0],
detection_batch.shape[0],
)
correct = np.zeros((npr, num_ious), dtype=bool)
if npr == 0:
if nl:
stats.append((correct, *np.zeros((2, 0)), true_batch[:, 4]))
# Gather matching stats for predictions and targets
for true_objs, predicted_objs in zip(targets, predictions):
if predicted_objs.shape[0] == 0:
if true_objs.shape[0]:
stats.append(
(
np.zeros((0, iou_thresholds.size), dtype=bool),
*np.zeros((2, 0)),
true_objs[:, 4],
)
)
continue
if nl:
correct = MeanAveragePrecision._match_detection_batch(
predictions=detection_batch,
targets=true_batch,
iou_levels=iou_levels,
if true_objs.shape[0]:
matches = cls._match_detection_batch(
predicted_objs, true_objs, iou_thresholds
)
stats.append(
(
correct,
detection_batch[:, conf_index],
detection_batch[:, class_index],
true_batch[:, class_index],
matches,
predicted_objs[:, 5],
predicted_objs[:, 4],
true_objs[:, 4],
)
)
stats = [np.concatenate(x, 0) for x in zip(*stats)]
# Compute average precisions if any matches exist
if stats:
concatenated_stats = [np.concatenate(items, 0) for items in zip(*stats)]
average_precisions = cls._average_precisions_per_class(*concatenated_stats)
map50 = average_precisions[:, 0].mean()
map75 = average_precisions[:, 5].mean()
map50_95 = average_precisions.mean()
else:
map50, map75, map50_95 = 0, 0, 0
average_precisions = []
if len(stats) and stats[0].any():
average_precisions = cls._average_precisions_per_class(*stats)
ap50, ap75, average_precisions = (
average_precisions[:, 0],
average_precisions[:, 5],
average_precisions.mean(1),
)
map50, map75, map = ap50.mean(), ap75.mean(), average_precisions.mean()
return cls(map=map, map50=map50, map75=map75, per_class_ap=average_precisions)
@staticmethod
def _match_detection_batch(
predictions: np.ndarray, targets: np.ndarray, iou_levels: np.ndarray
) -> np.ndarray:
"""
Args:
predictions (np.ndarray): batch prediction
targets (np.ndarray): batch target labels
iou_levels (np.ndarray): iou levels array contains different iou levels
Returns:
(np.ndarray): matched prediction with target lebels result
"""
correct = np.zeros((predictions.shape[0], iou_levels.shape[0])).astype(bool)
iou = box_iou_batch(targets[:, :4], predictions[:, :4])
correct_class = targets[:, 4:5] == predictions[:, 4]
for i in range(len(iou_levels)):
x = np.where((iou >= iou_levels[i]) & correct_class)
if x[0].shape[0]:
_X1 = np.concatenate(
[np.expand_dims(x[0], 1), np.expand_dims(x[1], 1)], axis=1
)
_x2 = iou[x[0], x[1]][:, None]
matches = np.concatenate([_X1, _x2], axis=1)
if x[0].shape[0] > 1:
matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
correct[matches[:, 1].astype(int), i] = True
correct[matches[:, 1].astype(int), i] = True
return correct
return cls(
map50_95=map50_95,
map50=map50,
map75=map75,
per_class_ap50_95=average_precisions,
)
@staticmethod
def compute_average_precision(recall: np.ndarray, precision: np.ndarray) -> float:
"""
Compute the average precision using 101-point interpolation (COCO), given the recall and precision curves.
Compute the average precision using 101-point interpolation (COCO), given
the recall and precision curves.
Args:
recall (np.ndarray): The recall curve.
@ -737,54 +746,98 @@ class MeanAveragePrecision:
average_precision = np.trapz(interpolated_precision, interpolated_recall_levels)
return average_precision
@staticmethod
def _match_detection_batch(
predictions: np.ndarray, targets: np.ndarray, iou_thresholds: np.ndarray
) -> np.ndarray:
"""
Match predictions with target labels based on IoU levels.
Args:
predictions (np.ndarray): Batch prediction. 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 (np.ndarray): Batch target labels. 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.
iou_thresholds (np.ndarray): Array contains different IoU thresholds.
Returns:
np.ndarray: Matched prediction with target labels result.
"""
num_predictions, num_iou_levels = predictions.shape[0], iou_thresholds.shape[0]
correct = np.zeros((num_predictions, num_iou_levels), dtype=bool)
iou = box_iou_batch(targets[:, :4], predictions[:, :4])
correct_class = targets[:, 4:5] == predictions[:, 4]
for i, iou_level in enumerate(iou_thresholds):
matched_indices = np.where((iou >= iou_level) & correct_class)
if matched_indices[0].shape[0]:
combined_indices = np.stack(matched_indices, axis=1)
iou_values = iou[matched_indices][:, None]
matches = np.hstack([combined_indices, iou_values])
if matched_indices[0].shape[0] > 1:
matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
correct[matches[:, 1].astype(int), i] = True
return correct
@staticmethod
def _average_precisions_per_class(
matches: np.ndarray,
prediction_confidence: np.ndarray,
prediction_class_ids: np.ndarray,
true_batch_class_ids: np.ndarray,
true_class_ids: np.ndarray,
eps: float = 1e-16,
) -> np.ndarray:
"""
Compute the average precision, given the recall and precision curves.
Source: https://github.com/rafaelpadilla/Object-Detection-Metrics.
Args:
matches (np.ndarray): True positives (nparray, nx1 or nx10).
prediction_confidence (np.ndarray): Objectness value from 0-1 (nparray).
prediction_class_ids (np.ndarray): Predicted object classes (nparray).
true_batch_class_ids (np.ndarray): True object classes (nparray).
Returns:
(np.ndarray): Average precision for different iou level array
"""
sorted_confidences = np.argsort(-prediction_confidence)
matches = matches[sorted_confidences]
prediction_class_ids = prediction_class_ids[sorted_confidences]
# Find unique classes
unique_classes, class_counts = np.unique(
true_batch_class_ids, return_counts=True
)
num_classes = unique_classes.shape[0] # number of classes, number of detections
Args:
matches (np.ndarray): True positives.
prediction_confidence (np.ndarray): Objectness value from 0-1.
prediction_class_ids (np.ndarray): Predicted object classes.
true_class_ids (np.ndarray): True object classes.
eps (float, optional): Small value to prevent division by zero.
Returns:
np.ndarray: Average precision for different IoU levels.
"""
sorted_indices = np.argsort(-prediction_confidence)
matches = matches[sorted_indices]
prediction_class_ids = prediction_class_ids[sorted_indices]
unique_classes, class_counts = np.unique(true_class_ids, return_counts=True)
num_classes = unique_classes.shape[0]
average_precisions = np.zeros((num_classes, matches.shape[1]))
for ci, c in enumerate(unique_classes):
valid = prediction_class_ids == c
num_targets = class_counts[ci] # number of labels
num_predictions = valid.sum() # number of predictions
if num_predictions == 0 or num_targets == 0:
for class_idx, class_id in enumerate(unique_classes):
is_class = prediction_class_ids == class_id
total_true = class_counts[class_idx]
total_prediction = is_class.sum()
if total_prediction == 0 or total_true == 0:
continue
fp_pool = (1 - matches[valid]).cumsum(0)
tp_pool = matches[valid].cumsum(0)
false_positives = (1 - matches[is_class]).cumsum(0)
true_positives = matches[is_class].cumsum(0)
recall = true_positives / (total_true + eps)
precision = true_positives / (true_positives + false_positives)
recall = tp_pool / (num_targets + eps)
precision = tp_pool / (tp_pool + fp_pool)
for j in range(matches.shape[1]):
for iou_level_idx in range(matches.shape[1]):
average_precisions[
ci, j
class_idx, iou_level_idx
] = MeanAveragePrecision.compute_average_precision(
recall[:, j], precision[:, j]
recall[:, iou_level_idx], precision[:, iou_level_idx]
)
return average_precisions

0
supervision/tracker/__init__.py Normal file
View File

0
supervision/tracker/byte_tracker/__init__.py Normal file
View File

55
supervision/tracker/byte_tracker/basetrack.py Normal file
View File

@ -0,0 +1,55 @@
from collections import OrderedDict
from enum import Enum
import numpy as np
class TrackState(Enum):
New = 0
Tracked = 1
Lost = 2
Removed = 3
class BaseTrack:
_count = 0
def __init__(self):
self.track_id = 0
self.is_activated = False
self.state = TrackState.New
self.history = OrderedDict()
self.features = []
self.curr_feature = None
self.score = 0
self.start_frame = 0
self.frame_id = 0
self.time_since_update = 0
# multi-camera
self.location = (np.inf, np.inf)
@property
def end_frame(self) -> int:
return self.frame_id
@staticmethod
def next_id() -> int:
BaseTrack._count += 1
return BaseTrack._count
def activate(self, *args):
raise NotImplementedError
def predict(self):
raise NotImplementedError
def update(self, *args, **kwargs):
raise NotImplementedError
def mark_lost(self):
self.state = TrackState.Lost
def mark_removed(self):
self.state = TrackState.Removed

472
supervision/tracker/byte_tracker/core.py Normal file
View File

@ -0,0 +1,472 @@
from typing import List, Tuple
import numpy as np
from supervision.detection.core import Detections
from supervision.tracker.byte_tracker import matching
from supervision.tracker.byte_tracker.basetrack import BaseTrack, TrackState
from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter
class STrack(BaseTrack):
shared_kalman = KalmanFilter()
def __init__(self, tlwh, score, class_ids):
# wait activate
self._tlwh = np.asarray(tlwh, dtype=np.float32)
self.kalman_filter = None
self.mean, self.covariance = None, None
self.is_activated = False
self.score = score
self.class_ids = class_ids
self.tracklet_len = 0
def predict(self):
mean_state = self.mean.copy()
if self.state != TrackState.Tracked:
mean_state[7] = 0
self.mean, self.covariance = self.kalman_filter.predict(
mean_state, self.covariance
)
@staticmethod
def multi_predict(stracks):
if len(stracks) > 0:
multi_mean = np.asarray([st.mean.copy() for st in stracks])
multi_covariance = np.asarray([st.covariance for st in stracks])
for i, st in enumerate(stracks):
if st.state != TrackState.Tracked:
multi_mean[i][7] = 0
multi_mean, multi_covariance = STrack.shared_kalman.multi_predict(
multi_mean, multi_covariance
)
for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
stracks[i].mean = mean
stracks[i].covariance = cov
def activate(self, kalman_filter, frame_id):
"""Start a new tracklet"""
self.kalman_filter = kalman_filter
self.track_id = self.next_id()
self.mean, self.covariance = self.kalman_filter.initiate(
self.tlwh_to_xyah(self._tlwh)
)
self.tracklet_len = 0
self.state = TrackState.Tracked
if frame_id == 1:
self.is_activated = True
self.frame_id = frame_id
self.start_frame = frame_id
def re_activate(self, new_track, frame_id, new_id=False):
self.mean, self.covariance = self.kalman_filter.update(
self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
)
self.tracklet_len = 0
self.state = TrackState.Tracked
self.is_activated = True
self.frame_id = frame_id
if new_id:
self.track_id = self.next_id()
self.score = new_track.score
def update(self, new_track, frame_id):
"""
Update a matched track
:type new_track: STrack
:type frame_id: int
:type update_feature: bool
:return:
"""
self.frame_id = frame_id
self.tracklet_len += 1
new_tlwh = new_track.tlwh
self.mean, self.covariance = self.kalman_filter.update(
self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh)
)
self.state = TrackState.Tracked
self.is_activated = True
self.score = new_track.score
@property
def tlwh(self):
"""Get current position in bounding box format `(top left x, top left y,
width, height)`.
"""
if self.mean is None:
return self._tlwh.copy()
ret = self.mean[:4].copy()
ret[2] *= ret[3]
ret[:2] -= ret[2:] / 2
return ret
@property
def tlbr(self):
"""Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
`(top left, bottom right)`.
"""
ret = self.tlwh.copy()
ret[2:] += ret[:2]
return ret
@staticmethod
def tlwh_to_xyah(tlwh):
"""Convert bounding box to format `(center x, center y, aspect ratio,
height)`, where the aspect ratio is `width / height`.
"""
ret = np.asarray(tlwh).copy()
ret[:2] += ret[2:] / 2
ret[2] /= ret[3]
return ret
def to_xyah(self):
return self.tlwh_to_xyah(self.tlwh)
@staticmethod
def tlbr_to_tlwh(tlbr):
ret = np.asarray(tlbr).copy()
ret[2:] -= ret[:2]
return ret
@staticmethod
def tlwh_to_tlbr(tlwh):
ret = np.asarray(tlwh).copy()
ret[2:] += ret[:2]
return ret
def __repr__(self):
return "OT_{}_({}-{})".format(self.track_id, self.start_frame, self.end_frame)
def detections2boxes(detections: Detections) -> np.ndarray:
"""
Convert Supervision Detections to numpy tensors for further computation.
Args:
detections (Detections): Detections/Targets in the format of sv.Detections.
Returns:
(np.ndarray): Detections as numpy tensors as in
`(x_min, y_min, x_max, y_max, confidence, class_id)` order.
"""
return np.hstack(
(
detections.xyxy,
detections.confidence[:, np.newaxis],
detections.class_id[:, np.newaxis],
)
)
class ByteTrack:
"""
Initialize the ByteTrack object.
Parameters:
track_thresh (float, optional): Detection confidence threshold
for track activation.
track_buffer (int, optional): Number of frames to buffer when a track is lost.
match_thresh (float, optional): Threshold for matching tracks with detections.
frame_rate (int, optional): The frame rate of the video.
"""
def __init__(
self,
track_thresh: float = 0.25,
track_buffer: int = 30,
match_thresh: float = 0.8,
frame_rate: int = 30,
):
self.track_thresh = track_thresh
self.match_thresh = match_thresh
self.frame_id = 0
self.det_thresh = self.track_thresh + 0.1
self.max_time_lost = int(frame_rate / 30.0 * track_buffer)
self.kalman_filter = KalmanFilter()
self.tracked_tracks: List[STrack] = []
self.lost_tracks: List[STrack] = []
self.removed_tracks: List[STrack] = []
def update_with_detections(self, detections: Detections) -> Detections:
"""
Updates the tracker with the provided detections and
returns the updated detection results.
Parameters:
detections: The new detections to update with.
Returns:
Detection: The updated detection results that now include tracking IDs.
Example:
```python
>>> import supervision as sv
>>> from ultralytics import YOLO
>>> model = YOLO(...)
>>> byte_tracker = sv.ByteTrack()
>>> annotator = sv.BoxAnnotator()
>>> def callback(frame: np.ndarray, index: int) -> np.ndarray:
... results = model(frame)[0]
... detections = sv.Detections.from_ultralytics(results)
... detections = byte_tracker.update_with_detections(detections)
... labels = [
... f"#{tracker_id} {model.model.names[class_id]} {confidence:0.2f}"
... for _, _, confidence, class_id, tracker_id
... in detections
... ]
... return annotator.annotate(scene=frame.copy(),
... detections=detections, labels=labels)
>>> sv.process_video(
... source_path='...',
... target_path='...',
... callback=callback
... )
```
"""
tracks = self.update_with_tensors(
tensors=detections2boxes(detections=detections)
)
detections = Detections.empty()
if len(tracks) > 0:
detections.xyxy = np.array(
[track.tlbr for track in tracks], dtype=np.float32
)
detections.class_id = np.array(
[int(t.class_ids) for t in tracks], dtype=int
)
detections.tracker_id = np.array(
[int(t.track_id) for t in tracks], dtype=int
)
detections.confidence = np.array(
[t.score for t in tracks], dtype=np.float32
)
return detections
def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
"""
Updates the tracker with the provided tensors and returns the updated tracks.
Parameters:
tensors: The new tensors to update with.
Returns:
List[STrack]: Updated tracks.
"""
self.frame_id += 1
activated_starcks = []
refind_stracks = []
lost_stracks = []
removed_stracks = []
class_ids = tensors[:, 5]
scores = tensors[:, 4]
bboxes = tensors[:, :4]
remain_inds = scores > self.track_thresh
inds_low = scores > 0.1
inds_high = scores < self.track_thresh
inds_second = np.logical_and(inds_low, inds_high)
dets_second = bboxes[inds_second]
dets = bboxes[remain_inds]
scores_keep = scores[remain_inds]
scores_second = scores[inds_second]
class_ids_keep = class_ids[remain_inds]
class_ids_second = class_ids[inds_second]
if len(dets) > 0:
"""Detections"""
detections = [
STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
for (tlbr, s, c) in zip(dets, scores_keep, class_ids_keep)
]
else:
detections = []
""" Add newly detected tracklets to tracked_stracks"""
unconfirmed = []
tracked_stracks = [] # type: list[STrack]
for track in self.tracked_tracks:
if not track.is_activated:
unconfirmed.append(track)
else:
tracked_stracks.append(track)
""" Step 2: First association, with high score detection boxes"""
strack_pool = joint_tracks(tracked_stracks, self.lost_tracks)
# Predict the current location with KF
STrack.multi_predict(strack_pool)
dists = matching.iou_distance(strack_pool, detections)
dists = matching.fuse_score(dists, detections)
matches, u_track, u_detection = matching.linear_assignment(
dists, thresh=self.match_thresh
)
for itracked, idet in matches:
track = strack_pool[itracked]
det = detections[idet]
if track.state == TrackState.Tracked:
track.update(detections[idet], self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
""" Step 3: Second association, with low score detection boxes"""
# association the untrack to the low score detections
if len(dets_second) > 0:
"""Detections"""
detections_second = [
STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
for (tlbr, s, c) in zip(dets_second, scores_second, class_ids_second)
]
else:
detections_second = []
r_tracked_stracks = [
strack_pool[i]
for i in u_track
if strack_pool[i].state == TrackState.Tracked
]
dists = matching.iou_distance(r_tracked_stracks, detections_second)
matches, u_track, u_detection_second = matching.linear_assignment(
dists, thresh=0.5
)
for itracked, idet in matches:
track = r_tracked_stracks[itracked]
det = detections_second[idet]
if track.state == TrackState.Tracked:
track.update(det, self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
refind_stracks.append(track)
for it in u_track:
track = r_tracked_stracks[it]
if not track.state == TrackState.Lost:
track.mark_lost()
lost_stracks.append(track)
"""Deal with unconfirmed tracks, usually tracks with only one beginning frame"""
detections = [detections[i] for i in u_detection]
dists = matching.iou_distance(unconfirmed, detections)
dists = matching.fuse_score(dists, detections)
matches, u_unconfirmed, u_detection = matching.linear_assignment(
dists, thresh=0.7
)
for itracked, idet in matches:
unconfirmed[itracked].update(detections[idet], self.frame_id)
activated_starcks.append(unconfirmed[itracked])
for it in u_unconfirmed:
track = unconfirmed[it]
track.mark_removed()
removed_stracks.append(track)
""" Step 4: Init new stracks"""
for inew in u_detection:
track = detections[inew]
if track.score < self.det_thresh:
continue
track.activate(self.kalman_filter, self.frame_id)
activated_starcks.append(track)
""" Step 5: Update state"""
for track in self.lost_tracks:
if self.frame_id - track.end_frame > self.max_time_lost:
track.mark_removed()
removed_stracks.append(track)
self.tracked_tracks = [
t for t in self.tracked_tracks if t.state == TrackState.Tracked
]
self.tracked_tracks = joint_tracks(self.tracked_tracks, activated_starcks)
self.tracked_tracks = joint_tracks(self.tracked_tracks, refind_stracks)
self.lost_tracks = sub_tracks(self.lost_tracks, self.tracked_tracks)
self.lost_tracks.extend(lost_stracks)
self.lost_tracks = sub_tracks(self.lost_tracks, self.removed_tracks)
self.removed_tracks.extend(removed_stracks)
self.tracked_tracks, self.lost_tracks = remove_duplicate_tracks(
self.tracked_tracks, self.lost_tracks
)
output_stracks = [track for track in self.tracked_tracks if track.is_activated]
return output_stracks
def joint_tracks(
track_list_a: List[STrack], track_list_b: List[STrack]
) -> List[STrack]:
"""
Joins two lists of tracks, ensuring that the resulting list does not
contain tracks with duplicate track_id values.
Parameters:
track_list_a: First list of tracks (with track_id attribute).
track_list_b: Second list of tracks (with track_id attribute).
Returns:
Combined list of tracks from track_list_a and track_list_b
without duplicate track_id values.
"""
seen_track_ids = set()
result = []
for track in track_list_a + track_list_b:
if track.track_id not in seen_track_ids:
seen_track_ids.add(track.track_id)
result.append(track)
return result
def sub_tracks(track_list_a: List, track_list_b: List) -> List[int]:
"""
Returns a list of tracks from track_list_a after removing any tracks
that share the same track_id with tracks in track_list_b.
Parameters:
track_list_a: List of tracks (with track_id attribute).
track_list_b: List of tracks (with track_id attribute) to
be subtracted from track_list_a.
Returns:
List of remaining tracks from track_list_a after subtraction.
"""
tracks = {track.track_id: track for track in track_list_a}
track_ids_b = {track.track_id for track in track_list_b}
for track_id in track_ids_b:
tracks.pop(track_id, None)
return list(tracks.values())
def remove_duplicate_tracks(tracks_a: List, tracks_b: List) -> Tuple[List, List]:
pairwise_distance = matching.iou_distance(tracks_a, tracks_b)
matching_pairs = np.where(pairwise_distance < 0.15)
duplicates_a, duplicates_b = set(), set()
for track_index_a, track_index_b in zip(*matching_pairs):
time_a = tracks_a[track_index_a].frame_id - tracks_a[track_index_a].start_frame
time_b = tracks_b[track_index_b].frame_id - tracks_b[track_index_b].start_frame
if time_a > time_b:
duplicates_b.add(track_index_b)
else:
duplicates_a.add(track_index_a)
result_a = [
track for index, track in enumerate(tracks_a) if index not in duplicates_a
]
result_b = [
track for index, track in enumerate(tracks_b) if index not in duplicates_b
]
return result_a, result_b

205
supervision/tracker/byte_tracker/kalman_filter.py Normal file
View File

@ -0,0 +1,205 @@
from typing import Tuple
import numpy as np
import scipy.linalg
class KalmanFilter:
"""
A simple Kalman filter for tracking bounding boxes in image space.
The 8-dimensional state space
x, y, a, h, vx, vy, va, vh
contains the bounding box center position (x, y), aspect ratio a, height h,
and their respective velocities.
Object motion follows a constant velocity model. The bounding box location
(x, y, a, h) is taken as direct observation of the state space (linear
observation model).
"""
def __init__(self):
ndim, dt = 4, 1.0
self._motion_mat = np.eye(2 * ndim, 2 * ndim)
for i in range(ndim):
self._motion_mat[i, ndim + i] = dt
self._update_mat = np.eye(ndim, 2 * ndim)
self._std_weight_position = 1.0 / 20
self._std_weight_velocity = 1.0 / 160
def initiate(self, measurement: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
"""
Create track from an unassociated measurement.
Args:
measurement (ndarray): Bounding box coordinates (x, y, a, h) with
center position (x, y), aspect ratio a, and height h.
Returns:
Tuple[ndarray, ndarray]: Returns the mean vector (8 dimensional) and
covariance matrix (8x8 dimensional) of the new track.
Unobserved velocities are initialized to 0 mean.
"""
mean_pos = measurement
mean_vel = np.zeros_like(mean_pos)
mean = np.r_[mean_pos, mean_vel]
std = [
2 * self._std_weight_position * measurement[3],
2 * self._std_weight_position * measurement[3],
1e-2,
2 * self._std_weight_position * measurement[3],
10 * self._std_weight_velocity * measurement[3],
10 * self._std_weight_velocity * measurement[3],
1e-5,
10 * self._std_weight_velocity * measurement[3],
]
covariance = np.diag(np.square(std))
return mean, covariance
def predict(
self, mean: np.ndarray, covariance: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
"""
Run Kalman filter prediction step.
Args:
mean (ndarray): The 8 dimensional mean vector of the object
state at the previous time step.
covariance (ndarray): The 8x8 dimensional covariance matrix of
the object state at the previous time step.
Returns:
Tuple[ndarray, ndarray]: Returns the mean vector and
covariance matrix of the predicted state.
Unobserved velocities are initialized to 0 mean.
"""
std_pos = [
self._std_weight_position * mean[3],
self._std_weight_position * mean[3],
1e-2,
self._std_weight_position * mean[3],
]
std_vel = [
self._std_weight_velocity * mean[3],
self._std_weight_velocity * mean[3],
1e-5,
self._std_weight_velocity * mean[3],
]
motion_cov = np.diag(np.square(np.r_[std_pos, std_vel]))
mean = np.dot(mean, self._motion_mat.T)
covariance = (
np.linalg.multi_dot((self._motion_mat, covariance, self._motion_mat.T))
+ motion_cov
)
return mean, covariance
def project(
self, mean: np.ndarray, covariance: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
"""
Project state distribution to measurement space.
Args:
mean (ndarray): The state's mean vector (8 dimensional array).
covariance (ndarray): The state's covariance matrix (8x8 dimensional).
Returns:
Tuple[ndarray, ndarray]: Returns the projected mean and
covariance matrix of the given state estimate.
"""
std = [
self._std_weight_position * mean[3],
self._std_weight_position * mean[3],
1e-1,
self._std_weight_position * mean[3],
]
innovation_cov = np.diag(np.square(std))
mean = np.dot(self._update_mat, mean)
covariance = np.linalg.multi_dot(
(self._update_mat, covariance, self._update_mat.T)
)
return mean, covariance + innovation_cov
def multi_predict(
self, mean: np.ndarray, covariance: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
"""
Run Kalman filter prediction step (Vectorized version).
Args:
mean (ndarray): The Nx8 dimensional mean matrix
of the object states at the previous time step.
covariance (ndarray): The Nx8x8 dimensional covariance matrices
of the object states at the previous time step.
Returns:
Tuple[ndarray, ndarray]: Returns the mean vector and
covariance matrix of the predicted state.
Unobserved velocities are initialized to 0 mean.
"""
std_pos = [
self._std_weight_position * mean[:, 3],
self._std_weight_position * mean[:, 3],
1e-2 * np.ones_like(mean[:, 3]),
self._std_weight_position * mean[:, 3],
]
std_vel = [
self._std_weight_velocity * mean[:, 3],
self._std_weight_velocity * mean[:, 3],
1e-5 * np.ones_like(mean[:, 3]),
self._std_weight_velocity * mean[:, 3],
]
sqr = np.square(np.r_[std_pos, std_vel]).T
motion_cov = []
for i in range(len(mean)):
motion_cov.append(np.diag(sqr[i]))
motion_cov = np.asarray(motion_cov)
mean = np.dot(mean, self._motion_mat.T)
left = np.dot(self._motion_mat, covariance).transpose((1, 0, 2))
covariance = np.dot(left, self._motion_mat.T) + motion_cov
return mean, covariance
def update(
self, mean: np.ndarray, covariance: np.ndarray, measurement: np.ndarray
) -> Tuple[np.ndarray, np.ndarray]:
"""
Run Kalman filter correction step.
Args:
mean (ndarray): The predicted state's mean vector (8 dimensional).
covariance (ndarray): The state's covariance matrix (8x8 dimensional).
measurement (ndarray): The 4-dimensional measurement vector (x, y, a, h),
where (x, y) is the center position, a the aspect ratio,
and h the height of the bounding box.
Returns:
Tuple[ndarray, ndarray]: Returns the measurement-corrected
state distribution.
"""
projected_mean, projected_cov = self.project(mean, covariance)
chol_factor, lower = scipy.linalg.cho_factor(
projected_cov, lower=True, check_finite=False
)
kalman_gain = scipy.linalg.cho_solve(
(chol_factor, lower),
np.dot(covariance, self._update_mat.T).T,
check_finite=False,
).T
innovation = measurement - projected_mean
new_mean = mean + np.dot(innovation, kalman_gain.T)
new_covariance = covariance - np.linalg.multi_dot(
(kalman_gain, projected_cov, kalman_gain.T)
)
return new_mean, new_covariance

64
supervision/tracker/byte_tracker/matching.py Normal file
View File

@ -0,0 +1,64 @@
from typing import List, Tuple
import numpy as np
from scipy.optimize import linear_sum_assignment
from supervision.detection.utils import box_iou_batch
def indices_to_matches(
cost_matrix: np.ndarray, indices: np.ndarray, thresh: float
) -> Tuple[np.ndarray, tuple, tuple]:
matched_cost = cost_matrix[tuple(zip(*indices))]
matched_mask = matched_cost <= thresh
matches = indices[matched_mask]
unmatched_a = tuple(set(range(cost_matrix.shape[0])) - set(matches[:, 0]))
unmatched_b = tuple(set(range(cost_matrix.shape[1])) - set(matches[:, 1]))
return matches, unmatched_a, unmatched_b
def linear_assignment(
cost_matrix: np.ndarray, thresh: float
) -> [np.ndarray, Tuple[int], Tuple[int, int]]:
if cost_matrix.size == 0:
return (
np.empty((0, 2), dtype=int),
tuple(range(cost_matrix.shape[0])),
tuple(range(cost_matrix.shape[1])),
)
cost_matrix[cost_matrix > thresh] = thresh + 1e-4
row_ind, col_ind = linear_sum_assignment(cost_matrix)
indices = np.column_stack((row_ind, col_ind))
return indices_to_matches(cost_matrix, indices, thresh)
def iou_distance(atracks: List, btracks: List) -> np.ndarray:
if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
len(btracks) > 0 and isinstance(btracks[0], np.ndarray)
):
atlbrs = atracks
btlbrs = btracks
else:
atlbrs = [track.tlbr for track in atracks]
btlbrs = [track.tlbr for track in btracks]
_ious = np.zeros((len(atlbrs), len(btlbrs)), dtype=np.float32)
if _ious.size != 0:
_ious = box_iou_batch(np.asarray(atlbrs), np.asarray(btlbrs))
cost_matrix = 1 - _ious
return cost_matrix
def fuse_score(cost_matrix: np.ndarray, detections: List) -> np.ndarray:
if cost_matrix.size == 0:
return cost_matrix
iou_sim = 1 - cost_matrix
det_scores = np.array([det.score for det in detections])
det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
fuse_sim = iou_sim * det_scores
fuse_cost = 1 - fuse_sim
return fuse_cost

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)

12
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:
@ -78,7 +81,8 @@ def plot_images_grid(
if len(images) > nrows * ncols:
raise ValueError(
"The number of images exceeds the grid size. Please increase the grid size or reduce the number of images."
"The number of images exceeds the grid size. Please increase the grid size"
" or reduce the number of images."
)
fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=size)

47
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)
```
"""
@ -103,7 +108,7 @@ def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
raise Exception(f"Could not open video at {source_path}")
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
if end is not None and end > total_frames:
raise Exception(f"Requested frames are outbound")
raise Exception("Requested frames are outbound")
start = max(start, 0)
end = min(end, total_frames) if end is not None else total_frames
video.set(cv2.CAP_PROP_POS_FRAMES, start)
@ -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,24 +163,28 @@ 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
>>> from supervision import process_video
>>> import supervision as sv
>>> def process_frame(scene: np.ndarray) -> np.ndarray:
>>> def callback(scene: np.ndarray, index: int) -> np.ndarray:
... ...
>>> process_video(
... source_path='source_video.mp4',
... target_path='target_video.mp4',
... callback=process_frame
... source_path='...',
... target_path='...',
... callback=callback
... )
```
"""

4
test/classification/test_core.py
View File

@ -1,8 +1,8 @@
from typing import List, Optional, Tuple
from contextlib import ExitStack as DoesNotRaise
from typing import Optional, Tuple
import pytest
import numpy as np
import pytest
from supervision.classification.core import Classifications

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

@ -1,13 +1,17 @@
from contextlib import ExitStack as DoesNotRaise
from typing import List, Tuple, Dict
from typing import Dict, List, Tuple
import numpy as np
import pytest
from supervision import Detections
from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes, \
group_coco_annotations_by_image_id, coco_annotations_to_detections, build_coco_class_index_mapping
import numpy as np
from supervision.dataset.formats.coco import (
build_coco_class_index_mapping,
classes_to_coco_categories,
coco_annotations_to_detections,
coco_categories_to_classes,
group_coco_annotations_by_image_id,
)
def mock_cock_coco_annotation(
@ -15,7 +19,7 @@ def mock_cock_coco_annotation(
image_id: int = 0,
category_id: int = 0,
bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
area: float = 0.0
area: float = 0.0,
) -> dict:
return {
"id": annotation_id,
@ -23,103 +27,52 @@ def mock_cock_coco_annotation(
"category_id": category_id,
"bbox": list(bbox),
"area": area,
"iscrowd": 0
"iscrowd": 0,
}
@pytest.mark.parametrize(
"coco_categories, expected_result, exception",
[
([], [], DoesNotRaise()), # empty coco categories
(
[],
[],
DoesNotRaise()
), # empty coco categories
(
[
{
"id": 0,
"name": "fashion-assistant",
"supercategory": "none"
}
],
[{"id": 0, "name": "fashion-assistant", "supercategory": "none"}],
["fashion-assistant"],
DoesNotRaise()
DoesNotRaise(),
), # single coco category with supercategory == "none"
(
[
{
"id": 0,
"name": "fashion-assistant",
"supercategory": "none"
},
{
"id": 1,
"name": "baseball cap",
"supercategory": "fashion-assistant"
}
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
],
[
"fashion-assistant", "baseball cap"
],
DoesNotRaise()
), # two coco categories; one with supercategory == "none" and one with supercategory != "none"
["fashion-assistant", "baseball cap"],
DoesNotRaise(),
), # two coco categories; one with supercategory == "none" and
# one with supercategory != "none"
(
[
{
"id": 0,
"name": "fashion-assistant",
"supercategory": "none"
},
{
"id": 1,
"name": "baseball cap",
"supercategory": "fashion-assistant"
},
{
"id": 2,
"name": "hoodie",
"supercategory": "fashion-assistant"
}
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
],
[
"fashion-assistant",
"baseball cap",
"hoodie"
],
DoesNotRaise()
), # three coco categories; one with supercategory == "none" and two with supercategory != "none"
["fashion-assistant", "baseball cap", "hoodie"],
DoesNotRaise(),
), # three coco categories; one with supercategory == "none" and
# two with supercategory != "none"
(
[
{
"id": 0,
"name": "fashion-assistant",
"supercategory": "none"
},
{
"id": 2,
"name": "hoodie",
"supercategory": "fashion-assistant"
},
{
"id": 1,
"name": "baseball cap",
"supercategory": "fashion-assistant"
}
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
],
[
"fashion-assistant",
"baseball cap",
"hoodie"
],
DoesNotRaise()
), # three coco categories; one with supercategory == "none" and two with supercategory != "none" (different order)
]
["fashion-assistant", "baseball cap", "hoodie"],
DoesNotRaise(),
), # three coco categories; one with supercategory == "none" and
# two with supercategory != "none" (different order)
],
)
def test_coco_categories_to_classes(
coco_categories: List[dict],
expected_result: List[str],
exception: Exception
coco_categories: List[dict], expected_result: List[str], exception: Exception
) -> None:
with exception:
result = coco_categories_to_classes(coco_categories=coco_categories)
@ -129,26 +82,14 @@ def test_coco_categories_to_classes(
@pytest.mark.parametrize(
"classes, exception",
[
(
[],
DoesNotRaise()
), # empty classes
(
[
"baseball cap"
],
DoesNotRaise()
), # single class
(
[
"baseball cap",
"hoodie"
],
DoesNotRaise()
), # two classes
]
([], DoesNotRaise()), # empty classes
(["baseball cap"], DoesNotRaise()), # single class
(["baseball cap", "hoodie"], DoesNotRaise()), # two classes
],
)
def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exception: Exception) -> None:
def test_classes_to_coco_categories_and_back_to_classes(
classes: List[str], exception: Exception
) -> None:
with exception:
coco_categories = classes_to_coco_categories(classes=classes)
result = coco_categories_to_classes(coco_categories=coco_categories)
@ -158,36 +99,36 @@ def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exce
@pytest.mark.parametrize(
"coco_annotations, expected_result, exception",
[
([], {}, DoesNotRaise()), # empty coco annotations
(
[],
{},
DoesNotRaise()
), # empty coco annotations
(
[
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
],
[mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)],
{
0: [
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
mock_cock_coco_annotation(
annotation_id=0, image_id=0, category_id=0
)
]
},
DoesNotRaise()
DoesNotRaise(),
), # single coco annotation
(
[
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0),
],
{
0: [
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
mock_cock_coco_annotation(
annotation_id=0, image_id=0, category_id=0
)
],
1: [
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
]
mock_cock_coco_annotation(
annotation_id=1, image_id=1, category_id=0
)
],
},
DoesNotRaise()
DoesNotRaise(),
), # two coco annotations
(
[
@ -201,29 +142,41 @@ def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exce
],
{
0: [
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
mock_cock_coco_annotation(
annotation_id=0, image_id=0, category_id=0
),
],
1: [
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
mock_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
mock_cock_coco_annotation(
annotation_id=1, image_id=1, category_id=1
),
mock_cock_coco_annotation(
annotation_id=2, image_id=1, category_id=2
),
],
2: [
mock_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
mock_cock_coco_annotation(
annotation_id=3, image_id=2, category_id=3
),
],
3: [
mock_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
]
mock_cock_coco_annotation(
annotation_id=4, image_id=3, category_id=1
),
mock_cock_coco_annotation(
annotation_id=5, image_id=3, category_id=2
),
mock_cock_coco_annotation(
annotation_id=5, image_id=3, category_id=3
),
],
},
DoesNotRaise()
DoesNotRaise(),
), # two coco annotations
]
],
)
def test_group_coco_annotations_by_image_id(
coco_annotations: List[dict],
expected_result: dict,
exception: Exception
coco_annotations: List[dict], expected_result: dict, exception: Exception
) -> None:
with exception:
result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations)
@ -238,187 +191,113 @@ def test_group_coco_annotations_by_image_id(
(1000, 1000),
False,
Detections.empty(),
DoesNotRaise()
DoesNotRaise(),
), # empty image annotations
(
[
mock_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100)
mock_cock_coco_annotation(
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
)
],
(1000, 1000),
False,
Detections(
xyxy=np.array([
[ 0, 0, 100, 100]
], dtype=np.float32),
class_id=np.array([
0
], dtype=int)
xyxy=np.array([[0, 0, 100, 100]], dtype=np.float32),
class_id=np.array([0], dtype=int),
),
DoesNotRaise()
DoesNotRaise(),
), # single image annotations
(
[
mock_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100),
mock_cock_coco_annotation(category_id=0, bbox=(100, 100, 100, 100), area=100 * 100),
mock_cock_coco_annotation(
category_id=0, bbox=(0, 0, 100, 100), area=100 * 100
),
mock_cock_coco_annotation(
category_id=0, bbox=(100, 100, 100, 100), area=100 * 100
),
],
(1000, 1000),
False,
Detections(
xyxy=np.array([
[ 0, 0, 100, 100],
[ 100, 100, 200, 200]
], dtype=np.float32),
class_id=np.array([
0, 0
], dtype=int)
xyxy=np.array(
[[0, 0, 100, 100], [100, 100, 200, 200]], dtype=np.float32
),
class_id=np.array([0, 0], dtype=int),
),
DoesNotRaise()
DoesNotRaise(),
), # two image annotations
]
],
)
def test_coco_annotations_to_detections(
image_annotations: List[dict],
resolution_wh: Tuple[int, int],
with_masks: bool,
expected_result: Detections,
exception: Exception
exception: Exception,
) -> None:
with exception:
result = coco_annotations_to_detections(
image_annotations=image_annotations,
resolution_wh=resolution_wh,
with_masks=with_masks
with_masks=with_masks,
)
assert result == expected_result
@pytest.mark.parametrize(
"coco_categories, target_classes, expected_result, exception",
[
([], [], {}, DoesNotRaise()), # empty coco categories
(
[],
[],
{},
DoesNotRaise()
), # empty coco categories
(
[
{
"id": 0,
"name": "fashion-assistant",
"supercategory": "none"
}
],
[
"fashion-assistant"
],
{
0: 0
},
DoesNotRaise()
[{"id": 0, "name": "fashion-assistant", "supercategory": "none"}],
["fashion-assistant"],
{0: 0},
DoesNotRaise(),
), # single coco category starting from 0
(
[
{
"id": 1,
"name": "fashion-assistant",
"supercategory": "none"
}
],
[
"fashion-assistant"
],
{
1: 0
},
DoesNotRaise()
[{"id": 1, "name": "fashion-assistant", "supercategory": "none"}],
["fashion-assistant"],
{1: 0},
DoesNotRaise(),
), # single coco category starting from 1
(
[
{
"id": 0,
"name": "fashion-assistant",
"supercategory": "none"
},
{
"id": 2,
"name": "hoodie",
"supercategory": "fashion-assistant"
},
{
"id": 1,
"name": "baseball cap",
"supercategory": "fashion-assistant"
}
{"id": 0, "name": "fashion-assistant", "supercategory": "none"},
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
],
[
"fashion-assistant",
"baseball cap",
"hoodie"
],
{
0: 0,
1: 1,
2: 2
},
DoesNotRaise()
["fashion-assistant", "baseball cap", "hoodie"],
{0: 0, 1: 1, 2: 2},
DoesNotRaise(),
), # three coco categories
(
[
{
"id": 2,
"name": "hoodie",
"supercategory": "fashion-assistant"
},
{
"id": 1,
"name": "baseball cap",
"supercategory": "fashion-assistant"
}
{"id": 2, "name": "hoodie", "supercategory": "fashion-assistant"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
],
[
"baseball cap",
"hoodie"
],
{
2: 1,
1: 0
},
DoesNotRaise()
["baseball cap", "hoodie"],
{2: 1, 1: 0},
DoesNotRaise(),
), # two coco categories
(
[
{
"id": 3,
"name": "hoodie",
"supercategory": "fashion-assistant"
},
{
"id": 1,
"name": "baseball cap",
"supercategory": "fashion-assistant"
}
{"id": 3, "name": "hoodie", "supercategory": "fashion-assistant"},
{"id": 1, "name": "baseball cap", "supercategory": "fashion-assistant"},
],
[
"baseball cap",
"hoodie"
],
{
3: 1,
1: 0
},
DoesNotRaise()
["baseball cap", "hoodie"],
{3: 1, 1: 0},
DoesNotRaise(),
), # two coco categories with missing category
]
],
)
def test_build_coco_class_index_mapping(
coco_categories: List[dict],
target_classes: List[str],
expected_result: Dict[int, int],
exception: Exception
exception: Exception,
) -> None:
with exception:
result = build_coco_class_index_mapping(
coco_categories=coco_categories,
target_classes=target_classes
coco_categories=coco_categories, target_classes=target_classes
)
assert result == expected_result

39
test/dataset/formats/test_pascal_voc.py
View File

@ -33,20 +33,25 @@ def are_xml_elements_equal(elem1, elem2):
"xyxy, name, polygon, expected_result, exception",
[
(
[0, 0, 10, 10],
np.array([0, 0, 10, 10]),
"test",
None,
ET.fromstring(
"""<object><name>test</name><bndbox><xmin>0</xmin><ymin>0</ymin><xmax>10</xmax><ymax>10</ymax></bndbox></object>"""
"""<object><name>test</name><bndbox><xmin>1</xmin><ymin>1</ymin>
<xmax>11</xmax><ymax>11</ymax></bndbox></object>"""
),
DoesNotRaise(),
),
(
[0, 0, 10, 10],
np.array([0, 0, 10, 10]),
"test",
[[0, 0], [10, 0], [10, 10], [0, 10]],
np.array([[0, 0], [10, 0], [10, 10], [0, 10]]),
ET.fromstring(
"""<object><name>test</name><bndbox><xmin>0</xmin><ymin>0</ymin><xmax>10</xmax><ymax>10</ymax></bndbox><polygon><x1>0</x1><y1>0</y1><x2>10</x2><y2>0</y2><x3>10</x3><y3>10</y3><x4>0</x4><y4>10</y4></polygon></object>"""
"""<object><name>test</name><bndbox><xmin>1</xmin><ymin>1</ymin>
<xmax>11</xmax><ymax>11</ymax>
</bndbox><polygon><x1>1</x1><y1>1</y1><x2>11</x2>
<y2>1</y2><x3>11</x3><y3>11</y3><x4>1</x4><y4>11</y4>
</polygon></object>"""
),
DoesNotRaise(),
),
@ -69,9 +74,10 @@ def test_object_to_pascal_voc(
[
(
ET.fromstring(
"""<polygon><x1>0</x1><y1>0</y1><x2>10</x2><y2>0</y2><x3>10</x3><y3>10</y3><x4>0</x4><y4>10</y4></polygon>"""
"""<polygon><x1>0</x1><y1>0</y1><x2>10</x2><y2>0</y2><x3>10</x3>
<y3>10</y3><x4>0</x4><y4>10</y4></polygon>"""
),
[[0, 0], [10, 0], [10, 10], [0, 10]],
np.array([[0, 0], [10, 0], [10, 10], [0, 10]]),
DoesNotRaise(),
)
],
@ -83,16 +89,27 @@ def test_parse_polygon_points(
):
with exception:
result = parse_polygon_points(polygon_element)
assert result == expected_result
assert np.array_equal(result, expected_result)
ONE_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>0</xmin><ymin>0</ymin><xmax>10</xmax><ymax>10</ymax></bndbox></object><object><name>test</name><bndbox><xmin>10</xmin><ymin>10</ymin><xmax>20</xmax><ymax>20</ymax></bndbox></object></annotation>"""
ONE_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>1</xmin>
<ymin>1</ymin><xmax>11</xmax><ymax>11</ymax>
</bndbox></object><object><name>test</name><bndbox><xmin>11</xmin><ymin>11</ymin>
<xmax>21</xmax><ymax>21</ymax></bndbox></object></annotation>"""
ONE_CLASS_ONE_BBOX = """<annotation><object><name>test</name><bndbox><xmin>0</xmin><ymin>0</ymin><xmax>10</xmax><ymax>10</ymax></bndbox></object></annotation>"""
ONE_CLASS_ONE_BBOX = """<annotation><object><name>test</name><bndbox>
<xmin>1</xmin><ymin>1</ymin><xmax>11</xmax><ymax>11</ymax></bndbox></object>
</annotation>"""
N_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>0</xmin><ymin>0</ymin><xmax>10</xmax><ymax>10</ymax></bndbox></object><object><name>test</name><bndbox><xmin>20</xmin><ymin>30</ymin><xmax>30</xmax><ymax>40</ymax></bndbox></object><object><name>test2</name><bndbox><xmin>10</xmin><ymin>10</ymin><xmax>20</xmax><ymax>20</ymax></bndbox></object></annotation>"""
N_CLASS_N_BBOX = """<annotation><object><name>test</name><bndbox><xmin>1</xmin>
<ymin>1</ymin><xmax>11</xmax><ymax>11</ymax>
</bndbox></object><object><name>test</name><bndbox>
<xmin>21</xmin><ymin>31</ymin><xmax>31</xmax><ymax>41</ymax></bndbox>
</object><object><name>test2</name><bndbox><xmin>
11</xmin><ymin>11</ymin><xmax>21</xmax><ymax>
21</ymax></bndbox></object></annotation>"""
NO_DETECTIONS = """<annotation></annotation>"""

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

@ -1,12 +1,16 @@
from contextlib import ExitStack as DoesNotRaise
from typing import List, Tuple, Optional
from typing import List, Optional, Tuple
import pytest
import numpy as np
import pytest
from supervision.dataset.formats.yolo import (
_image_name_to_annotation_name,
_with_mask,
object_to_yolo,
yolo_annotations_to_detections,
)
from supervision.detection.core import Detections
from supervision.dataset.formats.yolo import yolo_annotations_to_detections, _with_mask, _image_name_to_annotation_name, \
object_to_yolo
def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.array:
@ -17,223 +21,188 @@ def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.arra
# The result of _mock_simple_mask is a little different from the result produced by cv2.
def _arrays_almost_equal(arr1: np.ndarray, arr2: np.ndarray, threshold: float = 0.99) -> bool:
def _arrays_almost_equal(
arr1: np.ndarray, arr2: np.ndarray, threshold: float = 0.99
) -> bool:
equal_elements = np.equal(arr1, arr2)
proportion_equal = np.mean(equal_elements)
return proportion_equal >= threshold
@pytest.mark.parametrize(
'lines, expected_result, exception',
"lines, expected_result, exception",
[
([], False, DoesNotRaise()), # empty yolo annotation file
(
[],
["0 0.5 0.5 0.2 0.2"],
False,
DoesNotRaise()
), # empty yolo annotation file
(
[
'0 0.5 0.5 0.2 0.2'
],
False,
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with single line with box
(
[
'0 0.50 0.50 0.20 0.20',
'1 0.11 0.47 0.22 0.30'
],
False,
DoesNotRaise()
), # yolo annotation file with two lines with box
(
[
'0 0.5 0.5 0.2 0.2'
],
["0 0.50 0.50 0.20 0.20", "1 0.11 0.47 0.22 0.30"],
False,
DoesNotRaise()
),
DoesNotRaise(),
), # yolo annotation file with two lines with box
(["0 0.5 0.5 0.2 0.2"], False, DoesNotRaise()),
(
[
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
],
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6"],
True,
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with single line with polygon
(
[
'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'
],
["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"],
True,
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with two lines - one box and one polygon
]
],
)
def test_with_mask(lines: List[str], expected_result: Optional[bool], exception: Exception) -> None:
def test_with_mask(
lines: List[str], expected_result: Optional[bool], exception: Exception
) -> None:
with exception:
result = _with_mask(lines=lines)
assert result == expected_result
@pytest.mark.parametrize(
'lines, resolution_wh, with_masks, expected_result, exception',
"lines, resolution_wh, with_masks, expected_result, exception",
[
(
[],
(1000, 1000),
False,
Detections.empty(),
DoesNotRaise()
DoesNotRaise(),
), # empty yolo annotation file
(
[
'0 0.5 0.5 0.2 0.2'
],
["0 0.5 0.5 0.2 0.2"],
(1000, 1000),
False,
Detections(
xyxy=np.array([
[400, 400, 600, 600]
], dtype=np.float32),
class_id=np.array([0], dtype=int)
xyxy=np.array([[400, 400, 600, 600]], dtype=np.float32),
class_id=np.array([0], dtype=int),
),
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with single line with box
(
[
'0 0.50 0.50 0.20 0.20',
'1 0.11 0.47 0.22 0.30'
],
["0 0.50 0.50 0.20 0.20", "1 0.11 0.47 0.22 0.30"],
(1000, 1000),
False,
Detections(
xyxy=np.array([
[400, 400, 600, 600],
[ 0, 320, 220, 620]
], dtype=np.float32),
class_id=np.array([0, 1], dtype=int)
xyxy=np.array(
[[400, 400, 600, 600], [0, 320, 220, 620]], dtype=np.float32
),
class_id=np.array([0, 1], dtype=int),
),
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with two lines with box
(
[
'0 0.5 0.5 0.2 0.2'
],
["0 0.5 0.5 0.2 0.2"],
(1000, 1000),
True,
Detections(
xyxy=np.array([
[400, 400, 600, 600]
], dtype=np.float32),
xyxy=np.array([[400, 400, 600, 600]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array([
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
], dtype=bool)
mask=np.array(
[
_mock_simple_mask(
resolution_wh=(1000, 1000), box=[400, 400, 600, 600]
)
],
dtype=bool,
),
),
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with single line with box in with_masks mode
(
[
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
],
["0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6"],
(1000, 1000),
True,
Detections(
xyxy=np.array([
[400, 400, 600, 600]
], dtype=np.float32),
xyxy=np.array([[400, 400, 600, 600]], dtype=np.float32),
class_id=np.array([0], dtype=int),
mask=np.array([
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
], dtype=bool)
mask=np.array(
[
_mock_simple_mask(
resolution_wh=(1000, 1000), box=[400, 400, 600, 600]
)
],
dtype=bool,
),
),
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with single line with polygon
(
[
'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'
],
["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),
True,
Detections(
xyxy=np.array([
[400, 400, 600, 600],
[ 0, 320, 220, 620]
], dtype=np.float32),
xyxy=np.array(
[[400, 400, 600, 600], [0, 320, 220, 620]], dtype=np.float32
),
class_id=np.array([0, 1], dtype=int),
mask=np.array([
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600]),
_mock_simple_mask(resolution_wh=(1000, 1000), box=[ 0, 320, 220, 620])
], dtype=bool)
mask=np.array(
[
_mock_simple_mask(
resolution_wh=(1000, 1000), box=[400, 400, 600, 600]
),
_mock_simple_mask(
resolution_wh=(1000, 1000), box=[0, 320, 220, 620]
),
],
dtype=bool,
),
),
DoesNotRaise()
), # yolo annotation file with two lines - one box and one polygon in with_masks mode
DoesNotRaise(),
), # 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'
],
["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),
False,
Detections(
xyxy=np.array([
[400, 400, 600, 600],
[ 0, 320, 220, 620]
], dtype=np.float32),
class_id=np.array([0, 1], dtype=int)
xyxy=np.array(
[[400, 400, 600, 600], [0, 320, 220, 620]], dtype=np.float32
),
class_id=np.array([0, 1], dtype=int),
),
DoesNotRaise()
DoesNotRaise(),
), # yolo annotation file with two lines - one box and one polygon
]
],
)
def test_yolo_annotations_to_detections(
lines: List[str],
resolution_wh: Tuple[int, int],
with_masks: bool,
expected_result: Optional[Detections],
exception: Exception
exception: Exception,
) -> None:
with exception:
result = yolo_annotations_to_detections(
lines=lines,
resolution_wh=resolution_wh,
with_masks=with_masks)
lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
)
assert np.array_equal(result.xyxy, expected_result.xyxy)
assert np.array_equal(result.class_id, expected_result.class_id)
assert (result.mask is None and expected_result.mask is None) or _arrays_almost_equal(result.mask, expected_result.mask)
assert (
result.mask is None and expected_result.mask is None
) or _arrays_almost_equal(result.mask, expected_result.mask)
@pytest.mark.parametrize(
'image_name, expected_result, exception',
"image_name, expected_result, exception",
[
("image.png", "image.txt", DoesNotRaise()), # simple png image
("image.jpeg", "image.txt", DoesNotRaise()), # simple jpeg image
("image.jpg", "image.txt", DoesNotRaise()), # simple jpg image
(
'image.png',
'image.txt',
DoesNotRaise()
), # simple png image
(
'image.jpeg',
'image.txt',
DoesNotRaise()
), # simple jpeg image
(
'image.jpg',
'image.txt',
DoesNotRaise()
), # simple jpg image
(
'image.000.jpg',
'image.000.txt',
DoesNotRaise()
"image.000.jpg",
"image.000.txt",
DoesNotRaise(),
), # jpg image with multiple dots in name
]
],
)
def test_image_name_to_annotation_name(
image_name: str,
expected_result: Optional[str],
exception: Exception
image_name: str, expected_result: Optional[str], exception: Exception
) -> None:
with exception:
result = _image_name_to_annotation_name(image_name=image_name)
@ -241,64 +210,70 @@ def test_image_name_to_annotation_name(
@pytest.mark.parametrize(
'xyxy, class_id, image_shape, polygon, expected_result, exception',
"xyxy, class_id, image_shape, polygon, expected_result, exception",
[
(
np.array([100, 100, 200, 200], dtype=np.float32),
1,
(1000, 1000, 3),
None,
'1 0.15000 0.15000 0.10000 0.10000',
DoesNotRaise()
"1 0.15000 0.15000 0.10000 0.10000",
DoesNotRaise(),
), # square bounding box on square image
(
np.array([100, 100, 200, 200], dtype=np.float32),
1,
(800, 1000, 3),
None,
'1 0.15000 0.18750 0.10000 0.12500',
DoesNotRaise()
"1 0.15000 0.18750 0.10000 0.12500",
DoesNotRaise(),
), # square bounding box on horizontal image
(
np.array([100, 100, 200, 200], dtype=np.float32),
1,
(1000, 800, 3),
None,
'1 0.18750 0.15000 0.12500 0.10000',
DoesNotRaise()
"1 0.18750 0.15000 0.12500 0.10000",
DoesNotRaise(),
), # square bounding box on vertical image
(
np.array([100, 200, 200, 400], dtype=np.float32),
1,
(1000, 1000, 3),
None,
'1 0.15000 0.30000 0.10000 0.20000',
DoesNotRaise()
"1 0.15000 0.30000 0.10000 0.20000",
DoesNotRaise(),
), # horizontal bounding box on square image
(
np.array([200, 100, 400, 200], dtype=np.float32),
1,
(1000, 1000, 3),
None,
'1 0.30000 0.15000 0.20000 0.10000',
DoesNotRaise()
"1 0.30000 0.15000 0.20000 0.10000",
DoesNotRaise(),
), # vertical bounding box on square image
(
np.array([100, 100, 200, 200], dtype=np.float32),
1,
(1000, 1000, 3),
np.array([
[100, 100],
[200, 100],
[200, 200],
[100, 100]
], dtype=np.float32),
'1 0.10000 0.10000 0.20000 0.10000 0.20000 0.20000 0.10000 0.10000',
DoesNotRaise()
np.array(
[[100, 100], [200, 100], [200, 200], [100, 100]], dtype=np.float32
),
"1 0.10000 0.10000 0.20000 0.10000 0.20000 0.20000 0.10000 0.10000",
DoesNotRaise(),
), # square mask on square image
]
],
)
def test_object_to_yolo(xyxy: np.ndarray, class_id: int, image_shape: Tuple[int, int, int], polygon: Optional[np.ndarray], expected_result: Optional[str], exception: Exception) -> None:
def test_object_to_yolo(
xyxy: np.ndarray,
class_id: int,
image_shape: Tuple[int, int, int],
polygon: Optional[np.ndarray],
expected_result: Optional[str],
exception: Exception,
) -> None:
with exception:
result = object_to_yolo(xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon)
result = object_to_yolo(
xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon
)
assert result == expected_result

180
test/dataset/test_core.py
View File

@ -1,13 +1,11 @@
from contextlib import ExitStack as DoesNotRaise
from test.utils import mock_detections
from typing import List, Optional
import numpy as np
import pytest
from supervision import DetectionDataset
from contextlib import ExitStack as DoesNotRaise
import numpy as np
from test.utils import mock_detections
@pytest.mark.parametrize(
@ -16,159 +14,181 @@ from test.utils import mock_detections
(
[],
DetectionDataset(classes=[], images={}, annotations={}),
DoesNotRaise()
DoesNotRaise(),
), # empty dataset list
(
[
DetectionDataset(classes=[], images={}, annotations={})
],
[DetectionDataset(classes=[], images={}, annotations={})],
DetectionDataset(classes=[], images={}, annotations={}),
DoesNotRaise()
DoesNotRaise(),
), # single empty dataset
(
[
DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
DetectionDataset(classes=['dog', 'person'], images={}, annotations={})
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
],
DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
DoesNotRaise()
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
DoesNotRaise(),
), # two datasets; no images and annotations, the same classes
(
[
DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
DetectionDataset(classes=['cat'], images={}, annotations={})
DetectionDataset(classes=["dog", "person"], images={}, annotations={}),
DetectionDataset(classes=["cat"], images={}, annotations={}),
],
DetectionDataset(classes=['cat', 'dog', 'person'], images={}, annotations={}),
DoesNotRaise()
DetectionDataset(
classes=["cat", "dog", "person"], images={}, annotations={}
),
DoesNotRaise(),
), # two datasets; no images and annotations, different classes
(
[
DetectionDataset(
classes=['dog', 'person'],
classes=["dog", "person"],
images={
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
}
"image-1.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[0]
),
"image-2.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[1]
),
},
),
DetectionDataset(classes=[], images={}, annotations={}),
],
DetectionDataset(
classes=['dog', 'person'],
classes=["dog", "person"],
images={
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
}
"image-1.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
"image-2.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
},
),
DoesNotRaise()
DoesNotRaise(),
), # two datasets; images and annotations, the same classes
(
[
DetectionDataset(
classes=['dog', 'person'],
classes=["dog", "person"],
images={
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
}
"image-1.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[0]
),
"image-2.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[1]
),
},
),
DetectionDataset(classes=['cat'], images={}, annotations={}),
DetectionDataset(classes=["cat"], images={}, annotations={}),
],
DetectionDataset(
classes=['cat', 'dog', 'person'],
classes=["cat", "dog", "person"],
images={
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
}
"image-1.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
"image-2.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
},
),
DoesNotRaise()
DoesNotRaise(),
), # two datasets; images and annotations, different classes
(
[
DetectionDataset(
classes=['dog', 'person'],
classes=["dog", "person"],
images={
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
}
"image-1.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[0]
),
"image-2.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[1]
),
},
),
DetectionDataset(
classes=['cat'],
classes=["cat"],
images={
'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-3.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
}
"image-3.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[0]
),
},
),
],
DetectionDataset(
classes=['cat', 'dog', 'person'],
classes=["cat", "dog", "person"],
images={
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-3.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
}
"image-1.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
"image-2.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
"image-3.png": mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
},
),
DoesNotRaise()
DoesNotRaise(),
), # two datasets; images and annotations, different classes
(
[
DetectionDataset(
classes=['dog', 'person'],
classes=["dog", "person"],
images={
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-1.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
}
"image-1.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[0]
),
"image-2.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[1]
),
},
),
DetectionDataset(
classes=['dog', 'person'],
classes=["dog", "person"],
images={
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
"image-2.png": np.zeros((100, 100, 3), dtype=np.uint8),
"image-3.png": np.zeros((100, 100, 3), dtype=np.uint8),
},
annotations={
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
}
"image-2.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[0]
),
"image-3.png": mock_detections(
xyxy=[[0, 0, 10, 10]], class_id=[1]
),
},
),
],
None,
pytest.raises(ValueError)
)
]
pytest.raises(ValueError),
),
],
)
def test_dataset_merge(
dataset_list: List[DetectionDataset],
expected_result: Optional[DetectionDataset],
exception: Exception
exception: Exception,
) -> None:
with exception:
result = DetectionDataset.merge(dataset_list=dataset_list)

178
test/dataset/test_utils.py
View File

@ -1,34 +1,31 @@
from contextlib import ExitStack as DoesNotRaise
from typing import List, TypeVar, Optional, Tuple, Dict
from test.utils import mock_detections
from typing import Dict, List, Optional, Tuple, TypeVar
import pytest
from supervision import Detections
from supervision.dataset.utils import train_test_split, merge_class_lists, build_class_index_mapping, \
map_detections_class_id
from test.utils import mock_detections
from supervision.dataset.utils import (
build_class_index_mapping,
map_detections_class_id,
merge_class_lists,
train_test_split,
)
T = TypeVar("T")
@pytest.mark.parametrize(
'data, train_ratio, random_state, shuffle, expected_result, exception',
"data, train_ratio, random_state, shuffle, expected_result, exception",
[
(
[],
0.5,
None,
False,
([], []),
DoesNotRaise()
), # empty data
([], 0.5, None, False, ([], []), DoesNotRaise()), # empty data
(
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
0.5,
None,
False,
([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]),
DoesNotRaise()
DoesNotRaise(),
), # data with 10 numbers and 50% train split
(
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
@ -36,7 +33,7 @@ T = TypeVar("T")
None,
False,
([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], []),
DoesNotRaise()
DoesNotRaise(),
), # data with 10 numbers and 100% train split
(
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
@ -44,15 +41,15 @@ T = TypeVar("T")
None,
False,
([], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
DoesNotRaise()
DoesNotRaise(),
), # data with 10 numbers and 0% train split
(
['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'],
["a", "b", "c", "d", "e", "f", "g", "h", "i", "j"],
0.5,
None,
False,
(['a', 'b', 'c', 'd', 'e'], ['f', 'g', 'h', 'i', 'j']),
DoesNotRaise()
(["a", "b", "c", "d", "e"], ["f", "g", "h", "i", "j"]),
DoesNotRaise(),
), # data with 10 chars and 50% train split
(
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
@ -60,7 +57,7 @@ T = TypeVar("T")
23,
True,
([7, 8, 5, 6, 3], [2, 9, 0, 1, 4]),
DoesNotRaise()
DoesNotRaise(),
), # data with 10 numbers and 50% train split with 23 random seed
(
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
@ -68,9 +65,9 @@ T = TypeVar("T")
32,
True,
([4, 6, 0, 8, 9], [5, 7, 2, 3, 1]),
DoesNotRaise()
DoesNotRaise(),
), # data with 10 numbers and 50% train split with 23 random seed
]
],
)
def test_train_test_split(
data: List[T],
@ -78,155 +75,140 @@ def test_train_test_split(
random_state: int,
shuffle: bool,
expected_result: Optional[Tuple[List[T], List[T]]],
exception: Exception
exception: Exception,
) -> None:
with exception:
result = train_test_split(data=data, train_ratio=train_ratio, random_state=random_state, shuffle=shuffle)
result = train_test_split(
data=data,
train_ratio=train_ratio,
random_state=random_state,
shuffle=shuffle,
)
assert result == expected_result
@pytest.mark.parametrize(
'class_lists, expected_result, exception',
"class_lists, expected_result, exception",
[
([], [], DoesNotRaise()), # empty class lists
(
[],
[],
DoesNotRaise()
), # empty class lists
(
[
['dog', 'person']
],
['dog', 'person'],
DoesNotRaise()
[["dog", "person"]],
["dog", "person"],
DoesNotRaise(),
), # single class list; already alphabetically sorted
(
[
['person', 'dog']
],
['dog', 'person'],
DoesNotRaise()
[["person", "dog"]],
["dog", "person"],
DoesNotRaise(),
), # single class list; not alphabetically sorted
(
[
['dog', 'person'],
['dog', 'person']
],
['dog', 'person'],
DoesNotRaise()
[["dog", "person"], ["dog", "person"]],
["dog", "person"],
DoesNotRaise(),
), # two class lists; the same classes; already alphabetically sorted
(
[
['dog', 'person'],
['cat']
],
['cat', 'dog', 'person'],
DoesNotRaise()
[["dog", "person"], ["cat"]],
["cat", "dog", "person"],
DoesNotRaise(),
), # two class lists; different classes; already alphabetically sorted
]
],
)
def test_merge_class_maps(class_lists: List[List[str]], expected_result: List[str], exception: Exception) -> None:
def test_merge_class_maps(
class_lists: List[List[str]], expected_result: List[str], exception: Exception
) -> None:
with exception:
result = merge_class_lists(class_lists=class_lists)
assert result == expected_result
@pytest.mark.parametrize(
'source_classes, target_classes, expected_result, exception',
"source_classes, target_classes, expected_result, exception",
[
([], [], {}, DoesNotRaise()), # empty class lists
([], ["dog", "person"], {}, DoesNotRaise()), # empty source class list
(
[],
[],
{},
DoesNotRaise()
), # empty class lists
(
[],
['dog', 'person'],
{},
DoesNotRaise()
), # empty source class list
(
['dog', 'person'],
["dog", "person"],
[],
None,
pytest.raises(ValueError)
pytest.raises(ValueError),
), # empty target class list
(
['dog', 'person'],
['dog', 'person'],
["dog", "person"],
["dog", "person"],
{0: 0, 1: 1},
DoesNotRaise()
DoesNotRaise(),
), # same class lists
(
['dog', 'person'],
['person', 'dog'],
["dog", "person"],
["person", "dog"],
{0: 1, 1: 0},
DoesNotRaise()
DoesNotRaise(),
), # same class lists but not alphabetically sorted
(
['dog', 'person'],
['cat', 'dog', 'person'],
["dog", "person"],
["cat", "dog", "person"],
{0: 1, 1: 2},
DoesNotRaise()
DoesNotRaise(),
), # source class list is a subset of target class list
(
['dog', 'person'],
['cat', 'dog'],
["dog", "person"],
["cat", "dog"],
None,
pytest.raises(ValueError)
pytest.raises(ValueError),
), # source class list is not a subset of target class list
]
],
)
def test_build_class_index_mapping(
source_classes: List[str],
target_classes: List[str],
expected_result: Optional[Dict[int, int]],
exception: Exception
exception: Exception,
) -> None:
with exception:
result = build_class_index_mapping(source_classes=source_classes, target_classes=target_classes)
result = build_class_index_mapping(
source_classes=source_classes, target_classes=target_classes
)
assert result == expected_result
@pytest.mark.parametrize(
'source_to_target_mapping, detections, expected_result, exception',
"source_to_target_mapping, detections, expected_result, exception",
[
(
{},
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
None,
pytest.raises(ValueError)
pytest.raises(ValueError),
), # empty mapping
(
{0: 1},
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
DoesNotRaise()
DoesNotRaise(),
), # single mapping
(
{0: 1, 1: 2},
Detections.empty(),
Detections.empty(),
DoesNotRaise()
{0: 1, 1: 2},
Detections.empty(),
Detections.empty(),
DoesNotRaise(),
), # empty detections
(
{0: 1, 1: 2},
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
DoesNotRaise()
DoesNotRaise(),
), # multiple mappings
(
{0: 1, 1: 2},
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[0, 1]),
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[1, 2]),
DoesNotRaise()
DoesNotRaise(),
), # multiple mappings
(
{0: 1, 1: 2},
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
None,
pytest.raises(ValueError)
pytest.raises(ValueError),
), # class_id not in mapping
(
{0: 1, 1: 2},
@ -234,14 +216,16 @@ def test_build_class_index_mapping(
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1], confidence=[0.5]),
DoesNotRaise(),
), # confidence is not None
]
],
)
def test_map_detections_class_id(
source_to_target_mapping: Dict[int, int],
detections: Detections,
expected_result: Optional[Detections],
exception: Exception
exception: Exception,
) -> None:
with exception:
result = map_detections_class_id(source_to_target_mapping=source_to_target_mapping, detections=detections)
result = map_detections_class_id(
source_to_target_mapping=source_to_target_mapping, detections=detections
)
assert result == expected_result

215
test/detection/test_core.py
View File

@ -1,46 +1,44 @@
from contextlib import ExitStack as DoesNotRaise
from test.utils import mock_detections
from typing import List, Optional, Union
import numpy as np
import pytest
from supervision import Detections
from typing import Optional, Union, List
import numpy as np
from test.utils import mock_detections
PREDICTIONS = np.array([
[ 2254, 906, 2447, 1353, 0.90538, 0],
[ 2049, 1133, 2226, 1371, 0.59002, 56],
[ 727, 1224, 838, 1601, 0.51119, 39],
[ 808, 1214, 910, 1564, 0.45287, 39],
[ 6, 52, 1131, 2133, 0.45057, 72],
[ 299, 1225, 512, 1663, 0.45029, 39],
[ 529, 874, 645, 945, 0.31101, 39],
[ 8, 47, 1935, 2135, 0.28192, 72],
[ 2265, 813, 2328, 901, 0.2714, 62]
], dtype=np.float32)
PREDICTIONS = np.array(
[
[2254, 906, 2447, 1353, 0.90538, 0],
[2049, 1133, 2226, 1371, 0.59002, 56],
[727, 1224, 838, 1601, 0.51119, 39],
[808, 1214, 910, 1564, 0.45287, 39],
[6, 52, 1131, 2133, 0.45057, 72],
[299, 1225, 512, 1663, 0.45029, 39],
[529, 874, 645, 945, 0.31101, 39],
[8, 47, 1935, 2135, 0.28192, 72],
[2265, 813, 2328, 901, 0.2714, 62],
],
dtype=np.float32,
)
DETECTIONS = Detections(
xyxy=PREDICTIONS[:, :4],
confidence=PREDICTIONS[:, 4],
class_id=PREDICTIONS[:, 5].astype(int)
class_id=PREDICTIONS[:, 5].astype(int),
)
@pytest.mark.parametrize(
'detections, index, expected_result, exception',
"detections, index, expected_result, exception",
[
(
DETECTIONS,
DETECTIONS.class_id == 0,
mock_detections(
xyxy=[[2254, 906, 2447, 1353]],
confidence=[0.90538],
class_id=[0]
xyxy=[[2254, 906, 2447, 1353]], confidence=[0.90538], class_id=[0]
),
DoesNotRaise()
DoesNotRaise(),
), # take only detections with class_id = 0
(
DETECTIONS,
@ -49,109 +47,90 @@ DETECTIONS = Detections(
xyxy=[
[2254, 906, 2447, 1353],
[2049, 1133, 2226, 1371],
[727, 1224, 838, 1601]
[727, 1224, 838, 1601],
],
confidence=[0.90538, 0.59002, 0.51119],
class_id=[0, 56, 39]
class_id=[0, 56, 39],
),
DoesNotRaise()
DoesNotRaise(),
), # take only detections with confidence > 0.5
(
DETECTIONS,
np.array([True, True, True, True, True, True, True, True, True], dtype=bool),
np.array(
[True, True, True, True, True, True, True, True, True], dtype=bool
),
DETECTIONS,
DoesNotRaise()
DoesNotRaise(),
), # take all detections
(
DETECTIONS,
np.array([False, False, False, False, False, False, False, False, False], dtype=bool),
np.array(
[False, False, False, False, False, False, False, False, False],
dtype=bool,
),
Detections(
xyxy=np.empty((0, 4), dtype=np.float32),
confidence=np.array([], dtype=np.float32),
class_id=np.array([], dtype=int)
class_id=np.array([], dtype=int),
),
DoesNotRaise()
DoesNotRaise(),
), # take no detections
(
DETECTIONS,
[0, 2],
mock_detections(
xyxy=[
[2254, 906, 2447, 1353],
[727, 1224, 838, 1601]
],
xyxy=[[2254, 906, 2447, 1353], [727, 1224, 838, 1601]],
confidence=[0.90538, 0.51119],
class_id=[0, 39]
class_id=[0, 39],
),
DoesNotRaise()
DoesNotRaise(),
), # take only first and third detection using List[int] index
(
DETECTIONS,
np.array([0, 2]),
mock_detections(
xyxy=[
[2254, 906, 2447, 1353],
[727, 1224, 838, 1601]
],
xyxy=[[2254, 906, 2447, 1353], [727, 1224, 838, 1601]],
confidence=[0.90538, 0.51119],
class_id=[0, 39]
class_id=[0, 39],
),
DoesNotRaise()
DoesNotRaise(),
), # take only first and third detection using np.ndarray index
(
DETECTIONS,
0,
mock_detections(
xyxy=[[2254, 906, 2447, 1353]],
confidence=[0.90538],
class_id=[0]
xyxy=[[2254, 906, 2447, 1353]], confidence=[0.90538], class_id=[0]
),
DoesNotRaise()
DoesNotRaise(),
), # take only first detection by index
(
DETECTIONS,
slice(1, 3),
mock_detections(
xyxy=[
[2049, 1133, 2226, 1371],
[727, 1224, 838, 1601]
],
xyxy=[[2049, 1133, 2226, 1371], [727, 1224, 838, 1601]],
confidence=[0.59002, 0.51119],
class_id=[56, 39]
class_id=[56, 39],
),
DoesNotRaise()
DoesNotRaise(),
), # take only first detection by index slice (1, 3)
(DETECTIONS, 10, None, pytest.raises(IndexError)), # index out of range
(DETECTIONS, [0, 2, 10], None, pytest.raises(IndexError)), # index out of range
(DETECTIONS, np.array([0, 2, 10]), None, pytest.raises(IndexError)),
(
DETECTIONS,
10,
np.array(
[True, True, True, True, True, True, True, True, True, True, True]
),
None,
pytest.raises(IndexError)
), # index out of range
(
DETECTIONS,
[0, 2, 10],
None,
pytest.raises(IndexError)
), # index out of range
(
DETECTIONS,
np.array([0, 2, 10]),
None,
pytest.raises(IndexError)
pytest.raises(IndexError),
),
(
DETECTIONS,
np.array([True, True, True, True, True, True, True, True, True, True, True]),
None,
pytest.raises(IndexError)
)
]
],
)
def test_getitem(
detections: Detections,
index: Union[int, slice, List[int], np.ndarray],
expected_result: Optional[Detections],
exception: Exception
detections: Detections,
index: Union[int, slice, List[int], np.ndarray],
expected_result: Optional[Detections],
exception: Exception,
) -> None:
with exception:
result = detections[index]
@ -159,86 +138,54 @@ def test_getitem(
@pytest.mark.parametrize(
'detections_list, expected_result, exception',
"detections_list, expected_result, exception",
[
([], Detections.empty(), DoesNotRaise()), # empty detections list
(
[],
[Detections.empty()],
Detections.empty(),
DoesNotRaise()
), # empty detections list
(
[
Detections.empty()
],
Detections.empty(),
DoesNotRaise()
DoesNotRaise(),
), # single empty detections
(
[
mock_detections(xyxy=[[10, 10, 20, 20]])
],
[mock_detections(xyxy=[[10, 10, 20, 20]])],
mock_detections(xyxy=[[10, 10, 20, 20]]),
DoesNotRaise()
DoesNotRaise(),
), # single detection with xyxy field
(
[
mock_detections(xyxy=[[10, 10, 20, 20]]),
Detections.empty()
],
[mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()],
mock_detections(xyxy=[[10, 10, 20, 20]]),
DoesNotRaise()
DoesNotRaise(),
), # single detection with xyxy field + empty detection
(
[
mock_detections(xyxy=[[10, 10, 20, 20]]),
mock_detections(xyxy=[[20, 20, 30, 30]])
mock_detections(xyxy=[[20, 20, 30, 30]]),
],
mock_detections(
xyxy=[
[10, 10, 20, 20],
[20, 20, 30, 30]
]),
DoesNotRaise()
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
DoesNotRaise(),
), # two detections with xyxy field
(
[
mock_detections(
xyxy=[[10, 10, 20, 20]],
class_id=[0]),
mock_detections(
xyxy=[[20, 20, 30, 30]])
mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]),
mock_detections(xyxy=[[20, 20, 30, 30]]),
],
mock_detections(
xyxy=[
[10, 10, 20, 20],
[20, 20, 30, 30]
]),
DoesNotRaise()
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
DoesNotRaise(),
), # detection with xyxy, class_id fields + detection with xyxy field
(
(
[
mock_detections(
xyxy=[[10, 10, 20, 20]],
class_id=[0]),
mock_detections(
xyxy=[[20, 20, 30, 30]],
class_id=[1]),
mock_detections(xyxy=[[10, 10, 20, 20]], class_id=[0]),
mock_detections(xyxy=[[20, 20, 30, 30]], class_id=[1]),
],
mock_detections(
xyxy=[
[10, 10, 20, 20],
[20, 20, 30, 30]
],
class_id=[0, 1]
),
DoesNotRaise()
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]),
DoesNotRaise(),
), # two detections with xyxy, class_id fields
]
],
)
def test_merge(
detections_list: List[Detections],
expected_result: Optional[Detections],
exception: Exception
detections_list: List[Detections],
expected_result: Optional[Detections],
exception: Exception,
) -> None:
with exception:
result = Detections.merge(detections_list=detections_list)

302
test/detection/test_utils.py
View File

@ -1,12 +1,15 @@
from contextlib import ExitStack as DoesNotRaise
from typing import Optional, Tuple, List
import pytest
from typing import List, Optional, Tuple
import numpy as np
import pytest
from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area, \
process_roboflow_result
from supervision.detection.utils import (
clip_boxes,
filter_polygons_by_area,
non_max_suppression,
process_roboflow_result,
)
@pytest.mark.parametrize(
@ -16,116 +19,101 @@ from supervision.detection.utils import non_max_suppression, clip_boxes, filter_
np.empty(shape=(0, 5)),
0.5,
np.array([]),
DoesNotRaise()
DoesNotRaise(),
), # single box with no category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8]
]),
np.array([[10.0, 10.0, 40.0, 40.0, 0.8]]),
0.5,
np.array([
True
]),
DoesNotRaise()
np.array([True]),
DoesNotRaise(),
), # single box with no category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8, 0]
]),
np.array([[10.0, 10.0, 40.0, 40.0, 0.8, 0]]),
0.5,
np.array([
True
]),
DoesNotRaise()
np.array([True]),
DoesNotRaise(),
), # single box with category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8],
[15.0, 15.0, 40.0, 40.0, 0.9],
]),
np.array(
[
[10.0, 10.0, 40.0, 40.0, 0.8],
[15.0, 15.0, 40.0, 40.0, 0.9],
]
),
0.5,
np.array([
False,
True
]),
DoesNotRaise()
np.array([False, True]),
DoesNotRaise(),
), # two boxes with no category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
]),
np.array(
[
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
]
),
0.5,
np.array([
True,
True
]),
DoesNotRaise()
np.array([True, True]),
DoesNotRaise(),
), # two boxes with different category
(
np.array([
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
]),
np.array(
[
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
]
),
0.5,
np.array([
False,
True
]),
DoesNotRaise()
np.array([False, True]),
DoesNotRaise(),
), # two boxes with same category
(
np.array([
[0.0, 0.0, 30.0, 40.0, 0.8],
[5.0, 5.0, 35.0, 45.0, 0.9],
[10.0, 10.0, 40.0, 50.0, 0.85],
]),
np.array(
[
[0.0, 0.0, 30.0, 40.0, 0.8],
[5.0, 5.0, 35.0, 45.0, 0.9],
[10.0, 10.0, 40.0, 50.0, 0.85],
]
),
0.5,
np.array([
False,
True,
False
]),
DoesNotRaise()
np.array([False, True, False]),
DoesNotRaise(),
), # three boxes with no category
(
np.array([
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
]),
np.array(
[
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
]
),
0.5,
np.array([
True,
True,
True
]),
DoesNotRaise()
np.array([True, True, True]),
DoesNotRaise(),
), # three boxes with same category
(
np.array([
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
]),
np.array(
[
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
]
),
0.5,
np.array([
False,
True,
True
]),
DoesNotRaise()
np.array([False, True, True]),
DoesNotRaise(),
), # three boxes with different category
]
],
)
def test_non_max_suppression(
predictions: np.ndarray,
iou_threshold: float,
expected_result: Optional[np.ndarray],
exception: Exception
predictions: np.ndarray,
iou_threshold: float,
expected_result: Optional[np.ndarray],
exception: Exception,
) -> None:
with exception:
result = non_max_suppression(predictions=predictions, iou_threshold=iou_threshold)
result = non_max_suppression(
predictions=predictions, iou_threshold=iou_threshold
)
assert np.array_equal(result, expected_result)
@ -138,53 +126,37 @@ def test_non_max_suppression(
np.empty(shape=(0, 4)),
),
(
np.array([
[1.0, 1.0, 1279.0, 719.0]
]),
np.array([[1.0, 1.0, 1279.0, 719.0]]),
(1280, 720),
np.array([
[1.0, 1.0, 1279.0, 719.0]
]),
np.array([[1.0, 1.0, 1279.0, 719.0]]),
),
(
np.array([
[-1.0, 1.0, 1279.0, 719.0]
]),
np.array([[-1.0, 1.0, 1279.0, 719.0]]),
(1280, 720),
np.array([
[0.0, 1.0, 1279.0, 719.0]
]),
np.array([[0.0, 1.0, 1279.0, 719.0]]),
),
(
np.array([
[1.0, -1.0, 1279.0, 719.0]
]),
np.array([[1.0, -1.0, 1279.0, 719.0]]),
(1280, 720),
np.array([
[1.0, 0.0, 1279.0, 719.0]
]),
np.array([[1.0, 0.0, 1279.0, 719.0]]),
),
(
np.array([
[1.0, 1.0, 1281.0, 719.0]
]),
np.array([[1.0, 1.0, 1281.0, 719.0]]),
(1280, 720),
np.array([
[1.0, 1.0, 1280.0, 719.0]
]),
np.array([[1.0, 1.0, 1280.0, 719.0]]),
),
(
np.array([
[1.0, 1.0, 1279.0, 721.0]
]),
np.array([[1.0, 1.0, 1279.0, 721.0]]),
(1280, 720),
np.array([
[1.0, 1.0, 1279.0, 720.0]
]),
np.array([[1.0, 1.0, 1279.0, 720.0]]),
),
]
],
)
def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None:
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)
@ -197,83 +169,87 @@ def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int]
None,
None,
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
DoesNotRaise()
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()
DoesNotRaise(),
), # single polygon with min_area constraint
(
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
None,
50,
[],
DoesNotRaise()
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]])
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
],
200,
None,
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
DoesNotRaise()
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]])
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
],
None,
200,
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
DoesNotRaise()
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]])
np.array([[0, 0], [0, 20], [20, 20], [20, 0]]),
],
200,
200,
[],
DoesNotRaise()
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]])
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
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]])
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
]
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
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)
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)
@ -283,18 +259,10 @@ def test_filter_polygons_by_area(
"roboflow_result, class_list, expected_result, exception",
[
(
{
"predictions": [],
"image": {"width": 1000, "height": 1000}
},
{"predictions": [], "image": {"width": 1000, "height": 1000}},
["person", "car", "truck"],
(
np.empty((0, 4)),
np.empty(0),
np.empty(0),
None
),
DoesNotRaise()
(np.empty((0, 4)), np.empty(0), np.empty(0), None),
DoesNotRaise(),
), # empty result
(
{
@ -305,33 +273,35 @@ def test_filter_polygons_by_area(
"width": 50.0,
"height": 50.0,
"confidence": 0.9,
"class": "person"
"class": "person",
}
],
"image": {"width": 1000, "height": 1000}
"image": {"width": 1000, "height": 1000},
},
["person", "car", "truck"],
(
np.array([
[175.0, 275.0, 225.0, 325.0]
]),
np.array([[175.0, 275.0, 225.0, 325.0]]),
np.array([0.9]),
np.array([0]),
None
),
DoesNotRaise()
None,
),
DoesNotRaise(),
), # single bounding box
]
],
)
def test_process_roboflow_result(
roboflow_result: dict,
class_list: List[str],
expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
exception: Exception
roboflow_result: dict,
class_list: List[str],
expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
exception: Exception,
) -> None:
with exception:
result = process_roboflow_result(roboflow_result=roboflow_result, class_list=class_list)
result = process_roboflow_result(
roboflow_result=roboflow_result, class_list=class_list
)
assert np.array_equal(result[0], expected_result[0])
assert np.array_equal(result[1], expected_result[1])
assert np.array_equal(result[2], expected_result[2])
assert (result[3] is None and expected_result[3] is None) or (np.array_equal(result[3], expected_result[3]))
assert (result[3] is None and expected_result[3] is None) or (
np.array_equal(result[3], expected_result[3])
)

2
test/geometry/test_dataclasses.py
View File

@ -1,6 +1,6 @@
import pytest
from supervision.geometry.core import Vector, Point
from supervision.geometry.core import Point, Vector
@pytest.mark.parametrize(

61
test/metrics/test_detection.py
View File

@ -1,12 +1,16 @@
from contextlib import ExitStack as DoesNotRaise
from test.utils import assert_almost_equal, mock_detections
from typing import Optional, Union
import numpy as np
import pytest
from supervision.detection.core import Detections
from supervision.metrics.detection import ConfusionMatrix, detections_to_tensor, MeanAveragePrecision
from test.utils import mock_detections, assert_almost_equal
from supervision.metrics.detection import (
ConfusionMatrix,
MeanAveragePrecision,
detections_to_tensor,
)
CLASSES = np.arange(80)
NUM_CLASSES = len(CLASSES)
@ -147,15 +151,25 @@ BAD_CONF_MATRIX = worsen_ideal_conf_matrix(
DoesNotRaise(),
), # single detection; with confidence
(
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]], class_id=[0, 1], confidence=[0.5, 0.2]),
mock_detections(
xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]],
class_id=[0, 1],
confidence=[0.5, 0.2],
),
False,
np.array([[0, 0, 10, 10, 0], [0, 0, 20, 20, 1]], dtype=np.float32),
DoesNotRaise(),
), # multiple detections; no confidence
(
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]], class_id=[0, 1], confidence=[0.5, 0.2]),
mock_detections(
xyxy=[[0, 0, 10, 10], [0, 0, 20, 20]],
class_id=[0, 1],
confidence=[0.5, 0.2],
),
True,
np.array([[0, 0, 10, 10, 0, 0.5], [0, 0, 20, 20, 1, 0.2]], dtype=np.float32),
np.array(
[[0, 0, 10, 10, 0, 0.5], [0, 0, 20, 20, 1, 0.2]], dtype=np.float32
),
DoesNotRaise(),
), # multiple detections; with confidence
],
@ -164,18 +178,18 @@ def test_detections_to_tensor(
detections: Detections,
with_confidence: bool,
expected_result: Optional[np.ndarray],
exception: Exception
exception: Exception,
):
with exception:
result = detections_to_tensor(
detections=detections,
with_confidence=with_confidence
detections=detections, with_confidence=with_confidence
)
assert np.array_equal(result, expected_result)
@pytest.mark.parametrize(
"predictions, targets, classes, conf_threshold, iou_threshold, expected_result, exception",
"predictions, targets, classes, conf_threshold, iou_threshold, expected_result,"
" exception",
[
(
DETECTION_TENSORS,
@ -346,7 +360,8 @@ def test_from_tensors(
@pytest.mark.parametrize(
"predictions, targets, num_classes, conf_threshold, iou_threshold, expected_result, exception",
"predictions, targets, num_classes, conf_threshold, iou_threshold, expected_result,"
" exception",
[
(
DETECTION_TENSORS[0],
@ -403,40 +418,42 @@ def test_drop_extra_matches(
@pytest.mark.parametrize(
'recall, precision, expected_result, exception',
"recall, precision, expected_result, exception",
[
(
np.array([1.0]),
np.array([1.0]),
1.0,
DoesNotRaise()
DoesNotRaise(),
), # perfect recall and precision
(
np.array([0.0]),
np.array([0.0]),
0.0,
DoesNotRaise()
DoesNotRaise(),
), # no recall and precision
(
np.array([0.0, 0.2, 0.2, 0.8, 0.8, 1.0]),
np.array([0.7, 0.8, 0.4, 0.5, 0.1, 0.2]),
0.5,
DoesNotRaise()
DoesNotRaise(),
),
(
np.array([0.0, 0.5, 0.5, 1.0]),
np.array([0.75, 0.75, 0.75, 0.75]),
0.75,
DoesNotRaise()
)
]
DoesNotRaise(),
),
],
)
def test_compute_average_precision(
recall: np.ndarray,
precision: np.ndarray,
expected_result: float,
exception: Exception
recall: np.ndarray,
precision: np.ndarray,
expected_result: float,
exception: Exception,
) -> None:
with exception:
result = MeanAveragePrecision.compute_average_precision(recall=recall, precision=precision)
result = MeanAveragePrecision.compute_average_precision(
recall=recall, precision=precision
)
assert_almost_equal(result, expected_result, tolerance=0.01)