Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
800e39d64c | ||
|
|
926f935bad | ||
|
|
d3da6afa07 | ||
|
|
97ae155b1e | ||
|
|
24f013757f | ||
|
|
532245640b | ||
|
|
864764e665 | ||
|
|
ca348401d5 | ||
|
|
300ebea8ec | ||
|
|
5cc7da289e | ||
|
|
7a07517750 | ||
|
|
699fc01a0c | ||
|
|
43339e380a | ||
|
|
89a33cda8a | ||
|
|
e84302b43f | ||
|
|
93ab21dda7 | ||
|
|
effa1415b5 | ||
|
|
e9e075acaf | ||
|
|
61e9584f5a | ||
|
|
2a640d3ab6 | ||
|
|
2f478f4d06 | ||
|
|
0c20830be5 | ||
|
|
6c51c4a1b4 | ||
|
|
d6883caadc | ||
|
|
f1f2728e0d | ||
|
|
b6b7663ea7 | ||
|
|
3d074fa213 | ||
|
|
9fd267acb2 | ||
|
|
727451b382 | ||
|
|
5f432fbc93 | ||
|
|
ea4638f5f5 | ||
|
|
4b38634c63 | ||
|
|
04b7badcbe | ||
|
|
54524a4a1f | ||
|
|
3f55e7a79a | ||
|
|
7a782b1ac3 | ||
|
|
6016338a4f | ||
|
|
8d131ad3e8 | ||
|
|
ae55e9c908 | ||
|
|
5cb3c603cd | ||
|
|
56e0b77211 | ||
|
|
bc9fe6f690 | ||
|
|
fb211a6160 | ||
|
|
37b7224380 | ||
|
|
314454c7a7 | ||
|
|
ac39bdd487 | ||
|
|
fc1e899575 | ||
|
|
fccbe6baa2 | ||
|
|
d21154e882 | ||
|
|
1b3af74756 | ||
|
|
dc1083711a | ||
|
|
8df5c5b0c4 | ||
|
|
abbc0e44b0 | ||
|
|
950744111a | ||
|
|
20954a2996 | ||
|
|
22d82c306f | ||
|
|
399121b65e | ||
|
|
1b9c65bdf8 | ||
|
|
3a8c4d63ea | ||
|
|
bfe058c307 | ||
|
|
2c0905ba28 | ||
|
|
fd717ea749 | ||
|
|
40853350fd | ||
|
|
fdd8127293 | ||
|
|
21a480b483 | ||
|
|
1b013f85e5 | ||
|
|
36b48def89 | ||
|
|
b7f72b79de | ||
|
|
141428310a | ||
|
|
8f39df2fc0 | ||
|
|
5107dd359f | ||
|
|
d480ff5be2 | ||
|
|
2b876bdd64 | ||
|
|
476be527e5 | ||
|
|
f78ac90464 | ||
|
|
d4ea295f09 | ||
|
|
9c5c296fb6 | ||
|
|
85c4d5a644 | ||
|
|
ad06937653 | ||
|
|
e01e495c92 | ||
|
|
6784c78b4f | ||
|
|
4375e19245 | ||
|
|
a75a92708c | ||
|
|
c83e57c638 | ||
|
|
fcde1cea9c | ||
|
|
8614984f8f | ||
|
|
cc81ae723b | ||
|
|
f7a4a788ed | ||
|
|
35d9ee7bad | ||
|
|
abbde887dc | ||
|
|
4f5f87acb8 | ||
|
|
3f78150a1d | ||
|
|
183ba1e809 | ||
|
|
597c776c39 | ||
|
|
5291706bc1 | ||
|
|
a50e47e7c0 | ||
|
|
c9d85477f1 | ||
|
|
19271f362e | ||
|
|
9e1e9fc2fc | ||
|
|
629744dee1 | ||
|
|
c9d2b52900 | ||
|
|
43cca52a58 | ||
|
|
e1ee80cb37 | ||
|
|
6a8f3fc03b | ||
|
|
c4b33fe1d6 | ||
|
|
ebfaee22eb | ||
|
|
9c062aac29 | ||
|
|
3e4cb6df42 | ||
|
|
5fdbb70c60 | ||
|
|
9aee232d04 | ||
|
|
9adc953780 | ||
|
|
4ace93f9ab | ||
|
|
c1d3ab2944 | ||
|
|
2072eeac12 | ||
|
|
c74efc4908 | ||
|
|
08f399cf19 | ||
|
|
17020b23de | ||
|
|
cd929525f6 | ||
|
|
bb18de8dce | ||
|
|
dec96aec4d | ||
|
|
ee45c684d0 | ||
|
|
ea517642f0 | ||
|
|
cdc08c0455 | ||
|
|
a0c808cfdf | ||
|
|
f146e88ceb | ||
|
|
66d29015d6 | ||
|
|
46cb82797e | ||
|
|
194ff0aaad | ||
|
|
a579ed46c6 | ||
|
|
8918783702 | ||
|
|
6adb0f019d | ||
|
|
84ae12768a | ||
|
|
9cbec9fc94 | ||
|
|
ed12daf77d | ||
|
|
84cfa9609b | ||
|
|
f2eccae88e | ||
|
|
55c8295755 | ||
|
|
7849033da3 | ||
|
|
2e443c80c7 | ||
|
|
006caa0b9b | ||
|
|
34e45bed5e | ||
|
|
0d39de898a | ||
|
|
782e7e5a39 | ||
|
|
b40265db9e | ||
|
|
de1ceff065 | ||
|
|
5bf78e5505 | ||
|
|
f63336a824 | ||
|
|
4d82460358 | ||
|
|
540efc80ce | ||
|
|
beea21a97d | ||
|
|
484e3d4afd | ||
|
|
ab9728cca2 | ||
|
|
769b6b967c | ||
|
|
253113fa88 | ||
|
|
3b2902c87c | ||
|
|
9cff624e15 | ||
|
|
584c4fef82 | ||
|
|
282aa2c5d3 | ||
|
|
453d4df068 | ||
|
|
4d4b38b292 | ||
|
|
13e8afa996 | ||
|
|
dd4cae328b | ||
|
|
ed7103851a | ||
|
|
ca40b001dc | ||
|
|
de3ccc5c5d | ||
|
|
1b6d02898c | ||
|
|
5ac92bc43b | ||
|
|
86d9293770 | ||
|
|
0e943e673e | ||
|
|
234e376fca | ||
|
|
54b7cbeb70 | ||
|
|
d000292eb2 | ||
|
|
b551e6d410 | ||
|
|
138c271af0 | ||
|
|
d22940ce30 | ||
|
|
06da0e70d2 | ||
|
|
5d8b378f35 | ||
|
|
a4830389de | ||
|
|
d9a2efc64b | ||
|
|
f91b844b7d | ||
|
|
dc45d47e7e | ||
|
|
b5dedf09cd | ||
|
|
eb1ae1b790 | ||
|
|
f2f8a564e0 | ||
|
|
140a219091 | ||
|
|
86af9b9dce | ||
|
|
4366e4480f | ||
|
|
b66e9c8c32 | ||
|
|
d0d5626156 | ||
|
|
c661bf250d | ||
|
|
ccb32f8aa1 | ||
|
|
97b9edcbec | ||
|
|
6b52863c9f | ||
|
|
2fcc035413 | ||
|
|
7c8f16bf0f | ||
|
|
1870670c77 | ||
|
|
802347f4d8 | ||
|
|
920019c887 | ||
|
|
477b9a8f38 | ||
|
|
14f62deacf | ||
|
|
08eb513594 | ||
|
|
a1615fd275 | ||
|
|
ce979216f2 | ||
|
|
95f94319f8 | ||
|
|
98f60c24d3 | ||
|
|
1c043d6cc7 | ||
|
|
86aee5ed98 | ||
|
|
0d38b95da7 | ||
|
|
97892e3138 | ||
|
|
17a42bc728 | ||
|
|
be65ba7370 | ||
|
|
72d86ba180 | ||
|
|
fba83ec371 | ||
|
|
4a81a17ddb | ||
|
|
30d253d5cc | ||
|
|
544f9cb0e2 | ||
|
|
84928a2da9 | ||
|
|
08ae0b2fc7 | ||
|
|
0261bec2ce | ||
|
|
5a4cb1b302 | ||
|
|
91655f9609 | ||
|
|
42d4079b23 | ||
|
|
f03a47a5df | ||
|
|
c01efeaaca | ||
|
|
5489d405c6 | ||
|
|
5f97eac9a9 | ||
|
|
b474a47065 |
@@ -0,0 +1 @@
|
||||
*.ipynb linguist-vendored
|
||||
@@ -5,14 +5,15 @@ on:
|
||||
branches:
|
||||
- master
|
||||
- main
|
||||
- develop
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v2
|
||||
- uses: actions/setup-python@v2
|
||||
- uses: actions/checkout@v3
|
||||
- uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: 3.x
|
||||
- run: pip install mkdocs-material
|
||||
- run: pip install "mkdocstrings[python]"
|
||||
- run: mkdocs gh-deploy --force
|
||||
- run: mkdocs gh-deploy --force
|
||||
|
||||
@@ -16,17 +16,32 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.head_ref }}
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v2
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: 🦾 Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade virtualenv
|
||||
python -m pip install --upgrade pip
|
||||
pip install ".[dev]"
|
||||
- name: 🚀 Publish to PyPi
|
||||
env:
|
||||
PYPI_USERNAME: ${{ secrets.PYPI_USERNAME }}
|
||||
PYPI_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
|
||||
PYPI_TEST_PASSWORD: ${{ secrets.PYPI_TEST_PASSWORD }}
|
||||
virtualenv venv
|
||||
source venv/bin/activate
|
||||
python -m pip install --upgrade pip
|
||||
python -m pip install --upgrade poetry
|
||||
poetry install
|
||||
|
||||
- name: 🏗️ Build source and wheel distributions
|
||||
run: |
|
||||
make publish -e PYPI_USERNAME=$PYPI_USERNAME -e PYPI_PASSWORD=$PYPI_PASSWORD -e PYPI_TEST_PASSWORD=$PYPI_TEST_PASSWORD
|
||||
python -m pip install --upgrade build twine
|
||||
python -m build
|
||||
twine check --strict dist/*
|
||||
- name: 🚀 Publish to PyPi
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
user: ${{ secrets.PYPI_USERNAME }}
|
||||
password: ${{ secrets.PYPI_PASSWORD }}
|
||||
- name: 🚀 Publish to Test-PyPi
|
||||
uses: pypa/gh-action-pypi-publish@release/v1
|
||||
with:
|
||||
repository-url: https://test.pypi.org/legacy/
|
||||
user: ${{ secrets.PYPI_USERNAME }}
|
||||
password: ${{ secrets.PYPI_TEST_PASSWORD }}
|
||||
|
||||
@@ -2,21 +2,19 @@ name: Test WorkFlow
|
||||
|
||||
on:
|
||||
pull_request:
|
||||
branches: [main]
|
||||
branches: [main, develop]
|
||||
|
||||
jobs:
|
||||
build:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: [3.7, 3.8, 3.9, '3.10']
|
||||
python-version: ["3.8", "3.9", "3.10","3.11"]
|
||||
steps:
|
||||
- name: 🛎️ Checkout
|
||||
uses: actions/checkout@v3
|
||||
with:
|
||||
ref: ${{ github.head_ref }}
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v2
|
||||
uses: actions/setup-python@v4
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: 🦾 Install dependencies
|
||||
|
||||
+1
-1
@@ -2,7 +2,7 @@ cff-version: 1.2.0
|
||||
author: Roboflow
|
||||
message: If you use this software, please cite it as below.
|
||||
title: supervision
|
||||
version: 0.1.0
|
||||
version: 0.11.1
|
||||
date-released: 2023-01-19
|
||||
license: MIT
|
||||
repository-code: https://github.com/roboflow/supervision
|
||||
|
||||
@@ -16,7 +16,7 @@ check_code_quality:
|
||||
flake8 $(check_dirs) --count --max-line-length=88 --exit-zero --ignore=D --extend-ignore=E203,E501,W503 --statistics
|
||||
|
||||
publish:
|
||||
python setup.py sdist bdist_wheel
|
||||
poetry build
|
||||
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
|
||||
@@ -10,6 +10,269 @@
|
||||
|
||||
<br>
|
||||
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://pypistats.org/packages/supervision)
|
||||
[](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb)
|
||||
|
||||
</div>
|
||||
|
||||
## 👋 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! 🤝
|
||||
|
||||
## 💻 install
|
||||
|
||||
Pip install the supervision package in a
|
||||
[**3.11>=Python>=3.8**](https://www.python.org/) environment.
|
||||
|
||||
```bash
|
||||
pip install supervision[desktop]
|
||||
```
|
||||
|
||||
Read more about desktop, headless and local installation in our [guide](https://roboflow.github.io/supervision/).
|
||||
|
||||
## 🔥 quickstart
|
||||
|
||||
### [detections processing](https://roboflow.github.io/supervision/detection/core/)
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> model = YOLO('yolov8s.pt')
|
||||
>>> result = model(IMAGE)[0]
|
||||
>>> detections = sv.Detections.from_yolov8(result)
|
||||
|
||||
>>> len(detections)
|
||||
5
|
||||
```
|
||||
|
||||
<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
|
||||
|
||||
>>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
|
||||
>>> mask_generator = SamAutomaticMaskGenerator(sam)
|
||||
>>> sam_result = mask_generator.generate(IMAGE)
|
||||
>>> detections = sv.Detections.from_sam(sam_result=sam_result)
|
||||
```
|
||||
|
||||
- [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
|
||||
|
||||
>>> box_annotator = sv.BoxAnnotator()
|
||||
>>> annotated_frame = box_annotator.annotate(
|
||||
... scene=IMAGE,
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### [datasets processing](https://roboflow.github.io/supervision/dataset/core/)
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> len(dataset)
|
||||
1000
|
||||
```
|
||||
|
||||
<details close>
|
||||
<summary>👉 more dataset utils</summary>
|
||||
|
||||
- Load object detection / instance segmentation datasets in one of supported formats
|
||||
|
||||
```python
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_pascal_voc(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Loop over dataset entries
|
||||
|
||||
```python
|
||||
>>> for name, image, labels in dataset:
|
||||
... print(labels.xyxy)
|
||||
|
||||
array([[404. , 719. , 538. , 884.5 ],
|
||||
[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
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
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='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset.as_pascal_voc(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset.as_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Convert labels between supported formats
|
||||
|
||||
```python
|
||||
>>> sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... ).as_pascal_voc(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Load classification datasets in one of supported formats
|
||||
|
||||
```python
|
||||
>>> cs = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Save classification datasets in one of supported formats
|
||||
|
||||
```python
|
||||
>>> cs.as_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### [model evaluation](https://roboflow.github.io/supervision/metrics/detection/)
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(...)
|
||||
|
||||
>>> def callback(image: np.ndarray) -> sv.Detections:
|
||||
... ...
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.benchmark(
|
||||
... dataset = dataset,
|
||||
... callback = callback
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
|
||||
## 🎬 tutorials
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/ae1ca38e-40b7-4b35-8582-e8ea5de3806e" alt="Accelerate Image Annotation with SAM and Grounding DINO" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><strong>Accelerate Image Annotation with SAM and Grounding DINO</strong></a>
|
||||
<div><strong>Created: 20 Apr 2023</strong> | <strong>Updated: 20 Apr 2023</strong></div>
|
||||
<br/> Discover how to speed up your image annotation process using Grounding DINO and Segment Anything Model (SAM). Learn how to convert object detection datasets into instance segmentation datasets, and see the potential of using these models to automatically annotate your datasets for real-time detectors like YOLOv8... </p>
|
||||
|
||||
<br/>
|
||||
|
||||
<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>
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><strong>SAM - Segment Anything Model by Meta AI: Complete Guide</strong></a>
|
||||
<div><strong>Created: 11 Apr 2023</strong> | <strong>Updated: 11 Apr 2023</strong></div>
|
||||
<br/> Discover the incredible potential of Meta AI's Segment Anything Model (SAM)! We dive into SAM, an efficient and promptable model for image segmentation, which has revolutionized computer vision tasks. With over 1 billion masks on 11M licensed and privacy-respecting images, SAM's zero-shot performance is often competitive with or even superior to prior fully supervised results... </p>
|
||||
|
||||
## 📚 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.
|
||||
|
||||
## 🏆 contribution
|
||||
|
||||
We love your input! Please see our [contributing guide](https://github.com/roboflow/supervision/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!
|
||||
|
||||
<br>
|
||||
|
||||
<div align="center">
|
||||
|
||||
<div align="center">
|
||||
<a href="https://youtube.com/roboflow">
|
||||
<img
|
||||
@@ -54,77 +317,4 @@
|
||||
</a>
|
||||
</div>
|
||||
|
||||
<br>
|
||||
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://pypistats.org/packages/supervision)
|
||||
[](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
|
||||
|
||||
</div>
|
||||
|
||||
## 👋 hello
|
||||
|
||||
A set of easy-to-use utils that will come in handy in any computer vision project. **Supervision is still in
|
||||
pre-release stage. 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid
|
||||
and may change.**
|
||||
|
||||
## 💻 install
|
||||
|
||||
Pip install the supervision package in a
|
||||
[**3.10>=Python>=3.7**](https://www.python.org/) environment.
|
||||
|
||||
```bash
|
||||
pip install supervision
|
||||
```
|
||||
|
||||
<details close>
|
||||
<summary>Install from source</summary>
|
||||
|
||||
```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
|
||||
|
||||
# install
|
||||
pip install -e ".[dev]"
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 📖 documentation
|
||||
|
||||
Curious how Supervision can help you solve problems on your project? Visit our [documentation](https://roboflow.github.io/supervision) page!
|
||||
|
||||
## 🎬 tutorials
|
||||
|
||||
🔥 [Subscribe](https://www.youtube.com/@Roboflow) and stay up to date with the latest changes to the Supervision library.
|
||||
|
||||
<p align="center">
|
||||
<a href="https://youtu.be/l_kf9CfZ_8M">
|
||||
<img
|
||||
width="90%"
|
||||
src="https://user-images.githubusercontent.com/26109316/217950212-311de186-1862-4b4c-a86e-89cafd68b233.jpg"
|
||||
alt="count-objects-in-zone"
|
||||
>
|
||||
</a>
|
||||
</p>
|
||||
|
||||
<p align="center">
|
||||
<a href="https://youtu.be/OS5qI9YBkfk">
|
||||
<img
|
||||
width="90%"
|
||||
src="https://user-images.githubusercontent.com/26109316/213702005-ddd568f0-b902-46c2-9af9-e6ff33db23bf.jpg"
|
||||
alt="count-objects-crossing-line"
|
||||
>
|
||||
</a>
|
||||
</p>
|
||||
|
||||
## 🏆 contribution
|
||||
|
||||
We love your input! Please see our [contributing guide](https://github.com/roboflow/supervision/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!
|
||||
|
||||
Vendored
+1325
File diff suppressed because one or more lines are too long
+211
-28
@@ -1,29 +1,212 @@
|
||||
### 0.12.0 <small>July 24, 2023</small>
|
||||
|
||||
!!! warning
|
||||
|
||||
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).
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(...)
|
||||
|
||||
>>> model = YOLO(...)
|
||||
>>> def callback(image: np.ndarray) -> sv.Detections:
|
||||
... result = model(image)[0]
|
||||
... return sv.Detections.from_yolov8(result)
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.benchmark(
|
||||
... dataset = dataset,
|
||||
... callback = callback
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
|
||||
- Added [#173](https://github.com/roboflow/supervision/pull/173): [`Detections.from_mmdetection`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_mmdetection) to enable seamless integration with [MMDetection](https://github.com/open-mmlab/mmdetection) framework.
|
||||
|
||||
- Added [#130](https://github.com/roboflow/supervision/issues/130): ability to [install](https://roboflow.github.io/supervision/) package in `headless` or `desktop` mode.
|
||||
|
||||
- Changed [#180](https://github.com/roboflow/supervision/pull/180): packing method from `setup.py` to `pyproject.toml`.
|
||||
|
||||
- Fixed [#188](https://github.com/roboflow/supervision/issues/188): [`sv.DetectionDataset.from_cooc`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_coco) can't be loaded when there are images without annotations.
|
||||
|
||||
- Fixed [#226](https://github.com/roboflow/supervision/issues/226): [`sv.DetectionDataset.from_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_yolo) can't load background instances.
|
||||
|
||||
### 0.11.1 <small>June 29, 2023</small>
|
||||
|
||||
- Fix [#165](https://github.com/roboflow/supervision/pull/165): [`as_folder_structure`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.ClassificationDataset.as_folder_structure) fails to save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) when it is result of inference.
|
||||
|
||||
### 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.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
|
||||
>>> ds.as_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- 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
|
||||
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
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']
|
||||
```
|
||||
|
||||
- Added [#162](https://github.com/roboflow/supervision/pull/162): additional `start` and `end` arguments to [`sv.get_video_frames_generator`](https://roboflow.github.io/supervision/utils/video/#get_video_frames_generator) allowing to generate frames only for a selected part of the video.
|
||||
|
||||
- Fix [#157](https://github.com/roboflow/supervision/pull/157): incorrect loading of YOLO dataset class names from `data.yaml`.
|
||||
|
||||
### 0.10.0 <small>June 14, 2023</small>
|
||||
|
||||
- Added [#125](https://github.com/roboflow/supervision/pull/125): ability to load and save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) in a folder structure format.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> cs = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> cs.as_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- 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 [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.
|
||||
|
||||
### 0.9.0 <small>June 7, 2023</small>
|
||||
|
||||
- Added [#118](https://github.com/roboflow/supervision/pull/118): ability to select [`sv.Detections`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.__getitem__) by index, list of indexes or slice. Here is an example illustrating the new selection methods.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> detections = sv.Detections(...)
|
||||
>>> len(detections[0])
|
||||
1
|
||||
>>> len(detections[[0, 1]])
|
||||
2
|
||||
>>> len(detections[0:2])
|
||||
2
|
||||
```
|
||||
|
||||
- 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.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
|
||||
... for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=10):
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
|
||||
- Fixed [#106](https://github.com/roboflow/supervision/issues/106): inconvenient handling of [`sv.PolygonZone`](https://roboflow.github.io/supervision/detection/tools/polygon_zone/#polygonzone) coordinates. Now `sv.PolygonZone` accepts coordinates in the form of `[[x1, y1], [x2, y2], ...]` that can be both integers and floats.
|
||||
|
||||
### 0.8.0 <small>May 17, 2023</small>
|
||||
|
||||
- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: support for Dataset inheritance. Current `Dataset` got renamed to `DetectionDataset` and make it inherit from `BaseDataset`.
|
||||
- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: ability to save datasets in YOLO format using `DetectionDataset.as_yolo`.
|
||||
- Added [[#102](https://github.com/roboflow/supervision/pull/103)]: support for splitting `DetectionDataset`.
|
||||
- Changed [[#100](https://github.com/roboflow/supervision/pull/100)]: default value of `approximation_percentage` parameter from `0.75` to `0.0` in `DetectionDataset.as_yolo` and `DetectionDataset.as_pascal_voc`.
|
||||
- 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).
|
||||
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("yolov5")
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=f"{dataset.location}/train/images",
|
||||
... annotations_directory_path=f"{dataset.location}/train/labels",
|
||||
... data_yaml_path=f"{dataset.location}/data.yaml"
|
||||
... )
|
||||
|
||||
>>> ds.classes
|
||||
['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.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds = sv.DetectionDataset(...)
|
||||
>>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
|
||||
>>> len(train_ds), len(test_ds)
|
||||
(700, 300)
|
||||
```
|
||||
|
||||
- Changed [#100](https://github.com/roboflow/supervision/pull/100): default value of `approximation_percentage` parameter from `0.75` to `0.0` in `DetectionDataset.as_yolo` and `DetectionDataset.as_pascal_voc`.
|
||||
|
||||
### 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 [[#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.
|
||||
- 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 [#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)]: new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
|
||||
- Added [[#72](https://github.com/roboflow/supervision/pull/72)]: ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
|
||||
- Changed [[#70](https://github.com/roboflow/supervision/pull/70)]: order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
|
||||
- Changed [[#71](https://github.com/roboflow/supervision/pull/71)]: `generate_2d_mask` to `polygon_to_mask`.
|
||||
- Added [#71](https://github.com/roboflow/supervision/pull/71): initial `Dataset` support and ability to save `Detections` in Pascal VOC XML format.
|
||||
- Added [#71](https://github.com/roboflow/supervision/pull/71): new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
|
||||
- Added [#72](https://github.com/roboflow/supervision/pull/72): ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
|
||||
- Changed [#70](https://github.com/roboflow/supervision/pull/70): order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
|
||||
- Changed [#71](https://github.com/roboflow/supervision/pull/71): `generate_2d_mask` to `polygon_to_mask`.
|
||||
|
||||
### 0.5.2 <small>April 13, 2023</small>
|
||||
|
||||
- Fixed [[#63](https://github.com/roboflow/supervision/pull/63)]: `LineZone.trigger` function expects 4 values instead of 5.
|
||||
- Fixed [#63](https://github.com/roboflow/supervision/pull/63): `LineZone.trigger` function expects 4 values instead of 5.
|
||||
|
||||
### 0.5.1 <small>April 12, 2023</small>
|
||||
|
||||
@@ -32,28 +215,28 @@
|
||||
|
||||
### 0.5.0 <small>April 10, 2023</small>
|
||||
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.mask` to enable segmentation support.
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `MaskAnnotator` to allow easy `Detections.mask` annotation.
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
|
||||
- Changed [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.area` behaviour to work not only with boxes but also with masks.
|
||||
- Added [#58](https://github.com/roboflow/supervision/pull/58): `Detections.mask` to enable segmentation support.
|
||||
- Added [#58](https://github.com/roboflow/supervision/pull/58): `MaskAnnotator` to allow easy `Detections.mask` annotation.
|
||||
- Added [#58](https://github.com/roboflow/supervision/pull/58): `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
|
||||
- Changed [#58](https://github.com/roboflow/supervision/pull/58): `Detections.area` behaviour to work not only with boxes but also with masks.
|
||||
|
||||
### 0.4.0 <small>April 5, 2023</small>
|
||||
|
||||
- Added [[#46](https://github.com/roboflow/supervision/discussions/48)]: `Detections.empty` to allow easy creation of empty `Detections` objects.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `plot_images_grid` to allow easy plotting of multiple images on single plot.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
|
||||
- Changed [[#56](https://github.com/roboflow/supervision/pull/56)]: `show_frame_in_notebook` refactored and renamed to `plot_image`.
|
||||
- Added [#46](https://github.com/roboflow/supervision/discussions/48): `Detections.empty` to allow easy creation of empty `Detections` objects.
|
||||
- Added [#56](https://github.com/roboflow/supervision/pull/56): `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
|
||||
- Added [#56](https://github.com/roboflow/supervision/pull/56): `plot_images_grid` to allow easy plotting of multiple images on single plot.
|
||||
- Added [#56](https://github.com/roboflow/supervision/pull/56): initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
|
||||
- Changed [#56](https://github.com/roboflow/supervision/pull/56): `show_frame_in_notebook` refactored and renamed to `plot_image`.
|
||||
|
||||
### 0.3.2 <small>March 23, 2023</small>
|
||||
|
||||
- Changed [[#50](https://github.com/roboflow/supervision/issues/50)]: Allow `Detections.class_id` to be `None`.
|
||||
- Changed [#50](https://github.com/roboflow/supervision/issues/50): Allow `Detections.class_id` to be `None`.
|
||||
|
||||
### 0.3.1 <small>March 6, 2023</small>
|
||||
|
||||
- Fixed [[#41](https://github.com/roboflow/supervision/issues/41)]: `PolygonZone` throws an exception when the object touches the bottom edge of the image.
|
||||
- Fixed [[#42](https://github.com/roboflow/supervision/issues/42)]: `Detections.wth_nms` method throws an exception when `Detections` is empty.
|
||||
- Changed [[#36](https://github.com/roboflow/supervision/pull/36)]: `Detections.wth_nms` support class agnostic and non-class agnostic case.
|
||||
- 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,0 +1,3 @@
|
||||
## Classifications
|
||||
|
||||
:::supervision.classification.core.Classifications
|
||||
@@ -5,4 +5,8 @@
|
||||
|
||||
## DetectionDataset
|
||||
|
||||
:::supervision.dataset.core.DetectionDataset
|
||||
:::supervision.dataset.core.DetectionDataset
|
||||
|
||||
## ClassificationDataset
|
||||
|
||||
:::supervision.dataset.core.ClassificationDataset
|
||||
+54
-18
@@ -9,29 +9,65 @@
|
||||
</p>
|
||||
</div>
|
||||
|
||||
## 👋 Welcome
|
||||
## 👋 Hello
|
||||
|
||||
Supervision is a set of easy-to-use utilities that will come in handy in any computer vision project.
|
||||
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!
|
||||
|
||||
**Supervision** is still in
|
||||
pre-release stage 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid and may change.
|
||||
|
||||
## 💻 How to Install
|
||||
## 💻 Install
|
||||
|
||||
You can install `supervision` with pip in a
|
||||
[**3.10>=Python>=3.7**](https://www.python.org/) environment.
|
||||
[**3.11>=Python>=3.8**](https://www.python.org/) environment.
|
||||
|
||||
!!! example "Pip install method (recommended)"
|
||||
!!! example "pip install (recommended)"
|
||||
|
||||
```bash
|
||||
pip install supervision
|
||||
```
|
||||
=== "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.
|
||||
|
||||
!!! example "Git clone method (for development)"
|
||||
```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
|
||||
git https://github.com/roboflow/supervision.git
|
||||
cd supervision
|
||||
pip install -e '.[dev]'
|
||||
```
|
||||
See contributing section to know more about contributing to the project
|
||||
```bash
|
||||
pip install supervision[desktop]
|
||||
```
|
||||
|
||||
!!! example "git clone (for development)"
|
||||
|
||||
=== "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
|
||||
|
||||
# setup python environment and activate it
|
||||
poetry env use python 3.10
|
||||
poetry shell
|
||||
|
||||
# headless install
|
||||
poetry install
|
||||
|
||||
# desktop install
|
||||
poetry install --extras "desktop"
|
||||
```
|
||||
@@ -0,0 +1,8 @@
|
||||
!!! 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`.
|
||||
|
||||
## ConfusionMatrix
|
||||
|
||||
:::supervision.metrics.detection.ConfusionMatrix
|
||||
@@ -1,7 +0,0 @@
|
||||
## plot_image
|
||||
|
||||
:::supervision.notebook.utils.plot_image
|
||||
|
||||
## plot_images_grid
|
||||
|
||||
:::supervision.notebook.utils.plot_images_grid
|
||||
@@ -0,0 +1,3 @@
|
||||
## list_files_with_extensions
|
||||
|
||||
:::supervision.utils.file.list_files_with_extensions
|
||||
@@ -0,0 +1,7 @@
|
||||
## ImageSink
|
||||
|
||||
:::supervision.utils.image.ImageSink
|
||||
|
||||
## crop
|
||||
|
||||
:::supervision.utils.image.crop
|
||||
@@ -0,0 +1,7 @@
|
||||
## plot_image
|
||||
|
||||
:::supervision.utils.notebook.plot_image
|
||||
|
||||
## plot_images_grid
|
||||
|
||||
:::supervision.utils.notebook.plot_images_grid
|
||||
@@ -0,0 +1,15 @@
|
||||
## VideoInfo
|
||||
|
||||
:::supervision.utils.video.VideoInfo
|
||||
|
||||
## VideoSink
|
||||
|
||||
:::supervision.utils.video.VideoSink
|
||||
|
||||
## get_video_frames_generator
|
||||
|
||||
:::supervision.utils.video.get_video_frames_generator
|
||||
|
||||
## process_video
|
||||
|
||||
:::supervision.utils.video.process_video
|
||||
@@ -1,15 +0,0 @@
|
||||
## VideoInfo
|
||||
|
||||
:::supervision.video.VideoInfo
|
||||
|
||||
## VideoSink
|
||||
|
||||
:::supervision.video.VideoSink
|
||||
|
||||
## get_video_frames_generator
|
||||
|
||||
:::supervision.video.get_video_frames_generator
|
||||
|
||||
## process_video
|
||||
|
||||
:::supervision.video.process_video
|
||||
+9
-4
@@ -29,6 +29,8 @@ nav:
|
||||
- Quickstart:
|
||||
- Detections: quickstart/detections.md
|
||||
- API reference:
|
||||
- Classifications:
|
||||
- Core: classification/core.md
|
||||
- Detections:
|
||||
- Core: detection/core.md
|
||||
- Annotate: detection/annotate.md
|
||||
@@ -37,12 +39,15 @@ nav:
|
||||
- Polygon Zone: detection/tools/polygon_zone.md
|
||||
- Dataset:
|
||||
- Core: dataset/core.md
|
||||
- Metrics:
|
||||
- Detection Models: metrics/detection.md
|
||||
- Draw:
|
||||
- Utils: draw/utils.md
|
||||
- Notebook:
|
||||
- Utils: notebook/utils.md
|
||||
- Video:
|
||||
- Utils: video.md
|
||||
- Utils:
|
||||
- Video: utils/video.md
|
||||
- Image: utils/image.md
|
||||
- Notebook: utils/notebook.md
|
||||
- File: utils/file.md
|
||||
- Changelog: changelog.md
|
||||
|
||||
theme:
|
||||
|
||||
Generated
+3446
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,75 @@
|
||||
[tool.poetry]
|
||||
name = "supervision"
|
||||
version = "0.12.0"
|
||||
description = "A set of easy-to-use utils that will come in handy in any Computer Vision project"
|
||||
authors = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
maintainers = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
readme = "README.md"
|
||||
packages = [{include = "supervision"}]
|
||||
homepage = "https://github.com/roboflow/supervision"
|
||||
repository = "https://github.com/roboflow/supervision"
|
||||
documentation = "https://github.com/roboflow/supervision/blob/main/README.md"
|
||||
keywords = ["machine-learning", "deep-learning", "vision", "ML", "DL", "AI", "YOLOv5", "YOLOv8", "Roboflow"]
|
||||
|
||||
classifiers=[
|
||||
'Intended Audience :: Developers',
|
||||
'Intended Audience :: Science/Research',
|
||||
'License :: OSI Approved :: BSD License',
|
||||
'Programming Language :: Python :: 3',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
'Programming Language :: Python :: 3.11',
|
||||
'Programming Language :: Python :: 3 :: Only',
|
||||
'Topic :: Software Development',
|
||||
'Topic :: Scientific/Engineering',
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Scientific/Engineering :: Image Recognition",
|
||||
'Typing :: Typed',
|
||||
'Operating System :: Microsoft :: Windows',
|
||||
'Operating System :: POSIX',
|
||||
'Operating System :: Unix',
|
||||
'Operating System :: MacOS'
|
||||
]
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.8,<3.12.0"
|
||||
numpy = "^1.20.0"
|
||||
matplotlib = "^3.7.1"
|
||||
pyyaml = "^6.0"
|
||||
pillow = "^8.4.0"
|
||||
opencv-python = { version = "^4.8.0.74", optional = true }
|
||||
opencv-python-headless = "^4.8.0.74"
|
||||
|
||||
|
||||
[tool.poetry.extras]
|
||||
desktop = ["opencv-python"]
|
||||
|
||||
|
||||
[tool.poetry.group.dev.dependencies]
|
||||
twine = "^4.0.2"
|
||||
pytest = "^7.2.2"
|
||||
wheel = "^0.40.0"
|
||||
notebook = "^6.5.3"
|
||||
build = "^0.10.0"
|
||||
ruff = "^0.0.280"
|
||||
isort = "^5.12.0"
|
||||
black = "^23.7.0"
|
||||
mypy = "^1.4.1"
|
||||
pre-commit = "^3.3.3"
|
||||
flake8 = { version = "*", python = ">=3.8.1,<3.12.0" }
|
||||
|
||||
|
||||
[tool.poetry.group.docs.dependencies]
|
||||
mkdocs-material = "^9.1.4"
|
||||
mkdocstrings = {extras = ["python"], version = "^0.20.0"}
|
||||
|
||||
[tool.setuptools]
|
||||
include-package-data = false
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
exclude = ["docs*", "test*"]
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
build-backend = "poetry.core.masonry.api"
|
||||
@@ -1,68 +0,0 @@
|
||||
import setuptools
|
||||
from setuptools import find_packages
|
||||
import re
|
||||
from pathlib import Path
|
||||
|
||||
FILE = Path(__file__).resolve()
|
||||
PARENT = FILE.parent # root directory
|
||||
README = (PARENT / "README.md").read_text(encoding="utf-8")
|
||||
|
||||
|
||||
def get_version():
|
||||
file = PARENT / 'supervision/__init__.py'
|
||||
return re.search(r'^__version__ = [\'"]([^\'"]*)[\'"]', file.read_text(encoding="utf-8"), re.M)[1]
|
||||
|
||||
|
||||
setuptools.setup(
|
||||
name='supervision',
|
||||
version=get_version(),
|
||||
author='Piotr Skalski',
|
||||
author_email='piotr.skalski92@gmail.com',
|
||||
license='MIT',
|
||||
description='A set of easy-to-use utils that will come in handy in any Computer Vision project',
|
||||
long_description=README,
|
||||
long_description_content_type='text/markdown',
|
||||
url='https://github.com/roboflow/supervision',
|
||||
install_requires=[
|
||||
'numpy>=1.20.0',
|
||||
'opencv-python',
|
||||
'matplotlib',
|
||||
'pyyaml'
|
||||
],
|
||||
packages=find_packages(exclude=("tests",)),
|
||||
extras_require={
|
||||
'dev': [
|
||||
'flake8',
|
||||
'black==22.3.0',
|
||||
'isort',
|
||||
'twine',
|
||||
'pytest',
|
||||
'wheel',
|
||||
'notebook',
|
||||
'mkdocs-material',
|
||||
'mkdocstrings[python]'
|
||||
],
|
||||
},
|
||||
classifiers=[
|
||||
'Intended Audience :: Developers',
|
||||
'Intended Audience :: Science/Research',
|
||||
'License :: OSI Approved :: BSD License',
|
||||
'Programming Language :: Python :: 3',
|
||||
'Programming Language :: Python :: 3.7',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
'Programming Language :: Python :: 3 :: Only',
|
||||
'Topic :: Software Development',
|
||||
'Topic :: Scientific/Engineering',
|
||||
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
||||
"Topic :: Scientific/Engineering :: Image Recognition",
|
||||
'Typing :: Typed',
|
||||
'Operating System :: Microsoft :: Windows',
|
||||
'Operating System :: POSIX',
|
||||
'Operating System :: Unix',
|
||||
'Operating System :: MacOS'
|
||||
],
|
||||
keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLOv5, YOLOv8, Roboflow",
|
||||
python_requires='>=3.7',
|
||||
)
|
||||
+19
-5
@@ -1,6 +1,18 @@
|
||||
__version__ = "0.8.0"
|
||||
import importlib.metadata as importlib_metadata
|
||||
|
||||
from supervision.dataset.core import BaseDataset, DetectionDataset
|
||||
try:
|
||||
# This will read version from pyproject.toml
|
||||
__version__ = importlib_metadata.version(__package__ or __name__)
|
||||
except importlib_metadata.PackageNotFoundError:
|
||||
__version__ = "development"
|
||||
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.core import (
|
||||
BaseDataset,
|
||||
ClassificationDataset,
|
||||
DetectionDataset,
|
||||
)
|
||||
from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
|
||||
@@ -16,11 +28,13 @@ from supervision.detection.utils import (
|
||||
)
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
|
||||
from supervision.file import list_files_with_extensions
|
||||
from supervision.geometry.core import Point, Position, Rect
|
||||
from supervision.geometry.utils import get_polygon_center
|
||||
from supervision.notebook.utils import plot_image, plot_images_grid
|
||||
from supervision.video import (
|
||||
from supervision.metrics.detection import ConfusionMatrix
|
||||
from supervision.utils.file import list_files_with_extensions
|
||||
from supervision.utils.image import ImageSink, crop
|
||||
from supervision.utils.notebook import plot_image, plot_images_grid
|
||||
from supervision.utils.video import (
|
||||
VideoInfo,
|
||||
VideoSink,
|
||||
get_video_frames_generator,
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _validate_class_ids(class_id: Any, n: int) -> None:
|
||||
"""
|
||||
Ensure that class_id is a 1d np.ndarray with (n, ) shape.
|
||||
"""
|
||||
is_valid = isinstance(class_id, np.ndarray) and class_id.shape == (n,)
|
||||
if not is_valid:
|
||||
raise ValueError("class_id must be 1d np.ndarray with (n, ) shape")
|
||||
|
||||
|
||||
def _validate_confidence(confidence: Any, n: int) -> None:
|
||||
"""
|
||||
Ensure that confidence is a 1d np.ndarray with (n, ) shape.
|
||||
"""
|
||||
if confidence is not None:
|
||||
is_valid = isinstance(confidence, np.ndarray) and confidence.shape == (n,)
|
||||
if not is_valid:
|
||||
raise ValueError("confidence must be 1d np.ndarray with (n, ) shape")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Classifications:
|
||||
class_id: np.ndarray
|
||||
confidence: Optional[np.ndarray] = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""
|
||||
Validate the classification inputs.
|
||||
"""
|
||||
n = len(self.class_id)
|
||||
|
||||
_validate_class_ids(self.class_id, n)
|
||||
_validate_confidence(self.confidence, n)
|
||||
|
||||
@classmethod
|
||||
def from_yolov8(cls, yolov8_results) -> Classifications:
|
||||
"""
|
||||
Creates a Classifications instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
|
||||
|
||||
Args:
|
||||
yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
|
||||
|
||||
Returns:
|
||||
Detections: A new Classifications object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = YOLO('yolov8s-cls.pt')
|
||||
>>> result = model(image)[0]
|
||||
>>> classifications = sv.Classifications.from_yolov8(result)
|
||||
```
|
||||
"""
|
||||
confidence = yolov8_results.probs.data.cpu().numpy()
|
||||
return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
|
||||
|
||||
def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Retrieve the top k class IDs and confidences, ordered in descending order by confidence.
|
||||
|
||||
Args:
|
||||
k (int): The number of top class IDs and confidences to retrieve.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, np.ndarray]: A tuple containing the top k class IDs and confidences.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> classifications = sv.Classifications(...)
|
||||
|
||||
>>> classifications.get_top_k(1)
|
||||
|
||||
(array([1]), array([0.9]))
|
||||
```
|
||||
"""
|
||||
if self.confidence is None:
|
||||
raise ValueError("top_k could not be calculated, confidence is None")
|
||||
|
||||
order = np.argsort(self.confidence)[::-1]
|
||||
top_k_order = order[:k]
|
||||
top_k_class_id = self.class_id[top_k_order]
|
||||
top_k_confidence = self.confidence[top_k_order]
|
||||
|
||||
return top_k_class_id, top_k_confidence
|
||||
+335
-25
@@ -1,5 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
@@ -8,6 +9,11 @@ from typing import Dict, Iterator, List, Optional, Tuple
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.formats.coco import (
|
||||
load_coco_annotations,
|
||||
save_coco_annotations,
|
||||
)
|
||||
from supervision.dataset.formats.pascal_voc import (
|
||||
detections_to_pascal_voc,
|
||||
load_pascal_voc_annotations,
|
||||
@@ -17,9 +23,15 @@ from supervision.dataset.formats.yolo import (
|
||||
save_data_yaml,
|
||||
save_yolo_annotations,
|
||||
)
|
||||
from supervision.dataset.ultils import save_dataset_images, train_test_split
|
||||
from supervision.dataset.utils import (
|
||||
build_class_index_mapping,
|
||||
map_detections_class_id,
|
||||
merge_class_lists,
|
||||
save_dataset_images,
|
||||
train_test_split,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.file import list_files_with_extensions
|
||||
from supervision.utils.file import list_files_with_extensions
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -70,6 +82,21 @@ class DetectionDataset(BaseDataset):
|
||||
for image_name, image in self.images.items():
|
||||
yield image_name, image, self.annotations.get(image_name, None)
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, DetectionDataset):
|
||||
return False
|
||||
|
||||
if set(self.classes) != set(other.classes):
|
||||
return False
|
||||
|
||||
for key in self.images:
|
||||
if not np.array_equal(self.images[key], other.images[key]):
|
||||
return False
|
||||
if not self.annotations[key] == other.annotations[key]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
) -> Tuple[DetectionDataset, DetectionDataset]:
|
||||
@@ -77,9 +104,9 @@ class DetectionDataset(BaseDataset):
|
||||
Splits the dataset into two parts (training and testing) using the provided split_ratio.
|
||||
|
||||
Args:
|
||||
split_ratio (float, optional): The ratio of the training set to the entire dataset. Default is 0.8.
|
||||
random_state (int, optional): The seed for the random number generator. This is used for reproducibility. Default is None.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting. Default is True.
|
||||
split_ratio (float, optional): The ratio of the training set to the entire dataset.
|
||||
random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting.
|
||||
|
||||
Returns:
|
||||
Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets.
|
||||
@@ -88,9 +115,9 @@ class DetectionDataset(BaseDataset):
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> dataset = sv.DetectionDataset(...)
|
||||
>>> train_dataset, test_dataset = dataset.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
>>> len(train_dataset), len(test_dataset)
|
||||
>>> ds = sv.DetectionDataset(...)
|
||||
>>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
>>> len(train_ds), len(test_ds)
|
||||
(700, 300)
|
||||
```
|
||||
"""
|
||||
@@ -125,8 +152,7 @@ class DetectionDataset(BaseDataset):
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
|
||||
PASCAL VOC format, which consists of XML files. The method allows filtering the detections based on their area
|
||||
percentage.
|
||||
PASCAL VOC format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
@@ -134,10 +160,11 @@ class DetectionDataset(BaseDataset):
|
||||
annotations_directory_path (Optional[str]): The path to the directory where the annotations in
|
||||
PASCAL VOC format should be saved. If not provided, annotations will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). Argument is used only for segmentation datasets.
|
||||
"""
|
||||
if images_directory_path:
|
||||
images_path = Path(images_directory_path)
|
||||
@@ -195,12 +222,12 @@ class DetectionDataset(BaseDataset):
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("voc")
|
||||
|
||||
>>> train_dataset = sv.Dataset.from_yolo(
|
||||
>>> ds = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=f"{dataset.location}/train/images",
|
||||
... annotations_directory_path=f"{dataset.location}/train/labels"
|
||||
... )
|
||||
|
||||
>>> dataset.classes
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
"""
|
||||
@@ -267,13 +294,13 @@ class DetectionDataset(BaseDataset):
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("yolov5")
|
||||
|
||||
>>> train_dataset = sv.Dataset.from_yolo(
|
||||
>>> ds = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=f"{dataset.location}/train/images",
|
||||
... annotations_directory_path=f"{dataset.location}/train/labels",
|
||||
... data_yaml_path=f"{dataset.location}/data.yaml"
|
||||
... )
|
||||
|
||||
>>> dataset.classes
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
"""
|
||||
@@ -295,11 +322,8 @@ class DetectionDataset(BaseDataset):
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to YOLO (You Only Look Once) format. This method saves the images and their corresponding
|
||||
annotations in YOLO format, which is a simple text file that describes an object in the image. It also allows
|
||||
for the optional saving of a data.yaml file, used in YOLOv5, that contains metadata about the dataset.
|
||||
|
||||
The method allows filtering the detections based on their area percentage and offers an option for polygon approximation.
|
||||
Exports the dataset to YOLO format. This method saves the images and their corresponding
|
||||
annotations in YOLO format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
@@ -309,11 +333,12 @@ class DetectionDataset(BaseDataset):
|
||||
data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
|
||||
If not provided, the file will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
the image area for a detection to be included.
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). This is useful for simplifying the annotations.
|
||||
in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
|
||||
segmentation datasets.
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
@@ -330,3 +355,288 @@ class DetectionDataset(BaseDataset):
|
||||
)
|
||||
if data_yaml_path is not None:
|
||||
save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)
|
||||
|
||||
@classmethod
|
||||
def from_coco(
|
||||
cls,
|
||||
images_directory_path: str,
|
||||
annotations_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> DetectionDataset:
|
||||
"""
|
||||
Creates a Dataset instance from COCO formatted data.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_path (str): The path to the json annotation files.
|
||||
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
|
||||
|
||||
Returns:
|
||||
DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("coco")
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path=f"{dataset.location}/train",
|
||||
... annotations_path=f"{dataset.location}/train/_annotations.coco.json",
|
||||
... )
|
||||
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
"""
|
||||
classes, images, annotations = load_coco_annotations(
|
||||
images_directory_path=images_directory_path,
|
||||
annotations_path=annotations_path,
|
||||
force_masks=force_masks,
|
||||
)
|
||||
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
||||
|
||||
def as_coco(
|
||||
self,
|
||||
images_directory_path: Optional[str] = None,
|
||||
annotations_path: Optional[str] = None,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to COCO format. This method saves the images and their corresponding
|
||||
annotations in COCO format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_path (Optional[str]): The path to COCO annotation file.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
|
||||
segmentation datasets.
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
)
|
||||
if annotations_path is not None:
|
||||
save_coco_annotations(
|
||||
annotation_path=annotations_path,
|
||||
images=self.images,
|
||||
annotations=self.annotations,
|
||||
classes=self.classes,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def merge(cls, dataset_list: List[DetectionDataset]) -> DetectionDataset:
|
||||
"""
|
||||
Merge a list of `DetectionDataset` objects into a single `DetectionDataset` object.
|
||||
|
||||
This method takes a list of `DetectionDataset` objects and combines their respective fields (`classes`, `images`,
|
||||
`annotations`) into a single `DetectionDataset` object.
|
||||
|
||||
Args:
|
||||
dataset_list (List[DetectionDataset]): A list of `DetectionDataset` objects to merge.
|
||||
|
||||
Returns:
|
||||
(DetectionDataset): A single `DetectionDataset` object containing the merged data from the input list.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
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']
|
||||
```
|
||||
"""
|
||||
merged_images, merged_annotations = {}, {}
|
||||
class_lists = [dataset.classes for dataset in dataset_list]
|
||||
merged_classes = merge_class_lists(class_lists=class_lists)
|
||||
|
||||
for dataset in dataset_list:
|
||||
class_index_mapping = build_class_index_mapping(
|
||||
source_classes=dataset.classes, target_classes=merged_classes
|
||||
)
|
||||
for image_name, image, detections in dataset:
|
||||
if image_name in merged_annotations:
|
||||
raise ValueError(
|
||||
f"Image name {image_name} is not unique across datasets."
|
||||
)
|
||||
|
||||
merged_images[image_name] = image
|
||||
merged_annotations[image_name] = map_detections_class_id(
|
||||
source_to_target_mapping=class_index_mapping,
|
||||
detections=detections,
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=merged_classes, images=merged_images, annotations=merged_annotations
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClassificationDataset(BaseDataset):
|
||||
"""
|
||||
Dataclass containing information about a classification dataset.
|
||||
|
||||
Attributes:
|
||||
classes (List[str]): List containing dataset class names.
|
||||
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
images: Dict[str, np.ndarray]
|
||||
annotations: Dict[str, Classifications]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.images)
|
||||
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
) -> Tuple[ClassificationDataset, ClassificationDataset]:
|
||||
"""
|
||||
Splits the dataset into two parts (training and testing) using the provided split_ratio.
|
||||
|
||||
Args:
|
||||
split_ratio (float, optional): The ratio of the training set to the entire dataset.
|
||||
random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting.
|
||||
|
||||
Returns:
|
||||
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing the training and testing datasets.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> cd = sv.ClassificationDataset(...)
|
||||
>>> train_cd, test_cd = cd.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
>>> len(train_cd), len(test_cd)
|
||||
(700, 300)
|
||||
```
|
||||
"""
|
||||
image_names = list(self.images.keys())
|
||||
train_names, test_names = train_test_split(
|
||||
data=image_names,
|
||||
train_ratio=split_ratio,
|
||||
random_state=random_state,
|
||||
shuffle=shuffle,
|
||||
)
|
||||
|
||||
train_dataset = ClassificationDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in train_names},
|
||||
annotations={name: self.annotations[name] for name in train_names},
|
||||
)
|
||||
test_dataset = ClassificationDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in test_names},
|
||||
annotations={name: self.annotations[name] for name in test_names},
|
||||
)
|
||||
return train_dataset, test_dataset
|
||||
|
||||
def as_folder_structure(self, root_directory_path: str) -> None:
|
||||
"""
|
||||
Saves the dataset as a multi-class folder structure.
|
||||
|
||||
Args:
|
||||
root_directory_path (str): The path to the directory where the dataset will be saved.
|
||||
"""
|
||||
os.makedirs(root_directory_path, exist_ok=True)
|
||||
|
||||
for class_name in self.classes:
|
||||
os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
|
||||
|
||||
for image_name in self.images:
|
||||
classification = self.annotations[image_name]
|
||||
image = self.images[image_name]
|
||||
class_id = (
|
||||
classification.class_id[0]
|
||||
if classification.confidence is None
|
||||
else classification.get_top_k(1)[0][0]
|
||||
)
|
||||
class_name = self.classes[class_id]
|
||||
image_path = os.path.join(root_directory_path, class_name, image_name)
|
||||
cv2.imwrite(image_path, image)
|
||||
|
||||
@classmethod
|
||||
def from_folder_structure(cls, root_directory_path: str) -> ClassificationDataset:
|
||||
"""
|
||||
Load data from a multiclass folder structure into a ClassificationDataset.
|
||||
|
||||
Args:
|
||||
root_directory_path (str): The path to the dataset directory.
|
||||
|
||||
Returns:
|
||||
ClassificationDataset: The dataset.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("folder")
|
||||
|
||||
>>> cd = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path=f"{dataset.location}/train"
|
||||
... )
|
||||
```
|
||||
"""
|
||||
classes = os.listdir(root_directory_path)
|
||||
classes = sorted(set(classes))
|
||||
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
for class_name in classes:
|
||||
class_id = classes.index(class_name)
|
||||
|
||||
for image in os.listdir(os.path.join(root_directory_path, class_name)):
|
||||
image_dir = os.path.join(root_directory_path, class_name, image)
|
||||
images[image] = cv2.imread(image_dir)
|
||||
annotations[image] = Classifications(
|
||||
class_id=np.array([class_id]),
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=classes,
|
||||
images=images,
|
||||
annotations=annotations,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.utils import (
|
||||
approximate_mask_with_polygons,
|
||||
map_detections_class_id,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_mask
|
||||
from supervision.utils.file import read_json_file, save_json_file
|
||||
|
||||
|
||||
def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]:
|
||||
return [
|
||||
category["name"]
|
||||
for category in sorted(coco_categories, key=lambda category: category["id"])
|
||||
]
|
||||
|
||||
|
||||
def build_coco_class_index_mapping(
|
||||
coco_categories: List[dict], target_classes: List[str]
|
||||
) -> Dict[int, int]:
|
||||
source_class_to_index = {
|
||||
category["name"]: category["id"] for category in coco_categories
|
||||
}
|
||||
return {
|
||||
source_class_to_index[target_class_name]: target_class_index
|
||||
for target_class_index, target_class_name in enumerate(target_classes)
|
||||
}
|
||||
|
||||
|
||||
def classes_to_coco_categories(classes: List[str]) -> List[dict]:
|
||||
return [
|
||||
{
|
||||
"id": class_id,
|
||||
"name": class_name,
|
||||
"supercategory": "common-objects",
|
||||
}
|
||||
for class_id, class_name in enumerate(classes)
|
||||
]
|
||||
|
||||
|
||||
def group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict],
|
||||
) -> Dict[int, List[dict]]:
|
||||
annotations = {}
|
||||
for annotation in coco_annotations:
|
||||
image_id = annotation["image_id"]
|
||||
if image_id not in annotations:
|
||||
annotations[image_id] = []
|
||||
annotations[image_id].append(annotation)
|
||||
return annotations
|
||||
|
||||
|
||||
def _polygons_to_masks(
|
||||
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
|
||||
) -> np.ndarray:
|
||||
return np.array(
|
||||
[
|
||||
polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
|
||||
for polygon in polygons
|
||||
],
|
||||
dtype=bool,
|
||||
)
|
||||
|
||||
|
||||
def coco_annotations_to_detections(
|
||||
image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool
|
||||
) -> Detections:
|
||||
if not image_annotations:
|
||||
return Detections.empty()
|
||||
|
||||
class_ids = [
|
||||
image_annotation["category_id"] for image_annotation in image_annotations
|
||||
]
|
||||
xyxy = [image_annotation["bbox"] for image_annotation in image_annotations]
|
||||
xyxy = np.asarray(xyxy)
|
||||
xyxy[:, 2:4] += xyxy[:, 0:2]
|
||||
|
||||
if with_masks:
|
||||
polygons = [
|
||||
np.reshape(
|
||||
np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2)
|
||||
)
|
||||
for image_annotation in image_annotations
|
||||
]
|
||||
mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
|
||||
return Detections(
|
||||
class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
|
||||
)
|
||||
|
||||
return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
|
||||
|
||||
|
||||
def detections_to_coco_annotations(
|
||||
detections: Detections,
|
||||
image_id: int,
|
||||
annotation_id: int,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> Tuple[List[Dict], int]:
|
||||
coco_annotations = []
|
||||
for xyxy, mask, _, class_id, _ in detections:
|
||||
box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
|
||||
polygon = []
|
||||
if mask is not None:
|
||||
polygon = list(
|
||||
approximate_mask_with_polygons(
|
||||
mask=mask,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)[0].flatten()
|
||||
)
|
||||
coco_annotation = {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": int(class_id),
|
||||
"bbox": [xyxy[0], xyxy[1], box_width, box_height],
|
||||
"area": box_width * box_height,
|
||||
"segmentation": polygon,
|
||||
"iscrowd": 0,
|
||||
}
|
||||
coco_annotations.append(coco_annotation)
|
||||
annotation_id += 1
|
||||
return coco_annotations, annotation_id
|
||||
|
||||
|
||||
def load_coco_annotations(
|
||||
images_directory_path: str,
|
||||
annotations_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
coco_data = read_json_file(file_path=annotations_path)
|
||||
classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
|
||||
class_index_mapping = build_coco_class_index_mapping(
|
||||
coco_categories=coco_data["categories"], target_classes=classes
|
||||
)
|
||||
coco_images = coco_data["images"]
|
||||
coco_annotations_groups = group_coco_annotations_by_image_id(
|
||||
coco_annotations=coco_data["annotations"]
|
||||
)
|
||||
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
for coco_image in coco_images:
|
||||
image_name, image_width, image_height = (
|
||||
coco_image["file_name"],
|
||||
coco_image["width"],
|
||||
coco_image["height"],
|
||||
)
|
||||
image_annotations = coco_annotations_groups.get(coco_image["id"], [])
|
||||
image_path = os.path.join(images_directory_path, image_name)
|
||||
|
||||
image = cv2.imread(str(image_path))
|
||||
annotation = coco_annotations_to_detections(
|
||||
image_annotations=image_annotations,
|
||||
resolution_wh=(image_width, image_height),
|
||||
with_masks=force_masks,
|
||||
)
|
||||
annotation = map_detections_class_id(
|
||||
source_to_target_mapping=class_index_mapping,
|
||||
detections=annotation,
|
||||
)
|
||||
|
||||
images[image_name] = image
|
||||
annotations[image_name] = annotation
|
||||
|
||||
return classes, images, annotations
|
||||
|
||||
|
||||
def save_coco_annotations(
|
||||
annotation_path: str,
|
||||
images: Dict[str, np.ndarray],
|
||||
annotations: Dict[str, Detections],
|
||||
classes: List[str],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> None:
|
||||
Path(annotation_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
info = {}
|
||||
licenses = [
|
||||
{
|
||||
"id": 1,
|
||||
"url": "https://creativecommons.org/licenses/by/4.0/",
|
||||
"name": "CC BY 4.0",
|
||||
}
|
||||
]
|
||||
|
||||
coco_annotations = []
|
||||
coco_images = []
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
|
||||
image_id, annotation_id = 1, 1
|
||||
for image_name, image in images.items():
|
||||
image_height, image_width, _ = image.shape
|
||||
|
||||
coco_image = {
|
||||
"id": image_id,
|
||||
"license": 1,
|
||||
"file_name": image_name,
|
||||
"height": image_height,
|
||||
"width": image_width,
|
||||
"date_captured": datetime.now().strftime("%m/%d/%Y,%H:%M:%S"),
|
||||
}
|
||||
|
||||
coco_images.append(coco_image)
|
||||
detections = annotations[image_name]
|
||||
|
||||
coco_annotation, label_id = detections_to_coco_annotations(
|
||||
detections=detections,
|
||||
image_id=image_id,
|
||||
annotation_id=annotation_id,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
coco_annotations.extend(coco_annotation)
|
||||
image_id += 1
|
||||
|
||||
annotation_dict = {
|
||||
"info": info,
|
||||
"licenses": licenses,
|
||||
"categories": coco_categories,
|
||||
"images": coco_images,
|
||||
"annotations": coco_annotations,
|
||||
}
|
||||
save_json_file(annotation_dict, file_path=annotation_path)
|
||||
@@ -4,7 +4,7 @@ from xml.etree.ElementTree import Element, SubElement, parse, tostring
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.ultils import approximate_mask_with_polygons
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_xyxy
|
||||
|
||||
|
||||
@@ -4,12 +4,17 @@ from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
from supervision.dataset.ultils import approximate_mask_with_polygons
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
|
||||
from supervision.file import list_files_with_extensions, read_txt_file, save_text_file
|
||||
from supervision.utils.file import (
|
||||
list_files_with_extensions,
|
||||
read_txt_file,
|
||||
read_yaml_file,
|
||||
save_text_file,
|
||||
save_yaml_file,
|
||||
)
|
||||
|
||||
|
||||
def _parse_box(values: List[str]) -> np.ndarray:
|
||||
@@ -47,37 +52,15 @@ def _polygons_to_masks(
|
||||
)
|
||||
|
||||
|
||||
def _pair_image_with_annotation(
|
||||
image_paths: List[Union[str, Path]], annotation_paths: List[Union[str, Path]]
|
||||
) -> List[Tuple[Union[str, Path], Union[str, Path]]]:
|
||||
image_dict = {p.stem: p for p in image_paths}
|
||||
return [
|
||||
(image_dict[annotation_path.stem], annotation_path)
|
||||
for annotation_path in annotation_paths
|
||||
if annotation_path.stem in image_dict
|
||||
]
|
||||
|
||||
|
||||
def _with_mask(lines: List[str]) -> bool:
|
||||
return any([len(line.split()) > 5 for line in lines])
|
||||
|
||||
|
||||
def _extract_class_names(file_path: str) -> List[str]:
|
||||
names = []
|
||||
|
||||
with open(file_path, "r") as file:
|
||||
lines = file.readlines()
|
||||
for line in lines:
|
||||
if line.strip().startswith("names:"):
|
||||
start_index = lines.index(line) + 1
|
||||
break
|
||||
|
||||
for line in lines[start_index:]:
|
||||
if line.strip().startswith("-"):
|
||||
names.append(line.strip().replace("-", "").strip())
|
||||
else:
|
||||
break
|
||||
|
||||
data = read_yaml_file(file_path=file_path)
|
||||
names = data["names"]
|
||||
if isinstance(names, dict):
|
||||
names = [names[key] for key in sorted(names.keys())]
|
||||
return names
|
||||
|
||||
|
||||
@@ -146,14 +129,21 @@ def load_yolo_annotations(
|
||||
annotation_paths = list_files_with_extensions(
|
||||
directory=annotations_directory_path, extensions=["txt"]
|
||||
)
|
||||
path_pairs = _pair_image_with_annotation(
|
||||
image_paths=image_paths, annotation_paths=annotation_paths
|
||||
)
|
||||
|
||||
classes = _extract_class_names(file_path=data_yaml_path)
|
||||
images = {}
|
||||
annotations = {}
|
||||
for image_path, annotation_path in path_pairs:
|
||||
|
||||
for image_path in image_paths:
|
||||
image_name = Path(image_path).stem
|
||||
image = cv2.imread(str(image_path))
|
||||
|
||||
annotation_path = os.path.join(annotations_directory_path, f"{image_name}.txt")
|
||||
if not os.path.exists(annotation_path):
|
||||
images[image_path.name] = image
|
||||
annotations[image_path.name] = Detections.empty()
|
||||
continue
|
||||
|
||||
lines = read_txt_file(str(annotation_path))
|
||||
h, w, _ = image.shape
|
||||
resolution_wh = (w, h)
|
||||
@@ -252,5 +242,4 @@ def save_yolo_annotations(
|
||||
def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
|
||||
data = {"nc": len(classes), "names": classes}
|
||||
Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
with open(data_yaml_path, "w") as outfile:
|
||||
yaml.dump(data, outfile, default_flow_style=False)
|
||||
save_yaml_file(data=data, file_path=data_yaml_path)
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import copy
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
@@ -6,6 +7,7 @@ from typing import Dict, List, Optional, Tuple, TypeVar
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import (
|
||||
approximate_polygon,
|
||||
filter_polygons_by_area,
|
||||
@@ -43,6 +45,53 @@ def approximate_mask_with_polygons(
|
||||
]
|
||||
|
||||
|
||||
def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
|
||||
unique_classes = set()
|
||||
|
||||
for class_list in class_lists:
|
||||
for class_name in class_list:
|
||||
unique_classes.add(class_name.lower())
|
||||
|
||||
return sorted(list(unique_classes))
|
||||
|
||||
|
||||
def build_class_index_mapping(
|
||||
source_classes: List[str], target_classes: List[str]
|
||||
) -> Dict[int, int]:
|
||||
index_mapping = {}
|
||||
|
||||
for i, class_name in enumerate(source_classes):
|
||||
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."
|
||||
)
|
||||
corresponding_index = target_classes.index(class_name)
|
||||
index_mapping[i] = corresponding_index
|
||||
|
||||
return index_mapping
|
||||
|
||||
|
||||
def map_detections_class_id(
|
||||
source_to_target_mapping: Dict[int, int], detections: Detections
|
||||
) -> Detections:
|
||||
if detections.class_id is None:
|
||||
raise ValueError("Detections must have class_id attribute.")
|
||||
if set(np.unique(detections.class_id)) - set(source_to_target_mapping.keys()):
|
||||
raise ValueError(
|
||||
"Detections class_id must be a subset of source_to_target_mapping keys."
|
||||
)
|
||||
|
||||
detections_copy = copy.deepcopy(detections)
|
||||
|
||||
if len(detections) > 0:
|
||||
detections_copy.class_id = np.vectorize(source_to_target_mapping.get)(
|
||||
detections_copy.class_id
|
||||
)
|
||||
|
||||
return detections_copy
|
||||
|
||||
|
||||
def save_dataset_images(
|
||||
images_directory_path: str, images: Dict[str, np.ndarray]
|
||||
) -> None:
|
||||
@@ -1,13 +1,18 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import astuple, dataclass
|
||||
from typing import Any, Iterator, List, Optional, Tuple
|
||||
from typing import Any, Iterator, List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.utils import non_max_suppression, xywh_to_xyxy
|
||||
from supervision.detection.utils import (
|
||||
extract_yolov8_masks,
|
||||
non_max_suppression,
|
||||
process_roboflow_result,
|
||||
xywh_to_xyxy,
|
||||
)
|
||||
from supervision.geometry.core import Position
|
||||
from supervision.internal import deprecated
|
||||
|
||||
|
||||
def _validate_xyxy(xyxy: Any, n: int) -> None:
|
||||
@@ -191,6 +196,7 @@ class Detections:
|
||||
xyxy=yolov8_results.boxes.xyxy.cpu().numpy(),
|
||||
confidence=yolov8_results.boxes.conf.cpu().numpy(),
|
||||
class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
|
||||
mask=extract_yolov8_masks(yolov8_results),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -222,6 +228,35 @@ class Detections:
|
||||
class_id=yolo_nas_results.prediction.labels.astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_mmdetection(cls, mmdet_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a [mmdetection](https://github.com/open-mmlab/mmdetection) inference result.
|
||||
Also supported for [mmyolo](https://github.com/open-mmlab/mmyolo)
|
||||
|
||||
Args:
|
||||
mmdet_results (mmdet.structures.DetDataSample): The output Results instance from MMDetection
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
>>> from mmdet.apis import DetInferencer
|
||||
|
||||
>>> inferencer = DetInferencer(model_name, checkpoint, device)
|
||||
>>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output', return_datasample=True)["predictions"][0]
|
||||
>>> detections = sv.Detections.from_mmdet(mmdet_result)
|
||||
```
|
||||
"""
|
||||
return cls(
|
||||
xyxy=mmdet_results.pred_instances.bboxes.cpu().numpy(),
|
||||
confidence=mmdet_results.pred_instances.scores.cpu().numpy(),
|
||||
class_id=mmdet_results.pred_instances.labels.cpu().numpy().astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_transformers(cls, transformers_results: dict) -> Detections:
|
||||
"""
|
||||
@@ -307,27 +342,14 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_roboflow(roboflow_result, class_list)
|
||||
```
|
||||
"""
|
||||
xyxy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
|
||||
for prediction in roboflow_result["predictions"]:
|
||||
x = prediction["x"]
|
||||
y = prediction["y"]
|
||||
width = prediction["width"]
|
||||
height = prediction["height"]
|
||||
x_min = x - width / 2
|
||||
y_min = y - height / 2
|
||||
x_max = x_min + width
|
||||
y_max = y_min + height
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(class_list.index(prediction["class"]))
|
||||
confidence.append(prediction["confidence"])
|
||||
|
||||
xyxy, confidence, class_id, masks = process_roboflow_result(
|
||||
roboflow_result=roboflow_result, class_list=class_list
|
||||
)
|
||||
return Detections(
|
||||
xyxy=np.array(xyxy),
|
||||
confidence=np.array(confidence),
|
||||
class_id=np.array(class_id).astype(int),
|
||||
xyxy=xyxy,
|
||||
confidence=confidence,
|
||||
class_id=class_id,
|
||||
mask=masks,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@@ -361,20 +383,6 @@ class Detections:
|
||||
|
||||
return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"Dataset loading and saving is going to be executed by supervision.dataset.core.Dataset"
|
||||
)
|
||||
def from_coco_annotations(cls, coco_annotation: dict) -> Detections:
|
||||
xyxy, class_id = [], []
|
||||
|
||||
for annotation in coco_annotation:
|
||||
x_min, y_min, width, height = annotation["bbox"]
|
||||
xyxy.append([x_min, y_min, x_min + width, y_min + height])
|
||||
class_id.append(annotation["category_id"])
|
||||
|
||||
return cls(xyxy=np.array(xyxy), class_id=np.array(class_id))
|
||||
|
||||
@classmethod
|
||||
def empty(cls) -> Detections:
|
||||
"""
|
||||
@@ -469,23 +477,43 @@ class Detections:
|
||||
|
||||
raise ValueError(f"{anchor} is not supported.")
|
||||
|
||||
def __getitem__(self, index: np.ndarray) -> Detections:
|
||||
if isinstance(index, np.ndarray) and (
|
||||
index.dtype == bool or index.dtype == int
|
||||
):
|
||||
return Detections(
|
||||
xyxy=self.xyxy[index],
|
||||
mask=self.mask[index] if self.mask is not None else None,
|
||||
confidence=self.confidence[index]
|
||||
if self.confidence is not None
|
||||
else None,
|
||||
class_id=self.class_id[index] if self.class_id is not None else None,
|
||||
tracker_id=self.tracker_id[index]
|
||||
if self.tracker_id is not None
|
||||
else None,
|
||||
)
|
||||
raise TypeError(
|
||||
f"Detections.__getitem__ not supported for index of type {type(index)}."
|
||||
def __getitem__(
|
||||
self, index: Union[int, slice, List[int], np.ndarray]
|
||||
) -> Detections:
|
||||
"""
|
||||
Get a subset of the Detections object.
|
||||
|
||||
Args:
|
||||
index (Union[int, slice, List[int], np.ndarray]): The index or indices of the subset of the Detections
|
||||
|
||||
Returns:
|
||||
(Detections): A subset of the Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> first_detection = detections[0]
|
||||
|
||||
>>> first_10_detections = detections[0:10]
|
||||
|
||||
>>> some_detections = detections[[0, 2, 4]]
|
||||
|
||||
>>> class_0_detections = detections[detections.class_id == 0]
|
||||
|
||||
>>> high_confidence_detections = detections[detections.confidence > 0.5]
|
||||
```
|
||||
"""
|
||||
if isinstance(index, int):
|
||||
index = [index]
|
||||
return Detections(
|
||||
xyxy=self.xyxy[index],
|
||||
mask=self.mask[index] if self.mask is not None else None,
|
||||
confidence=self.confidence[index] if self.confidence is not None else None,
|
||||
class_id=self.class_id[index] if self.class_id is not None else None,
|
||||
tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
|
||||
)
|
||||
|
||||
@property
|
||||
|
||||
@@ -17,7 +17,7 @@ class PolygonZone:
|
||||
A class for defining a polygon-shaped zone within a frame for detecting objects.
|
||||
|
||||
Attributes:
|
||||
polygon (np.ndarray): A numpy array defining the polygon vertices
|
||||
polygon (np.ndarray): A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points.
|
||||
frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height)
|
||||
triggering_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER)
|
||||
current_count (int): The current count of detected objects within the zone
|
||||
@@ -30,7 +30,7 @@ class PolygonZone:
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
triggering_position: Position = Position.BOTTOM_CENTER,
|
||||
):
|
||||
self.polygon = polygon
|
||||
self.polygon = polygon.astype(int)
|
||||
self.frame_resolution_wh = frame_resolution_wh
|
||||
self.triggering_position = triggering_position
|
||||
self.current_count = 0
|
||||
|
||||
@@ -258,3 +258,83 @@ def approximate_polygon(
|
||||
break
|
||||
|
||||
return np.squeeze(approximated_points, axis=1)
|
||||
|
||||
|
||||
def extract_yolov8_masks(yolov8_results) -> Optional[np.ndarray]:
|
||||
if not yolov8_results.masks:
|
||||
return None
|
||||
|
||||
orig_shape = yolov8_results.orig_shape
|
||||
inference_shape = tuple(yolov8_results.masks.data.shape[1:])
|
||||
|
||||
gain = 0
|
||||
pad = (0, 0)
|
||||
|
||||
if inference_shape != orig_shape:
|
||||
gain = min(
|
||||
inference_shape[0] / orig_shape[0],
|
||||
inference_shape[1] / orig_shape[1],
|
||||
)
|
||||
pad = (
|
||||
(inference_shape[1] - orig_shape[1] * gain) / 2,
|
||||
(inference_shape[0] - orig_shape[0] * gain) / 2,
|
||||
)
|
||||
|
||||
top, left = int(pad[1]), int(pad[0])
|
||||
bottom, right = int(inference_shape[0] - pad[1]), int(inference_shape[1] - pad[0])
|
||||
|
||||
mask_maps = []
|
||||
masks = yolov8_results.masks.data.cpu().numpy()
|
||||
for i in range(masks.shape[0]):
|
||||
mask = masks[i]
|
||||
mask = mask[top:bottom, left:right]
|
||||
mask = cv2.resize(mask, (orig_shape[1], orig_shape[0]))
|
||||
mask_maps.append(mask)
|
||||
|
||||
return np.asarray(mask_maps, dtype=bool)
|
||||
|
||||
|
||||
def process_roboflow_result(
|
||||
roboflow_result: dict, class_list: List[str]
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
|
||||
if not roboflow_result["predictions"]:
|
||||
return np.empty((0, 4)), np.empty(0), np.empty(0), None
|
||||
|
||||
xyxy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
masks = []
|
||||
|
||||
image_width = int(roboflow_result["image"]["width"])
|
||||
image_height = int(roboflow_result["image"]["height"])
|
||||
|
||||
for prediction in roboflow_result["predictions"]:
|
||||
x = prediction["x"]
|
||||
y = prediction["y"]
|
||||
width = prediction["width"]
|
||||
height = prediction["height"]
|
||||
x_min = x - width / 2
|
||||
y_min = y - height / 2
|
||||
x_max = x_min + width
|
||||
y_max = y_min + height
|
||||
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(class_list.index(prediction["class"]))
|
||||
confidence.append(prediction["confidence"])
|
||||
|
||||
if "points" not in prediction:
|
||||
continue
|
||||
|
||||
polygon = np.array(
|
||||
[[point["x"], point["y"]] for point in prediction["points"]], dtype=int
|
||||
)
|
||||
|
||||
mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height))
|
||||
masks.append(mask)
|
||||
|
||||
xyxy = np.array(xyxy)
|
||||
confidence = np.array(confidence)
|
||||
class_id = np.array(class_id).astype(int)
|
||||
masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
|
||||
|
||||
return xyxy, confidence, class_id, masks
|
||||
|
||||
@@ -1,68 +0,0 @@
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
|
||||
def list_files_with_extensions(
|
||||
directory: Union[str, Path], extensions: Optional[List[str]] = None
|
||||
) -> List[Path]:
|
||||
"""
|
||||
List files in a directory with specified extensions or all files if no extensions are provided.
|
||||
|
||||
Args:
|
||||
directory (Union[str, Path]): The directory path as a string or Path object.
|
||||
extensions (Optional[List[str]]): A list of file extensions to filter. Default is None, which lists all files.
|
||||
|
||||
Returns:
|
||||
(List[Path]): A list of Path objects for the matching files.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> from supervision import list_files_with_extensions
|
||||
|
||||
>>> # List all files in the directory
|
||||
>>> files = list_files_with_extensions(directory='my_directory')
|
||||
|
||||
>>> # List only files with '.txt' and '.md' extensions
|
||||
>>> files = list_files_with_extensions(directory='my_directory', extensions=['txt', 'md'])
|
||||
```
|
||||
"""
|
||||
directory = Path(directory)
|
||||
files_with_extensions = []
|
||||
|
||||
if extensions is not None:
|
||||
for ext in extensions:
|
||||
files_with_extensions.extend(directory.glob(f"*.{ext}"))
|
||||
else:
|
||||
files_with_extensions.extend(directory.glob("*"))
|
||||
|
||||
return files_with_extensions
|
||||
|
||||
|
||||
def read_txt_file(file_path: str) -> List[str]:
|
||||
"""
|
||||
Read a text file and return a list of strings without newline characters.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the text file.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of strings representing the lines in the text file.
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
lines = file.readlines()
|
||||
lines = [line.rstrip("\n") for line in lines]
|
||||
|
||||
return lines
|
||||
|
||||
|
||||
def save_text_file(lines: List[str], file_path: str):
|
||||
"""
|
||||
Write a list of strings to a text file, each string on a new line.
|
||||
|
||||
Args:
|
||||
lines (List[str]): The list of strings to be written to the file.
|
||||
file_path (str): The path to the text file.
|
||||
"""
|
||||
with open(file_path, "w") as file:
|
||||
for line in lines:
|
||||
file.write(line + "\n")
|
||||
@@ -0,0 +1,458 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, List, Optional, Tuple
|
||||
|
||||
import matplotlib
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.core import DetectionDataset
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import box_iou_batch
|
||||
|
||||
|
||||
@dataclass
|
||||
class ConfusionMatrix:
|
||||
"""
|
||||
Confusion matrix for object detection tasks.
|
||||
|
||||
Attributes:
|
||||
matrix (np.ndarray): An 2D `np.ndarray` of shape `(len(classes) + 1, len(classes) + 1)` containing the number of `TP`, `FP`, `FN` and `TN` for each class.
|
||||
classes (List[str]): Model class names.
|
||||
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded from the matrix.
|
||||
iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`.
|
||||
"""
|
||||
|
||||
matrix: np.ndarray
|
||||
classes: List[str]
|
||||
conf_threshold: float
|
||||
iou_threshold: float
|
||||
|
||||
@classmethod
|
||||
def from_detections(
|
||||
cls,
|
||||
predictions: List[Detections],
|
||||
targets: List[Detections],
|
||||
classes: List[str],
|
||||
conf_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> ConfusionMatrix:
|
||||
"""
|
||||
Calculate confusion matrix based on predicted and ground-truth detections.
|
||||
|
||||
Args:
|
||||
targets (List[Detections]): Detections objects from ground-truth.
|
||||
predictions (List[Detections]): Detections objects predicted by the model.
|
||||
classes (List[str]): Model class names.
|
||||
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
|
||||
iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`.
|
||||
|
||||
Returns:
|
||||
ConfusionMatrix: New instance of ConfusionMatrix.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> targets = [
|
||||
... sv.Detections(...),
|
||||
... sv.Detections(...)
|
||||
... ]
|
||||
|
||||
>>> predictions = [
|
||||
... sv.Detections(...),
|
||||
... sv.Detections(...)
|
||||
... ]
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.from_detections(
|
||||
... predictions=predictions,
|
||||
... targets=target,
|
||||
... classes=['person', ...]
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
"""
|
||||
|
||||
prediction_tensors = []
|
||||
target_tensors = []
|
||||
for prediction, target in zip(predictions, targets):
|
||||
prediction_tensors.append(
|
||||
ConfusionMatrix.detections_to_tensor(prediction, with_confidence=True)
|
||||
)
|
||||
target_tensors.append(
|
||||
ConfusionMatrix.detections_to_tensor(target, with_confidence=False)
|
||||
)
|
||||
return cls.from_tensors(
|
||||
predictions=prediction_tensors,
|
||||
targets=target_tensors,
|
||||
classes=classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def detections_to_tensor(
|
||||
detections: Detections, with_confidence: bool = False
|
||||
) -> np.ndarray:
|
||||
if detections.class_id is None:
|
||||
raise ValueError(
|
||||
"ConfusionMatrix can only be calculated for Detections with class_id"
|
||||
)
|
||||
|
||||
arrays_to_concat = [detections.xyxy, np.expand_dims(detections.class_id, 1)]
|
||||
|
||||
if with_confidence:
|
||||
if detections.confidence is None:
|
||||
raise ValueError(
|
||||
"ConfusionMatrix can only be calculated for Detections with confidence"
|
||||
)
|
||||
arrays_to_concat.append(np.expand_dims(detections.confidence, 1))
|
||||
|
||||
return np.concatenate(arrays_to_concat, axis=1)
|
||||
|
||||
@classmethod
|
||||
def from_tensors(
|
||||
cls,
|
||||
predictions: List[np.ndarray],
|
||||
targets: List[np.ndarray],
|
||||
classes: List[str],
|
||||
conf_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> ConfusionMatrix:
|
||||
"""
|
||||
Calculate confusion matrix based on predicted and ground-truth detections.
|
||||
|
||||
Args:
|
||||
predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format.
|
||||
targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format.
|
||||
classes (List[str]): Model class names.
|
||||
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
|
||||
iou_threshold (float): Detection iou threshold between `0` and `1`. Detections with lower iou will be classified as `FP`.
|
||||
|
||||
Returns:
|
||||
ConfusionMatrix: New instance of ConfusionMatrix.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> targets = (
|
||||
... [
|
||||
... array(
|
||||
... [
|
||||
... [0.0, 0.0, 3.0, 3.0, 1],
|
||||
... [2.0, 2.0, 5.0, 5.0, 1],
|
||||
... [6.0, 1.0, 8.0, 3.0, 2],
|
||||
... ]
|
||||
... ),
|
||||
... array([1.0, 1.0, 2.0, 2.0, 2]),
|
||||
... ]
|
||||
... )
|
||||
|
||||
>>> predictions = [
|
||||
... array(
|
||||
... [
|
||||
... [0.0, 0.0, 3.0, 3.0, 1, 0.9],
|
||||
... [0.1, 0.1, 3.0, 3.0, 0, 0.9],
|
||||
... [6.0, 1.0, 8.0, 3.0, 1, 0.8],
|
||||
... [1.0, 6.0, 2.0, 7.0, 1, 0.8],
|
||||
... ]
|
||||
... ),
|
||||
... array([[1.0, 1.0, 2.0, 2.0, 2, 0.8]])
|
||||
... ]
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.from_tensors(
|
||||
... predictions=predictions,
|
||||
... targets=targets,
|
||||
... classes=['person', ...]
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
"""
|
||||
cls._validate_input_tensors(predictions, targets)
|
||||
|
||||
num_classes = len(classes)
|
||||
matrix = np.zeros((num_classes + 1, num_classes + 1))
|
||||
for true_batch, detection_batch in zip(targets, predictions):
|
||||
matrix += cls.evaluate_detection_batch(
|
||||
predictions=detection_batch,
|
||||
targets=true_batch,
|
||||
num_classes=num_classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
return cls(
|
||||
matrix=matrix,
|
||||
classes=classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _validate_input_tensors(
|
||||
cls, predictions: List[np.ndarray], targets: List[np.ndarray]
|
||||
):
|
||||
"""
|
||||
Checks for shape consistency of input tensors.
|
||||
"""
|
||||
if len(predictions) != len(targets):
|
||||
raise ValueError(
|
||||
f"Number of predictions ({len(predictions)}) and targets ({len(targets)}) must be equal."
|
||||
)
|
||||
if len(predictions) > 0:
|
||||
if not isinstance(predictions[0], np.ndarray) or not isinstance(
|
||||
targets[0], np.ndarray
|
||||
):
|
||||
raise ValueError(
|
||||
f"Predictions and targets must be lists of numpy arrays. Got {type(predictions[0])} and {type(targets[0])} instead."
|
||||
)
|
||||
if predictions[0].shape[1] != 6:
|
||||
raise ValueError(
|
||||
f"Predictions must have shape (N, 6). Got {predictions[0].shape} instead."
|
||||
)
|
||||
if targets[0].shape[1] != 5:
|
||||
raise ValueError(
|
||||
f"Targets must have shape (N, 5). Got {targets[0].shape} instead."
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def evaluate_detection_batch(
|
||||
predictions: np.ndarray,
|
||||
targets: np.ndarray,
|
||||
num_classes: int,
|
||||
conf_threshold: float,
|
||||
iou_threshold: float,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Calculate confusion matrix for a batch of detections for a single image.
|
||||
|
||||
Args:
|
||||
predictions (List[np.ndarray]): Each element of the list describes a single image and has `shape = (M, 6)` where `M` is the number of detected objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class, conf)` format.
|
||||
targets (List[np.ndarray]): Each element of the list describes a single image and has `shape = (N, 5)` where `N` is the number of ground-truth objects. Each row is expected to be in `(x_min, y_min, x_max, y_max, class)` format.
|
||||
num_classes (int): Number of classes.
|
||||
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
|
||||
iou_threshold (float): Detection iou threshold between `0` and `1`. Detections with lower iou will be classified as `FP`.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Confusion matrix based on a single image.
|
||||
"""
|
||||
result_matrix = np.zeros((num_classes + 1, num_classes + 1))
|
||||
|
||||
conf_idx = 5
|
||||
confidence = predictions[:, conf_idx]
|
||||
detection_batch_filtered = predictions[confidence > conf_threshold]
|
||||
|
||||
class_id_idx = 4
|
||||
true_classes = np.array(targets[:, class_id_idx], dtype=np.int16)
|
||||
detection_classes = np.array(
|
||||
detection_batch_filtered[:, class_id_idx], dtype=np.int16
|
||||
)
|
||||
true_boxes = targets[:, :class_id_idx]
|
||||
detection_boxes = detection_batch_filtered[:, :class_id_idx]
|
||||
|
||||
iou_batch = box_iou_batch(
|
||||
boxes_true=true_boxes, boxes_detection=detection_boxes
|
||||
)
|
||||
matched_idx = np.asarray(iou_batch > iou_threshold).nonzero()
|
||||
|
||||
if matched_idx[0].shape[0]:
|
||||
matches = np.stack(
|
||||
(matched_idx[0], matched_idx[1], iou_batch[matched_idx]), axis=1
|
||||
)
|
||||
matches = ConfusionMatrix._drop_extra_matches(matches=matches)
|
||||
else:
|
||||
matches = np.zeros((0, 3))
|
||||
|
||||
matched_true_idx, matched_detection_idx, _ = matches.transpose().astype(
|
||||
np.int16
|
||||
)
|
||||
|
||||
for i, true_class_value in enumerate(true_classes):
|
||||
j = matched_true_idx == i
|
||||
if matches.shape[0] > 0 and sum(j) == 1:
|
||||
result_matrix[
|
||||
true_class_value, detection_classes[matched_detection_idx[j]]
|
||||
] += 1 # TP
|
||||
else:
|
||||
result_matrix[true_class_value, num_classes] += 1 # FN
|
||||
|
||||
for i, detection_class_value in enumerate(detection_classes):
|
||||
if not any(matched_detection_idx == i):
|
||||
result_matrix[num_classes, detection_class_value] += 1 # FP
|
||||
|
||||
return result_matrix
|
||||
|
||||
@staticmethod
|
||||
def _drop_extra_matches(matches: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Deduplicate matches. If there are multiple matches for the same true or predicted box,
|
||||
only the one with the highest IoU is kept.
|
||||
"""
|
||||
if matches.shape[0] > 0:
|
||||
matches = matches[matches[:, 2].argsort()[::-1]]
|
||||
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
|
||||
matches = matches[matches[:, 2].argsort()[::-1]]
|
||||
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
|
||||
return matches
|
||||
|
||||
@classmethod
|
||||
def benchmark(
|
||||
cls,
|
||||
dataset: DetectionDataset,
|
||||
callback: Callable[[np.ndarray], Detections],
|
||||
conf_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.5,
|
||||
) -> ConfusionMatrix:
|
||||
"""
|
||||
Create confusion matrix from dataset and callback function.
|
||||
|
||||
Args:
|
||||
dataset (DetectionDataset): Object detection dataset used for evaluation.
|
||||
callback (Callable[[np.ndarray], Detections]): Function that takes an image as input and returns Detections object.
|
||||
conf_threshold (float): Detection confidence threshold between `0` and `1`. Detections with lower confidence will be excluded.
|
||||
iou_threshold (float): Detection IoU threshold between `0` and `1`. Detections with lower IoU will be classified as `FP`.
|
||||
|
||||
Returns:
|
||||
ConfusionMatrix: New instance of ConfusionMatrix.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(...)
|
||||
|
||||
>>> model = YOLO(...)
|
||||
>>> def callback(image: np.ndarray) -> sv.Detections:
|
||||
... result = model(image)[0]
|
||||
... return sv.Detections.from_yolov8(result)
|
||||
|
||||
>>> confusion_matrix = sv.ConfusionMatrix.benchmark(
|
||||
... dataset = dataset,
|
||||
... callback = callback
|
||||
... )
|
||||
|
||||
>>> confusion_matrix.matrix
|
||||
array([
|
||||
[0., 0., 0., 0.],
|
||||
[0., 1., 0., 1.],
|
||||
[0., 1., 1., 0.],
|
||||
[1., 1., 0., 0.]
|
||||
])
|
||||
```
|
||||
"""
|
||||
predictions, targets = [], []
|
||||
for img_name, img in dataset.images.items():
|
||||
predictions_batch = callback(img)
|
||||
predictions.append(predictions_batch)
|
||||
targets_batch = dataset.annotations[img_name]
|
||||
targets.append(targets_batch)
|
||||
return cls.from_detections(
|
||||
predictions=predictions,
|
||||
targets=targets,
|
||||
classes=dataset.classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
def plot(
|
||||
self,
|
||||
save_path: Optional[str] = None,
|
||||
title: Optional[str] = None,
|
||||
classes: Optional[List[str]] = None,
|
||||
normalize: bool = False,
|
||||
fig_size: Tuple[int, int] = (12, 10),
|
||||
) -> matplotlib.figure.Figure:
|
||||
"""
|
||||
Create confusion matrix plot and save it at selected location.
|
||||
|
||||
Args:
|
||||
save_path (Optional[str]): Path to save the plot. If not provided, plot will be displayed.
|
||||
title (Optional[str]): Title of the plot.
|
||||
classes (Optional[List[str]]): List of classes to be displayed on the plot. If not provided, all classes will be displayed.
|
||||
normalize (bool): If True, normalize the confusion matrix.
|
||||
fig_size (Tuple[int, int]): Size of the plot.
|
||||
|
||||
Returns:
|
||||
matplotlib.figure.Figure: Confusion matrix plot.
|
||||
"""
|
||||
|
||||
array = self.matrix.copy()
|
||||
|
||||
if normalize:
|
||||
eps = 1e-8
|
||||
array = array / (array.sum(0).reshape(1, -1) + eps)
|
||||
|
||||
array[array < 0.005] = np.nan
|
||||
|
||||
fig, ax = plt.subplots(figsize=fig_size, tight_layout=True, facecolor="white")
|
||||
|
||||
class_names = classes if classes is not None else self.classes
|
||||
use_labels_for_ticks = class_names is not None and (0 < len(class_names) < 99)
|
||||
if use_labels_for_ticks:
|
||||
x_tick_labels = class_names + ["FN"]
|
||||
y_tick_labels = class_names + ["FP"]
|
||||
num_ticks = len(x_tick_labels)
|
||||
else:
|
||||
x_tick_labels = None
|
||||
y_tick_labels = None
|
||||
num_ticks = len(array)
|
||||
im = ax.imshow(array, cmap="Blues")
|
||||
|
||||
cbar = ax.figure.colorbar(im, ax=ax)
|
||||
cbar.mappable.set_clim(vmin=0, vmax=np.nanmax(array))
|
||||
|
||||
if x_tick_labels is None:
|
||||
tick_interval = 2
|
||||
else:
|
||||
tick_interval = 1
|
||||
ax.set_xticks(np.arange(0, num_ticks, tick_interval), labels=x_tick_labels)
|
||||
ax.set_yticks(np.arange(0, num_ticks, tick_interval), labels=y_tick_labels)
|
||||
|
||||
plt.setp(ax.get_xticklabels(), rotation=90, ha="right", rotation_mode="default")
|
||||
|
||||
labelsize = 10 if num_ticks < 50 else 8
|
||||
ax.tick_params(axis="both", which="both", labelsize=labelsize)
|
||||
|
||||
if num_ticks < 30:
|
||||
for i in range(array.shape[0]):
|
||||
for j in range(array.shape[1]):
|
||||
n_preds = array[i, j]
|
||||
if not np.isnan(n_preds):
|
||||
ax.text(
|
||||
j,
|
||||
i,
|
||||
f"{n_preds:.2f}" if normalize else f"{n_preds:.0f}",
|
||||
ha="center",
|
||||
va="center",
|
||||
color="black"
|
||||
if n_preds < 0.5 * np.nanmax(array)
|
||||
else "white",
|
||||
)
|
||||
|
||||
if title:
|
||||
ax.set_title(title, fontsize=20)
|
||||
|
||||
ax.set_xlabel("Predicted")
|
||||
ax.set_ylabel("True")
|
||||
ax.set_facecolor("white")
|
||||
if save_path:
|
||||
fig.savefig(
|
||||
save_path, dpi=250, facecolor=fig.get_facecolor(), transparent=True
|
||||
)
|
||||
return fig
|
||||
@@ -0,0 +1,140 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
|
||||
class NumpyJsonEncoder(json.JSONEncoder):
|
||||
def default(self, obj):
|
||||
if isinstance(obj, np.integer):
|
||||
return int(obj)
|
||||
if isinstance(obj, np.floating):
|
||||
return float(obj)
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
return super(NumpyJsonEncoder, self).default(obj)
|
||||
|
||||
|
||||
def list_files_with_extensions(
|
||||
directory: Union[str, Path], extensions: Optional[List[str]] = None
|
||||
) -> List[Path]:
|
||||
"""
|
||||
List files in a directory with specified extensions or all files if no extensions are provided.
|
||||
|
||||
Args:
|
||||
directory (Union[str, Path]): The directory path as a string or Path object.
|
||||
extensions (Optional[List[str]]): A list of file extensions to filter. Default is None, which lists all files.
|
||||
|
||||
Returns:
|
||||
(List[Path]): A list of Path objects for the matching files.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> # List all files in the directory
|
||||
>>> files = sv.list_files_with_extensions(directory='my_directory')
|
||||
|
||||
>>> # List only files with '.txt' and '.md' extensions
|
||||
>>> files = sv.list_files_with_extensions(directory='my_directory', extensions=['txt', 'md'])
|
||||
```
|
||||
"""
|
||||
directory = Path(directory)
|
||||
files_with_extensions = []
|
||||
|
||||
if extensions is not None:
|
||||
for ext in extensions:
|
||||
files_with_extensions.extend(directory.glob(f"*.{ext}"))
|
||||
else:
|
||||
files_with_extensions.extend(directory.glob("*"))
|
||||
|
||||
return files_with_extensions
|
||||
|
||||
|
||||
def read_txt_file(file_path: str) -> List[str]:
|
||||
"""
|
||||
Read a text file and return a list of strings without newline characters.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the text file.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of strings representing the lines in the text file.
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
lines = file.readlines()
|
||||
lines = [line.rstrip("\n") for line in lines]
|
||||
|
||||
return lines
|
||||
|
||||
|
||||
def save_text_file(lines: List[str], file_path: str):
|
||||
"""
|
||||
Write a list of strings to a text file, each string on a new line.
|
||||
|
||||
Args:
|
||||
lines (List[str]): The list of strings to be written to the file.
|
||||
file_path (str): The path to the text file.
|
||||
"""
|
||||
with open(file_path, "w") as file:
|
||||
for line in lines:
|
||||
file.write(line + "\n")
|
||||
|
||||
|
||||
def read_json_file(file_path: str) -> dict:
|
||||
"""
|
||||
Read a json file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the json file.
|
||||
|
||||
Returns:
|
||||
dict: A dict of annotations information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = json.load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_json_file(data: dict, file_path: str, indent: int = 3) -> None:
|
||||
"""
|
||||
Write a dict to a json file.
|
||||
|
||||
Args:
|
||||
indent:
|
||||
data (dict): dict with unique keys and value as pair.
|
||||
file_path (str): The path to the json file.
|
||||
"""
|
||||
with open(file_path, "w") as fp:
|
||||
json.dump(data, fp, cls=NumpyJsonEncoder, indent=indent)
|
||||
|
||||
|
||||
def read_yaml_file(file_path: str) -> dict:
|
||||
"""
|
||||
Read a yaml file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the yaml file.
|
||||
|
||||
Returns:
|
||||
dict: A dict of content information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = yaml.safe_load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_yaml_file(data: dict, file_path: str) -> None:
|
||||
"""
|
||||
Save a dict to a json file.
|
||||
|
||||
Args:
|
||||
indent:
|
||||
data (dict): dict with unique keys and value as pair.
|
||||
file_path (str): The path to the json file.
|
||||
"""
|
||||
|
||||
with open(file_path, "w") as outfile:
|
||||
yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)
|
||||
@@ -0,0 +1,93 @@
|
||||
import os
|
||||
import shutil
|
||||
from typing import Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Crops the given image based on the given bounding box.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The image to be cropped, represented as a numpy array.
|
||||
xyxy (np.ndarray): A numpy array containing the bounding box coordinates in the format (x1, y1, x2, y2).
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The cropped image as a numpy array.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> detection = sv.Detections(...)
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
|
||||
... for xyxy in detection.xyxy:
|
||||
... cropped_image = sv.crop(image=image, xyxy=xyxy)
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
"""
|
||||
|
||||
xyxy = np.round(xyxy).astype(int)
|
||||
x1, y1, x2, y2 = xyxy
|
||||
cropped_img = image[y1:y2, x1:x2]
|
||||
return cropped_img
|
||||
|
||||
|
||||
class ImageSink:
|
||||
def __init__(
|
||||
self,
|
||||
target_dir_path: str,
|
||||
overwrite: bool = False,
|
||||
image_name_pattern: str = "image_{:05d}.png",
|
||||
):
|
||||
"""
|
||||
Initialize a context manager for saving images.
|
||||
|
||||
Args:
|
||||
target_dir_path (str): The target directory where images will be saved.
|
||||
overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
|
||||
image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
|
||||
... for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
"""
|
||||
self.target_dir_path = target_dir_path
|
||||
self.overwrite = overwrite
|
||||
self.image_name_pattern = image_name_pattern
|
||||
self.image_count = 0
|
||||
|
||||
def __enter__(self):
|
||||
if os.path.exists(self.target_dir_path):
|
||||
if self.overwrite:
|
||||
shutil.rmtree(self.target_dir_path)
|
||||
os.makedirs(self.target_dir_path)
|
||||
else:
|
||||
os.makedirs(self.target_dir_path)
|
||||
|
||||
return self
|
||||
|
||||
def save_image(self, image: np.ndarray, image_name: Optional[str] = None):
|
||||
"""
|
||||
Save a given image in the target directory.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The image to be saved.
|
||||
image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
|
||||
"""
|
||||
if image_name is None:
|
||||
image_name = self.image_name_pattern.format(self.image_count)
|
||||
|
||||
image_path = os.path.join(self.target_dir_path, image_name)
|
||||
cv2.imwrite(image_path, image)
|
||||
self.image_count += 1
|
||||
|
||||
def __exit__(self, exc_type, exc_value, exc_traceback):
|
||||
pass
|
||||
@@ -20,9 +20,9 @@ class VideoInfo:
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> from supervision import VideoInfo
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> video_info = VideoInfo.from_video_path(video_path='video.mp4')
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='video.mp4')
|
||||
|
||||
>>> video_info
|
||||
VideoInfo(width=3840, height=2160, fps=25, total_frames=538)
|
||||
@@ -63,15 +63,15 @@ class VideoSink:
|
||||
target_path (str): The path to the output file where the video will be saved.
|
||||
video_info (VideoInfo): Information about the video resolution, fps, and total frame count.
|
||||
|
||||
Examples:
|
||||
Example:
|
||||
```python
|
||||
>>> from supervision import VideoInfo, VideoSink
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> video_info = VideoInfo.from_video_path(video_path='source_video.mp4')
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4')
|
||||
|
||||
>>> with VideoSink(target_path='target_video.mp4', video_info=video_info) as s:
|
||||
... frame = ...
|
||||
... s.write_frame(frame=frame)
|
||||
>>> with sv.VideoSink(target_path='target_video.mp4', video_info=video_info) as sink:
|
||||
... for frame in get_video_frames_generator(source_path='source_video.mp4', stride=2):
|
||||
... sink.write_frame(frame=frame)
|
||||
```
|
||||
"""
|
||||
|
||||
@@ -93,35 +93,58 @@ class VideoSink:
|
||||
def write_frame(self, frame: np.ndarray):
|
||||
self.__writer.write(frame)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
def __exit__(self, exc_type, exc_value, exc_traceback):
|
||||
self.__writer.release()
|
||||
|
||||
|
||||
def get_video_frames_generator(source_path: str) -> Generator[np.ndarray, None, None]:
|
||||
def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
|
||||
video = cv2.VideoCapture(source_path)
|
||||
if not video.isOpened():
|
||||
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")
|
||||
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)
|
||||
return video, start, end
|
||||
|
||||
|
||||
def get_video_frames_generator(
|
||||
source_path: str, stride: int = 1, start: int = 0, end: Optional[int] = None
|
||||
) -> Generator[np.ndarray, None, None]:
|
||||
"""
|
||||
Get a generator that yields the frames of the video.
|
||||
|
||||
Args:
|
||||
source_path (str): The path of the video file.
|
||||
stride (int): Indicates the interval at which frames are returned, skipping stride - 1 frames between each.
|
||||
start (int): Indicates the starting position from which video should generate frames
|
||||
end (Optional[int]): Indicates the ending position at which video should stop generating frames. If None, video will be read to the end.
|
||||
|
||||
Returns:
|
||||
(Generator[np.ndarray, None, None]): A generator that yields the frames of the video.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> from supervision import get_video_frames_generator
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> for frame in get_video_frames_generator(source_path='source_video.mp4'):
|
||||
>>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'):
|
||||
... ...
|
||||
```
|
||||
"""
|
||||
video = cv2.VideoCapture(source_path)
|
||||
if not video.isOpened():
|
||||
raise Exception(f"Could not open video at {source_path}")
|
||||
success, frame = video.read()
|
||||
while success:
|
||||
yield frame
|
||||
video, start, end = _validate_and_setup_video(source_path, start, end)
|
||||
frame_position = start
|
||||
while True:
|
||||
success, frame = video.read()
|
||||
if not success or frame_position >= end:
|
||||
break
|
||||
yield frame
|
||||
for _ in range(stride - 1):
|
||||
success = video.grab()
|
||||
if not success:
|
||||
break
|
||||
frame_position += stride
|
||||
video.release()
|
||||
|
||||
|
||||
@@ -0,0 +1,63 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
|
||||
import pytest
|
||||
import numpy as np
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"class_id, confidence, k, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.array([0, 1, 2, 3, 4]),
|
||||
np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
|
||||
5,
|
||||
(np.array([2, 4, 3, 1, 0]), np.array([0.9, 0.5, 0.4, 0.2, 0.1])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences
|
||||
(
|
||||
np.array([5, 1, 2, 3, 4]),
|
||||
np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
|
||||
1,
|
||||
(np.array([2]), np.array([0.9])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences, retrieve where k = 1
|
||||
(
|
||||
np.array([4, 1, 2, 3, 6, 5]),
|
||||
np.array([0.8, 0.2, 0.9, 0.4, 0.5, 0.1]),
|
||||
2,
|
||||
(np.array([2, 4]), np.array([0.9, 0.8])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences, retrieve where k = 3
|
||||
(
|
||||
np.array([0, 1, 2, 3, 4]),
|
||||
np.array([]),
|
||||
5,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id with 5 numbers and 0 confidences
|
||||
(
|
||||
[0, 1, 2, 3, 4],
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
5,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id with 5 numbers and 4 confidences
|
||||
],
|
||||
)
|
||||
def test_top_k(
|
||||
class_id: np.ndarray,
|
||||
confidence: Optional[np.ndarray],
|
||||
k: int,
|
||||
expected_result: Optional[Tuple[np.ndarray, np.ndarray]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = Classifications(
|
||||
class_id=np.array(class_id), confidence=np.array(confidence)
|
||||
).get_top_k(k)
|
||||
|
||||
assert np.array_equal(result[0], expected_result[0])
|
||||
assert np.array_equal(result[1], expected_result[1])
|
||||
@@ -0,0 +1,424 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, Tuple, Dict
|
||||
|
||||
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
|
||||
|
||||
|
||||
def mock_cock_coco_annotation(
|
||||
annotation_id: int = 0,
|
||||
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
|
||||
) -> dict:
|
||||
return {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": category_id,
|
||||
"bbox": list(bbox),
|
||||
"area": area,
|
||||
"iscrowd": 0
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_categories, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # empty coco categories
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
}
|
||||
],
|
||||
["fashion-assistant"],
|
||||
DoesNotRaise()
|
||||
), # single coco category with supercategory == "none"
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"name": "baseball cap",
|
||||
"supercategory": "fashion-assistant"
|
||||
}
|
||||
],
|
||||
[
|
||||
"fashion-assistant", "baseball cap"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # two coco categories; one with supercategory == "none" and one with supercategory != "none"
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"name": "baseball cap",
|
||||
"supercategory": "fashion-assistant"
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"name": "hoodie",
|
||||
"supercategory": "fashion-assistant"
|
||||
}
|
||||
],
|
||||
[
|
||||
"fashion-assistant",
|
||||
"baseball cap",
|
||||
"hoodie"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # three coco categories; one with supercategory == "none" and two with supercategory != "none"
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"name": "hoodie",
|
||||
"supercategory": "fashion-assistant"
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"name": "baseball cap",
|
||||
"supercategory": "fashion-assistant"
|
||||
}
|
||||
],
|
||||
[
|
||||
"fashion-assistant",
|
||||
"baseball cap",
|
||||
"hoodie"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # three coco categories; one with supercategory == "none" and two with supercategory != "none" (different order)
|
||||
]
|
||||
)
|
||||
def test_coco_categories_to_classes(
|
||||
coco_categories: List[dict],
|
||||
expected_result: List[str],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = coco_categories_to_classes(coco_categories=coco_categories)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"classes, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
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:
|
||||
with exception:
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
result = coco_categories_to_classes(coco_categories=coco_categories)
|
||||
assert result == classes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_annotations, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty coco annotations
|
||||
(
|
||||
[
|
||||
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)
|
||||
]
|
||||
},
|
||||
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)
|
||||
],
|
||||
{
|
||||
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)
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # two coco annotations
|
||||
(
|
||||
[
|
||||
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=1),
|
||||
mock_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
|
||||
mock_cock_coco_annotation(annotation_id=3, image_id=2, 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),
|
||||
],
|
||||
{
|
||||
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),
|
||||
],
|
||||
2: [
|
||||
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),
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # two coco annotations
|
||||
]
|
||||
)
|
||||
def test_group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict],
|
||||
expected_result: dict,
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image_annotations, resolution_wh, with_masks, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
), # empty image annotations
|
||||
(
|
||||
[
|
||||
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)
|
||||
),
|
||||
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),
|
||||
],
|
||||
(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)
|
||||
),
|
||||
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
|
||||
) -> None:
|
||||
with exception:
|
||||
result = coco_annotations_to_detections(
|
||||
image_annotations=image_annotations,
|
||||
resolution_wh=resolution_wh,
|
||||
with_masks=with_masks
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_categories, target_classes, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
[],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty coco categories
|
||||
(
|
||||
[
|
||||
{
|
||||
"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()
|
||||
), # 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"
|
||||
}
|
||||
],
|
||||
[
|
||||
"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"
|
||||
}
|
||||
],
|
||||
[
|
||||
"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"
|
||||
}
|
||||
],
|
||||
[
|
||||
"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
|
||||
) -> None:
|
||||
with exception:
|
||||
result = build_coco_class_index_mapping(
|
||||
coco_categories=coco_categories,
|
||||
target_classes=target_classes
|
||||
)
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,175 @@
|
||||
from typing import List, Optional
|
||||
|
||||
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(
|
||||
"dataset_list, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
DoesNotRaise()
|
||||
), # empty dataset list
|
||||
(
|
||||
[
|
||||
DetectionDataset(classes=[], images={}, annotations={})
|
||||
],
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
DoesNotRaise()
|
||||
), # single empty dataset
|
||||
(
|
||||
[
|
||||
DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
|
||||
DetectionDataset(classes=['dog', 'person'], images={}, annotations={})
|
||||
],
|
||||
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=['cat', 'dog', 'person'], images={}, annotations={}),
|
||||
DoesNotRaise()
|
||||
), # two datasets; no images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two datasets; images and annotations, the same classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(classes=['cat'], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two datasets; images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(
|
||||
classes=['cat'],
|
||||
images={
|
||||
'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]),
|
||||
}
|
||||
),
|
||||
],
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two datasets; images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(
|
||||
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),
|
||||
},
|
||||
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]),
|
||||
}
|
||||
),
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
)
|
||||
]
|
||||
)
|
||||
def test_dataset_merge(
|
||||
dataset_list: List[DetectionDataset],
|
||||
expected_result: Optional[DetectionDataset],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = DetectionDataset.merge(dataset_list=dataset_list)
|
||||
assert result == expected_result
|
||||
+167
-2
@@ -1,9 +1,12 @@
|
||||
from typing import List, TypeVar, Optional, Tuple
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, TypeVar, Optional, Tuple, Dict
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision.dataset.ultils import train_test_split
|
||||
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
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
@@ -80,3 +83,165 @@ def test_train_test_split(
|
||||
with exception:
|
||||
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',
|
||||
[
|
||||
(
|
||||
[],
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # empty class lists
|
||||
(
|
||||
[
|
||||
['dog', 'person']
|
||||
],
|
||||
['dog', 'person'],
|
||||
DoesNotRaise()
|
||||
), # single class list; already alphabetically sorted
|
||||
(
|
||||
[
|
||||
['person', 'dog']
|
||||
],
|
||||
['dog', 'person'],
|
||||
DoesNotRaise()
|
||||
), # single class list; not alphabetically sorted
|
||||
(
|
||||
[
|
||||
['dog', 'person'],
|
||||
['dog', 'person']
|
||||
],
|
||||
['dog', 'person'],
|
||||
DoesNotRaise()
|
||||
), # two class lists; the same classes; already alphabetically sorted
|
||||
(
|
||||
[
|
||||
['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:
|
||||
with exception:
|
||||
result = merge_class_lists(class_lists=class_lists)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'source_classes, target_classes, expected_result, exception',
|
||||
[
|
||||
(
|
||||
[],
|
||||
[],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty class lists
|
||||
(
|
||||
[],
|
||||
['dog', 'person'],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty source class list
|
||||
(
|
||||
['dog', 'person'],
|
||||
[],
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
), # empty target class list
|
||||
(
|
||||
['dog', 'person'],
|
||||
['dog', 'person'],
|
||||
{0: 0, 1: 1},
|
||||
DoesNotRaise()
|
||||
), # same class lists
|
||||
(
|
||||
['dog', 'person'],
|
||||
['person', 'dog'],
|
||||
{0: 1, 1: 0},
|
||||
DoesNotRaise()
|
||||
), # same class lists but not alphabetically sorted
|
||||
(
|
||||
['dog', 'person'],
|
||||
['cat', 'dog', 'person'],
|
||||
{0: 1, 1: 2},
|
||||
DoesNotRaise()
|
||||
), # source class list is a subset of target class list
|
||||
(
|
||||
['dog', 'person'],
|
||||
['cat', 'dog'],
|
||||
None,
|
||||
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
|
||||
) -> None:
|
||||
with exception:
|
||||
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',
|
||||
[
|
||||
(
|
||||
{},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
None,
|
||||
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()
|
||||
), # single mapping
|
||||
(
|
||||
{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()
|
||||
), # 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()
|
||||
), # multiple mappings
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
), # class_id not in mapping
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
|
||||
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
|
||||
) -> None:
|
||||
with exception:
|
||||
result = map_detections_class_id(source_to_target_mapping=source_to_target_mapping, detections=detections)
|
||||
assert result == expected_result
|
||||
|
||||
+87
-43
@@ -8,6 +8,7 @@ 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],
|
||||
@@ -28,60 +29,30 @@ DETECTIONS = 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],
|
||||
], dtype=np.float32)
|
||||
|
||||
|
||||
def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None) -> Detections:
|
||||
return Detections(
|
||||
xyxy = np.array(xyxy, dtype=np.float32),
|
||||
confidence = confidence if confidence is None else np.array(confidence, dtype=np.float32),
|
||||
class_id = class_id if class_id is None else np.array(class_id, dtype=int),
|
||||
tracker_id = tracker_id if tracker_id is None else np.array(tracker_id, dtype=int),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'detections, index, expected_result, exception',
|
||||
[
|
||||
(
|
||||
DETECTIONS,
|
||||
DETECTIONS.class_id == 0,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2254, 906, 2447, 1353]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.90538
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[[2254, 906, 2447, 1353]],
|
||||
confidence=[0.90538],
|
||||
class_id=[0]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only detections with class_id = 0
|
||||
(
|
||||
DETECTIONS,
|
||||
DETECTIONS.confidence > 0.5,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 2254, 906, 2447, 1353],
|
||||
[ 2049, 1133, 2226, 1371],
|
||||
[ 727, 1224, 838, 1601]
|
||||
], dtype=np.float32),
|
||||
confidence=np.array([
|
||||
0.90538,
|
||||
0.59002,
|
||||
0.51119
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0,
|
||||
56,
|
||||
39
|
||||
], dtype=int)
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2254, 906, 2447, 1353],
|
||||
[2049, 1133, 2226, 1371],
|
||||
[727, 1224, 838, 1601]
|
||||
],
|
||||
confidence=[0.90538, 0.59002, 0.51119],
|
||||
class_id=[0, 56, 39]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only detections with confidence > 0.5
|
||||
@@ -101,11 +72,84 @@ def mock_detections(xyxy, confidence = None, class_id = None, tracker_id = None)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take no detections
|
||||
(
|
||||
DETECTIONS,
|
||||
[0, 2],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2254, 906, 2447, 1353],
|
||||
[727, 1224, 838, 1601]
|
||||
],
|
||||
confidence=[0.90538, 0.51119],
|
||||
class_id=[0, 39]
|
||||
),
|
||||
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]
|
||||
],
|
||||
confidence=[0.90538, 0.51119],
|
||||
class_id=[0, 39]
|
||||
),
|
||||
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]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only first detection by index
|
||||
(
|
||||
DETECTIONS,
|
||||
slice(1, 3),
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2049, 1133, 2226, 1371],
|
||||
[727, 1224, 838, 1601]
|
||||
],
|
||||
confidence=[0.59002, 0.51119],
|
||||
class_id=[56, 39]
|
||||
),
|
||||
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,
|
||||
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, np.ndarray],
|
||||
index: Union[int, slice, List[int], np.ndarray],
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
|
||||
@@ -5,7 +5,8 @@ import pytest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area
|
||||
from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area, \
|
||||
process_roboflow_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -276,3 +277,61 @@ def test_filter_polygons_by_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)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"roboflow_result, class_list, expected_result, exception",
|
||||
[
|
||||
(
|
||||
{
|
||||
"predictions": [],
|
||||
"image": {"width": 1000, "height": 1000}
|
||||
},
|
||||
["person", "car", "truck"],
|
||||
(
|
||||
np.empty((0, 4)),
|
||||
np.empty(0),
|
||||
np.empty(0),
|
||||
None
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # empty result
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class": "person"
|
||||
}
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000}
|
||||
},
|
||||
["person", "car", "truck"],
|
||||
(
|
||||
np.array([
|
||||
[175.0, 275.0, 225.0, 325.0]
|
||||
]),
|
||||
np.array([0.9]),
|
||||
np.array([0]),
|
||||
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
|
||||
) -> None:
|
||||
with exception:
|
||||
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]))
|
||||
|
||||
@@ -0,0 +1,402 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Optional, Union
|
||||
|
||||
import numpy as np
|
||||
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)
|
||||
|
||||
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,
|
||||
)
|
||||
|
||||
TARGET_TENSORS = [
|
||||
np.array(
|
||||
[
|
||||
[2254, 906, 2447, 1353, 0],
|
||||
[2049, 1133, 2226, 1371, 56],
|
||||
[727, 1224, 838, 1601, 39],
|
||||
[808, 1214, 910, 1564, 39],
|
||||
[6, 52, 1131, 2133, 72],
|
||||
[299, 1225, 512, 1663, 39],
|
||||
[529, 874, 645, 945, 39],
|
||||
[8, 47, 1935, 2135, 72],
|
||||
[2265, 813, 2328, 901, 62],
|
||||
]
|
||||
)
|
||||
]
|
||||
|
||||
DETECTIONS = Detections(
|
||||
xyxy=PREDICTIONS[:, :4],
|
||||
confidence=PREDICTIONS[:, 4],
|
||||
class_id=PREDICTIONS[:, 5].astype(int),
|
||||
)
|
||||
CERTAIN_DETECTIONS = Detections(
|
||||
xyxy=PREDICTIONS[:, :4],
|
||||
confidence=np.ones(len(PREDICTIONS)),
|
||||
class_id=PREDICTIONS[:, 5].astype(int),
|
||||
)
|
||||
|
||||
DETECTION_TENSORS = [
|
||||
np.concatenate(
|
||||
[
|
||||
det.xyxy,
|
||||
np.expand_dims(det.class_id, 1),
|
||||
np.expand_dims(det.confidence, 1),
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
for det in [DETECTIONS]
|
||||
]
|
||||
CERTAIN_DETECTION_TENSORS = [
|
||||
np.concatenate(
|
||||
[
|
||||
det.xyxy,
|
||||
np.expand_dims(det.class_id, 1),
|
||||
np.ones((len(det), 1)),
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
for det in [DETECTIONS]
|
||||
]
|
||||
|
||||
IDEAL_MATCHES = np.stack(
|
||||
[
|
||||
np.arange(len(PREDICTIONS)),
|
||||
np.arange(len(PREDICTIONS)),
|
||||
np.ones(len(PREDICTIONS)),
|
||||
],
|
||||
axis=1,
|
||||
)
|
||||
|
||||
|
||||
def create_empty_conf_matrix(num_classes: int, do_add_dummy_class: bool = True):
|
||||
if do_add_dummy_class:
|
||||
num_classes += 1
|
||||
return np.zeros((num_classes, num_classes))
|
||||
|
||||
|
||||
def update_ideal_conf_matrix(conf_matrix: np.ndarray, class_ids: np.ndarray):
|
||||
for class_id, count in zip(*np.unique(class_ids, return_counts=True)):
|
||||
class_id = int(class_id)
|
||||
conf_matrix[class_id, class_id] += count
|
||||
return conf_matrix
|
||||
|
||||
|
||||
def worsen_ideal_conf_matrix(
|
||||
conf_matrix: np.ndarray, class_ids: Union[np.ndarray, list]
|
||||
):
|
||||
for class_id in class_ids:
|
||||
class_id = int(class_id)
|
||||
conf_matrix[class_id, class_id] -= 1
|
||||
conf_matrix[class_id, 80] += 1
|
||||
return conf_matrix
|
||||
|
||||
|
||||
IDEAL_CONF_MATRIX = create_empty_conf_matrix(NUM_CLASSES)
|
||||
IDEAL_CONF_MATRIX = update_ideal_conf_matrix(IDEAL_CONF_MATRIX, PREDICTIONS[:, 5])
|
||||
|
||||
GOOD_CONF_MATRIX = worsen_ideal_conf_matrix(IDEAL_CONF_MATRIX.copy(), [62, 72])
|
||||
|
||||
BAD_CONF_MATRIX = worsen_ideal_conf_matrix(
|
||||
IDEAL_CONF_MATRIX.copy(), [62, 72, 72, 39, 39, 39, 39, 56]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, with_confidence, expected_result, exception",
|
||||
[
|
||||
(
|
||||
Detections.empty(),
|
||||
False,
|
||||
np.empty((0, 5), dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # empty detections; no confidence
|
||||
(
|
||||
Detections.empty(),
|
||||
True,
|
||||
np.empty((0, 6), dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # empty detections; with confidence
|
||||
(
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
|
||||
False,
|
||||
np.array([[0, 0, 10, 10, 0]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # single detection; no confidence
|
||||
(
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
|
||||
True,
|
||||
np.array([[0, 0, 10, 10, 0, 0.5]], dtype=np.float32),
|
||||
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]),
|
||||
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]),
|
||||
True,
|
||||
np.array([[0, 0, 10, 10, 0, 0.5], [0, 0, 20, 20, 1, 0.2]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; with confidence
|
||||
],
|
||||
)
|
||||
def test_detections_to_tensor(
|
||||
detections: Detections,
|
||||
with_confidence: bool,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception
|
||||
):
|
||||
with exception:
|
||||
result = ConfusionMatrix.detections_to_tensor(
|
||||
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",
|
||||
[
|
||||
(
|
||||
DETECTION_TENSORS,
|
||||
TARGET_TENSORS,
|
||||
CLASSES,
|
||||
0.2,
|
||||
0.5,
|
||||
IDEAL_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
[],
|
||||
[],
|
||||
CLASSES,
|
||||
0.2,
|
||||
0.5,
|
||||
create_empty_conf_matrix(NUM_CLASSES),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
DETECTION_TENSORS,
|
||||
TARGET_TENSORS,
|
||||
CLASSES,
|
||||
0.3,
|
||||
0.5,
|
||||
GOOD_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
DETECTION_TENSORS,
|
||||
TARGET_TENSORS,
|
||||
CLASSES,
|
||||
0.6,
|
||||
0.5,
|
||||
BAD_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0, 0.9], # correct detection of [0]
|
||||
[
|
||||
0.1,
|
||||
0.1,
|
||||
3.0,
|
||||
3.0,
|
||||
0,
|
||||
0.9,
|
||||
], # additional detection of [0] - FP
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
1,
|
||||
0.8,
|
||||
], # correct detection with incorrect class
|
||||
[1.0, 6.0, 2.0, 7.0, 1, 0.8], # incorrect detection - FP
|
||||
[
|
||||
1.0,
|
||||
2.0,
|
||||
2.0,
|
||||
4.0,
|
||||
1,
|
||||
0.8,
|
||||
], # incorrect detection with low IoU - FP
|
||||
]
|
||||
)
|
||||
],
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0], # [0] detected
|
||||
[2.0, 2.0, 5.0, 5.0, 1], # [1] undetected - FN
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
2,
|
||||
], # [2] correct detection with incorrect class
|
||||
]
|
||||
)
|
||||
],
|
||||
CLASSES[:3],
|
||||
0.6,
|
||||
0.5,
|
||||
np.array([[1, 0, 0, 0], [0, 0, 0, 1], [0, 1, 0, 0], [1, 2, 0, 0]]),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0, 0.9], # correct detection of [0]
|
||||
[
|
||||
0.1,
|
||||
0.1,
|
||||
3.0,
|
||||
3.0,
|
||||
0,
|
||||
0.9,
|
||||
], # additional detection of [0] - FP
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
1,
|
||||
0.8,
|
||||
], # correct detection with incorrect class
|
||||
[1.0, 6.0, 2.0, 7.0, 1, 0.8], # incorrect detection - FP
|
||||
[
|
||||
1.0,
|
||||
2.0,
|
||||
2.0,
|
||||
4.0,
|
||||
1,
|
||||
0.8,
|
||||
], # incorrect detection with low IoU - FP
|
||||
]
|
||||
)
|
||||
],
|
||||
[
|
||||
np.array(
|
||||
[
|
||||
[0.0, 0.0, 3.0, 3.0, 0], # [0] detected
|
||||
[2.0, 2.0, 5.0, 5.0, 1], # [1] undetected - FN
|
||||
[
|
||||
6.0,
|
||||
1.0,
|
||||
8.0,
|
||||
3.0,
|
||||
2,
|
||||
], # [2] correct detection with incorrect class
|
||||
]
|
||||
)
|
||||
],
|
||||
CLASSES[:3],
|
||||
0.6,
|
||||
1.0,
|
||||
np.array([[0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 1], [2, 3, 0, 0]]),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_from_tensors(
|
||||
predictions,
|
||||
targets,
|
||||
classes,
|
||||
conf_threshold,
|
||||
iou_threshold,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = ConfusionMatrix.from_tensors(
|
||||
predictions=predictions,
|
||||
targets=targets,
|
||||
classes=classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
assert result.matrix.diagonal().sum() == expected_result.diagonal().sum()
|
||||
assert np.array_equal(result.matrix, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, targets, num_classes, conf_threshold, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
DETECTION_TENSORS[0],
|
||||
CERTAIN_DETECTION_TENSORS[0],
|
||||
NUM_CLASSES,
|
||||
0.2,
|
||||
0.5,
|
||||
IDEAL_CONF_MATRIX,
|
||||
DoesNotRaise(),
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_evaluate_detection_batch(
|
||||
predictions,
|
||||
targets,
|
||||
num_classes,
|
||||
conf_threshold,
|
||||
iou_threshold,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = ConfusionMatrix.evaluate_detection_batch(
|
||||
predictions=predictions,
|
||||
targets=targets,
|
||||
num_classes=num_classes,
|
||||
conf_threshold=conf_threshold,
|
||||
iou_threshold=iou_threshold,
|
||||
)
|
||||
|
||||
assert result.diagonal().sum() == result.sum()
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"matches, expected_result, exception",
|
||||
[
|
||||
(
|
||||
IDEAL_MATCHES,
|
||||
IDEAL_MATCHES,
|
||||
DoesNotRaise(),
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_drop_extra_matches(
|
||||
matches,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = ConfusionMatrix._drop_extra_matches(matches)
|
||||
|
||||
assert np.array_equal(result, expected_result)
|
||||
@@ -0,0 +1,23 @@
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
|
||||
def mock_detections(
|
||||
xyxy: List[List[float]],
|
||||
confidence: List[float] = None,
|
||||
class_id: List[int] = None,
|
||||
tracker_id: List[int] = None,
|
||||
) -> Detections:
|
||||
return Detections(
|
||||
xyxy=np.array(xyxy, dtype=np.float32),
|
||||
confidence=confidence
|
||||
if confidence is None
|
||||
else np.array(confidence, dtype=np.float32),
|
||||
class_id=class_id if class_id is None else np.array(class_id, dtype=int),
|
||||
tracker_id=tracker_id
|
||||
if tracker_id is None
|
||||
else np.array(tracker_id, dtype=int),
|
||||
)
|
||||
Reference in New Issue
Block a user