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

feat: 🚀 pre-commit and initial formatted files 

Signed-off-by: Onuralp SEZER <thunderbirdtr@gmail.com>
This commit is contained in:
Onuralp SEZER 2023-07-25 13:02:26 +03:00
parent 8662671e6b
commit 57e17c34a0
No known key found for this signature in database
GPG Key ID: CF0835DFDF14CA38
43 changed files with 1193 additions and 1193 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?

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

@ -29,4 +29,4 @@ jobs:
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."

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

@ -0,0 +1,86 @@
default_language_version:
python: python3.8
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/executablebooks/mdformat
rev: 0.7.16
hooks:
- id: mdformat
name: MD formatting
additional_dependencies:
- mdformat-gfm
- mdformat-black
exclude: "README.md|CITATION|CONTRIBUTING.md"
- repo: https://github.com/PyCQA/flake8
rev: 6.0.0
hooks:
- id: flake8
name: PEP8 Max Line Length Check
args:
- --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics
- repo: https://github.com/PyCQA/flake8
rev: 6.0.0
hooks:
- id: flake8
name: PEP8
args:
- --count --max-line-length=88 --select=E9,F63,F7,F82 --show-source --statistics
- 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/autoflake
rev: v2.2.0
hooks:
- id: autoflake
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.0.280
hooks:
- id: ruff
args: [--fix, --exit-non-zero-on-fix]

2
CITATION.cff
View File

@ -9,4 +9,4 @@ repository-code: https://github.com/roboflow/supervision
keywords:
- computer vision
- image processing
- video processing
- video processing

8
CONTRIBUTING.md
View File

@ -37,15 +37,15 @@ All pull requests will be reviewed by the maintainers of the project. We will pr
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.
[`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

44
README.md
View File

@ -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
```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/)
@ -126,7 +126,7 @@ Read more about desktop, headless and local installation in our [guide](https://
... annotations_path='...'
... )
```
- Loop over dataset entries
```python
@ -137,17 +137,17 @@ 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
```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
```python
@ -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
```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,7 +202,7 @@ Read more about desktop, headless and local installation in our [guide](https://
... annotations_directory_path='...'
... )
```
- Load classification datasets in one of supported formats
```python
@ -251,9 +251,9 @@ array([
<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 +263,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

2
demo.ipynb vendored
View File

@ -1322,4 +1322,4 @@
}
}
]
}
}

40
docs/changelog.md
View File

@ -2,7 +2,9 @@
!!! warning
With the `supervision-0.12.0` release, we are terminating official support for Python 3.7.
```
With the `supervision-0.12.0` release, we are terminating official support for Python 3.7.
```
- Added [#177](https://github.com/roboflow/supervision/pull/177): initial support for object detection model benchmarking with [`sv.ConfusionMatrix`](https://roboflow.github.io/supervision/metrics/detection/#confusionmatrix).
@ -47,7 +49,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 +65,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 +111,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 +134,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 +152,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 +176,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 +193,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 +222,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 +230,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 👋

8
docs/dataset/core.md
View File

@ -1,7 +1,9 @@
!!! warning
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`.
```
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 +11,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

94
docs/index.md
View File

@ -11,63 +11,67 @@
## 👋 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.
````
=== "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.
```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.
```bash
pip install supervision
```
```bash
pip install supervision[desktop]
```
=== "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.
```bash
pip install supervision[desktop]
```
````
!!! example "git clone (for development)"
=== "virtualenv"
````
=== "virtualenv"
```bash
# 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
```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
# setup python environment and activate it
python3 -m venv venv
source venv/bin/activate
# headless install
poetry install
# desktop install
poetry install --extras "desktop"
```
# 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"
```
````

6
docs/metrics/detection.md
View File

@ -1,7 +1,9 @@
!!! warning
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`.
```
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

368
docs/quickstart/detections.md
View File

@ -1,10 +1,9 @@
## 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
@ -13,72 +12,79 @@ Allows you to select detections that belong only to one selected class.
=== "After"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
detections = detections[detections.class_id == 0]
```
detections = sv.Detections(...)
detections = detections[detections.class_id == 0]
```
<div class="result" markdown>
<div class="result" markdown>
![by-specific-class](https://media.roboflow.com/open-source/supervision/supervision-detection-by-specific-class.png){ align=center width="800" }
![by-specific-class](https://media.roboflow.com/open-source/supervision/supervision-detection-by-specific-class.png){ align=center width="800" }
</div>
</div>
````
=== "Before"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
detections = detections[detections.class_id == 0]
```
detections = sv.Detections(...)
detections = detections[detections.class_id == 0]
```
<div class="result" markdown>
<div class="result" markdown>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
````
### 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)]
```
````
```python
import numpy as np
import supervision as sv
<div class="result" markdown>
selected_classes = [0, 2, 3]
detections = sv.Detections(...)
detections = detections[np.isin(detections.class_id, selected_classes)]
```
![by-set-of-classes](https://media.roboflow.com/open-source/supervision/supervision-detection-by-set-of-classes.png){ align=center width="800" }
<div class="result" markdown>
</div>
![by-set-of-classes](https://media.roboflow.com/open-source/supervision/supervision-detection-by-set-of-classes.png){ align=center width="800" }
</div>
````
=== "Before"
```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)]
```
````
```python
import numpy as np
import supervision as sv
<div class="result" markdown>
class_id = [0, 2, 3]
detections = sv.Detections(...)
detections = detections[np.isin(detections.class_id, class_id)]
```
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
<div class="result" markdown>
</div>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
````
### by confidence
@ -86,192 +92,212 @@ Allows you to select detections with specific confidence value, for example high
=== "After"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
detections = detections[detections.confidence > 0.5]
```
detections = sv.Detections(...)
detections = detections[detections.confidence > 0.5]
```
<div class="result" markdown>
<div class="result" markdown>
![by-set-of-classes](https://media.roboflow.com/open-source/supervision/supervision-detection-by-confidence.png){ align=center width="800" }
![by-set-of-classes](https://media.roboflow.com/open-source/supervision/supervision-detection-by-confidence.png){ align=center width="800" }
</div>
</div>
````
=== "Before"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
detections = detections[detections.confidence > 0.5]
```
detections = sv.Detections(...)
detections = detections[detections.confidence > 0.5]
```
<div class="result" markdown>
<div class="result" markdown>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
</div>
````
### 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"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
detections = detections[detections.area > 1000]
```
detections = sv.Detections(...)
detections = detections[detections.area > 1000]
```
<div class="result" markdown>
<div class="result" markdown>
![by-area](https://media.roboflow.com/open-source/supervision/supervision-detection-by-area.png){ align=center width="800" }
![by-area](https://media.roboflow.com/open-source/supervision/supervision-detection-by-area.png){ align=center width="800" }
</div>
</div>
````
=== "Before"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
detections = detections[detections.area > 1000]
```
detections = sv.Detections(...)
detections = detections[detections.area > 1000]
```
<div class="result" markdown>
<div class="result" markdown>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
</div>
````
### 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"
```python
import supervision as sv
````
```python
import supervision as sv
image = ...
height, width, channels = image.shape
image_area = height * width
detections = sv.Detections(...)
detections = detections[(detections.area / image_area) < 0.8]
```
image = ...
height, width, channels = image.shape
image_area = height * width
<div class="result" markdown>
detections = sv.Detections(...)
detections = detections[(detections.area / image_area) < 0.8]
```
![by-relative-area](https://media.roboflow.com/open-source/supervision/supervision-detection-by-relative-area.png?updatedAt=1683207183434){ align=center width="800" }
<div class="result" markdown>
</div>
![by-relative-area](https://media.roboflow.com/open-source/supervision/supervision-detection-by-relative-area.png?updatedAt=1683207183434){ align=center width="800" }
</div>
````
=== "Before"
```python
import supervision as sv
````
```python
import supervision as sv
image = ...
height, width, channels = image.shape
image_area = height * width
detections = sv.Detections(...)
detections = detections[(detections.area / image_area) < 0.8]
```
image = ...
height, width, channels = image.shape
image_area = height * width
<div class="result" markdown>
detections = sv.Detections(...)
detections = detections[(detections.area / image_area) < 0.8]
```
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
<div class="result" markdown>
</div>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
````
### 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"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
detections = detections[(w > 200) & (h > 200)]
```
detections = sv.Detections(...)
w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
detections = detections[(w > 200) & (h > 200)]
```
<div class="result" markdown>
<div class="result" markdown>
![by-box-dimensions](https://media.roboflow.com/open-source/supervision/supervision-detection-by-box-dimensions.png){ align=center width="800" }
![by-box-dimensions](https://media.roboflow.com/open-source/supervision/supervision-detection-by-box-dimensions.png){ align=center width="800" }
</div>
</div>
````
=== "Before"
```python
import supervision as sv
````
```python
import supervision as sv
detections = sv.Detections(...)
w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
detections = detections[(w > 200) & (h > 200)]
```
detections = sv.Detections(...)
w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
detections = detections[(w > 200) & (h > 200)]
```
<div class="result" markdown>
<div class="result" markdown>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
</div>
````
### 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"
```python
import supervision as sv
````
```python
import supervision as sv
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[mask]
```
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[mask]
```
<div class="result" markdown>
<div class="result" markdown>
![by-polygon-zone](https://media.roboflow.com/open-source/supervision/supervision-detection-by-polygon-zone.png?updatedAt=1683211380445){ align=center width="800" }
![by-polygon-zone](https://media.roboflow.com/open-source/supervision/supervision-detection-by-polygon-zone.png?updatedAt=1683211380445){ align=center width="800" }
</div>
</div>
````
=== "Before"
```python
import supervision as sv
````
```python
import supervision as sv
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[mask]
```
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[mask]
```
<div class="result" markdown>
<div class="result" markdown>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
</div>
</div>
````
### by mixed conditions
@ -279,34 +305,38 @@ zone. In the example below you can see how to filter out all detections located
=== "After"
```python
import supervision as sv
````
```python
import supervision as sv
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[(detections.confidence > 0.7) & mask]
```
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[(detections.confidence > 0.7) & mask]
```
<div class="result" markdown>
<div class="result" markdown>
![by-mixed-conditions](https://media.roboflow.com/open-source/supervision/supervision-detection-by-mixed-conditions.png){ align=center width="800" }
![by-mixed-conditions](https://media.roboflow.com/open-source/supervision/supervision-detection-by-mixed-conditions.png){ align=center width="800" }
</div>
</div>
````
=== "Before"
```python
import supervision as sv
````
```python
import supervision as sv
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[mask]
```
zone = sv.PolygonZone(...)
detections = sv.Detections(...)
mask = zone.trigger(detections=detections)
detections = detections[mask]
```
<div class="result" markdown>
<div class="result" markdown>
![original](https://media.roboflow.com/open-source/supervision/supervision-detection-original.png){ align=center width="800" }
![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;
}
}

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

2
mkdocs.yml
View File

@ -85,4 +85,4 @@ markdown_extensions:
- pymdownx.tabbed:
alternate_style: true
- toc:
permalink: true
permalink: true

91
pyproject.toml
View File

@ -64,6 +64,97 @@ flake8 = { version = "*", python = ">=3.8.1,<3.12.0" }
mkdocs-material = "^9.1.4"
mkdocstrings = {extras = ["python"], version = "^0.20.0"}
[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
)/
'''
[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

6
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
@ -126,9 +126,7 @@ def load_yolo_annotations(
image_paths = list_files_with_extensions(
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
)
annotation_paths = list_files_with_extensions(
directory=annotations_directory_path, extensions=["txt"]
)
list_files_with_extensions(directory=annotations_directory_path, extensions=["txt"])
classes = _extract_class_names(file_path=data_yaml_path)
images = {}

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

10
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 (
@ -437,7 +436,8 @@ 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(l):
return all(x is not None for x in l)
xyxy = np.vstack(xyxy)
mask = np.vstack(mask) if all_not_none(mask) else None
@ -561,7 +561,7 @@ class Detections:
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)))
@ -571,8 +571,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(

2
supervision/detection/utils.py
View File

@ -65,7 +65,7 @@ def non_max_suppression(
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

12
supervision/metrics/detection.py
View File

@ -112,7 +112,8 @@ class ConfusionMatrix:
if with_confidence:
if detections.confidence is None:
raise ValueError(
"ConfusionMatrix can only be calculated for Detections with confidence"
"ConfusionMatrix can only be calculated for Detections with"
" confidence"
)
arrays_to_concat.append(np.expand_dims(detections.confidence, 1))
@ -212,18 +213,21 @@ class ConfusionMatrix:
"""
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 targets"
f" ({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."
"Predictions and targets must be lists of numpy arrays. Got"
f" {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."
"Predictions must have shape (N, 6). Got"
f" {predictions[0].shape} instead."
)
if targets[0].shape[1] != 5:
raise ValueError(

3
supervision/utils/notebook.py
View File

@ -78,7 +78,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)

2
supervision/utils/video.py
View File

@ -103,7 +103,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)

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

378
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,49 @@ 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()
["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()
["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()
["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 +79,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 +96,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 +139,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 +188,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

280
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,187 @@ 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()
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 +209,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)

296
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,85 @@ 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()
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()
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 +257,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 +271,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(

29
test/metrics/test_detection.py
View File

@ -1,4 +1,5 @@
from contextlib import ExitStack as DoesNotRaise
from test.utils import mock_detections
from typing import Optional, Union
import numpy as np
@ -6,7 +7,6 @@ import pytest
from supervision.detection.core import Detections
from supervision.metrics.detection import ConfusionMatrix
from test.utils import mock_detections
CLASSES = np.arange(80)
NUM_CLASSES = len(CLASSES)
@ -147,15 +147,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 +174,18 @@ def test_detections_to_tensor(
detections: Detections,
with_confidence: bool,
expected_result: Optional[np.ndarray],
exception: Exception
exception: Exception,
):
with exception:
result = ConfusionMatrix.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 +356,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],