Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9adc953780 | ||
|
|
4ace93f9ab | ||
|
|
c1d3ab2944 | ||
|
|
2072eeac12 | ||
|
|
c74efc4908 | ||
|
|
08f399cf19 | ||
|
|
17020b23de | ||
|
|
cd929525f6 | ||
|
|
bb18de8dce | ||
|
|
dec96aec4d | ||
|
|
ee45c684d0 | ||
|
|
ea517642f0 | ||
|
|
cdc08c0455 | ||
|
|
a0c808cfdf | ||
|
|
f146e88ceb | ||
|
|
66d29015d6 | ||
|
|
46cb82797e | ||
|
|
194ff0aaad | ||
|
|
a579ed46c6 | ||
|
|
8918783702 | ||
|
|
6adb0f019d | ||
|
|
84ae12768a | ||
|
|
9cbec9fc94 | ||
|
|
ed12daf77d | ||
|
|
84cfa9609b | ||
|
|
f2eccae88e | ||
|
|
55c8295755 | ||
|
|
7849033da3 | ||
|
|
2e443c80c7 | ||
|
|
006caa0b9b | ||
|
|
34e45bed5e | ||
|
|
0d39de898a | ||
|
|
782e7e5a39 | ||
|
|
b40265db9e | ||
|
|
de1ceff065 | ||
|
|
5bf78e5505 | ||
|
|
f63336a824 | ||
|
|
4d82460358 | ||
|
|
540efc80ce | ||
|
|
beea21a97d | ||
|
|
484e3d4afd | ||
|
|
ab9728cca2 | ||
|
|
769b6b967c | ||
|
|
253113fa88 | ||
|
|
3b2902c87c | ||
|
|
9cff624e15 | ||
|
|
584c4fef82 | ||
|
|
282aa2c5d3 | ||
|
|
453d4df068 | ||
|
|
4d4b38b292 | ||
|
|
13e8afa996 | ||
|
|
dd4cae328b | ||
|
|
ed7103851a | ||
|
|
ca40b001dc | ||
|
|
de3ccc5c5d | ||
|
|
1b6d02898c | ||
|
|
5ac92bc43b | ||
|
|
86d9293770 | ||
|
|
0e943e673e | ||
|
|
234e376fca | ||
|
|
54b7cbeb70 | ||
|
|
d000292eb2 | ||
|
|
b551e6d410 | ||
|
|
138c271af0 | ||
|
|
d22940ce30 | ||
|
|
06da0e70d2 | ||
|
|
5d8b378f35 | ||
|
|
a4830389de | ||
|
|
d9a2efc64b | ||
|
|
f91b844b7d | ||
|
|
dc45d47e7e | ||
|
|
b5dedf09cd | ||
|
|
eb1ae1b790 | ||
|
|
f2f8a564e0 | ||
|
|
140a219091 | ||
|
|
86af9b9dce | ||
|
|
4366e4480f | ||
|
|
b66e9c8c32 | ||
|
|
d0d5626156 | ||
|
|
c661bf250d | ||
|
|
ccb32f8aa1 | ||
|
|
97b9edcbec | ||
|
|
6b52863c9f | ||
|
|
2fcc035413 | ||
|
|
7c8f16bf0f | ||
|
|
1870670c77 | ||
|
|
802347f4d8 | ||
|
|
920019c887 | ||
|
|
477b9a8f38 | ||
|
|
14f62deacf | ||
|
|
08eb513594 | ||
|
|
a1615fd275 | ||
|
|
ce979216f2 | ||
|
|
95f94319f8 | ||
|
|
98f60c24d3 | ||
|
|
1c043d6cc7 | ||
|
|
86aee5ed98 | ||
|
|
0d38b95da7 | ||
|
|
97892e3138 | ||
|
|
17a42bc728 | ||
|
|
be65ba7370 | ||
|
|
72d86ba180 | ||
|
|
fba83ec371 | ||
|
|
4a81a17ddb | ||
|
|
30d253d5cc | ||
|
|
544f9cb0e2 | ||
|
|
84928a2da9 | ||
|
|
08ae0b2fc7 | ||
|
|
0261bec2ce | ||
|
|
5a4cb1b302 | ||
|
|
91655f9609 | ||
|
|
42d4079b23 | ||
|
|
f03a47a5df | ||
|
|
c01efeaaca | ||
|
|
5489d405c6 | ||
|
|
5f97eac9a9 | ||
|
|
30f105f267 | ||
|
|
a34e64866e | ||
|
|
c02f69428e | ||
|
|
55f4f00c0c | ||
|
|
91100e8393 | ||
|
|
b474a47065 | ||
|
|
83d357d973 | ||
|
|
b536c6c889 | ||
|
|
df579c866c | ||
|
|
f224bceb6f | ||
|
|
bd3503e7d5 | ||
|
|
f63f5acb2f | ||
|
|
2f5b1b0f0d | ||
|
|
3af32800d3 | ||
|
|
8222bee9dd | ||
|
|
06ea737336 | ||
|
|
f616bda40f | ||
|
|
0e835f1b0a | ||
|
|
4c024214b2 | ||
|
|
6a15093277 | ||
|
|
abcf453119 | ||
|
|
ba2926688f | ||
|
|
3da433f90f | ||
|
|
a911cbd554 | ||
|
|
d10618a8b3 | ||
|
|
eaf5b7e95f | ||
|
|
4ae8ac1db0 | ||
|
|
74efca7c22 | ||
|
|
3063dfb1f5 | ||
|
|
6c81d21a03 | ||
|
|
f4b91d0175 | ||
|
|
45fdb1e805 | ||
|
|
8c0a294ca1 | ||
|
|
cf7f003133 | ||
|
|
1bebbb14a6 | ||
|
|
3d0aec03e0 | ||
|
|
306d7bd27e | ||
|
|
43fc5d6205 | ||
|
|
d13238bc42 | ||
|
|
77243f1752 | ||
|
|
e16cca4349 | ||
|
|
e592a44bea | ||
|
|
c2833bc8ca | ||
|
|
feb08d10b2 | ||
|
|
74bd403dc0 | ||
|
|
6a013382d3 | ||
|
|
92290979c7 | ||
|
|
696b6f9276 | ||
|
|
29450139f9 | ||
|
|
ea400feb6a | ||
|
|
05495aca75 | ||
|
|
a3bea9ae37 | ||
|
|
651690f838 | ||
|
|
d87761f75d | ||
|
|
b20e4a6c19 | ||
|
|
0ef2ab7273 | ||
|
|
d1bb169593 | ||
|
|
c8f3f59a45 | ||
|
|
12523afe9a | ||
|
|
68412500fa | ||
|
|
49d4441419 | ||
|
|
0394a3f545 | ||
|
|
2bcd1e1dfd | ||
|
|
e62937c8ce | ||
|
|
b37aca52e2 | ||
|
|
790d2f6c5b | ||
|
|
0cb722e91e | ||
|
|
55cd04d0aa | ||
|
|
58181e6a7a | ||
|
|
ab00d2db0d | ||
|
|
6a76c96df5 | ||
|
|
d4a6a1a85b | ||
|
|
7c78feb4b8 | ||
|
|
73a7983719 | ||
|
|
1d6c2bc080 | ||
|
|
889461574f | ||
|
|
f95426540b | ||
|
|
c16599269c | ||
|
|
888edfbc67 | ||
|
|
7497f8d96a | ||
|
|
b34ecd53cf | ||
|
|
e9649acd90 | ||
|
|
2cf7a9b620 | ||
|
|
e091f9c959 | ||
|
|
f12504b0dd | ||
|
|
dba4d9f8ad | ||
|
|
ddc8a8ccce | ||
|
|
c1ccff474f | ||
|
|
d02bac75c4 | ||
|
|
f3c4a966c4 | ||
|
|
5bc436b1c5 | ||
|
|
c37590ef3a | ||
|
|
a14bb94b46 | ||
|
|
aab7e238e4 | ||
|
|
34dba17641 | ||
|
|
ef5ee71f0b | ||
|
|
bed0e41cb7 | ||
|
|
e439b29f07 | ||
|
|
511375f24b | ||
|
|
ff2788945d | ||
|
|
5b9eef363d | ||
|
|
bcc4b25eb9 | ||
|
|
d5d2035048 | ||
|
|
14795a6626 | ||
|
|
bc12a8ea0b | ||
|
|
2be4981c76 | ||
|
|
fe68327b75 | ||
|
|
141809ae92 | ||
|
|
ac61598ad4 | ||
|
|
1fa4a1c72e | ||
|
|
2df4261355 | ||
|
|
a88eff868e | ||
|
|
8bf686e23b | ||
|
|
b2f6520b87 | ||
|
|
420dae9e70 | ||
|
|
550c913dbb | ||
|
|
4dcca5309c | ||
|
|
d30cc180d0 | ||
|
|
ace8301b12 | ||
|
|
b91bc91897 | ||
|
|
e04c531c05 | ||
|
|
63fd5a3f8a | ||
|
|
38f52a04ae | ||
|
|
65090f612c | ||
|
|
9401210aa5 | ||
|
|
59d8655014 | ||
|
|
4dde484141 | ||
|
|
e6f250dbfb | ||
|
|
e6090be9eb | ||
|
|
25c19b4422 | ||
|
|
ac1658268c | ||
|
|
8e7228313b | ||
|
|
a27377aaf3 | ||
|
|
eefc5ddf7d | ||
|
|
96d8badf1a | ||
|
|
c0ae62a460 | ||
|
|
1c2d946527 | ||
|
|
95b44db2bd | ||
|
|
ba732330a5 | ||
|
|
3ea95b3966 | ||
|
|
c2a5161ff9 | ||
|
|
c3eec02b6c | ||
|
|
8c68999324 | ||
|
|
4a43f4ec69 | ||
|
|
78d267e169 | ||
|
|
7780308533 | ||
|
|
2e76e3db0f | ||
|
|
a0eb6c3237 | ||
|
|
fd35ad9653 | ||
|
|
14d23eb9df | ||
|
|
602d3af009 | ||
|
|
b8b08cacd3 | ||
|
|
c8868c3412 | ||
|
|
52a8193b4a | ||
|
|
51a1655bd4 | ||
|
|
38bf4b0f54 | ||
|
|
05705bd4f2 | ||
|
|
940bcb8bc5 | ||
|
|
a64a252ae3 | ||
|
|
4afc650eb2 | ||
|
|
8d091a7a00 | ||
|
|
ccafa3339d | ||
|
|
807785a820 | ||
|
|
66f81ceeb9 | ||
|
|
b7c1791236 | ||
|
|
0ac184b01d | ||
|
|
307ec8a042 | ||
|
|
d5d7187fe4 | ||
|
|
0e3851506c | ||
|
|
301f1df386 | ||
|
|
ceabb83ac6 | ||
|
|
dfb757b8cd | ||
|
|
ed0fd7b768 | ||
|
|
5b2f8e4dcc | ||
|
|
465830f3f6 | ||
|
|
ae84cdeb02 | ||
|
|
4fbec12b06 | ||
|
|
b5e33c5b30 | ||
|
|
5a3df85fc6 | ||
|
|
61a2f4e0e6 | ||
|
|
90a26747b6 | ||
|
|
507f08f27f | ||
|
|
fcf30b79d4 | ||
|
|
a7b26abe0e | ||
|
|
9590f3845e | ||
|
|
8d8bd3b997 | ||
|
|
6b11a1fdad | ||
|
|
cb485878cc | ||
|
|
8e70a3cd46 | ||
|
|
123dd77ba4 | ||
|
|
7687d4a873 | ||
|
|
d941fe256d | ||
|
|
d4a2adde9e | ||
|
|
8b229427e5 | ||
|
|
b4641245c1 | ||
|
|
bfcc3e40b2 | ||
|
|
1561bd0027 | ||
|
|
8d14a4d93b | ||
|
|
bd889ae3cc | ||
|
|
02c20e81a9 | ||
|
|
c236bb71d1 | ||
|
|
22fa6f0a3a | ||
|
|
c2ac59f992 |
@@ -0,0 +1 @@
|
||||
*.ipynb linguist-vendored
|
||||
@@ -0,0 +1,65 @@
|
||||
name: 🐞 Bug Report
|
||||
# title: " "
|
||||
description: Problems with Supervision
|
||||
labels: [bug]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thank you for submitting a Supervision 🐞 Bug Report!
|
||||
|
||||
- type: checkboxes
|
||||
attributes:
|
||||
label: Search before asking
|
||||
description: >
|
||||
Please search the [issues](https://github.com/roboflow/supervision/issues) to see if a similar bug report already exists.
|
||||
options:
|
||||
- label: >
|
||||
I have searched the Supervision [issues](https://github.com/roboflow/supervision/issues) and found no similar bug report.
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Bug
|
||||
description: Provide console output with error messages and/or screenshots of the bug.
|
||||
placeholder: |
|
||||
💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Environment
|
||||
description: Please specify the software and hardware you used to produce the bug.
|
||||
placeholder: |
|
||||
- Supervision: 0.1.0
|
||||
- OS: Ubuntu 20.04
|
||||
- Python: 3.8.10
|
||||
validations:
|
||||
required: false
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Minimal Reproducible Example
|
||||
description: >
|
||||
When asking a question, people will be better able to provide help if you provide code that they can easily understand and use to **reproduce** the problem.
|
||||
This is referred to by community members as creating a [minimal reproducible example](https://stackoverflow.com/help/minimal-reproducible-example).
|
||||
placeholder: |
|
||||
```
|
||||
# Code to reproduce your issue here
|
||||
```
|
||||
validations:
|
||||
required: false
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Additional
|
||||
description: Anything else you would like to share?
|
||||
|
||||
- type: checkboxes
|
||||
attributes:
|
||||
label: Are you willing to submit a PR?
|
||||
description: >
|
||||
(Optional) We encourage you to submit a [Pull Request](https://github.com/roboflow/supervision/pulls) (PR) to help improve Supervision for everyone, especially if you have a good understanding of how to implement a fix or feature.
|
||||
options:
|
||||
- label: Yes I'd like to help by submitting a PR!
|
||||
@@ -0,0 +1,49 @@
|
||||
name: 🤩 Feature Request
|
||||
description: Suggest a Supervision idea
|
||||
# title: " "
|
||||
labels: [enhancement]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thank you for submitting a Supervision 🤩 Feature Request!
|
||||
|
||||
- type: checkboxes
|
||||
attributes:
|
||||
label: Search before asking
|
||||
description: >
|
||||
Please search the [issues](https://github.com/roboflow/supervision/issues) to see if a similar feature request already exists.
|
||||
options:
|
||||
- label: >
|
||||
I have searched the Supervision [issues](https://github.com/roboflow/supervision/issues) and found no similar feature requests.
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Description
|
||||
description: A short description of your feature.
|
||||
placeholder: |
|
||||
What new feature would you like to see in Supervision?
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Use case
|
||||
description: |
|
||||
Describe the use case of your feature request. It will help us understand and prioritize the feature request.
|
||||
placeholder: |
|
||||
How would this feature be used, and who would use it?
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Additional
|
||||
description: Anything else you would like to share?
|
||||
|
||||
- type: checkboxes
|
||||
attributes:
|
||||
label: Are you willing to submit a PR?
|
||||
description: >
|
||||
(Optional) We encourage you to submit a [Pull Request](https://github.com/roboflow/supervision/pulls) (PR) to help improve Supervision for everyone, especially if you have a good understanding of how to implement a fix or feature.
|
||||
options:
|
||||
- label: Yes I'd like to help by submitting a PR!
|
||||
@@ -0,0 +1,33 @@
|
||||
name: ❓ Question
|
||||
description: Ask a Supervision question
|
||||
# title: " "
|
||||
labels: [question]
|
||||
body:
|
||||
- type: markdown
|
||||
attributes:
|
||||
value: |
|
||||
Thank you for asking a Supervision ❓ Question!
|
||||
|
||||
- type: checkboxes
|
||||
attributes:
|
||||
label: Search before asking
|
||||
description: >
|
||||
Please search the [issues](https://github.com/roboflow/supervision/issues) to see if a similar feature request already exists.
|
||||
options:
|
||||
- label: >
|
||||
I have searched the Supervision [issues](https://github.com/roboflow/supervision/issues) and found no similar feature requests.
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Question
|
||||
description: What is your question?
|
||||
placeholder: |
|
||||
💡 ProTip! Include as much information as possible (screenshots, logs, tracebacks etc.) to receive the most helpful response.
|
||||
validations:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
attributes:
|
||||
label: Additional
|
||||
description: Anything else you would like to share?
|
||||
@@ -0,0 +1,18 @@
|
||||
name: Docs WorkFlow
|
||||
|
||||
on:
|
||||
push:
|
||||
branches:
|
||||
- master
|
||||
- main
|
||||
jobs:
|
||||
deploy:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v2
|
||||
- uses: actions/setup-python@v2
|
||||
with:
|
||||
python-version: 3.x
|
||||
- run: pip install mkdocs-material
|
||||
- run: pip install "mkdocstrings[python]"
|
||||
- run: mkdocs gh-deploy --force
|
||||
@@ -9,7 +9,7 @@ jobs:
|
||||
runs-on: ubuntu-latest
|
||||
strategy:
|
||||
matrix:
|
||||
python-version: [3.7, 3.8, 3.9]
|
||||
python-version: [3.7, 3.8, 3.9, '3.10']
|
||||
steps:
|
||||
- name: 🛎️ Checkout
|
||||
uses: actions/checkout@v3
|
||||
@@ -22,7 +22,11 @@ jobs:
|
||||
- name: 🦾 Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install ".[dev]"
|
||||
pip install -e .
|
||||
pip install pytest
|
||||
pip install isort
|
||||
pip install flake8
|
||||
pip install "black==22.3.0"
|
||||
- name: 🧹 Lint with flake8
|
||||
run: |
|
||||
make check_code_quality
|
||||
|
||||
@@ -1,3 +1,5 @@
|
||||
name: Welcome WorkFlow
|
||||
|
||||
on:
|
||||
issues:
|
||||
types: [opened]
|
||||
|
||||
@@ -0,0 +1,12 @@
|
||||
cff-version: 1.2.0
|
||||
author: Roboflow
|
||||
message: If you use this software, please cite it as below.
|
||||
title: supervision
|
||||
version: 0.11.0
|
||||
date-released: 2023-01-19
|
||||
license: MIT
|
||||
repository-code: https://github.com/roboflow/supervision
|
||||
keywords:
|
||||
- computer vision
|
||||
- image processing
|
||||
- video processing
|
||||
@@ -0,0 +1,51 @@
|
||||
# Contributing to Supervision 🛠️
|
||||
|
||||
Thank you for your interest in contributing to Supervision!
|
||||
|
||||
We are actively improving this library to reduce the amount of work you need to do to solve common computer vision problems.
|
||||
|
||||
## Contribution Guidelines
|
||||
|
||||
We welcome contributions to:
|
||||
|
||||
1. Add a new feature to the library (guidance below).
|
||||
2. Improve our documentation and add examples to make it clear how to leverage the supervision library.
|
||||
3. Report bugs and issues in the project.
|
||||
4. Submit a request for a new feature.
|
||||
5. Improve our test coverage.
|
||||
|
||||
### Contributing Features
|
||||
|
||||
Supervision is designed to provide generic utilities to solve problems. Thus, we focus on contributions that can have an impact on a wide range of projects.
|
||||
|
||||
For example, counting objects that cross a line anywhere on an image is a common problem in computer vision, but counting objects that cross a line 75% of the way through is less useful.
|
||||
|
||||
Before you contribute a new feature, consider submitting an Issue to discuss the feature so the community can weigh in and assist.
|
||||
|
||||
## How to Contribute Changes
|
||||
|
||||
First, fork this repository to your own GitHub account. Create a new branch that describes your changes (i.e. `line-counter-docs`). Push your changes to the branch on your fork and then submit a pull request to this repository.
|
||||
|
||||
When creating new functions, please ensure you have the following:
|
||||
|
||||
1. Docstrings for the function and all parameters.
|
||||
2. Unit tests for the function.
|
||||
3. Examples in the documentation for the function.
|
||||
4. Created an entry in our docs to autogenerate the documentation for the function.
|
||||
|
||||
All pull requests will be reviewed by the maintainers of the project. We will provide feedback and ask for changes if necessary.
|
||||
|
||||
PRs must pass all tests and linting requirements before they can be merged.
|
||||
|
||||
## 🧹 code quality
|
||||
|
||||
We provide two handy commands inside the `Makefile`, namely:
|
||||
|
||||
- `make style` to format the code
|
||||
- `make check_code_quality` to check code quality (PEP8 basically)
|
||||
|
||||
So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
|
||||
|
||||
## 🧪 tests
|
||||
|
||||
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
|
||||
@@ -2,70 +2,41 @@
|
||||
<p>
|
||||
<a align="center" href="" target="_blank">
|
||||
<img
|
||||
width="850"
|
||||
src="https://media.roboflow.com/open-source/supervision/roboflow-supervision-banner.png?ik-sdk-version=javascript-1.4.3&updatedAt=1674062891088"
|
||||
width="100%"
|
||||
src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529"
|
||||
>
|
||||
</a>
|
||||
</p>
|
||||
|
||||
<br>
|
||||
|
||||
<div align="center">
|
||||
<a href="https://youtube.com/roboflow">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/youtube.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949634652"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/roboflow-app.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949746649"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://www.linkedin.com/company/roboflow-ai/">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/linkedin.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633691"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://docs.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/knowledge.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949634511"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://disuss.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/forum.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633584"
|
||||
width="3%"
|
||||
/>
|
||||
<img src="https://github.com/SkalskiP/SkalskiP/blob/master/icons/transparent.png" width="3%"/>
|
||||
<a href="https://blog.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/blog.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633605"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
</a>
|
||||
</div>
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://pypistats.org/packages/supervision)
|
||||
[](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb)
|
||||
|
||||
</div>
|
||||
|
||||
## 👋 hello
|
||||
|
||||
A set of easy-to-use utils that will come in handy in any Computer Vision project. **Supervision** is still in
|
||||
pre-release stage. 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid
|
||||
and may change.
|
||||
**We write your reusable computer vision tools.** Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us! 🤝
|
||||
|
||||
## 💻 install
|
||||
|
||||
```console
|
||||
Pip install the supervision package in a
|
||||
[**3.11>=Python>=3.7**](https://www.python.org/) environment.
|
||||
|
||||
```bash
|
||||
pip install supervision
|
||||
```
|
||||
|
||||
<details close>
|
||||
<summary>👉 install from source</summary>
|
||||
|
||||
```bash
|
||||
# clone repository and navigate to root directory
|
||||
git clone git@github.com:roboflow/supervision.git
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision
|
||||
|
||||
# setup python environment and activate it
|
||||
@@ -76,15 +47,270 @@ source venv/bin/activate
|
||||
pip install -e ".[dev]"
|
||||
```
|
||||
|
||||
## 🧹 Code Quality
|
||||
</details>
|
||||
|
||||
We provide two handy commands inside the `Makefile`, namely:
|
||||
## 🔥 quickstart
|
||||
|
||||
- `make style` to format the code
|
||||
- `make check_code_quality` to check code quality (PEP8 basically)
|
||||
### [detections processing](https://roboflow.github.io/supervision/detection/core/)
|
||||
|
||||
So far, **there is no types checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
## 🧪 Tests
|
||||
>>> model = YOLO('yolov8s.pt')
|
||||
>>> result = model(IMAGE)[0]
|
||||
>>> detections = sv.Detections.from_yolov8(result)
|
||||
|
||||
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
|
||||
>>> len(detections)
|
||||
5
|
||||
```
|
||||
|
||||
<details close>
|
||||
<summary>👉 more detections utils</summary>
|
||||
|
||||
- Easily switch inference pipeline between supported object detection / instance segmentation models
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
|
||||
|
||||
>>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
|
||||
>>> mask_generator = SamAutomaticMaskGenerator(sam)
|
||||
>>> sam_result = mask_generator.generate(IMAGE)
|
||||
>>> detections = sv.Detections.from_sam(sam_result=sam_result)
|
||||
```
|
||||
|
||||
- [Advanced filtering](https://roboflow.github.io/supervision/quickstart/detections/)
|
||||
|
||||
```python
|
||||
>>> detections = detections[detections.class_id == 0]
|
||||
>>> detections = detections[detections.confidence > 0.5]
|
||||
>>> detections = detections[detections.area > 1000]
|
||||
```
|
||||
|
||||
- Image annotation
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> box_annotator = sv.BoxAnnotator()
|
||||
>>> annotated_frame = box_annotator.annotate(
|
||||
... scene=IMAGE,
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
### [datasets processing](https://roboflow.github.io/supervision/dataset/core/)
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> len(dataset)
|
||||
1000
|
||||
```
|
||||
|
||||
<details close>
|
||||
<summary>👉 more dataset utils</summary>
|
||||
|
||||
- Load object detection / instance segmentation datasets in one of supported formats
|
||||
|
||||
```python
|
||||
>>> dataset = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_pascal_voc(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Loop over dataset entries
|
||||
|
||||
```python
|
||||
>>> for name, image, labels in dataset:
|
||||
... print(labels.xyxy)
|
||||
|
||||
array([[404. , 719. , 538. , 884.5 ],
|
||||
[155. , 497. , 404. , 833.5 ],
|
||||
[ 20.154999, 347.825 , 416.125 , 915.895 ]], dtype=float32)
|
||||
```
|
||||
|
||||
- Split dataset for training, testing and validation
|
||||
|
||||
```python
|
||||
>>> train_dataset, test_dataset = dataset.split(split_ratio=0.7)
|
||||
>>> test_dataset, valid_dataset = test_dataset.split(split_ratio=0.5)
|
||||
|
||||
>>> len(train_dataset), len(test_dataset), len(valid_dataset)
|
||||
(700, 150, 150)
|
||||
```
|
||||
|
||||
- Merge multiple datasets together
|
||||
|
||||
```python
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
100
|
||||
>>> ds_1.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> ds_2 = sv.DetectionDataset(...)
|
||||
>>> len(ds_2)
|
||||
200
|
||||
>>> ds_2.classes
|
||||
['cat']
|
||||
|
||||
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
|
||||
>>> len(ds_merged)
|
||||
300
|
||||
>>> ds_merged.classes
|
||||
['cat', 'dog', 'person']
|
||||
```
|
||||
|
||||
- Save object detection / instance segmentation datasets in one of supported formats
|
||||
|
||||
```python
|
||||
>>> dataset.as_yolo(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset.as_pascal_voc(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> dataset.as_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Convert labels between supported formats
|
||||
|
||||
```python
|
||||
>>> sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...',
|
||||
... data_yaml_path='...'
|
||||
... ).as_pascal_voc(
|
||||
... images_directory_path='...',
|
||||
... annotations_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Load classification datasets in one of supported formats
|
||||
|
||||
```python
|
||||
>>> cs = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Save classification datasets in one of supported formats
|
||||
|
||||
```python
|
||||
>>> cs.as_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
</details>
|
||||
|
||||
## 🎬 tutorials
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/ae1ca38e-40b7-4b35-8582-e8ea5de3806e" alt="Accelerate Image Annotation with SAM and Grounding DINO" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><strong>Accelerate Image Annotation with SAM and Grounding DINO</strong></a>
|
||||
<div><strong>Created: 20 Apr 2023</strong> | <strong>Updated: 20 Apr 2023</strong></div>
|
||||
<br/> Discover how to speed up your image annotation process using Grounding DINO and Segment Anything Model (SAM). Learn how to convert object detection datasets into instance segmentation datasets, and see the potential of using these models to automatically annotate your datasets for real-time detectors like YOLOv8... </p>
|
||||
|
||||
<br/>
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/6913ff11-53c6-4341-8d90-eaff3023c3fd" alt="SAM - Segment Anything Model by Meta AI: Complete Guide" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><strong>SAM - Segment Anything Model by Meta AI: Complete Guide</strong></a>
|
||||
<div><strong>Created: 11 Apr 2023</strong> | <strong>Updated: 11 Apr 2023</strong></div>
|
||||
<br/> Discover the incredible potential of Meta AI's Segment Anything Model (SAM)! We dive into SAM, an efficient and promptable model for image segmentation, which has revolutionized computer vision tasks. With over 1 billion masks on 11M licensed and privacy-respecting images, SAM's zero-shot performance is often competitive with or even superior to prior fully supervised results... </p>
|
||||
|
||||
## 📚 documentation
|
||||
|
||||
Curious how Supervision can help you solve problems on your project? Visit our [documentation](https://roboflow.github.io/supervision) page!
|
||||
|
||||
## 💜 built with supervision
|
||||
|
||||
You built something cool using supervision? [Let us know!](https://github.com/roboflow/supervision/discussions/categories/built-with-supervision)
|
||||
|
||||
## 🏆 contribution
|
||||
|
||||
We love your input! Please see our [contributing guide](https://github.com/roboflow/supervision/blob/main/CONTRIBUTING.md) to get started. Thank you 🙏 to all our contributors!
|
||||
|
||||
<br>
|
||||
|
||||
<div align="center">
|
||||
|
||||
<div align="center">
|
||||
<a href="https://youtube.com/roboflow">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/youtube.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949634652"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://raw.githubusercontent.com/ultralytics/assets/main/social/logo-transparent.png" width="3%"/>
|
||||
<a href="https://roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/roboflow-app.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949746649"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://raw.githubusercontent.com/ultralytics/assets/main/social/logo-transparent.png" width="3%"/>
|
||||
<a href="https://www.linkedin.com/company/roboflow-ai/">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/linkedin.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633691"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://raw.githubusercontent.com/ultralytics/assets/main/social/logo-transparent.png" width="3%"/>
|
||||
<a href="https://docs.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/knowledge.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949634511"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
<img src="https://raw.githubusercontent.com/ultralytics/assets/main/social/logo-transparent.png" width="3%"/>
|
||||
<a href="https://disuss.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/forum.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633584"
|
||||
width="3%"
|
||||
/>
|
||||
<img src="https://raw.githubusercontent.com/ultralytics/assets/main/social/logo-transparent.png" width="3%"/>
|
||||
<a href="https://blog.roboflow.com">
|
||||
<img
|
||||
src="https://media.roboflow.com/notebooks/template/icons/purple/blog.png?ik-sdk-version=javascript-1.4.3&updatedAt=1672949633605"
|
||||
width="3%"
|
||||
/>
|
||||
</a>
|
||||
</a>
|
||||
</div>
|
||||
|
||||
</div>
|
||||
|
||||
Vendored
+1325
File diff suppressed because one or more lines are too long
@@ -0,0 +1,29 @@
|
||||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<!-- Generator: Adobe Illustrator 25.3.1, SVG Export Plug-In . SVG Version: 6.00 Build 0) -->
|
||||
<svg version="1.1" id="Layer_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
|
||||
viewBox="0 0 600 600" style="enable-background:new 0 0 600 600;" xml:space="preserve">
|
||||
<style type="text/css">
|
||||
.st0{fill-rule:evenodd;clip-rule:evenodd;fill:#6706CE;}
|
||||
.st1{fill-rule:evenodd;clip-rule:evenodd;fill:#FFFFFF;}
|
||||
</style>
|
||||
<g>
|
||||
<circle class="st0" cx="300" cy="300" r="300"/>
|
||||
</g>
|
||||
<g>
|
||||
<path class="st1" d="M279.34,267.67c27.7-21.58,15.55-51.11-30.24-45.4c-53.03,6.6-120.62,58.45-137.46,120.7
|
||||
c-1.38,5.09-2.45,10.27-3.19,15.5c-0.73,5.15-1.12,10.35-1.17,15.55c-0.05,5.15,0.25,10.3,0.89,15.4
|
||||
c0.64,5.11,1.64,10.17,2.98,15.14c1.36,5.03,3.06,9.97,5.09,14.76c2.09,4.92,4.51,9.68,7.24,14.26c2.83,4.76,5.98,9.31,9.4,13.64
|
||||
c11.67,14.78,26.65,27.16,42.63,37.02c3.46,2.14,7.18,4.21,11.25,4.41c5.94,0.28,10.44-4.44,10.68-10.23
|
||||
c0.09-2.18-0.46-4.24-1.46-6c-13.42-23.67-24.52-58.81-22.24-80.99C179.43,336.29,235.28,301.98,279.34,267.67z"/>
|
||||
<path class="st1" d="M500.28,363.84c-2.9-6.91-10.16-10.81-17.51-10.15c-5.02,0.45-9.17,3.13-12.55,6.68
|
||||
c-5.49,4.97-9.31,11.89-13.43,17.91c-4.86,7.08-9.94,14.02-15.41,20.63c-10.8,13.04-23.24,24.79-37.76,33.61
|
||||
c-24.63,14.95-47.17,22.12-74.27,8.31c-19.83-15.86-13.13-34.6,3.29-49.87c28.08-26.1,61.93-44.42,88.52-72.44
|
||||
c25-26.34,47.05-59.11,52.24-95.78c1.46-10.27,1.48-20.74-0.17-31c-7.8-48.31-51.48-78.97-95.59-92.01
|
||||
C321.47,83.13,259.6,88,206.48,112.56c-21.86,10.11-42.23,23.65-59.94,39.98c-7.85,7.24-15.24,15.01-22.1,23.2
|
||||
c-11.1,13.27-23.73,28.77-25.89,46.57c-1.33,10.99,8.58,18.59,18.99,17.18c6.98-0.94,13.16-5.15,18.3-9.74
|
||||
c5.97-5.33,11.45-11.25,17.21-16.82c17.38-16.82,35.94-32.47,58.01-42.79c28.11-13.15,59.61-19.55,90.6-16.32
|
||||
c31.43,3.27,69.67,17.44,81.15,49.98c11.97,33.97-13.73,67.36-37.18,89.06c-38.29,35.44-110.99,71.63-102.98,134.19
|
||||
c4.29,33.51,28.4,62.19,59.24,74.96c53.66,22.23,111.74-4.84,148.89-46.18c12.46-13.24,30.38-35.82,46.68-69.46
|
||||
C500.61,379.84,503.23,370.87,500.28,363.84z"/>
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 2.1 KiB |
@@ -0,0 +1,29 @@
|
||||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<!-- Generator: Adobe Illustrator 25.3.1, SVG Export Plug-In . SVG Version: 6.00 Build 0) -->
|
||||
<svg version="1.1" id="Layer_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
|
||||
viewBox="0 0 600 600" style="enable-background:new 0 0 600 600;" xml:space="preserve">
|
||||
<style type="text/css">
|
||||
.st0{fill-rule:evenodd;clip-rule:evenodd;fill:#FFFFFF;}
|
||||
.st1{fill-rule:evenodd;clip-rule:evenodd;fill:#6706CE;}
|
||||
</style>
|
||||
<g>
|
||||
<circle class="st0" cx="300" cy="300" r="300"/>
|
||||
</g>
|
||||
<g>
|
||||
<path class="st1" d="M279.34,267.67c27.7-21.58,15.55-51.11-30.24-45.4c-53.03,6.6-120.62,58.45-137.46,120.7
|
||||
c-1.38,5.09-2.45,10.27-3.19,15.5c-0.73,5.15-1.12,10.35-1.17,15.55c-0.05,5.15,0.25,10.3,0.89,15.4
|
||||
c0.64,5.11,1.64,10.17,2.98,15.14c1.36,5.03,3.06,9.97,5.09,14.76c2.09,4.92,4.51,9.68,7.24,14.26c2.83,4.76,5.98,9.31,9.4,13.64
|
||||
c11.67,14.78,26.65,27.16,42.63,37.02c3.46,2.14,7.18,4.21,11.25,4.41c5.94,0.28,10.44-4.44,10.68-10.23
|
||||
c0.09-2.18-0.46-4.24-1.46-6c-13.42-23.67-24.52-58.81-22.24-80.99C179.43,336.29,235.28,301.98,279.34,267.67z"/>
|
||||
<path class="st1" d="M500.28,363.84c-2.9-6.91-10.16-10.81-17.51-10.15c-5.02,0.45-9.17,3.13-12.55,6.68
|
||||
c-5.49,4.97-9.31,11.89-13.43,17.91c-4.86,7.08-9.94,14.02-15.41,20.63c-10.8,13.04-23.24,24.79-37.76,33.61
|
||||
c-24.63,14.95-47.17,22.12-74.27,8.31c-19.83-15.86-13.13-34.6,3.29-49.87c28.08-26.1,61.93-44.42,88.52-72.44
|
||||
c25-26.34,47.05-59.11,52.24-95.78c1.46-10.27,1.48-20.74-0.17-31c-7.8-48.31-51.48-78.97-95.59-92.01
|
||||
C321.47,83.13,259.6,88,206.48,112.56c-21.86,10.11-42.23,23.65-59.94,39.98c-7.85,7.24-15.24,15.01-22.1,23.2
|
||||
c-11.1,13.27-23.73,28.77-25.89,46.57c-1.33,10.99,8.58,18.59,18.99,17.18c6.98-0.94,13.16-5.15,18.3-9.74
|
||||
c5.97-5.33,11.45-11.25,17.21-16.82c17.38-16.82,35.94-32.47,58.01-42.79c28.11-13.15,59.61-19.55,90.6-16.32
|
||||
c31.43,3.27,69.67,17.44,81.15,49.98c11.97,33.97-13.73,67.36-37.18,89.06c-38.29,35.44-110.99,71.63-102.98,134.19
|
||||
c4.29,33.51,28.4,62.19,59.24,74.96c53.66,22.23,111.74-4.84,148.89-46.18c12.46-13.24,30.38-35.82,46.68-69.46
|
||||
C500.61,379.84,503.23,370.87,500.28,363.84z"/>
|
||||
</g>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 2.1 KiB |
@@ -0,0 +1,23 @@
|
||||
<?xml version="1.0" encoding="utf-8"?>
|
||||
<!-- Generator: Adobe Illustrator 25.3.1, SVG Export Plug-In . SVG Version: 6.00 Build 0) -->
|
||||
<svg version="1.1" id="Layer_1" xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" x="0px" y="0px"
|
||||
viewBox="0 0 600 600" style="enable-background:new 0 0 600 600;" xml:space="preserve">
|
||||
<style type="text/css">
|
||||
.st0{fill-rule:evenodd;clip-rule:evenodd;fill:#FFFFFF;}
|
||||
</style>
|
||||
<path class="st0" d="M300,0C134.31,0,0,134.31,0,300s134.31,300,300,300s300-134.31,300-300S465.69,0,300,0z M195.98,472.42
|
||||
c1,1.76,1.55,3.82,1.46,6c-0.24,5.79-4.74,10.51-10.68,10.23c-4.07-0.19-7.78-2.27-11.25-4.41c-15.98-9.86-30.96-22.24-42.63-37.02
|
||||
c-3.42-4.33-6.57-8.89-9.4-13.64c-2.72-4.58-5.15-9.34-7.24-14.26c-2.03-4.79-3.74-9.72-5.09-14.76
|
||||
c-1.34-4.97-2.33-10.03-2.98-15.14c-0.64-5.1-0.94-10.25-0.89-15.4c0.05-5.2,0.45-10.4,1.17-15.55c0.74-5.23,1.81-10.41,3.19-15.5
|
||||
c16.84-62.24,84.43-114.09,137.46-120.7c45.79-5.71,57.94,23.81,30.24,45.4c-44.06,34.32-99.92,68.63-105.6,123.77
|
||||
C171.46,413.61,182.55,448.75,195.98,472.42z M497.44,386.39c-16.3,33.65-34.22,56.22-46.68,69.46
|
||||
c-37.14,41.34-95.23,68.4-148.89,46.18c-30.83-12.77-54.94-41.45-59.24-74.96c-8.02-62.55,64.69-98.74,102.98-134.19
|
||||
c23.45-21.71,49.16-55.1,37.18-89.06c-11.47-32.55-49.72-46.71-81.15-49.98c-30.99-3.23-62.49,3.17-90.6,16.32
|
||||
c-22.07,10.32-40.62,25.97-58.01,42.79c-5.75,5.57-11.24,11.49-17.21,16.82c-5.14,4.59-11.31,8.79-18.29,9.74
|
||||
c-10.41,1.41-20.32-6.18-18.99-17.18c2.16-17.8,14.79-33.29,25.89-46.57c6.86-8.19,14.25-15.96,22.1-23.2
|
||||
c17.71-16.34,38.08-29.87,59.94-39.98C259.6,88,321.47,83.13,377.64,99.73c44.11,13.04,87.79,43.7,95.59,92.01
|
||||
c1.65,10.25,1.63,20.72,0.17,31c-5.19,36.67-27.24,69.44-52.24,95.78c-26.59,28.02-60.44,46.35-88.52,72.44
|
||||
c-16.43,15.27-23.13,34.01-3.29,49.87c27.1,13.81,49.64,6.64,74.27-8.31c14.52-8.82,26.96-20.57,37.76-33.61
|
||||
c5.48-6.61,10.56-13.55,15.41-20.63c4.13-6.02,7.94-12.94,13.43-17.91c3.38-3.55,7.53-6.23,12.55-6.68
|
||||
c7.35-0.66,14.61,3.23,17.51,10.15C503.23,370.87,500.61,379.84,497.44,386.39z"/>
|
||||
</svg>
|
||||
|
After Width: | Height: | Size: 2.0 KiB |
@@ -0,0 +1,208 @@
|
||||
### 0.11.0 <small>June 28, 2023</small>
|
||||
|
||||
- Added [[#150](https://github.com/roboflow/supervision/pull/150)]: ability to load and save [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) in COCO format using [`as_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_coco) and [`from_coco`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.from_coco) methods.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
|
||||
>>> ds.as_coco(
|
||||
... images_directory_path='...',
|
||||
... annotations_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Added [[#158](https://github.com/roboflow/supervision/pull/158)]: ability to marge multiple [`sv.DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) together using [`merge`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.merge) method.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
100
|
||||
>>> ds_1.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> ds_2 = sv.DetectionDataset(...)
|
||||
>>> len(ds_2)
|
||||
200
|
||||
>>> ds_2.classes
|
||||
['cat']
|
||||
|
||||
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
|
||||
>>> len(ds_merged)
|
||||
300
|
||||
>>> ds_merged.classes
|
||||
['cat', 'dog', 'person']
|
||||
```
|
||||
|
||||
- Added [[#162](https://github.com/roboflow/supervision/pull/162)]: additional `start` and `end` arguments to [`sv.get_video_frames_generator`](https://roboflow.github.io/supervision/utils/video/#get_video_frames_generator) allowing to generate frames only for a selected part of the video.
|
||||
|
||||
- Fix [[#157](https://github.com/roboflow/supervision/pull/157)]: incorrect loading of YOLO dataset class names from `data.yaml`.
|
||||
|
||||
### 0.10.0 <small>June 14, 2023</small>
|
||||
|
||||
- Added [[#125](https://github.com/roboflow/supervision/pull/125)]: ability to load and save [`sv.ClassificationDataset`](https://roboflow.github.io/supervision/dataset/core/#classificationdataset) in a folder structure format.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> cs = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
|
||||
>>> cs.as_folder_structure(
|
||||
... root_directory_path='...'
|
||||
... )
|
||||
```
|
||||
|
||||
- Added [[#125](https://github.com/roboflow/supervision/pull/125)]: support for [`sv.ClassificationDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.ClassificationDataset.split) allowing to divide `sv.ClassificationDataset` into two parts.
|
||||
|
||||
- Added [[#110](https://github.com/roboflow/supervision/pull/110)]: ability to extract masks from Roboflow API results using [`sv.Detections.from_roboflow`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_roboflow).
|
||||
|
||||
- Added [[commit hash](https://github.com/roboflow/supervision/commit/d000292eb2f2342544e0947b65528082e60fb8d6)]: Supervision Quickstart [notebook](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb) where you can learn more about Detection, Dataset and Video APIs.
|
||||
|
||||
- Changed [[#135](https://github.com/roboflow/supervision/pull/135)]: `sv.get_video_frames_generator` documentation to better describe actual behavior.
|
||||
|
||||
### 0.9.0 <small>June 7, 2023</small>
|
||||
|
||||
- Added [[#118](https://github.com/roboflow/supervision/pull/118)]: ability to select [`sv.Detections`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.__getitem__) by index, list of indexes or slice. Here is an example illustrating the new selection methods.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> detections = sv.Detections(...)
|
||||
>>> len(detections[0])
|
||||
1
|
||||
>>> len(detections[[0, 1]])
|
||||
2
|
||||
>>> len(detections[0:2])
|
||||
2
|
||||
```
|
||||
|
||||
- Added [[#101](https://github.com/roboflow/supervision/pull/101)]: ability to extract masks from YOLOv8 result using [`sv.Detections.from_yolov8`](https://roboflow.github.io/supervision/detection/core/#supervision.detection.core.Detections.from_yolov8). Here is an example illustrating how to extract boolean masks from the result of the YOLOv8 model inference.
|
||||
|
||||
- Added [[#122](https://github.com/roboflow/supervision/pull/122)]: ability to crop image using [`sv.crop`](https://roboflow.github.io/supervision/utils/image/#crop). Here is an example showing how to get a separate crop for each detection in `sv.Detections`.
|
||||
|
||||
- Added [[#120](https://github.com/roboflow/supervision/pull/120)]: ability to conveniently save multiple images into directory using [`sv.ImageSink`](https://roboflow.github.io/supervision/utils/image/#imagesink). Here is an example showing how to save every tenth video frame as a separate image.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
|
||||
... for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=10):
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
|
||||
- Fixed [[#106](https://github.com/roboflow/supervision/issues/106)]: inconvenient handling of [`sv.PolygonZone`](https://roboflow.github.io/supervision/detection/tools/polygon_zone/#polygonzone) coordinates. Now `sv.PolygonZone` accepts coordinates in the form of `[[x1, y1], [x2, y2], ...]` that can be both integers and floats.
|
||||
|
||||
### 0.8.0 <small>May 17, 2023</small>
|
||||
|
||||
- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: support for dataset inheritance. The current `Dataset` got renamed to `DetectionDataset`. Now [`DetectionDataset`](https://roboflow.github.io/supervision/dataset/core/#detectiondataset) inherits from `BaseDataset`. This change was made to enforce the future consistency of APIs of different types of computer vision datasets.
|
||||
- Added [[#100](https://github.com/roboflow/supervision/pull/100)]: ability to save datasets in YOLO format using [`DetectionDataset.as_yolo`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.as_yolo).
|
||||
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("yolov5")
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=f"{dataset.location}/train/images",
|
||||
... annotations_directory_path=f"{dataset.location}/train/labels",
|
||||
... data_yaml_path=f"{dataset.location}/data.yaml"
|
||||
... )
|
||||
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
|
||||
- Added [[#102](https://github.com/roboflow/supervision/pull/103)]: support for [`DetectionDataset.split`](https://roboflow.github.io/supervision/dataset/core/#supervision.dataset.core.DetectionDataset.split) allowing to divide `DetectionDataset` into two parts.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds = sv.DetectionDataset(...)
|
||||
>>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
|
||||
>>> len(train_ds), len(test_ds)
|
||||
(700, 300)
|
||||
```
|
||||
|
||||
- Changed [[#100](https://github.com/roboflow/supervision/pull/100)]: default value of `approximation_percentage` parameter from `0.75` to `0.0` in `DetectionDataset.as_yolo` and `DetectionDataset.as_pascal_voc`.
|
||||
|
||||
### 0.7.0 <small>May 11, 2023</small>
|
||||
|
||||
- Added [[#91](https://github.com/roboflow/supervision/pull/91)]: `Detections.from_yolo_nas` to enable seamless integration with [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) model.
|
||||
- Added [[#86](https://github.com/roboflow/supervision/pull/86)]: ability to load datasets in YOLO format using `Dataset.from_yolo`.
|
||||
- Added [[#84](https://github.com/roboflow/supervision/pull/84)]: `Detections.merge` to merge multiple `Detections` objects together.
|
||||
- Fixed [[#81](https://github.com/roboflow/supervision/pull/81)]: `LineZoneAnnotator.annotate` does not return annotated frame.
|
||||
- Changed [[#44](https://github.com/roboflow/supervision/pull/44)]: `LineZoneAnnotator.annotate` to allow for custom text for the in and out tags.
|
||||
|
||||
### 0.6.0 <small>April 19, 2023</small>
|
||||
|
||||
- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: initial `Dataset` support and ability to save `Detections` in Pascal VOC XML format.
|
||||
- Added [[#71](https://github.com/roboflow/supervision/pull/71)]: new `mask_to_polygons`, `filter_polygons_by_area`, `polygon_to_xyxy` and `approximate_polygon` utilities.
|
||||
- Added [[#72](https://github.com/roboflow/supervision/pull/72)]: ability to load Pascal VOC XML **object detections** dataset as `Dataset`.
|
||||
- Changed [[#70](https://github.com/roboflow/supervision/pull/70)]: order of `Detections` attributes to make it consistent with order of objects in `__iter__` tuple.
|
||||
- Changed [[#71](https://github.com/roboflow/supervision/pull/71)]: `generate_2d_mask` to `polygon_to_mask`.
|
||||
|
||||
### 0.5.2 <small>April 13, 2023</small>
|
||||
|
||||
- Fixed [[#63](https://github.com/roboflow/supervision/pull/63)]: `LineZone.trigger` function expects 4 values instead of 5.
|
||||
|
||||
### 0.5.1 <small>April 12, 2023</small>
|
||||
|
||||
- Fixed `Detections.__getitem__` method did not return mask for selected item.
|
||||
- Fixed `Detections.area` crashed for mask detections.
|
||||
|
||||
### 0.5.0 <small>April 10, 2023</small>
|
||||
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.mask` to enable segmentation support.
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `MaskAnnotator` to allow easy `Detections.mask` annotation.
|
||||
- Added [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.from_sam` to enable native Segment Anything Model (SAM) support.
|
||||
- Changed [[#58](https://github.com/roboflow/supervision/pull/58)]: `Detections.area` behaviour to work not only with boxes but also with masks.
|
||||
|
||||
### 0.4.0 <small>April 5, 2023</small>
|
||||
|
||||
- Added [[#46](https://github.com/roboflow/supervision/discussions/48)]: `Detections.empty` to allow easy creation of empty `Detections` objects.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `Detections.from_roboflow` to allow easy creation of `Detections` objects from Roboflow API inference results.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: `plot_images_grid` to allow easy plotting of multiple images on single plot.
|
||||
- Added [[#56](https://github.com/roboflow/supervision/pull/56)]: initial support for Pascal VOC XML format with `detections_to_voc_xml` method.
|
||||
- Changed [[#56](https://github.com/roboflow/supervision/pull/56)]: `show_frame_in_notebook` refactored and renamed to `plot_image`.
|
||||
|
||||
### 0.3.2 <small>March 23, 2023</small>
|
||||
|
||||
- Changed [[#50](https://github.com/roboflow/supervision/issues/50)]: Allow `Detections.class_id` to be `None`.
|
||||
|
||||
### 0.3.1 <small>March 6, 2023</small>
|
||||
|
||||
- Fixed [[#41](https://github.com/roboflow/supervision/issues/41)]: `PolygonZone` throws an exception when the object touches the bottom edge of the image.
|
||||
- Fixed [[#42](https://github.com/roboflow/supervision/issues/42)]: `Detections.wth_nms` method throws an exception when `Detections` is empty.
|
||||
- Changed [[#36](https://github.com/roboflow/supervision/pull/36)]: `Detections.wth_nms` support class agnostic and non-class agnostic case.
|
||||
|
||||
### 0.3.0 <small>March 6, 2023</small>
|
||||
|
||||
- Changed: Allow `Detections.confidence` to be `None`.
|
||||
- Added: `Detections.from_transformers` and `Detections.from_detectron2` to enable seamless integration with Transformers and Detectron2 models.
|
||||
- Added: `Detections.area` to dynamically calculate bounding box area.
|
||||
- Added: `Detections.wth_nms` to filter out double detections with NMS. Initial - only class agnostic - implementation.
|
||||
|
||||
### 0.2.0 <small>February 2, 2023</small>
|
||||
|
||||
- Added: Advanced `Detections` filtering with pandas-like API.
|
||||
- Added: `Detections.from_yolov5` and `Detections.from_yolov8` to enable seamless integration with YOLOv5 and YOLOv8 models.
|
||||
|
||||
### 0.1.0 <small>January 19, 2023</small>
|
||||
|
||||
Say hello to Supervision 👋
|
||||
@@ -0,0 +1,3 @@
|
||||
## Classifications
|
||||
|
||||
:::supervision.classification.core.Classifications
|
||||
@@ -0,0 +1,12 @@
|
||||
!!! warning
|
||||
|
||||
Dataset API is still fluid and may change. If you use Dataset API in your project until further notice, freeze the
|
||||
`supervision` version in your `requirements.txt` or `setup.py`.
|
||||
|
||||
## DetectionDataset
|
||||
|
||||
:::supervision.dataset.core.DetectionDataset
|
||||
|
||||
## ClassificationDataset
|
||||
|
||||
:::supervision.dataset.core.ClassificationDataset
|
||||
@@ -0,0 +1,7 @@
|
||||
## BoxAnnotator
|
||||
|
||||
:::supervision.detection.annotate.BoxAnnotator
|
||||
|
||||
## MaskAnnotator
|
||||
|
||||
:::supervision.detection.annotate.MaskAnnotator
|
||||
@@ -0,0 +1,3 @@
|
||||
## Detections
|
||||
|
||||
:::supervision.detection.core.Detections
|
||||
@@ -0,0 +1,7 @@
|
||||
## PolygonZone
|
||||
|
||||
:::supervision.detection.tools.polygon_zone.PolygonZone
|
||||
|
||||
## PolygonZoneAnnotator
|
||||
|
||||
:::supervision.detection.tools.polygon_zone.PolygonZoneAnnotator
|
||||
@@ -0,0 +1,27 @@
|
||||
## box_iou_batch
|
||||
|
||||
:::supervision.detection.utils.box_iou_batch
|
||||
|
||||
## non_max_suppression
|
||||
|
||||
:::supervision.detection.utils.non_max_suppression
|
||||
|
||||
## polygon_to_mask
|
||||
|
||||
:::supervision.detection.utils.polygon_to_mask
|
||||
|
||||
## mask_to_xyxy
|
||||
|
||||
:::supervision.detection.utils.mask_to_xyxy
|
||||
|
||||
## mask_to_polygons
|
||||
|
||||
:::supervision.detection.utils.mask_to_polygons
|
||||
|
||||
## polygon_to_xyxy
|
||||
|
||||
:::supervision.detection.utils.polygon_to_xyxy
|
||||
|
||||
## filter_polygons_by_area
|
||||
|
||||
:::supervision.detection.utils.filter_polygons_by_area
|
||||
@@ -0,0 +1,19 @@
|
||||
## draw_line
|
||||
|
||||
:::supervision.draw.utils.draw_line
|
||||
|
||||
## draw_rectangle
|
||||
|
||||
:::supervision.draw.utils.draw_rectangle
|
||||
|
||||
## draw_filled_rectangle
|
||||
|
||||
:::supervision.draw.utils.draw_filled_rectangle
|
||||
|
||||
## draw_polygon
|
||||
|
||||
:::supervision.draw.utils.draw_polygon
|
||||
|
||||
## draw_text
|
||||
|
||||
:::supervision.draw.utils.draw_text
|
||||
@@ -0,0 +1,37 @@
|
||||
<div align="center">
|
||||
<p>
|
||||
<a align="center" href="" target="_blank">
|
||||
<img
|
||||
width="850"
|
||||
src="https://media.roboflow.com/open-source/supervision/roboflow-supervision-banner.png?ik-sdk-version=javascript-1.4.3&updatedAt=1674062891088"
|
||||
>
|
||||
</a>
|
||||
</p>
|
||||
</div>
|
||||
|
||||
## 👋 Welcome
|
||||
|
||||
Supervision is a set of easy-to-use utilities that will come in handy in any computer vision project.
|
||||
|
||||
**Supervision** is still in
|
||||
pre-release stage 🚧 Keep your eyes open for potential bugs and be aware that at this stage our API is still fluid and may change.
|
||||
|
||||
## 💻 How to Install
|
||||
|
||||
You can install `supervision` with pip in a
|
||||
[**3.11>=Python>=3.7**](https://www.python.org/) environment.
|
||||
|
||||
!!! example "Pip install method (recommended)"
|
||||
|
||||
```bash
|
||||
pip install supervision
|
||||
```
|
||||
|
||||
!!! example "Git clone method (for development)"
|
||||
|
||||
```bash
|
||||
git https://github.com/roboflow/supervision.git
|
||||
cd supervision
|
||||
pip install -e '.[dev]'
|
||||
```
|
||||
See contributing section to know more about contributing to the project
|
||||
@@ -0,0 +1,312 @@
|
||||
## advanced filtering
|
||||
|
||||
|
||||
The advanced filtering capabilities of the `Detections` class offer users a versatile and efficient way to narrow down
|
||||
and refine object detections. This section outlines various filtering methods, including filtering by specific class
|
||||
or a set of classes, confidence, object area, bounding box area, relative area, box dimensions, and designated zones.
|
||||
Each method is demonstrated with concise code examples to provide users with a clear understanding of how to implement
|
||||
the filters in their applications.
|
||||
|
||||
### by specific class
|
||||
|
||||
Allows you to select detections that belong only to one selected class.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[detections.class_id == 0]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[detections.class_id == 0]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
|
||||
### by set of classes
|
||||
|
||||
Allows you to select detections that belong only to selected set of classes.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
selected_classes = [0, 2, 3]
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[np.isin(detections.class_id, selected_classes)]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import numpy as np
|
||||
import supervision as sv
|
||||
|
||||
class_id = [0, 2, 3]
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[np.isin(detections.class_id, class_id)]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
### by confidence
|
||||
|
||||
Allows you to select detections with specific confidence value, for example higher than selected threshold.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[detections.confidence > 0.5]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[detections.confidence > 0.5]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
### by area
|
||||
|
||||
Allows you to select detections based on their size. We define the area as the number of pixels occupied by the
|
||||
detection in the image. In the example below, we have sifted out the detections that are too small.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[detections.area > 1000]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[detections.area > 1000]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
### by relative area
|
||||
|
||||
Allows you to select detections based on their size in relation to the size of whole image. Sometimes the concept of
|
||||
detection size changes depending on the image. Detection occupying 10000 square px can be large on a 1280x720 image
|
||||
but small on a 3840x2160 image. In such cases, we can filter out detections based on the percentage of the image area
|
||||
occupied by them. In the example below, we remove too large detections.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
height, width, channels = image.shape
|
||||
image_area = height * width
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[(detections.area / image_area) < 0.8]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
image = ...
|
||||
height, width, channels = image.shape
|
||||
image_area = height * width
|
||||
|
||||
detections = sv.Detections(...)
|
||||
detections = detections[(detections.area / image_area) < 0.8]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
### by box dimensions
|
||||
|
||||
Allows you to select detections based on their dimensions. The size of the bounding box, as well as its coordinates,
|
||||
can be criteria for rejecting detection. Implementing such filtering requires a bit of custom code but is relatively
|
||||
simple and fast.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
|
||||
h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
|
||||
detections = detections[(w > 200) & (h > 200)]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
detections = sv.Detections(...)
|
||||
w = detections.xyxy[:, 2] - detections.xyxy[:, 0]
|
||||
h = detections.xyxy[:, 3] - detections.xyxy[:, 1]
|
||||
detections = detections[(w > 200) & (h > 200)]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
### by `PolygonZone`
|
||||
|
||||
Allows you to use `Detections` in combination with `PolygonZone` to weed out bounding boxes that are in and out of the
|
||||
zone. In the example below you can see how to filter out all detections located in the lower part of the image.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
zone = sv.PolygonZone(...)
|
||||
detections = sv.Detections(...)
|
||||
mask = zone.trigger(detections=detections)
|
||||
detections = detections[mask]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
zone = sv.PolygonZone(...)
|
||||
detections = sv.Detections(...)
|
||||
mask = zone.trigger(detections=detections)
|
||||
detections = detections[mask]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
### by mixed conditions
|
||||
|
||||
`Detections`' greatest strength, however, is that you can build arbitrarily complex logical conditions by simply combining separate conditions using `&` or `|`.
|
||||
|
||||
=== "After"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
zone = sv.PolygonZone(...)
|
||||
detections = sv.Detections(...)
|
||||
mask = zone.trigger(detections=detections)
|
||||
detections = detections[(detections.confidence > 0.7) & mask]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Before"
|
||||
|
||||
```python
|
||||
import supervision as sv
|
||||
|
||||
zone = sv.PolygonZone(...)
|
||||
detections = sv.Detections(...)
|
||||
mask = zone.trigger(detections=detections)
|
||||
detections = detections[mask]
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
@@ -0,0 +1,4 @@
|
||||
:root {
|
||||
--md-primary-fg-color: #8315F9;
|
||||
--md-accent-fg-color: #00FFCE;
|
||||
}
|
||||
@@ -0,0 +1,3 @@
|
||||
## list_files_with_extensions
|
||||
|
||||
:::supervision.utils.file.list_files_with_extensions
|
||||
@@ -0,0 +1,7 @@
|
||||
## ImageSink
|
||||
|
||||
:::supervision.utils.image.ImageSink
|
||||
|
||||
## crop
|
||||
|
||||
:::supervision.utils.image.crop
|
||||
@@ -0,0 +1,7 @@
|
||||
## plot_image
|
||||
|
||||
:::supervision.utils.notebook.plot_image
|
||||
|
||||
## plot_images_grid
|
||||
|
||||
:::supervision.utils.notebook.plot_images_grid
|
||||
@@ -0,0 +1,15 @@
|
||||
## VideoInfo
|
||||
|
||||
:::supervision.utils.video.VideoInfo
|
||||
|
||||
## VideoSink
|
||||
|
||||
:::supervision.utils.video.VideoSink
|
||||
|
||||
## get_video_frames_generator
|
||||
|
||||
:::supervision.utils.video.get_video_frames_generator
|
||||
|
||||
## process_video
|
||||
|
||||
:::supervision.utils.video.process_video
|
||||
+86
@@ -0,0 +1,86 @@
|
||||
site_name: Supervision
|
||||
site_url: https://roboflow.github.io/supervision
|
||||
site_author: Roboflow
|
||||
site_description: A set of easy-to-use utils that will come in handy in any Computer Vision project
|
||||
repo_name: roboflow/supervision
|
||||
repo_url: https://github.com/roboflow/supervision
|
||||
edit_uri: https://github.com/roboflow/supervision/tree/main/docs
|
||||
copyright: Roboflow 2023. All rights reserved.
|
||||
|
||||
extra:
|
||||
social:
|
||||
- icon: fontawesome/brands/github
|
||||
link: https://github.com/roboflow
|
||||
- icon: fontawesome/brands/youtube
|
||||
link: https://www.youtube.com/roboflow
|
||||
- icon: fontawesome/brands/linkedin
|
||||
link: https://www.linkedin.com/company/roboflow-ai/mycompany/
|
||||
- icon: fontawesome/brands/twitter
|
||||
link: https://twitter.com/roboflow
|
||||
analytics:
|
||||
provider: google
|
||||
property: G-P7ZG0Y19G5
|
||||
|
||||
extra_css:
|
||||
- stylesheets/extra.css
|
||||
|
||||
nav:
|
||||
- Home: index.md
|
||||
- Quickstart:
|
||||
- Detections: quickstart/detections.md
|
||||
- API reference:
|
||||
- Classifications:
|
||||
- Core: classification/core.md
|
||||
- Detections:
|
||||
- Core: detection/core.md
|
||||
- Annotate: detection/annotate.md
|
||||
- Utils: detection/utils.md
|
||||
- Tools:
|
||||
- Polygon Zone: detection/tools/polygon_zone.md
|
||||
- Dataset:
|
||||
- Core: dataset/core.md
|
||||
- Draw:
|
||||
- Utils: draw/utils.md
|
||||
- Utils:
|
||||
- Video: utils/video.md
|
||||
- Image: utils/image.md
|
||||
- Notebook: utils/notebook.md
|
||||
- File: utils/file.md
|
||||
- Changelog: changelog.md
|
||||
|
||||
theme:
|
||||
name: 'material'
|
||||
logo: https://media.roboflow.com/open-source/supervision/supervision-lenny.png
|
||||
favicon: https://media.roboflow.com/open-source/supervision/supervision-lenny.png
|
||||
palette:
|
||||
# Palette for light mode
|
||||
- scheme: default
|
||||
primary: 'custom'
|
||||
toggle:
|
||||
icon: material/brightness-7
|
||||
name: Switch to dark mode
|
||||
|
||||
# Palette toggle for dark mode
|
||||
- scheme: slate
|
||||
primary: 'custom'
|
||||
toggle:
|
||||
icon: material/brightness-4
|
||||
name: Switch to light mode
|
||||
font:
|
||||
text: Roboto
|
||||
code: Roboto Mono
|
||||
|
||||
plugins:
|
||||
- mkdocstrings
|
||||
- search
|
||||
|
||||
markdown_extensions:
|
||||
- admonition
|
||||
- pymdownx.details
|
||||
- pymdownx.superfences
|
||||
- attr_list
|
||||
- md_in_html
|
||||
- pymdownx.tabbed:
|
||||
alternate_style: true
|
||||
- toc:
|
||||
permalink: true
|
||||
@@ -24,8 +24,10 @@ setuptools.setup(
|
||||
long_description_content_type='text/markdown',
|
||||
url='https://github.com/roboflow/supervision',
|
||||
install_requires=[
|
||||
'numpy',
|
||||
'opencv-python'
|
||||
'numpy>=1.20.0',
|
||||
'opencv-python',
|
||||
'matplotlib',
|
||||
'pyyaml'
|
||||
],
|
||||
packages=find_packages(exclude=("tests",)),
|
||||
extras_require={
|
||||
@@ -36,18 +38,23 @@ setuptools.setup(
|
||||
'twine',
|
||||
'pytest',
|
||||
'wheel',
|
||||
'notebook'
|
||||
'notebook',
|
||||
'mkdocs-material',
|
||||
'mkdocstrings[python]'
|
||||
],
|
||||
},
|
||||
classifiers=[
|
||||
'Development Status :: 4 - Beta',
|
||||
'Intended Audience :: Developers',
|
||||
'Intended Audience :: Education',
|
||||
'Intended Audience :: Science/Research',
|
||||
'License :: OSI Approved :: BSD License',
|
||||
'License :: OSI Approved :: MIT License',
|
||||
'Programming Language :: Python :: 3',
|
||||
'Programming Language :: Python :: 3.7',
|
||||
'Programming Language :: Python :: 3.8',
|
||||
'Programming Language :: Python :: 3.9',
|
||||
'Programming Language :: Python :: 3.10',
|
||||
'Programming Language :: Python :: 3.11',
|
||||
'Programming Language :: Python :: 3 :: Only',
|
||||
'Topic :: Software Development',
|
||||
'Topic :: Scientific/Engineering',
|
||||
@@ -59,6 +66,6 @@ setuptools.setup(
|
||||
'Operating System :: Unix',
|
||||
'Operating System :: MacOS'
|
||||
],
|
||||
keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLOv5, YOLOv8, Roboflow",
|
||||
keywords="machine-learning, deep-learning, vision, ML, DL, AI, YOLOv5, YOLOv8, SAM, Roboflow",
|
||||
python_requires='>=3.7',
|
||||
)
|
||||
|
||||
+34
-1
@@ -1 +1,34 @@
|
||||
__version__ = "0.0.2"
|
||||
__version__ = "0.11.0"
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.core import (
|
||||
BaseDataset,
|
||||
ClassificationDataset,
|
||||
DetectionDataset,
|
||||
)
|
||||
from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
|
||||
from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
|
||||
from supervision.detection.utils import (
|
||||
box_iou_batch,
|
||||
filter_polygons_by_area,
|
||||
mask_to_polygons,
|
||||
mask_to_xyxy,
|
||||
non_max_suppression,
|
||||
polygon_to_mask,
|
||||
polygon_to_xyxy,
|
||||
)
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
|
||||
from supervision.geometry.core import Point, Position, Rect
|
||||
from supervision.geometry.utils import get_polygon_center
|
||||
from supervision.utils.file import list_files_with_extensions
|
||||
from supervision.utils.image import ImageSink, crop
|
||||
from supervision.utils.notebook import plot_image, plot_images_grid
|
||||
from supervision.utils.video import (
|
||||
VideoInfo,
|
||||
VideoSink,
|
||||
get_video_frames_generator,
|
||||
process_video,
|
||||
)
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Optional, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def _validate_class_ids(class_id: Any, n: int) -> None:
|
||||
"""
|
||||
Ensure that class_id is a 1d np.ndarray with (n, ) shape.
|
||||
"""
|
||||
is_valid = isinstance(class_id, np.ndarray) and class_id.shape == (n,)
|
||||
if not is_valid:
|
||||
raise ValueError("class_id must be 1d np.ndarray with (n, ) shape")
|
||||
|
||||
|
||||
def _validate_confidence(confidence: Any, n: int) -> None:
|
||||
"""
|
||||
Ensure that confidence is a 1d np.ndarray with (n, ) shape.
|
||||
"""
|
||||
if confidence is not None:
|
||||
is_valid = isinstance(confidence, np.ndarray) and confidence.shape == (n,)
|
||||
if not is_valid:
|
||||
raise ValueError("confidence must be 1d np.ndarray with (n, ) shape")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Classifications:
|
||||
class_id: np.ndarray
|
||||
confidence: Optional[np.ndarray] = None
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""
|
||||
Validate the classification inputs.
|
||||
"""
|
||||
n = len(self.class_id)
|
||||
|
||||
_validate_class_ids(self.class_id, n)
|
||||
_validate_confidence(self.confidence, n)
|
||||
|
||||
@classmethod
|
||||
def from_yolov8(cls, yolov8_results) -> Classifications:
|
||||
"""
|
||||
Creates a Classifications instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
|
||||
|
||||
Args:
|
||||
yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
|
||||
|
||||
Returns:
|
||||
Detections: A new Classifications object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = YOLO('yolov8s-cls.pt')
|
||||
>>> result = model(image)[0]
|
||||
>>> classifications = sv.Classifications.from_yolov8(result)
|
||||
```
|
||||
"""
|
||||
confidence = yolov8_results.probs.data.cpu().numpy()
|
||||
return cls(class_id=np.arange(confidence.shape[0]), confidence=confidence)
|
||||
|
||||
def get_top_k(self, k: int) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Retrieve the top k class IDs and confidences, ordered in descending order by confidence.
|
||||
|
||||
Args:
|
||||
k (int): The number of top class IDs and confidences to retrieve.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, np.ndarray]: A tuple containing the top k class IDs and confidences.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> classifications = sv.Classifications(...)
|
||||
|
||||
>>> classifications.get_top_k(1)
|
||||
|
||||
(array([1]), array([0.9]))
|
||||
```
|
||||
"""
|
||||
if self.confidence is None:
|
||||
raise ValueError("top_k could not be calculated, confidence is None")
|
||||
|
||||
order = np.argsort(self.confidence)[::-1]
|
||||
top_k_order = order[:k]
|
||||
top_k_class_id = self.class_id[top_k_order]
|
||||
top_k_confidence = self.confidence[top_k_order]
|
||||
|
||||
return top_k_class_id, top_k_confidence
|
||||
@@ -0,0 +1,642 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Dict, Iterator, List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.formats.coco import (
|
||||
load_coco_annotations,
|
||||
save_coco_annotations,
|
||||
)
|
||||
from supervision.dataset.formats.pascal_voc import (
|
||||
detections_to_pascal_voc,
|
||||
load_pascal_voc_annotations,
|
||||
)
|
||||
from supervision.dataset.formats.yolo import (
|
||||
load_yolo_annotations,
|
||||
save_data_yaml,
|
||||
save_yolo_annotations,
|
||||
)
|
||||
from supervision.dataset.utils import (
|
||||
build_class_index_mapping,
|
||||
map_detections_class_id,
|
||||
merge_class_lists,
|
||||
save_dataset_images,
|
||||
train_test_split,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.utils.file import list_files_with_extensions
|
||||
|
||||
|
||||
@dataclass
|
||||
class BaseDataset(ABC):
|
||||
@abstractmethod
|
||||
def __len__(self) -> int:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
) -> Tuple[BaseDataset, BaseDataset]:
|
||||
pass
|
||||
|
||||
|
||||
@dataclass
|
||||
class DetectionDataset(BaseDataset):
|
||||
"""
|
||||
Dataclass containing information about object detection dataset.
|
||||
|
||||
Attributes:
|
||||
classes (List[str]): List containing dataset class names.
|
||||
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
images: Dict[str, np.ndarray]
|
||||
annotations: Dict[str, Detections]
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""
|
||||
Return the number of images in the dataset.
|
||||
|
||||
Returns:
|
||||
int: The number of images.
|
||||
"""
|
||||
return len(self.images)
|
||||
|
||||
def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Detections]]:
|
||||
"""
|
||||
Iterate over the images and annotations in the dataset.
|
||||
|
||||
Yields:
|
||||
Iterator[Tuple[str, np.ndarray, Detections]]: An iterator that yields tuples containing the image name,
|
||||
the image data, and its corresponding annotation.
|
||||
"""
|
||||
for image_name, image in self.images.items():
|
||||
yield image_name, image, self.annotations.get(image_name, None)
|
||||
|
||||
def __eq__(self, other):
|
||||
if not isinstance(other, DetectionDataset):
|
||||
return False
|
||||
|
||||
if set(self.classes) != set(other.classes):
|
||||
return False
|
||||
|
||||
for key in self.images:
|
||||
if not np.array_equal(self.images[key], other.images[key]):
|
||||
return False
|
||||
if not self.annotations[key] == other.annotations[key]:
|
||||
return False
|
||||
|
||||
return True
|
||||
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
) -> Tuple[DetectionDataset, DetectionDataset]:
|
||||
"""
|
||||
Splits the dataset into two parts (training and testing) using the provided split_ratio.
|
||||
|
||||
Args:
|
||||
split_ratio (float, optional): The ratio of the training set to the entire dataset.
|
||||
random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting.
|
||||
|
||||
Returns:
|
||||
Tuple[DetectionDataset, DetectionDataset]: A tuple containing the training and testing datasets.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds = sv.DetectionDataset(...)
|
||||
>>> train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
>>> len(train_ds), len(test_ds)
|
||||
(700, 300)
|
||||
```
|
||||
"""
|
||||
|
||||
image_names = list(self.images.keys())
|
||||
train_names, test_names = train_test_split(
|
||||
data=image_names,
|
||||
train_ratio=split_ratio,
|
||||
random_state=random_state,
|
||||
shuffle=shuffle,
|
||||
)
|
||||
|
||||
train_dataset = DetectionDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in train_names},
|
||||
annotations={name: self.annotations[name] for name in train_names},
|
||||
)
|
||||
test_dataset = DetectionDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in test_names},
|
||||
annotations={name: self.annotations[name] for name in test_names},
|
||||
)
|
||||
return train_dataset, test_dataset
|
||||
|
||||
def as_pascal_voc(
|
||||
self,
|
||||
images_directory_path: Optional[str] = None,
|
||||
annotations_directory_path: Optional[str] = None,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to PASCAL VOC format. This method saves the images and their corresponding annotations in
|
||||
PASCAL VOC format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_directory_path (Optional[str]): The path to the directory where the annotations in
|
||||
PASCAL VOC format should be saved. If not provided, annotations will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). Argument is used only for segmentation datasets.
|
||||
"""
|
||||
if images_directory_path:
|
||||
images_path = Path(images_directory_path)
|
||||
images_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if annotations_directory_path:
|
||||
annotations_path = Path(annotations_directory_path)
|
||||
annotations_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for image_name, image in self.images.items():
|
||||
detections = self.annotations[image_name]
|
||||
|
||||
if images_directory_path:
|
||||
cv2.imwrite(str(images_path / image_name), image)
|
||||
|
||||
if annotations_directory_path:
|
||||
annotation_name = Path(image_name).stem
|
||||
pascal_voc_xml = detections_to_pascal_voc(
|
||||
detections=detections,
|
||||
classes=self.classes,
|
||||
filename=image_name,
|
||||
image_shape=image.shape,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
with open(annotations_path / f"{annotation_name}.xml", "w") as f:
|
||||
f.write(pascal_voc_xml)
|
||||
|
||||
@classmethod
|
||||
def from_pascal_voc(
|
||||
cls, images_directory_path: str, annotations_directory_path: str
|
||||
) -> DetectionDataset:
|
||||
"""
|
||||
Creates a Dataset instance from PASCAL VOC formatted data.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_directory_path (str): The path to the directory containing the PASCAL VOC XML annotations.
|
||||
|
||||
Returns:
|
||||
DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("voc")
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=f"{dataset.location}/train/images",
|
||||
... annotations_directory_path=f"{dataset.location}/train/labels"
|
||||
... )
|
||||
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
"""
|
||||
image_paths = list_files_with_extensions(
|
||||
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
|
||||
)
|
||||
annotation_paths = list_files_with_extensions(
|
||||
directory=annotations_directory_path, extensions=["xml"]
|
||||
)
|
||||
|
||||
raw_annotations: List[Tuple[str, Detections, List[str]]] = [
|
||||
load_pascal_voc_annotations(annotation_path=str(annotation_path))
|
||||
for annotation_path in annotation_paths
|
||||
]
|
||||
|
||||
classes = []
|
||||
for annotation in raw_annotations:
|
||||
classes.extend(annotation[2])
|
||||
classes = list(set(classes))
|
||||
|
||||
for annotation in raw_annotations:
|
||||
class_id = [classes.index(class_name) for class_name in annotation[2]]
|
||||
annotation[1].class_id = np.array(class_id)
|
||||
|
||||
images = {
|
||||
image_path.name: cv2.imread(str(image_path)) for image_path in image_paths
|
||||
}
|
||||
|
||||
annotations = {
|
||||
image_name: detections for image_name, detections, _ in raw_annotations
|
||||
}
|
||||
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
||||
|
||||
@classmethod
|
||||
def from_yolo(
|
||||
cls,
|
||||
images_directory_path: str,
|
||||
annotations_directory_path: str,
|
||||
data_yaml_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> DetectionDataset:
|
||||
"""
|
||||
Creates a Dataset instance from YOLO formatted data.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
|
||||
data_yaml_path (str): The path to the data YAML file containing class information.
|
||||
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
|
||||
|
||||
Returns:
|
||||
DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("yolov5")
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_yolo(
|
||||
... images_directory_path=f"{dataset.location}/train/images",
|
||||
... annotations_directory_path=f"{dataset.location}/train/labels",
|
||||
... data_yaml_path=f"{dataset.location}/data.yaml"
|
||||
... )
|
||||
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
"""
|
||||
classes, images, annotations = load_yolo_annotations(
|
||||
images_directory_path=images_directory_path,
|
||||
annotations_directory_path=annotations_directory_path,
|
||||
data_yaml_path=data_yaml_path,
|
||||
force_masks=force_masks,
|
||||
)
|
||||
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
||||
|
||||
def as_yolo(
|
||||
self,
|
||||
images_directory_path: Optional[str] = None,
|
||||
annotations_directory_path: Optional[str] = None,
|
||||
data_yaml_path: Optional[str] = None,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to YOLO format. This method saves the images and their corresponding
|
||||
annotations in YOLO format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_directory_path (Optional[str]): The path to the directory where the annotations in
|
||||
YOLO format should be saved. If not provided, annotations will not be saved.
|
||||
data_yaml_path (Optional[str]): The path where the data.yaml file should be saved.
|
||||
If not provided, the file will not be saved.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
|
||||
segmentation datasets.
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
)
|
||||
if annotations_directory_path is not None:
|
||||
save_yolo_annotations(
|
||||
annotations_directory_path=annotations_directory_path,
|
||||
images=self.images,
|
||||
annotations=self.annotations,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
if data_yaml_path is not None:
|
||||
save_data_yaml(data_yaml_path=data_yaml_path, classes=self.classes)
|
||||
|
||||
@classmethod
|
||||
def from_coco(
|
||||
cls,
|
||||
images_directory_path: str,
|
||||
annotations_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> DetectionDataset:
|
||||
"""
|
||||
Creates a Dataset instance from COCO formatted data.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_path (str): The path to the json annotation files.
|
||||
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
|
||||
|
||||
Returns:
|
||||
DetectionDataset: A DetectionDataset instance containing the loaded images and annotations.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("coco")
|
||||
|
||||
>>> ds = sv.DetectionDataset.from_coco(
|
||||
... images_directory_path=f"{dataset.location}/train",
|
||||
... annotations_path=f"{dataset.location}/train/_annotations.coco.json",
|
||||
... )
|
||||
|
||||
>>> ds.classes
|
||||
['dog', 'person']
|
||||
```
|
||||
"""
|
||||
classes, images, annotations = load_coco_annotations(
|
||||
images_directory_path=images_directory_path,
|
||||
annotations_path=annotations_path,
|
||||
force_masks=force_masks,
|
||||
)
|
||||
return DetectionDataset(classes=classes, images=images, annotations=annotations)
|
||||
|
||||
def as_coco(
|
||||
self,
|
||||
images_directory_path: Optional[str] = None,
|
||||
annotations_path: Optional[str] = None,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.0,
|
||||
) -> None:
|
||||
"""
|
||||
Exports the dataset to COCO format. This method saves the images and their corresponding
|
||||
annotations in COCO format.
|
||||
|
||||
Args:
|
||||
images_directory_path (Optional[str]): The path to the directory where the images should be saved.
|
||||
If not provided, images will not be saved.
|
||||
annotations_path (Optional[str]): The path to COCO annotation file.
|
||||
min_image_area_percentage (float): The minimum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
max_image_area_percentage (float): The maximum percentage of detection area relative to
|
||||
the image area for a detection to be included. Argument is used only for segmentation datasets.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon,
|
||||
in the range [0, 1). This is useful for simplifying the annotations. Argument is used only for
|
||||
segmentation datasets.
|
||||
"""
|
||||
if images_directory_path is not None:
|
||||
save_dataset_images(
|
||||
images_directory_path=images_directory_path, images=self.images
|
||||
)
|
||||
if annotations_path is not None:
|
||||
save_coco_annotations(
|
||||
annotation_path=annotations_path,
|
||||
images=self.images,
|
||||
annotations=self.annotations,
|
||||
classes=self.classes,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def merge(cls, dataset_list: List[DetectionDataset]) -> DetectionDataset:
|
||||
"""
|
||||
Merge a list of `DetectionDataset` objects into a single `DetectionDataset` object.
|
||||
|
||||
This method takes a list of `DetectionDataset` objects and combines their respective fields (`classes`, `images`,
|
||||
`annotations`) into a single `DetectionDataset` object.
|
||||
|
||||
Args:
|
||||
dataset_list (List[DetectionDataset]): A list of `DetectionDataset` objects to merge.
|
||||
|
||||
Returns:
|
||||
(DetectionDataset): A single `DetectionDataset` object containing the merged data from the input list.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> ds_1 = sv.DetectionDataset(...)
|
||||
>>> len(ds_1)
|
||||
100
|
||||
>>> ds_1.classes
|
||||
['dog', 'person']
|
||||
|
||||
>>> ds_2 = sv.DetectionDataset(...)
|
||||
>>> len(ds_2)
|
||||
200
|
||||
>>> ds_2.classes
|
||||
['cat']
|
||||
|
||||
>>> ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
|
||||
>>> len(ds_merged)
|
||||
300
|
||||
>>> ds_merged.classes
|
||||
['cat', 'dog', 'person']
|
||||
```
|
||||
"""
|
||||
merged_images, merged_annotations = {}, {}
|
||||
class_lists = [dataset.classes for dataset in dataset_list]
|
||||
merged_classes = merge_class_lists(class_lists=class_lists)
|
||||
|
||||
for dataset in dataset_list:
|
||||
class_index_mapping = build_class_index_mapping(
|
||||
source_classes=dataset.classes, target_classes=merged_classes
|
||||
)
|
||||
for image_name, image, detections in dataset:
|
||||
if image_name in merged_annotations:
|
||||
raise ValueError(
|
||||
f"Image name {image_name} is not unique across datasets."
|
||||
)
|
||||
|
||||
merged_images[image_name] = image
|
||||
merged_annotations[image_name] = map_detections_class_id(
|
||||
source_to_target_mapping=class_index_mapping,
|
||||
detections=detections,
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=merged_classes, images=merged_images, annotations=merged_annotations
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ClassificationDataset(BaseDataset):
|
||||
"""
|
||||
Dataclass containing information about a classification dataset.
|
||||
|
||||
Attributes:
|
||||
classes (List[str]): List containing dataset class names.
|
||||
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
|
||||
annotations (Dict[str, Detections]): Dictionary mapping image name to annotations.
|
||||
"""
|
||||
|
||||
classes: List[str]
|
||||
images: Dict[str, np.ndarray]
|
||||
annotations: Dict[str, Classifications]
|
||||
|
||||
def __len__(self) -> int:
|
||||
return len(self.images)
|
||||
|
||||
def split(
|
||||
self, split_ratio=0.8, random_state=None, shuffle: bool = True
|
||||
) -> Tuple[ClassificationDataset, ClassificationDataset]:
|
||||
"""
|
||||
Splits the dataset into two parts (training and testing) using the provided split_ratio.
|
||||
|
||||
Args:
|
||||
split_ratio (float, optional): The ratio of the training set to the entire dataset.
|
||||
random_state (int, optional): The seed for the random number generator. This is used for reproducibility.
|
||||
shuffle (bool, optional): Whether to shuffle the data before splitting.
|
||||
|
||||
Returns:
|
||||
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing the training and testing datasets.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> cd = sv.ClassificationDataset(...)
|
||||
>>> train_cd, test_cd = cd.split(split_ratio=0.7, random_state=42, shuffle=True)
|
||||
>>> len(train_cd), len(test_cd)
|
||||
(700, 300)
|
||||
```
|
||||
"""
|
||||
image_names = list(self.images.keys())
|
||||
train_names, test_names = train_test_split(
|
||||
data=image_names,
|
||||
train_ratio=split_ratio,
|
||||
random_state=random_state,
|
||||
shuffle=shuffle,
|
||||
)
|
||||
|
||||
train_dataset = ClassificationDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in train_names},
|
||||
annotations={name: self.annotations[name] for name in train_names},
|
||||
)
|
||||
test_dataset = ClassificationDataset(
|
||||
classes=self.classes,
|
||||
images={name: self.images[name] for name in test_names},
|
||||
annotations={name: self.annotations[name] for name in test_names},
|
||||
)
|
||||
return train_dataset, test_dataset
|
||||
|
||||
def as_folder_structure(self, root_directory_path: str) -> None:
|
||||
"""
|
||||
Saves the dataset as a multi-class folder structure.
|
||||
|
||||
Args:
|
||||
root_directory_path (str): The path to the directory where the dataset will be saved.
|
||||
"""
|
||||
os.makedirs(root_directory_path, exist_ok=True)
|
||||
|
||||
for class_name in self.classes:
|
||||
os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
|
||||
|
||||
for image_name in self.images:
|
||||
classification = self.annotations[image_name]
|
||||
image = self.images[image_name]
|
||||
class_id = (
|
||||
classification.class_id[0]
|
||||
if classification.confidence is None
|
||||
else classification.get_top_k(1)[0]
|
||||
)
|
||||
class_name = self.classes[class_id]
|
||||
image_path = os.path.join(root_directory_path, class_name, image_name)
|
||||
cv2.imwrite(image_path, image)
|
||||
|
||||
@classmethod
|
||||
def from_folder_structure(cls, root_directory_path: str) -> ClassificationDataset:
|
||||
"""
|
||||
Load data from a multiclass folder structure into a ClassificationDataset.
|
||||
|
||||
Args:
|
||||
root_directory_path (str): The path to the dataset directory.
|
||||
|
||||
Returns:
|
||||
ClassificationDataset: The dataset.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import roboflow
|
||||
>>> from roboflow import Roboflow
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow.login()
|
||||
|
||||
>>> rf = Roboflow()
|
||||
|
||||
>>> project = rf.workspace(WORKSPACE_ID).project(PROJECT_ID)
|
||||
>>> dataset = project.version(PROJECT_VERSION).download("folder")
|
||||
|
||||
>>> cd = sv.ClassificationDataset.from_folder_structure(
|
||||
... root_directory_path=f"{dataset.location}/train"
|
||||
... )
|
||||
```
|
||||
"""
|
||||
classes = os.listdir(root_directory_path)
|
||||
classes = sorted(set(classes))
|
||||
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
for class_name in classes:
|
||||
class_id = classes.index(class_name)
|
||||
|
||||
for image in os.listdir(os.path.join(root_directory_path, class_name)):
|
||||
image_dir = os.path.join(root_directory_path, class_name, image)
|
||||
images[image] = cv2.imread(image_dir)
|
||||
annotations[image] = Classifications(
|
||||
class_id=np.array([class_id]),
|
||||
)
|
||||
|
||||
return cls(
|
||||
classes=classes,
|
||||
images=images,
|
||||
annotations=annotations,
|
||||
)
|
||||
@@ -0,0 +1,216 @@
|
||||
import os
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_mask
|
||||
from supervision.utils.file import read_json_file, save_json_file
|
||||
|
||||
|
||||
def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]:
|
||||
return [
|
||||
category["name"]
|
||||
for category in sorted(coco_categories, key=lambda category: category["id"])
|
||||
]
|
||||
|
||||
|
||||
def classes_to_coco_categories(classes: List[str]) -> List[dict]:
|
||||
return [
|
||||
{
|
||||
"id": class_id,
|
||||
"name": class_name,
|
||||
"supercategory": "common-objects",
|
||||
}
|
||||
for class_id, class_name in enumerate(classes)
|
||||
]
|
||||
|
||||
|
||||
def group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict],
|
||||
) -> Dict[int, List[dict]]:
|
||||
annotations = {}
|
||||
for annotation in coco_annotations:
|
||||
image_id = annotation["image_id"]
|
||||
if image_id not in annotations:
|
||||
annotations[image_id] = []
|
||||
annotations[image_id].append(annotation)
|
||||
return annotations
|
||||
|
||||
|
||||
def _polygons_to_masks(
|
||||
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
|
||||
) -> np.ndarray:
|
||||
return np.array(
|
||||
[
|
||||
polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
|
||||
for polygon in polygons
|
||||
],
|
||||
dtype=bool,
|
||||
)
|
||||
|
||||
|
||||
def coco_annotations_to_detections(
|
||||
image_annotations: List[dict], resolution_wh: Tuple[int, int], with_masks: bool
|
||||
) -> Detections:
|
||||
if not image_annotations:
|
||||
return Detections.empty()
|
||||
|
||||
class_ids = [
|
||||
image_annotation["category_id"] for image_annotation in image_annotations
|
||||
]
|
||||
xyxy = [image_annotation["bbox"] for image_annotation in image_annotations]
|
||||
xyxy = np.asarray(xyxy)
|
||||
xyxy[:, 2:4] += xyxy[:, 0:2]
|
||||
|
||||
if with_masks:
|
||||
polygons = [
|
||||
np.reshape(
|
||||
np.asarray(image_annotation["segmentation"], dtype=np.int32), (-1, 2)
|
||||
)
|
||||
for image_annotation in image_annotations
|
||||
]
|
||||
mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
|
||||
return Detections(
|
||||
class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
|
||||
)
|
||||
|
||||
return Detections(xyxy=xyxy, class_id=np.asarray(class_ids, dtype=int))
|
||||
|
||||
|
||||
def detections_to_coco_annotations(
|
||||
detections: Detections,
|
||||
image_id: int,
|
||||
annotation_id: int,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> Tuple[List[Dict], int]:
|
||||
coco_annotations = []
|
||||
for xyxy, mask, _, class_id, _ in detections:
|
||||
box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
|
||||
polygon = []
|
||||
if mask is not None:
|
||||
polygon = list(
|
||||
approximate_mask_with_polygons(
|
||||
mask=mask,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)[0].flatten()
|
||||
)
|
||||
coco_annotation = {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": int(class_id),
|
||||
"bbox": [xyxy[0], xyxy[1], box_width, box_height],
|
||||
"area": box_width * box_height,
|
||||
"segmentation": polygon,
|
||||
"iscrowd": 0,
|
||||
}
|
||||
coco_annotations.append(coco_annotation)
|
||||
annotation_id += 1
|
||||
return coco_annotations, annotation_id
|
||||
|
||||
|
||||
def load_coco_annotations(
|
||||
images_directory_path: str,
|
||||
annotations_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
coco_data = read_json_file(file_path=annotations_path)
|
||||
classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
|
||||
coco_images = coco_data["images"]
|
||||
coco_annotations_groups = group_coco_annotations_by_image_id(
|
||||
coco_annotations=coco_data["annotations"]
|
||||
)
|
||||
|
||||
images = {}
|
||||
annotations = {}
|
||||
|
||||
for coco_image in coco_images:
|
||||
image_name, image_width, image_height = (
|
||||
coco_image["file_name"],
|
||||
coco_image["width"],
|
||||
coco_image["height"],
|
||||
)
|
||||
image_annotations = coco_annotations_groups.get(coco_image["id"], [])
|
||||
image_path = os.path.join(images_directory_path, image_name)
|
||||
|
||||
image = cv2.imread(str(image_path))
|
||||
annotation = coco_annotations_to_detections(
|
||||
image_annotations=image_annotations,
|
||||
resolution_wh=(image_width, image_height),
|
||||
with_masks=force_masks,
|
||||
)
|
||||
|
||||
images[image_name] = image
|
||||
annotations[image_name] = annotation
|
||||
|
||||
return classes, images, annotations
|
||||
|
||||
|
||||
def save_coco_annotations(
|
||||
annotation_path: str,
|
||||
images: Dict[str, np.ndarray],
|
||||
annotations: Dict[str, Detections],
|
||||
classes: List[str],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> None:
|
||||
Path(annotation_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
info = {}
|
||||
licenses = [
|
||||
{
|
||||
"id": 1,
|
||||
"url": "https://creativecommons.org/licenses/by/4.0/",
|
||||
"name": "CC BY 4.0",
|
||||
}
|
||||
]
|
||||
|
||||
coco_annotations = []
|
||||
coco_images = []
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
|
||||
image_id = 0
|
||||
annotation_id = 0
|
||||
for image_name, image in images.items():
|
||||
image_height, image_width, _ = image.shape
|
||||
|
||||
coco_image = {
|
||||
"id": image_id,
|
||||
"license": 1,
|
||||
"file_name": image_name,
|
||||
"height": image_height,
|
||||
"width": image_width,
|
||||
"date_captured": datetime.now().strftime("%m/%d/%Y,%H:%M:%S"),
|
||||
}
|
||||
|
||||
coco_images.append(coco_image)
|
||||
detections = annotations[image_name]
|
||||
|
||||
coco_annotation, label_id = detections_to_coco_annotations(
|
||||
detections=detections,
|
||||
image_id=image_id,
|
||||
annotation_id=annotation_id,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
|
||||
coco_annotations.extend(coco_annotation)
|
||||
image_id += 1
|
||||
|
||||
annotation_dict = {
|
||||
"info": info,
|
||||
"licenses": licenses,
|
||||
"categories": coco_categories,
|
||||
"images": coco_images,
|
||||
"annotations": coco_annotations,
|
||||
}
|
||||
save_json_file(annotation_dict, file_path=annotation_path)
|
||||
@@ -0,0 +1,156 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from xml.dom.minidom import parseString
|
||||
from xml.etree.ElementTree import Element, SubElement, parse, tostring
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_xyxy
|
||||
|
||||
|
||||
def object_to_pascal_voc(
|
||||
xyxy: np.ndarray, name: str, polygon: Optional[np.ndarray] = None
|
||||
) -> Element:
|
||||
root = Element("object")
|
||||
|
||||
object_name = SubElement(root, "name")
|
||||
object_name.text = name
|
||||
|
||||
bndbox = SubElement(root, "bndbox")
|
||||
xmin = SubElement(bndbox, "xmin")
|
||||
xmin.text = str(int(xyxy[0]))
|
||||
ymin = SubElement(bndbox, "ymin")
|
||||
ymin.text = str(int(xyxy[1]))
|
||||
xmax = SubElement(bndbox, "xmax")
|
||||
xmax.text = str(int(xyxy[2]))
|
||||
ymax = SubElement(bndbox, "ymax")
|
||||
ymax.text = str(int(xyxy[3]))
|
||||
|
||||
if polygon is not None:
|
||||
object_polygon = SubElement(root, "polygon")
|
||||
for index, point in enumerate(polygon, start=1):
|
||||
x_coordinate, y_coordinate = point
|
||||
x = SubElement(object_polygon, f"x{index}")
|
||||
x.text = str(x_coordinate)
|
||||
y = SubElement(object_polygon, f"y{index}")
|
||||
y.text = str(y_coordinate)
|
||||
|
||||
return root
|
||||
|
||||
|
||||
def detections_to_pascal_voc(
|
||||
detections: Detections,
|
||||
classes: List[str],
|
||||
filename: str,
|
||||
image_shape: Tuple[int, int, int],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> str:
|
||||
"""
|
||||
Converts Detections object to Pascal VOC XML format.
|
||||
|
||||
Args:
|
||||
detections (Detections): A Detections object containing bounding boxes, class ids, and other relevant information.
|
||||
classes (List[str]): A list of class names corresponding to the class ids in the Detections object.
|
||||
filename (str): The name of the image file associated with the detections.
|
||||
image_shape (Tuple[int, int, int]): The shape of the image file associated with the detections.
|
||||
min_image_area_percentage (float): Minimum detection area relative to area of image associated with it.
|
||||
max_image_area_percentage (float): Maximum detection area relative to area of image associated with it.
|
||||
approximation_percentage (float): The percentage of polygon points to be removed from the input polygon, in the range [0, 1).
|
||||
Returns:
|
||||
str: An XML string in Pascal VOC format representing the detections.
|
||||
"""
|
||||
height, width, depth = image_shape
|
||||
|
||||
# Create root element
|
||||
annotation = Element("annotation")
|
||||
|
||||
# Add folder element
|
||||
folder = SubElement(annotation, "folder")
|
||||
folder.text = "VOC"
|
||||
|
||||
# Add filename element
|
||||
file_name = SubElement(annotation, "filename")
|
||||
file_name.text = filename
|
||||
|
||||
# Add source element
|
||||
source = SubElement(annotation, "source")
|
||||
database = SubElement(source, "database")
|
||||
database.text = "roboflow.ai"
|
||||
|
||||
# Add size element
|
||||
size = SubElement(annotation, "size")
|
||||
w = SubElement(size, "width")
|
||||
w.text = str(width)
|
||||
h = SubElement(size, "height")
|
||||
h.text = str(height)
|
||||
d = SubElement(size, "depth")
|
||||
d.text = str(depth)
|
||||
|
||||
# Add segmented element
|
||||
segmented = SubElement(annotation, "segmented")
|
||||
segmented.text = "0"
|
||||
|
||||
# Add object elements
|
||||
for xyxy, mask, _, class_id, _ in detections:
|
||||
name = classes[class_id]
|
||||
if mask is not None:
|
||||
polygons = approximate_mask_with_polygons(
|
||||
mask=mask,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
for polygon in polygons:
|
||||
xyxy = polygon_to_xyxy(polygon=polygon)
|
||||
next_object = object_to_pascal_voc(
|
||||
xyxy=xyxy, name=name, polygon=polygon
|
||||
)
|
||||
annotation.append(next_object)
|
||||
else:
|
||||
next_object = object_to_pascal_voc(xyxy=xyxy, name=name)
|
||||
annotation.append(next_object)
|
||||
|
||||
# Generate XML string
|
||||
xml_string = parseString(tostring(annotation)).toprettyxml(indent=" ")
|
||||
|
||||
return xml_string
|
||||
|
||||
|
||||
def load_pascal_voc_annotations(
|
||||
annotation_path: str,
|
||||
) -> Tuple[str, Detections, List[str]]:
|
||||
"""
|
||||
Loads PASCAL VOC XML annotations and returns the image name, a Detections instance, and a list of class names.
|
||||
|
||||
Args:
|
||||
annotation_path (str): The path to the PASCAL VOC XML annotations file.
|
||||
|
||||
Returns:
|
||||
Tuple[str, Detections, List[str]]: A tuple containing the image name, a Detections instance, and a list of class names of objects in the detections.
|
||||
"""
|
||||
tree = parse(annotation_path)
|
||||
root = tree.getroot()
|
||||
|
||||
image_name = root.find("filename").text
|
||||
|
||||
xyxy = []
|
||||
class_names = []
|
||||
for obj in root.findall("object"):
|
||||
class_name = obj.find("name").text
|
||||
class_names.append(class_name)
|
||||
|
||||
bbox = obj.find("bndbox")
|
||||
x1 = int(bbox.find("xmin").text)
|
||||
y1 = int(bbox.find("ymin").text)
|
||||
x2 = int(bbox.find("xmax").text)
|
||||
y2 = int(bbox.find("ymax").text)
|
||||
|
||||
xyxy.append([x1, y1, x2, y2])
|
||||
|
||||
xyxy = np.array(xyxy)
|
||||
detections = Detections(xyxy=xyxy)
|
||||
|
||||
return image_name, detections, class_names
|
||||
@@ -0,0 +1,246 @@
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.dataset.utils import approximate_mask_with_polygons
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
|
||||
from supervision.utils.file import (
|
||||
list_files_with_extensions,
|
||||
read_txt_file,
|
||||
read_yaml_file,
|
||||
save_text_file,
|
||||
save_yaml_file,
|
||||
)
|
||||
|
||||
|
||||
def _parse_box(values: List[str]) -> np.ndarray:
|
||||
x_center, y_center, width, height = values
|
||||
return np.array(
|
||||
[
|
||||
float(x_center) - float(width) / 2,
|
||||
float(y_center) - float(height) / 2,
|
||||
float(x_center) + float(width) / 2,
|
||||
float(y_center) + float(height) / 2,
|
||||
],
|
||||
dtype=np.float32,
|
||||
)
|
||||
|
||||
|
||||
def _box_to_polygon(box: np.ndarray) -> np.ndarray:
|
||||
return np.array(
|
||||
[[box[0], box[1]], [box[2], box[1]], [box[2], box[3]], [box[0], box[3]]]
|
||||
)
|
||||
|
||||
|
||||
def _parse_polygon(values: List[str]) -> np.ndarray:
|
||||
return np.array(values, dtype=np.float32).reshape(-1, 2)
|
||||
|
||||
|
||||
def _polygons_to_masks(
|
||||
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
|
||||
) -> np.ndarray:
|
||||
return np.array(
|
||||
[
|
||||
polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
|
||||
for polygon in polygons
|
||||
],
|
||||
dtype=bool,
|
||||
)
|
||||
|
||||
|
||||
def _pair_image_with_annotation(
|
||||
image_paths: List[Union[str, Path]], annotation_paths: List[Union[str, Path]]
|
||||
) -> List[Tuple[Union[str, Path], Union[str, Path]]]:
|
||||
image_dict = {p.stem: p for p in image_paths}
|
||||
return [
|
||||
(image_dict[annotation_path.stem], annotation_path)
|
||||
for annotation_path in annotation_paths
|
||||
if annotation_path.stem in image_dict
|
||||
]
|
||||
|
||||
|
||||
def _with_mask(lines: List[str]) -> bool:
|
||||
return any([len(line.split()) > 5 for line in lines])
|
||||
|
||||
|
||||
def _extract_class_names(file_path: str) -> List[str]:
|
||||
data = read_yaml_file(file_path=file_path)
|
||||
return data["names"]
|
||||
|
||||
|
||||
def _image_name_to_annotation_name(image_name: str) -> str:
|
||||
base_name, _ = os.path.splitext(image_name)
|
||||
return base_name + ".txt"
|
||||
|
||||
|
||||
def yolo_annotations_to_detections(
|
||||
lines: List[str], resolution_wh: Tuple[int, int], with_masks: bool
|
||||
) -> Detections:
|
||||
if len(lines) == 0:
|
||||
return Detections.empty()
|
||||
|
||||
class_id, relative_xyxy, relative_polygon = [], [], []
|
||||
w, h = resolution_wh
|
||||
for line in lines:
|
||||
values = line.split()
|
||||
class_id.append(int(values[0]))
|
||||
if len(values) == 5:
|
||||
box = _parse_box(values=values[1:])
|
||||
relative_xyxy.append(box)
|
||||
if with_masks:
|
||||
relative_polygon.append(_box_to_polygon(box=box))
|
||||
elif len(values) > 5:
|
||||
polygon = _parse_polygon(values=values[1:])
|
||||
relative_xyxy.append(polygon_to_xyxy(polygon=polygon))
|
||||
if with_masks:
|
||||
relative_polygon.append(polygon)
|
||||
|
||||
class_id = np.array(class_id, dtype=int)
|
||||
relative_xyxy = np.array(relative_xyxy, dtype=np.float32)
|
||||
xyxy = relative_xyxy * np.array([w, h, w, h], dtype=np.float32)
|
||||
|
||||
if not with_masks:
|
||||
return Detections(class_id=class_id, xyxy=xyxy)
|
||||
|
||||
polygons = [
|
||||
(polygon * np.array(resolution_wh)).astype(int) for polygon in relative_polygon
|
||||
]
|
||||
mask = _polygons_to_masks(polygons=polygons, resolution_wh=resolution_wh)
|
||||
return Detections(class_id=class_id, xyxy=xyxy, mask=mask)
|
||||
|
||||
|
||||
def load_yolo_annotations(
|
||||
images_directory_path: str,
|
||||
annotations_directory_path: str,
|
||||
data_yaml_path: str,
|
||||
force_masks: bool = False,
|
||||
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
|
||||
"""
|
||||
Loads YOLO annotations and returns class names, images, and their corresponding detections.
|
||||
|
||||
Args:
|
||||
images_directory_path (str): The path to the directory containing the images.
|
||||
annotations_directory_path (str): The path to the directory containing the YOLO annotation files.
|
||||
data_yaml_path (str): The path to the data YAML file containing class information.
|
||||
force_masks (bool, optional): If True, forces masks to be loaded for all annotations, regardless of whether they are present.
|
||||
|
||||
Returns:
|
||||
Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple containing a list of class names, a dictionary with image names as keys and images as values, and a dictionary with image names as keys and corresponding Detections instances as values.
|
||||
"""
|
||||
image_paths = list_files_with_extensions(
|
||||
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
|
||||
)
|
||||
annotation_paths = list_files_with_extensions(
|
||||
directory=annotations_directory_path, extensions=["txt"]
|
||||
)
|
||||
path_pairs = _pair_image_with_annotation(
|
||||
image_paths=image_paths, annotation_paths=annotation_paths
|
||||
)
|
||||
classes = _extract_class_names(file_path=data_yaml_path)
|
||||
images = {}
|
||||
annotations = {}
|
||||
for image_path, annotation_path in path_pairs:
|
||||
image = cv2.imread(str(image_path))
|
||||
lines = read_txt_file(str(annotation_path))
|
||||
h, w, _ = image.shape
|
||||
resolution_wh = (w, h)
|
||||
|
||||
with_masks = _with_mask(lines=lines)
|
||||
with_masks = force_masks if force_masks else with_masks
|
||||
annotation = yolo_annotations_to_detections(
|
||||
lines=lines, resolution_wh=resolution_wh, with_masks=with_masks
|
||||
)
|
||||
|
||||
images[image_path.name] = image
|
||||
annotations[image_path.name] = annotation
|
||||
return classes, images, annotations
|
||||
|
||||
|
||||
def object_to_yolo(
|
||||
xyxy: np.ndarray,
|
||||
class_id: int,
|
||||
image_shape: Tuple[int, int, int],
|
||||
polygon: Optional[np.ndarray] = None,
|
||||
) -> str:
|
||||
h, w, _ = image_shape
|
||||
if polygon is None:
|
||||
xyxy_relative = xyxy / np.array([w, h, w, h], dtype=np.float32)
|
||||
x_min, y_min, x_max, y_max = xyxy_relative
|
||||
x_center = (x_min + x_max) / 2
|
||||
y_center = (y_min + y_max) / 2
|
||||
width = x_max - x_min
|
||||
height = y_max - y_min
|
||||
return f"{int(class_id)} {x_center:.5f} {y_center:.5f} {width:.5f} {height:.5f}"
|
||||
else:
|
||||
polygon_relative = polygon / np.array([w, h], dtype=np.float32)
|
||||
polygon_relative = polygon_relative.reshape(-1)
|
||||
polygon_parsed = " ".join([f"{value:.5f}" for value in polygon_relative])
|
||||
return f"{int(class_id)} {polygon_parsed}"
|
||||
|
||||
|
||||
def detections_to_yolo_annotations(
|
||||
detections: Detections,
|
||||
image_shape: Tuple[int, int, int],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> List[str]:
|
||||
annotation = []
|
||||
for xyxy, mask, _, class_id, _ in detections:
|
||||
if mask is not None:
|
||||
polygons = approximate_mask_with_polygons(
|
||||
mask=mask,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
for polygon in polygons:
|
||||
xyxy = polygon_to_xyxy(polygon=polygon)
|
||||
next_object = object_to_yolo(
|
||||
xyxy=xyxy,
|
||||
class_id=class_id,
|
||||
image_shape=image_shape,
|
||||
polygon=polygon,
|
||||
)
|
||||
annotation.append(next_object)
|
||||
else:
|
||||
next_object = object_to_yolo(
|
||||
xyxy=xyxy, class_id=class_id, image_shape=image_shape
|
||||
)
|
||||
annotation.append(next_object)
|
||||
return annotation
|
||||
|
||||
|
||||
def save_yolo_annotations(
|
||||
annotations_directory_path: str,
|
||||
images: Dict[str, np.ndarray],
|
||||
annotations: Dict[str, Detections],
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> None:
|
||||
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
for image_name, image in images.items():
|
||||
detections = annotations[image_name]
|
||||
yolo_annotations_name = _image_name_to_annotation_name(image_name=image_name)
|
||||
yolo_annotations_path = os.path.join(
|
||||
annotations_directory_path, yolo_annotations_name
|
||||
)
|
||||
lines = detections_to_yolo_annotations(
|
||||
detections=detections,
|
||||
image_shape=image.shape,
|
||||
min_image_area_percentage=min_image_area_percentage,
|
||||
max_image_area_percentage=max_image_area_percentage,
|
||||
approximation_percentage=approximation_percentage,
|
||||
)
|
||||
save_text_file(lines=lines, file_path=yolo_annotations_path)
|
||||
|
||||
|
||||
def save_data_yaml(data_yaml_path: str, classes: List[str]) -> None:
|
||||
data = {"nc": len(classes), "names": classes}
|
||||
Path(data_yaml_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
save_yaml_file(data=data, file_path=data_yaml_path)
|
||||
@@ -0,0 +1,128 @@
|
||||
import copy
|
||||
import os
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Dict, List, Optional, Tuple, TypeVar
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.utils import (
|
||||
approximate_polygon,
|
||||
filter_polygons_by_area,
|
||||
mask_to_polygons,
|
||||
)
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
def approximate_mask_with_polygons(
|
||||
mask: np.ndarray,
|
||||
min_image_area_percentage: float = 0.0,
|
||||
max_image_area_percentage: float = 1.0,
|
||||
approximation_percentage: float = 0.75,
|
||||
) -> List[np.ndarray]:
|
||||
height, width = mask.shape
|
||||
image_area = height * width
|
||||
minimum_detection_area = min_image_area_percentage * image_area
|
||||
maximum_detection_area = max_image_area_percentage * image_area
|
||||
|
||||
polygons = mask_to_polygons(mask=mask)
|
||||
if len(polygons) == 1:
|
||||
polygons = filter_polygons_by_area(
|
||||
polygons=polygons, min_area=None, max_area=maximum_detection_area
|
||||
)
|
||||
else:
|
||||
polygons = filter_polygons_by_area(
|
||||
polygons=polygons,
|
||||
min_area=minimum_detection_area,
|
||||
max_area=maximum_detection_area,
|
||||
)
|
||||
return [
|
||||
approximate_polygon(polygon=polygon, percentage=approximation_percentage)
|
||||
for polygon in polygons
|
||||
]
|
||||
|
||||
|
||||
def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
|
||||
unique_classes = set()
|
||||
|
||||
for class_list in class_lists:
|
||||
for class_name in class_list:
|
||||
unique_classes.add(class_name.lower())
|
||||
|
||||
return sorted(list(unique_classes))
|
||||
|
||||
|
||||
def build_class_index_mapping(
|
||||
source_classes: List[str], target_classes: List[str]
|
||||
) -> Dict[int, int]:
|
||||
index_mapping = {}
|
||||
|
||||
for i, class_name in enumerate(source_classes):
|
||||
if class_name not in target_classes:
|
||||
raise ValueError(
|
||||
f"Class {class_name} not found in target classes. "
|
||||
f"source_classes must be a subset of target_classes."
|
||||
)
|
||||
corresponding_index = target_classes.index(class_name)
|
||||
index_mapping[i] = corresponding_index
|
||||
|
||||
return index_mapping
|
||||
|
||||
|
||||
def map_detections_class_id(
|
||||
source_to_target_mapping: Dict[int, int], detections: Detections
|
||||
) -> Detections:
|
||||
if detections.class_id is None:
|
||||
raise ValueError("Detections must have class_id attribute.")
|
||||
if set(np.unique(detections.class_id)) - set(source_to_target_mapping.keys()):
|
||||
raise ValueError(
|
||||
"Detections class_id must be a subset of source_to_target_mapping keys."
|
||||
)
|
||||
|
||||
detections_copy = copy.deepcopy(detections)
|
||||
detections_copy.class_id = np.vectorize(source_to_target_mapping.get)(
|
||||
detections_copy.class_id
|
||||
)
|
||||
|
||||
return detections_copy
|
||||
|
||||
|
||||
def save_dataset_images(
|
||||
images_directory_path: str, images: Dict[str, np.ndarray]
|
||||
) -> None:
|
||||
Path(images_directory_path).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
for image_name, image in images.items():
|
||||
target_image_path = os.path.join(images_directory_path, image_name)
|
||||
cv2.imwrite(target_image_path, image)
|
||||
|
||||
|
||||
def train_test_split(
|
||||
data: List[T],
|
||||
train_ratio: float = 0.8,
|
||||
random_state: Optional[int] = None,
|
||||
shuffle: bool = True,
|
||||
) -> Tuple[List[T], List[T]]:
|
||||
"""
|
||||
Splits the data into two parts using the provided train_ratio.
|
||||
|
||||
Args:
|
||||
data (List[T]): The data to split.
|
||||
train_ratio (float): The ratio of the training set to the entire dataset.
|
||||
random_state (Optional[int]): The seed for the random number generator.
|
||||
shuffle (bool): Whether to shuffle the data before splitting.
|
||||
|
||||
Returns:
|
||||
Tuple[List[T], List[T]]: The split data.
|
||||
"""
|
||||
if random_state is not None:
|
||||
random.seed(random_state)
|
||||
|
||||
if shuffle:
|
||||
random.shuffle(data)
|
||||
|
||||
split_index = int(len(data) * train_ratio)
|
||||
return data[:split_index], data[split_index:]
|
||||
@@ -0,0 +1,195 @@
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
|
||||
|
||||
class BoxAnnotator:
|
||||
"""
|
||||
A class for drawing bounding boxes on an image using detections provided.
|
||||
|
||||
Attributes:
|
||||
color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette
|
||||
thickness (int): The thickness of the bounding box lines, default is 2
|
||||
text_color (Color): The color of the text on the bounding box, default is white
|
||||
text_scale (float): The scale of the text on the bounding box, default is 0.5
|
||||
text_thickness (int): The thickness of the text on the bounding box, default is 1
|
||||
text_padding (int): The padding around the text on the bounding box, default is 5
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
text_color: Color = Color.black(),
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
):
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.text_color: Color = text_color
|
||||
self.text_scale: float = text_scale
|
||||
self.text_thickness: int = text_thickness
|
||||
self.text_padding: int = text_padding
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
labels: Optional[List[str]] = None,
|
||||
skip_label: bool = False,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draws bounding boxes on the frame using the detections provided.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the bounding boxes will be drawn
|
||||
detections (Detections): The detections for which the bounding boxes will be drawn
|
||||
labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If `labels` are not provided, corresponding `class_id` will be used as label.
|
||||
skip_label (bool): Is set to `True`, skips bounding box label annotation.
|
||||
Returns:
|
||||
np.ndarray: The image with the bounding boxes drawn on it
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> classes = ['person', ...]
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> box_annotator = sv.BoxAnnotator()
|
||||
>>> labels = [
|
||||
... f"{classes[class_id]} {confidence:0.2f}"
|
||||
... for _, _, confidence, class_id, _
|
||||
... in detections
|
||||
... ]
|
||||
>>> annotated_frame = box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections,
|
||||
... labels=labels
|
||||
... )
|
||||
```
|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for i in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[i].astype(int)
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
idx = class_id if class_id is not None else i
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
if skip_label:
|
||||
continue
|
||||
|
||||
text = (
|
||||
f"{class_id}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[i]
|
||||
)
|
||||
|
||||
text_width, text_height = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
thickness=self.text_thickness,
|
||||
)[0]
|
||||
|
||||
text_x = x1 + self.text_padding
|
||||
text_y = y1 - self.text_padding
|
||||
|
||||
text_background_x1 = x1
|
||||
text_background_y1 = y1 - 2 * self.text_padding - text_height
|
||||
|
||||
text_background_x2 = x1 + 2 * self.text_padding + text_width
|
||||
text_background_y2 = y1
|
||||
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(text_background_x1, text_background_y1),
|
||||
pt2=(text_background_x2, text_background_y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=cv2.FILLED,
|
||||
)
|
||||
cv2.putText(
|
||||
img=scene,
|
||||
text=text,
|
||||
org=(text_x, text_y),
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
color=self.text_color.as_rgb(),
|
||||
thickness=self.text_thickness,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class MaskAnnotator:
|
||||
"""
|
||||
A class for overlaying masks on an image using detections provided.
|
||||
|
||||
Attributes:
|
||||
color (Union[Color, ColorPalette]): The color to fill the mask, can be a single color or a color palette
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
):
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
|
||||
def annotate(
|
||||
self, scene: np.ndarray, detections: Detections, opacity: float = 0.5
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Overlays the masks on the given image based on the provided detections, with a specified opacity.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the masks will be overlaid
|
||||
detections (Detections): The detections for which the masks will be overlaid
|
||||
opacity (float): The opacity of the masks, between 0 and 1, default is 0.5
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the masks overlaid
|
||||
"""
|
||||
if detections.mask is None:
|
||||
return scene
|
||||
|
||||
for i in np.flip(np.argsort(detections.area)):
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
idx = class_id if class_id is not None else i
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
|
||||
mask = detections.mask[i]
|
||||
colored_mask = np.zeros_like(scene, dtype=np.uint8)
|
||||
colored_mask[:] = color.as_bgr()
|
||||
|
||||
scene = np.where(
|
||||
np.expand_dims(mask, axis=-1),
|
||||
np.uint8(opacity * colored_mask + (1 - opacity) * scene),
|
||||
scene,
|
||||
)
|
||||
|
||||
return scene
|
||||
@@ -0,0 +1,553 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import astuple, dataclass
|
||||
from typing import Any, Iterator, List, Optional, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.utils import (
|
||||
extract_yolov8_masks,
|
||||
non_max_suppression,
|
||||
process_roboflow_result,
|
||||
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, W, H) shape")
|
||||
|
||||
|
||||
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")
|
||||
|
||||
|
||||
@dataclass
|
||||
class Detections:
|
||||
"""
|
||||
Data class containing information about the detections in a video frame.
|
||||
Attributes:
|
||||
xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]`
|
||||
mask: (Optional[np.ndarray]): An array of shape `(n, W, H)` containing the segmentation masks.
|
||||
confidence (Optional[np.ndarray]): An array of shape `(n,)` containing the confidence scores of the detections.
|
||||
class_id (Optional[np.ndarray]): An array of shape `(n,)` containing the class ids of the detections.
|
||||
tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections.
|
||||
"""
|
||||
|
||||
xyxy: np.ndarray
|
||||
mask: np.Optional[np.ndarray] = None
|
||||
confidence: Optional[np.ndarray] = None
|
||||
class_id: Optional[np.ndarray] = None
|
||||
tracker_id: Optional[np.ndarray] = None
|
||||
|
||||
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)
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
Returns the number of detections in the Detections object.
|
||||
"""
|
||||
return len(self.xyxy)
|
||||
|
||||
def __iter__(
|
||||
self,
|
||||
) -> Iterator[
|
||||
Tuple[
|
||||
np.ndarray,
|
||||
Optional[np.ndarray],
|
||||
Optional[float],
|
||||
Optional[int],
|
||||
Optional[int],
|
||||
]
|
||||
]:
|
||||
"""
|
||||
Iterates over the Detections object and yield a tuple of `(xyxy, mask, confidence, class_id, tracker_id)` for each detection.
|
||||
"""
|
||||
for i in range(len(self.xyxy)):
|
||||
yield (
|
||||
self.xyxy[i],
|
||||
self.mask[i] if self.mask is not None else None,
|
||||
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,
|
||||
)
|
||||
|
||||
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),
|
||||
]
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_yolov5(cls, yolov5_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a [YOLOv5](https://github.com/ultralytics/yolov5) inference result.
|
||||
|
||||
Args:
|
||||
yolov5_results (yolov5.models.common.Detections): The output Detections instance from YOLOv5
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import torch
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
|
||||
>>> result = model(image)
|
||||
>>> detections = sv.Detections.from_yolov5(result)
|
||||
```
|
||||
"""
|
||||
yolov5_detections_predictions = yolov5_results.pred[0].cpu().cpu().numpy()
|
||||
return cls(
|
||||
xyxy=yolov5_detections_predictions[:, :4],
|
||||
confidence=yolov5_detections_predictions[:, 4],
|
||||
class_id=yolov5_detections_predictions[:, 5].astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_yolov8(cls, yolov8_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a [YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
|
||||
|
||||
Args:
|
||||
yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from ultralytics import YOLO
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = YOLO('yolov8s.pt')
|
||||
>>> result = model(image)[0]
|
||||
>>> detections = sv.Detections.from_yolov8(result)
|
||||
```
|
||||
"""
|
||||
return cls(
|
||||
xyxy=yolov8_results.boxes.xyxy.cpu().numpy(),
|
||||
confidence=yolov8_results.boxes.conf.cpu().numpy(),
|
||||
class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
|
||||
mask=extract_yolov8_masks(yolov8_results),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_yolo_nas(cls, yolo_nas_results) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from a [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) inference result.
|
||||
|
||||
Args:
|
||||
yolo_nas_results (super_gradients.training.models.prediction_results.ImageDetectionPrediction): The output Results instance from YOLO-NAS
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from super_gradients.training import models
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> model = models.get('yolo_nas_l', pretrained_weights="coco")
|
||||
>>> result = list(model.predict(image, conf=0.35))[0]
|
||||
>>> detections = sv.Detections.from_yolo_nas(result)
|
||||
```
|
||||
"""
|
||||
return cls(
|
||||
xyxy=yolo_nas_results.prediction.bboxes_xyxy,
|
||||
confidence=yolo_nas_results.prediction.confidence,
|
||||
class_id=yolo_nas_results.prediction.labels.astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_transformers(cls, transformers_results: dict) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from object detection [transformer](https://github.com/huggingface/transformers) inference result.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
"""
|
||||
return cls(
|
||||
xyxy=transformers_results["boxes"].cpu().numpy(),
|
||||
confidence=transformers_results["scores"].cpu().numpy(),
|
||||
class_id=transformers_results["labels"].cpu().numpy().astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_detectron2(cls, detectron2_results) -> Detections:
|
||||
"""
|
||||
Create a Detections object from the [Detectron2](https://github.com/facebookresearch/detectron2) inference result.
|
||||
|
||||
Args:
|
||||
detectron2_results: The output of a Detectron2 model containing instances with prediction data.
|
||||
|
||||
Returns:
|
||||
(Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> from detectron2.engine import DefaultPredictor
|
||||
>>> from detectron2.config import get_cfg
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread(SOURCE_IMAGE_PATH)
|
||||
>>> cfg = get_cfg()
|
||||
>>> cfg.merge_from_file("path/to/config.yaml")
|
||||
>>> cfg.MODEL.WEIGHTS = "path/to/model_weights.pth"
|
||||
>>> predictor = DefaultPredictor(cfg)
|
||||
>>> result = predictor(image)
|
||||
|
||||
>>> detections = sv.Detections.from_detectron2(result)
|
||||
```
|
||||
"""
|
||||
return cls(
|
||||
xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
|
||||
confidence=detectron2_results["instances"].scores.cpu().numpy(),
|
||||
class_id=detectron2_results["instances"]
|
||||
.pred_classes.cpu()
|
||||
.numpy()
|
||||
.astype(int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_roboflow(cls, roboflow_result: dict, class_list: List[str]) -> Detections:
|
||||
"""
|
||||
Create a Detections object from the [Roboflow](https://roboflow.com/) API inference result.
|
||||
|
||||
Args:
|
||||
roboflow_result (dict): The result from the Roboflow API containing predictions.
|
||||
class_list (List[str]): A list of class names corresponding to the class IDs in the API result.
|
||||
|
||||
Returns:
|
||||
(Detections): A Detections object containing the bounding boxes, class IDs, and confidences of the predictions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> roboflow_result = {
|
||||
... "predictions": [
|
||||
... {
|
||||
... "x": 0.5,
|
||||
... "y": 0.5,
|
||||
... "width": 0.2,
|
||||
... "height": 0.3,
|
||||
... "class": "person",
|
||||
... "confidence": 0.9
|
||||
... },
|
||||
... # ... more predictions ...
|
||||
... ]
|
||||
... }
|
||||
>>> class_list = ["person", "car", "dog"]
|
||||
|
||||
>>> detections = sv.Detections.from_roboflow(roboflow_result, class_list)
|
||||
```
|
||||
"""
|
||||
xyxy, confidence, class_id, masks = process_roboflow_result(
|
||||
roboflow_result=roboflow_result, class_list=class_list
|
||||
)
|
||||
return Detections(
|
||||
xyxy=xyxy,
|
||||
confidence=confidence,
|
||||
class_id=class_id,
|
||||
mask=masks,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_sam(cls, sam_result: List[dict]) -> Detections:
|
||||
"""
|
||||
Creates a Detections instance from [Segment Anything Model](https://github.com/facebookresearch/segment-anything) inference result.
|
||||
|
||||
Args:
|
||||
sam_result (List[dict]): The output Results instance from SAM
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
|
||||
|
||||
>>> sam = sam_model_registry[MODEL_TYPE](checkpoint=CHECKPOINT_PATH).to(device=DEVICE)
|
||||
>>> mask_generator = SamAutomaticMaskGenerator(sam)
|
||||
>>> sam_result = mask_generator.generate(IMAGE)
|
||||
>>> detections = sv.Detections.from_sam(sam_result=sam_result)
|
||||
```
|
||||
"""
|
||||
sorted_generated_masks = sorted(
|
||||
sam_result, key=lambda x: x["area"], reverse=True
|
||||
)
|
||||
|
||||
xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
|
||||
mask = np.array([mask["segmentation"] for mask in sorted_generated_masks])
|
||||
|
||||
return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
|
||||
|
||||
@classmethod
|
||||
def empty(cls) -> Detections:
|
||||
"""
|
||||
Create an empty Detections object with no bounding boxes, confidences, or class IDs.
|
||||
|
||||
Returns:
|
||||
(Detections): An empty Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from supervision import Detections
|
||||
|
||||
>>> empty_detections = Detections.empty()
|
||||
```
|
||||
"""
|
||||
return cls(
|
||||
xyxy=np.empty((0, 4), dtype=np.float32),
|
||||
confidence=np.array([], dtype=np.float32),
|
||||
class_id=np.array([], dtype=int),
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def merge(cls, detections_list: List[Detections]) -> Detections:
|
||||
"""
|
||||
Merge a list of Detections objects into a single Detections object.
|
||||
|
||||
This method takes a list of Detections objects and combines their respective fields (`xyxy`, `mask`,
|
||||
`confidence`, `class_id`, and `tracker_id`) into a single Detections object. If all elements in a field are not
|
||||
`None`, the corresponding field will be stacked. Otherwise, the field will be set to `None`.
|
||||
|
||||
Args:
|
||||
detections_list (List[Detections]): A list of Detections objects to merge.
|
||||
|
||||
Returns:
|
||||
(Detections): A single Detections object containing the merged data from the input list.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> from supervision import Detections
|
||||
|
||||
>>> detections_1 = Detections(...)
|
||||
>>> detections_2 = Detections(...)
|
||||
|
||||
>>> merged_detections = Detections.merge([detections_1, detections_2])
|
||||
```
|
||||
"""
|
||||
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)
|
||||
]
|
||||
|
||||
all_not_none = lambda l: all(x is not None for x in l)
|
||||
|
||||
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
|
||||
|
||||
return cls(
|
||||
xyxy=xyxy,
|
||||
mask=mask,
|
||||
confidence=confidence,
|
||||
class_id=class_id,
|
||||
tracker_id=tracker_id,
|
||||
)
|
||||
|
||||
def get_anchor_coordinates(self, anchor: Position) -> np.ndarray:
|
||||
"""
|
||||
Returns the bounding box coordinates for a specific anchor.
|
||||
|
||||
Args:
|
||||
anchor (Position): Position of bounding box anchor for which to return the coordinates.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of shape `(n, 2)` containing the bounding box anchor coordinates in format `[x, y]`.
|
||||
"""
|
||||
if anchor == Position.CENTER:
|
||||
return np.array(
|
||||
[
|
||||
(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2,
|
||||
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
|
||||
]
|
||||
).transpose()
|
||||
elif anchor == Position.BOTTOM_CENTER:
|
||||
return np.array(
|
||||
[(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 3]]
|
||||
).transpose()
|
||||
|
||||
raise ValueError(f"{anchor} is not supported.")
|
||||
|
||||
def __getitem__(
|
||||
self, index: Union[int, slice, List[int], np.ndarray]
|
||||
) -> Detections:
|
||||
"""
|
||||
Get a subset of the Detections object.
|
||||
|
||||
Args:
|
||||
index (Union[int, slice, List[int], np.ndarray]): The index or indices of the subset of the Detections
|
||||
|
||||
Returns:
|
||||
(Detections): A subset of the Detections object.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> first_detection = detections[0]
|
||||
|
||||
>>> first_10_detections = detections[0:10]
|
||||
|
||||
>>> some_detections = detections[[0, 2, 4]]
|
||||
|
||||
>>> class_0_detections = detections[detections.class_id == 0]
|
||||
|
||||
>>> high_confidence_detections = detections[detections.confidence > 0.5]
|
||||
```
|
||||
"""
|
||||
if isinstance(index, int):
|
||||
index = [index]
|
||||
return Detections(
|
||||
xyxy=self.xyxy[index],
|
||||
mask=self.mask[index] if self.mask is not None else None,
|
||||
confidence=self.confidence[index] if self.confidence is not None else None,
|
||||
class_id=self.class_id[index] if self.class_id is not None else None,
|
||||
tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
|
||||
)
|
||||
|
||||
@property
|
||||
def area(self) -> np.ndarray:
|
||||
"""
|
||||
Calculate the area of each detection in the set of object detections. If masks field is defined property
|
||||
returns are of each mask. If only box is given property return area of each box.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of floats containing the area of each detection in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections.
|
||||
"""
|
||||
if self.mask is not None:
|
||||
return np.array([np.sum(mask) for mask in self.mask])
|
||||
else:
|
||||
return self.box_area
|
||||
|
||||
@property
|
||||
def box_area(self) -> np.ndarray:
|
||||
"""
|
||||
Calculate the area of each bounding box in the set of object detections.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of floats containing the area of each bounding box in the format of `(area_1, area_2, ..., area_n)`, where n is the number of detections.
|
||||
"""
|
||||
return (self.xyxy[:, 3] - self.xyxy[:, 1]) * (self.xyxy[:, 2] - self.xyxy[:, 0])
|
||||
|
||||
def with_nms(
|
||||
self, threshold: float = 0.5, class_agnostic: bool = False
|
||||
) -> Detections:
|
||||
"""
|
||||
Perform non-maximum suppression on the current set of object detections.
|
||||
|
||||
Args:
|
||||
threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression. Defaults to 0.5.
|
||||
class_agnostic (bool, optional): Whether to perform class-agnostic non-maximum suppression. If True, the class_id of each detection will be ignored. Defaults to False.
|
||||
|
||||
Returns:
|
||||
Detections: A new Detections object containing the subset of detections after non-maximum suppression.
|
||||
|
||||
Raises:
|
||||
AssertionError: If `confidence` is None and class_agnostic is False. If `class_id` is None and class_agnostic is False.
|
||||
"""
|
||||
if len(self) == 0:
|
||||
return self
|
||||
|
||||
assert (
|
||||
self.confidence is not None
|
||||
), f"Detections confidence must be given for NMS to be executed."
|
||||
|
||||
if class_agnostic:
|
||||
predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
|
||||
indices = non_max_suppression(
|
||||
predictions=predictions, iou_threshold=threshold
|
||||
)
|
||||
return self[indices]
|
||||
|
||||
assert self.class_id is not None, (
|
||||
f"Detections class_id must be given for NMS to be executed. If you intended to perform class agnostic "
|
||||
f"NMS set class_agnostic=True."
|
||||
)
|
||||
|
||||
predictions = np.hstack(
|
||||
(self.xyxy, self.confidence.reshape(-1, 1), self.class_id.reshape(-1, 1))
|
||||
)
|
||||
indices = non_max_suppression(predictions=predictions, iou_threshold=threshold)
|
||||
return self[indices]
|
||||
@@ -1,33 +1,41 @@
|
||||
from typing import Dict
|
||||
from typing import Dict, Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color
|
||||
from supervision.geometry.dataclasses import Point, Rect, Vector
|
||||
from supervision.tools.detections import Detections
|
||||
from supervision.geometry.core import Point, Rect, Vector
|
||||
|
||||
|
||||
class LineCounter:
|
||||
class LineZone:
|
||||
"""
|
||||
Count the number of objects that cross a line.
|
||||
"""
|
||||
|
||||
def __init__(self, start: Point, end: Point):
|
||||
"""
|
||||
Initialize a LineCounter object.
|
||||
|
||||
:param start: Point : The starting point of the line.
|
||||
:param end: Point : The ending point of the line.
|
||||
Attributes:
|
||||
start (Point): The starting point of the line.
|
||||
end (Point): The ending point of the line.
|
||||
|
||||
"""
|
||||
self.vector = Vector(start=start, end=end)
|
||||
self.tracker_state: Dict[str, bool] = {}
|
||||
self.in_count: int = 0
|
||||
self.out_count: int = 0
|
||||
|
||||
def update(self, detections: Detections):
|
||||
def trigger(self, detections: Detections):
|
||||
"""
|
||||
Update the in_count and out_count for the detections that cross the line.
|
||||
|
||||
:param detections: Detections : The detections for which to update the counts.
|
||||
Attributes:
|
||||
detections (Detections): The detections for which to update the counts.
|
||||
|
||||
"""
|
||||
for xyxy, confidence, class_id, tracker_id in detections:
|
||||
for xyxy, _, confidence, class_id, tracker_id in detections:
|
||||
# handle detections with no tracker_id
|
||||
if tracker_id is None:
|
||||
continue
|
||||
@@ -63,7 +71,7 @@ class LineCounter:
|
||||
self.out_count += 1
|
||||
|
||||
|
||||
class LineCounterAnnotator:
|
||||
class LineZoneAnnotator:
|
||||
def __init__(
|
||||
self,
|
||||
thickness: float = 2,
|
||||
@@ -73,17 +81,21 @@ class LineCounterAnnotator:
|
||||
text_scale: float = 0.5,
|
||||
text_offset: float = 1.5,
|
||||
text_padding: int = 10,
|
||||
custom_in_text: Optional[str] = None,
|
||||
custom_out_text: Optional[str] = None,
|
||||
):
|
||||
"""
|
||||
Initialize the LineCounterAnnotator object with default values.
|
||||
|
||||
:param thickness: float : The thickness of the line that will be drawn.
|
||||
:param color: Color : The color of the line that will be drawn.
|
||||
:param text_thickness: float : The thickness of the text that will be drawn.
|
||||
:param text_color: Color : The color of the text that will be drawn.
|
||||
:param text_scale: float : The scale of the text that will be drawn.
|
||||
:param text_offset: float : The offset of the text that will be drawn.
|
||||
:param text_padding: int : The padding of the text that will be drawn.
|
||||
Attributes:
|
||||
thickness (float): The thickness of the line that will be drawn.
|
||||
color (Color): The color of the line that will be drawn.
|
||||
text_thickness (float): The thickness of the text that will be drawn.
|
||||
text_color (Color): The color of the text that will be drawn.
|
||||
text_scale (float): The scale of the text that will be drawn.
|
||||
text_offset (float): The offset of the text that will be drawn.
|
||||
text_padding (int): The padding of the text that will be drawn.
|
||||
|
||||
"""
|
||||
self.thickness: float = thickness
|
||||
self.color: Color = color
|
||||
@@ -92,14 +104,20 @@ class LineCounterAnnotator:
|
||||
self.text_scale: float = text_scale
|
||||
self.text_offset: float = text_offset
|
||||
self.text_padding: int = text_padding
|
||||
self.custom_in_text: str = custom_in_text
|
||||
self.custom_out_text: str = custom_out_text
|
||||
|
||||
def annotate(self, frame: np.ndarray, line_counter: LineCounter) -> np.ndarray:
|
||||
def annotate(self, frame: np.ndarray, line_counter: LineZone) -> np.ndarray:
|
||||
"""
|
||||
Draws the line on the frame using the line_counter provided.
|
||||
|
||||
:param frame: np.ndarray : The image on which the line will be drawn
|
||||
:param line_counter: LineCounter : The line counter that will be used to draw the line
|
||||
:return: np.ndarray : The image with the line drawn on it
|
||||
Attributes:
|
||||
frame (np.ndarray): The image on which the line will be drawn.
|
||||
line_counter (LineCounter): The line counter that will be used to draw the line.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the line drawn on it.
|
||||
|
||||
"""
|
||||
cv2.line(
|
||||
frame,
|
||||
@@ -127,8 +145,16 @@ class LineCounterAnnotator:
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
|
||||
in_text = f"in: {line_counter.in_count}"
|
||||
out_text = f"out: {line_counter.out_count}"
|
||||
in_text = (
|
||||
f"{self.custom_in_text}: {line_counter.in_count}"
|
||||
if self.custom_in_text is not None
|
||||
else f"in: {line_counter.in_count}"
|
||||
)
|
||||
out_text = (
|
||||
f"{self.custom_out_text}: {line_counter.out_count}"
|
||||
if self.custom_out_text is not None
|
||||
else f"out: {line_counter.out_count}"
|
||||
)
|
||||
|
||||
(in_text_width, in_text_height), _ = cv2.getTextSize(
|
||||
in_text, cv2.FONT_HERSHEY_SIMPLEX, self.text_scale, self.text_thickness
|
||||
@@ -205,3 +231,4 @@ class LineCounterAnnotator:
|
||||
self.text_thickness,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
return frame
|
||||
@@ -0,0 +1,132 @@
|
||||
from dataclasses import replace
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision import Detections
|
||||
from supervision.detection.utils import clip_boxes, polygon_to_mask
|
||||
from supervision.draw.color import Color
|
||||
from supervision.draw.utils import draw_polygon, draw_text
|
||||
from supervision.geometry.core import Position
|
||||
from supervision.geometry.utils import get_polygon_center
|
||||
|
||||
|
||||
class PolygonZone:
|
||||
"""
|
||||
A class for defining a polygon-shaped zone within a frame for detecting objects.
|
||||
|
||||
Attributes:
|
||||
polygon (np.ndarray): A polygon represented by a numpy array of shape `(N, 2)`, containing the `x`, `y` coordinates of the points.
|
||||
frame_resolution_wh (Tuple[int, int]): The frame resolution (width, height)
|
||||
triggering_position (Position): The position within the bounding box that triggers the zone (default: Position.BOTTOM_CENTER)
|
||||
current_count (int): The current count of detected objects within the zone
|
||||
mask (np.ndarray): The 2D bool mask for the polygon zone
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
polygon: np.ndarray,
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
triggering_position: Position = Position.BOTTOM_CENTER,
|
||||
):
|
||||
self.polygon = polygon.astype(int)
|
||||
self.frame_resolution_wh = frame_resolution_wh
|
||||
self.triggering_position = triggering_position
|
||||
self.current_count = 0
|
||||
|
||||
width, height = frame_resolution_wh
|
||||
self.mask = polygon_to_mask(
|
||||
polygon=polygon, resolution_wh=(width + 1, height + 1)
|
||||
)
|
||||
|
||||
def trigger(self, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Determines if the detections are within the polygon zone.
|
||||
|
||||
Parameters:
|
||||
detections (Detections): The detections to be checked against the polygon zone
|
||||
|
||||
Returns:
|
||||
np.ndarray: A boolean numpy array indicating if each detection is within the polygon zone
|
||||
"""
|
||||
|
||||
clipped_xyxy = clip_boxes(
|
||||
boxes_xyxy=detections.xyxy, frame_resolution_wh=self.frame_resolution_wh
|
||||
)
|
||||
clipped_detections = replace(detections, xyxy=clipped_xyxy)
|
||||
clipped_anchors = np.ceil(
|
||||
clipped_detections.get_anchor_coordinates(anchor=self.triggering_position)
|
||||
).astype(int)
|
||||
is_in_zone = self.mask[clipped_anchors[:, 1], clipped_anchors[:, 0]]
|
||||
self.current_count = np.sum(is_in_zone)
|
||||
return is_in_zone.astype(bool)
|
||||
|
||||
|
||||
class PolygonZoneAnnotator:
|
||||
"""
|
||||
A class for annotating a polygon-shaped zone within a frame with a count of detected objects.
|
||||
|
||||
Attributes:
|
||||
zone (PolygonZone): The polygon zone to be annotated
|
||||
color (Color): The color to draw the polygon lines
|
||||
thickness (int): The thickness of the polygon lines, default is 2
|
||||
text_color (Color): The color of the text on the polygon, default is black
|
||||
text_scale (float): The scale of the text on the polygon, default is 0.5
|
||||
text_thickness (int): The thickness of the text on the polygon, default is 1
|
||||
text_padding (int): The padding around the text on the polygon, default is 10
|
||||
font (int): The font type for the text on the polygon, default is cv2.FONT_HERSHEY_SIMPLEX
|
||||
center (Tuple[int, int]): The center of the polygon for text placement
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
zone: PolygonZone,
|
||||
color: Color,
|
||||
thickness: int = 2,
|
||||
text_color: Color = Color.black(),
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
):
|
||||
self.zone = zone
|
||||
self.color = color
|
||||
self.thickness = thickness
|
||||
self.text_color = text_color
|
||||
self.text_scale = text_scale
|
||||
self.text_thickness = text_thickness
|
||||
self.text_padding = text_padding
|
||||
self.font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
self.center = get_polygon_center(polygon=zone.polygon)
|
||||
|
||||
def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray:
|
||||
"""
|
||||
Annotates the polygon zone within a frame with a count of detected objects.
|
||||
|
||||
Parameters:
|
||||
scene (np.ndarray): The image on which the polygon zone will be annotated
|
||||
label (Optional[str]): An optional label for the count of detected objects within the polygon zone (default: None)
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the polygon zone and count of detected objects
|
||||
"""
|
||||
annotated_frame = draw_polygon(
|
||||
scene=scene,
|
||||
polygon=self.zone.polygon,
|
||||
color=self.color,
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
annotated_frame = draw_text(
|
||||
scene=annotated_frame,
|
||||
text=str(self.zone.current_count) if label is None else label,
|
||||
text_anchor=self.center,
|
||||
background_color=self.color,
|
||||
text_color=self.text_color,
|
||||
text_scale=self.text_scale,
|
||||
text_thickness=self.text_thickness,
|
||||
text_padding=self.text_padding,
|
||||
text_font=self.font,
|
||||
)
|
||||
|
||||
return annotated_frame
|
||||
@@ -0,0 +1,340 @@
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
MIN_POLYGON_POINT_COUNT = 3
|
||||
|
||||
|
||||
def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray:
|
||||
"""Generate a mask from a polygon.
|
||||
|
||||
Args:
|
||||
polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices.
|
||||
resolution_wh (Tuple[int, int]): The width and height of the desired resolution.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The generated 2D mask, where the polygon is marked with `1`'s and the rest is filled with `0`'s.
|
||||
"""
|
||||
width, height = resolution_wh
|
||||
mask = np.zeros((height, width), dtype=np.uint8)
|
||||
cv2.fillPoly(mask, [polygon], color=1)
|
||||
return mask
|
||||
|
||||
|
||||
def box_iou_batch(boxes_true: np.ndarray, boxes_detection: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Compute Intersection over Union (IoU) of two sets of bounding boxes - `boxes_true` and `boxes_detection`. Both sets
|
||||
of boxes are expected to be in `(x_min, y_min, x_max, y_max)` format.
|
||||
|
||||
Args:
|
||||
boxes_true (np.ndarray): 2D `np.ndarray` representing ground-truth boxes. `shape = (N, 4)` where `N` is number of true objects.
|
||||
boxes_detection (np.ndarray): 2D `np.ndarray` representing detection boxes. `shape = (M, 4)` where `M` is number of detected objects.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Pairwise IoU of boxes from `boxes_true` and `boxes_detection`. `shape = (N, M)` where `N` is number of true objects and `M` is number of detected objects.
|
||||
"""
|
||||
|
||||
def box_area(box):
|
||||
return (box[2] - box[0]) * (box[3] - box[1])
|
||||
|
||||
area_true = box_area(boxes_true.T)
|
||||
area_detection = box_area(boxes_detection.T)
|
||||
|
||||
top_left = np.maximum(boxes_true[:, None, :2], boxes_detection[:, :2])
|
||||
bottom_right = np.minimum(boxes_true[:, None, 2:], boxes_detection[:, 2:])
|
||||
|
||||
area_inter = np.prod(np.clip(bottom_right - top_left, a_min=0, a_max=None), 2)
|
||||
return area_inter / (area_true[:, None] + area_detection - area_inter)
|
||||
|
||||
|
||||
def non_max_suppression(
|
||||
predictions: np.ndarray, iou_threshold: float = 0.5
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Perform Non-Maximum Suppression (NMS) on object detection predictions.
|
||||
|
||||
Args:
|
||||
predictions (np.ndarray): An array of object detection predictions in the format of `(x_min, y_min, x_max, y_max, score)` or `(x_min, y_min, x_max, y_max, score, class)`.
|
||||
iou_threshold (float, optional): The intersection-over-union threshold to use for non-maximum suppression.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A boolean array indicating which predictions to keep after non-maximum suppression.
|
||||
|
||||
Raises:
|
||||
AssertionError: If `iou_threshold` is not within the closed range from `0` to `1`.
|
||||
"""
|
||||
assert 0 <= iou_threshold <= 1, (
|
||||
f"Value of `iou_threshold` must be in the closed range from 0 to 1, "
|
||||
f"{iou_threshold} given."
|
||||
)
|
||||
rows, columns = predictions.shape
|
||||
|
||||
# add column #5 - category filled with zeros for agnostic nms
|
||||
if columns == 5:
|
||||
predictions = np.c_[predictions, np.zeros(rows)]
|
||||
|
||||
# sort predictions column #4 - score
|
||||
sort_index = np.flip(predictions[:, 4].argsort())
|
||||
predictions = predictions[sort_index]
|
||||
|
||||
boxes = predictions[:, :4]
|
||||
categories = predictions[:, 5]
|
||||
ious = box_iou_batch(boxes, boxes)
|
||||
ious = ious - np.eye(rows)
|
||||
|
||||
keep = np.ones(rows, dtype=bool)
|
||||
|
||||
for index, (iou, category) in enumerate(zip(ious, categories)):
|
||||
if not keep[index]:
|
||||
continue
|
||||
|
||||
# drop detections with iou > iou_threshold and same category as current detections
|
||||
condition = (iou > iou_threshold) & (categories == category)
|
||||
keep = keep & ~condition
|
||||
|
||||
return keep[sort_index.argsort()]
|
||||
|
||||
|
||||
def clip_boxes(
|
||||
boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int]
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Clips bounding boxes coordinates to fit within the frame resolution.
|
||||
|
||||
Args:
|
||||
boxes_xyxy (np.ndarray): A numpy array of shape `(N, 4)` where each row corresponds to a bounding box in
|
||||
the format `(x_min, y_min, x_max, y_max)`.
|
||||
frame_resolution_wh (Tuple[int, int]): A tuple of the form `(width, height)` representing the resolution of the
|
||||
frame.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A numpy array of shape `(N, 4)` where each row corresponds to a bounding box with coordinates
|
||||
clipped to fit within the frame resolution.
|
||||
"""
|
||||
result = np.copy(boxes_xyxy)
|
||||
width, height = frame_resolution_wh
|
||||
result[:, [0, 2]] = result[:, [0, 2]].clip(0, width)
|
||||
result[:, [1, 3]] = result[:, [1, 3]].clip(0, height)
|
||||
return result
|
||||
|
||||
|
||||
def xywh_to_xyxy(boxes_xywh: np.ndarray) -> np.ndarray:
|
||||
xyxy = boxes_xywh.copy()
|
||||
xyxy[:, 2] = boxes_xywh[:, 0] + boxes_xywh[:, 2]
|
||||
xyxy[:, 3] = boxes_xywh[:, 1] + boxes_xywh[:, 3]
|
||||
return xyxy
|
||||
|
||||
|
||||
def mask_to_xyxy(masks: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Converts a 3D `np.array` of 2D bool masks into a 2D `np.array` of bounding boxes.
|
||||
|
||||
Parameters:
|
||||
masks (np.ndarray): A 3D `np.array` of shape `(N, W, H)` containing 2D bool masks
|
||||
|
||||
Returns:
|
||||
np.ndarray: A 2D `np.array` of shape `(N, 4)` containing the bounding boxes `(x_min, y_min, x_max, y_max)` for each mask
|
||||
"""
|
||||
n = masks.shape[0]
|
||||
bboxes = np.zeros((n, 4), dtype=int)
|
||||
|
||||
for i, mask in enumerate(masks):
|
||||
rows, cols = np.where(mask)
|
||||
|
||||
if len(rows) > 0 and len(cols) > 0:
|
||||
x_min, x_max = np.min(cols), np.max(cols)
|
||||
y_min, y_max = np.min(rows), np.max(rows)
|
||||
bboxes[i, :] = [x_min, y_min, x_max, y_max]
|
||||
|
||||
return bboxes
|
||||
|
||||
|
||||
def mask_to_polygons(mask: np.ndarray) -> List[np.ndarray]:
|
||||
"""
|
||||
Converts a binary mask to a list of polygons.
|
||||
|
||||
Parameters:
|
||||
mask (np.ndarray): A binary mask represented as a 2D NumPy array of shape `(H, W)`,
|
||||
where H and W are the height and width of the mask, respectively.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points. Polygons with fewer points than `MIN_POLYGON_POINT_COUNT = 3`
|
||||
are excluded from the output.
|
||||
"""
|
||||
|
||||
contours, _ = cv2.findContours(
|
||||
mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE
|
||||
)
|
||||
return [
|
||||
np.squeeze(contour, axis=1)
|
||||
for contour in contours
|
||||
if contour.shape[0] >= MIN_POLYGON_POINT_COUNT
|
||||
]
|
||||
|
||||
|
||||
def filter_polygons_by_area(
|
||||
polygons: List[np.ndarray],
|
||||
min_area: Optional[float] = None,
|
||||
max_area: Optional[float] = None,
|
||||
) -> List[np.ndarray]:
|
||||
"""
|
||||
Filters a list of polygons based on their area.
|
||||
|
||||
Parameters:
|
||||
polygons (List[np.ndarray]): A list of polygons, where each polygon is represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points.
|
||||
min_area (Optional[float]): The minimum area threshold. Only polygons with an area greater than or equal to this value
|
||||
will be included in the output. If set to None, no minimum area constraint will be applied.
|
||||
max_area (Optional[float]): The maximum area threshold. Only polygons with an area less than or equal to this value
|
||||
will be included in the output. If set to None, no maximum area constraint will be applied.
|
||||
|
||||
Returns:
|
||||
List[np.ndarray]: A new list of polygons containing only those with areas within the specified thresholds.
|
||||
"""
|
||||
if min_area is None and max_area is None:
|
||||
return polygons
|
||||
ares = [cv2.contourArea(polygon) for polygon in polygons]
|
||||
return [
|
||||
polygon
|
||||
for polygon, area in zip(polygons, ares)
|
||||
if (min_area is None or area >= min_area)
|
||||
and (max_area is None or area <= max_area)
|
||||
]
|
||||
|
||||
|
||||
def polygon_to_xyxy(polygon: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Converts a polygon represented by a NumPy array into a bounding box.
|
||||
|
||||
Parameters:
|
||||
polygon (np.ndarray): A polygon represented by a NumPy array of shape `(N, 2)`,
|
||||
containing the `x`, `y` coordinates of the points.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A 1D NumPy array containing the bounding box `(x_min, y_min, x_max, y_max)` of the input polygon.
|
||||
"""
|
||||
x_min, y_min = np.min(polygon, axis=0)
|
||||
x_max, y_max = np.max(polygon, axis=0)
|
||||
return np.array([x_min, y_min, x_max, y_max])
|
||||
|
||||
|
||||
def approximate_polygon(
|
||||
polygon: np.ndarray, percentage: float, epsilon_step: float = 0.05
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Approximates a given polygon by reducing a certain percentage of points.
|
||||
|
||||
This function uses the Ramer-Douglas-Peucker algorithm to simplify the input polygon by reducing the number of points
|
||||
while preserving the general shape.
|
||||
|
||||
Parameters:
|
||||
polygon (np.ndarray): A 2D NumPy array of shape `(N, 2)` containing the `x`, `y` coordinates of the input polygon's points.
|
||||
percentage (float): The percentage of points to be removed from the input polygon, in the range `[0, 1)`.
|
||||
epsilon_step (float): Approximation accuracy step. Epsilon is the maximum distance between the original curve and its approximation.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A new 2D NumPy array of shape `(M, 2)`, where `M <= N * (1 - percentage)`, containing the `x`, `y` coordinates of the
|
||||
approximated polygon's points.
|
||||
"""
|
||||
|
||||
if percentage < 0 or percentage >= 1:
|
||||
raise ValueError("Percentage must be in the range [0, 1).")
|
||||
|
||||
target_points = max(int(len(polygon) * (1 - percentage)), 3)
|
||||
|
||||
if len(polygon) <= target_points:
|
||||
return polygon
|
||||
|
||||
epsilon = 0
|
||||
approximated_points = polygon
|
||||
while True:
|
||||
epsilon += epsilon_step
|
||||
new_approximated_points = cv2.approxPolyDP(polygon, epsilon, closed=True)
|
||||
if len(new_approximated_points) > target_points:
|
||||
approximated_points = new_approximated_points
|
||||
else:
|
||||
break
|
||||
|
||||
return np.squeeze(approximated_points, axis=1)
|
||||
|
||||
|
||||
def extract_yolov8_masks(yolov8_results) -> Optional[np.ndarray]:
|
||||
if not yolov8_results.masks:
|
||||
return None
|
||||
|
||||
orig_shape = yolov8_results.orig_shape
|
||||
inference_shape = tuple(yolov8_results.masks.data.shape[1:])
|
||||
|
||||
gain = 0
|
||||
pad = (0, 0)
|
||||
|
||||
if inference_shape != orig_shape:
|
||||
gain = min(
|
||||
inference_shape[0] / orig_shape[0],
|
||||
inference_shape[1] / orig_shape[1],
|
||||
)
|
||||
pad = (
|
||||
(inference_shape[1] - orig_shape[1] * gain) / 2,
|
||||
(inference_shape[0] - orig_shape[0] * gain) / 2,
|
||||
)
|
||||
|
||||
top, left = int(pad[1]), int(pad[0])
|
||||
bottom, right = int(inference_shape[0] - pad[1]), int(inference_shape[1] - pad[0])
|
||||
|
||||
mask_maps = []
|
||||
masks = yolov8_results.masks.data.cpu().numpy()
|
||||
for i in range(masks.shape[0]):
|
||||
mask = masks[i]
|
||||
mask = mask[top:bottom, left:right]
|
||||
mask = cv2.resize(mask, (orig_shape[1], orig_shape[0]))
|
||||
mask_maps.append(mask)
|
||||
|
||||
return np.asarray(mask_maps, dtype=bool)
|
||||
|
||||
|
||||
def process_roboflow_result(
|
||||
roboflow_result: dict, class_list: List[str]
|
||||
) -> Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]]:
|
||||
if not roboflow_result["predictions"]:
|
||||
return np.empty((0, 4)), np.empty(0), np.empty(0), None
|
||||
|
||||
xyxy = []
|
||||
confidence = []
|
||||
class_id = []
|
||||
masks = []
|
||||
|
||||
image_width = int(roboflow_result["image"]["width"])
|
||||
image_height = int(roboflow_result["image"]["height"])
|
||||
|
||||
for prediction in roboflow_result["predictions"]:
|
||||
x = prediction["x"]
|
||||
y = prediction["y"]
|
||||
width = prediction["width"]
|
||||
height = prediction["height"]
|
||||
x_min = x - width / 2
|
||||
y_min = y - height / 2
|
||||
x_max = x_min + width
|
||||
y_max = y_min + height
|
||||
|
||||
xyxy.append([x_min, y_min, x_max, y_max])
|
||||
class_id.append(class_list.index(prediction["class"]))
|
||||
confidence.append(prediction["confidence"])
|
||||
|
||||
if "points" not in prediction:
|
||||
continue
|
||||
|
||||
polygon = np.array(
|
||||
[[point["x"], point["y"]] for point in prediction["points"]], dtype=int
|
||||
)
|
||||
|
||||
mask = polygon_to_mask(polygon, resolution_wh=(image_width, image_height))
|
||||
masks.append(mask)
|
||||
|
||||
xyxy = np.array(xyxy)
|
||||
confidence = np.array(confidence)
|
||||
class_id = np.array(class_id).astype(int)
|
||||
masks = np.array(masks, dtype=bool) if len(masks) > 0 else None
|
||||
|
||||
return xyxy, confidence, class_id, masks
|
||||
@@ -1,6 +1,6 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
from dataclasses import dataclass
|
||||
from typing import List, Tuple
|
||||
|
||||
DEFAULT_COLOR_PALETTE = [
|
||||
@@ -94,11 +94,11 @@ class Color:
|
||||
|
||||
@dataclass
|
||||
class ColorPalette:
|
||||
colors: List[Color] = field(
|
||||
default_factory=lambda: [
|
||||
Color.from_hex(color_hex) for color_hex in DEFAULT_COLOR_PALETTE
|
||||
]
|
||||
)
|
||||
colors: List[Color]
|
||||
|
||||
@classmethod
|
||||
def default(cls) -> ColorPalette:
|
||||
return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, color_hex_list: List[str]):
|
||||
|
||||
+114
-17
@@ -1,8 +1,10 @@
|
||||
from typing import Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.commons.dataclasses import Point, Rect
|
||||
from supervision.draw.color import Color
|
||||
from supervision.geometry.core import Point, Rect
|
||||
|
||||
|
||||
def draw_line(
|
||||
@@ -11,12 +13,15 @@ def draw_line(
|
||||
"""
|
||||
Draws a line on a given scene.
|
||||
|
||||
:param scene: np.ndarray : The scene on which the line will be drawn
|
||||
:param start: Point : The starting point of the line
|
||||
:param end: Point : The end point of the line
|
||||
:param color: Color : The color of the line
|
||||
:param thickness: int : The thickness of the line
|
||||
:return: np.ndarray : The scene with the line drawn on it
|
||||
Parameters:
|
||||
scene (np.ndarray): The scene on which the line will be drawn
|
||||
start (Point): The starting point of the line
|
||||
end (Point): The end point of the line
|
||||
color (Color): The color of the line
|
||||
thickness (int): The thickness of the line
|
||||
|
||||
Returns:
|
||||
np.ndarray: The scene with the line drawn on it
|
||||
"""
|
||||
cv2.line(
|
||||
scene,
|
||||
@@ -34,11 +39,14 @@ def draw_rectangle(
|
||||
"""
|
||||
Draws a rectangle on an image.
|
||||
|
||||
:param scene: np.ndarray : The image on which to draw the rectangle.
|
||||
:param rect: Rect : The rectangle to draw.
|
||||
:param color: Color : The color of the rectangle.
|
||||
:param thickness: int : The thickness of the rectangle border.
|
||||
:return: np.ndarray : The image with the rectangle drawn on it.
|
||||
Parameters:
|
||||
scene (np.ndarray): The scene on which the rectangle will be drawn
|
||||
rect (Rect): The rectangle to be drawn
|
||||
color (Color): The color of the rectangle
|
||||
thickness (int): The thickness of the rectangle border
|
||||
|
||||
Returns:
|
||||
np.ndarray: The scene with the rectangle drawn on it
|
||||
"""
|
||||
cv2.rectangle(
|
||||
scene,
|
||||
@@ -52,12 +60,15 @@ def draw_rectangle(
|
||||
|
||||
def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.ndarray:
|
||||
"""
|
||||
Draws a filled rectangle on the given scene.
|
||||
Draws a filled rectangle on an image.
|
||||
|
||||
:param scene: np.ndarray : The scene on which to draw the rectangle.
|
||||
:param rect: Rect : The rectangle to be drawn.
|
||||
:param color: Color : The color of the rectangle.
|
||||
:return: np.ndarray : The updated scene with the filled rectangle drawn on it.
|
||||
Parameters:
|
||||
scene (np.ndarray): The scene on which the rectangle will be drawn
|
||||
rect (Rect): The rectangle to be drawn
|
||||
color (Color): The color of the rectangle
|
||||
|
||||
Returns:
|
||||
np.ndarray: The scene with the rectangle drawn on it
|
||||
"""
|
||||
cv2.rectangle(
|
||||
scene,
|
||||
@@ -67,3 +78,89 @@ def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.nda
|
||||
-1,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
def draw_polygon(
|
||||
scene: np.ndarray, polygon: np.ndarray, color: Color, thickness: int = 2
|
||||
) -> np.ndarray:
|
||||
"""Draw a polygon on a scene.
|
||||
|
||||
Parameters:
|
||||
scene (np.ndarray): The scene to draw the polygon on.
|
||||
polygon (np.ndarray): The polygon to be drawn, given as a list of vertices.
|
||||
color (Color): The color of the polygon.
|
||||
thickness (int, optional): The thickness of the polygon lines, by default 2.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The scene with the polygon drawn on it.
|
||||
"""
|
||||
cv2.polylines(
|
||||
scene, [polygon], isClosed=True, color=color.as_bgr(), thickness=thickness
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
def draw_text(
|
||||
scene: np.ndarray,
|
||||
text: str,
|
||||
text_anchor: Point,
|
||||
text_color: Color = Color.black(),
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
text_font: int = cv2.FONT_HERSHEY_SIMPLEX,
|
||||
background_color: Optional[Color] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draw text with background on a scene.
|
||||
|
||||
Parameters:
|
||||
scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene.
|
||||
text (str): The text to be drawn.
|
||||
text_anchor (Point): The anchor point for the text, represented as a Point object with x and y attributes.
|
||||
text_color (Color, optional): The color of the text. Defaults to black.
|
||||
text_scale (float, optional): The scale of the text. Defaults to 0.5.
|
||||
text_thickness (int, optional): The thickness of the text. Defaults to 1.
|
||||
text_padding (int, optional): The amount of padding to add around the text when drawing a rectangle in the background. Defaults to 10.
|
||||
text_font (int, optional): The font to use for the text. Defaults to cv2.FONT_HERSHEY_SIMPLEX.
|
||||
background_color (Color, optional): The color of the background rectangle, if one is to be drawn. Defaults to None.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The input scene with the text drawn on it.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> scene = np.zeros((100, 100, 3), dtype=np.uint8)
|
||||
>>> text_anchor = Point(x=50, y=50)
|
||||
>>> scene = draw_text(scene=scene, text="Hello, world!", text_anchor=text_anchor)
|
||||
```
|
||||
"""
|
||||
text_width, text_height = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=text_font,
|
||||
fontScale=text_scale,
|
||||
thickness=text_thickness,
|
||||
)[0]
|
||||
text_rect = Rect(
|
||||
x=text_anchor.x - text_width // 2,
|
||||
y=text_anchor.y - text_height // 2,
|
||||
width=text_width,
|
||||
height=text_height,
|
||||
).pad(text_padding)
|
||||
|
||||
if background_color is not None:
|
||||
scene = draw_filled_rectangle(
|
||||
scene=scene, rect=text_rect, color=background_color
|
||||
)
|
||||
|
||||
cv2.putText(
|
||||
img=scene,
|
||||
text=text,
|
||||
org=(text_anchor.x - text_width // 2, text_anchor.y + text_height // 2),
|
||||
fontFace=text_font,
|
||||
fontScale=text_scale,
|
||||
color=text_color.as_bgr(),
|
||||
thickness=text_thickness,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
|
||||
@@ -1,9 +1,19 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Tuple
|
||||
|
||||
|
||||
class Position(Enum):
|
||||
CENTER = "CENTER"
|
||||
BOTTOM_CENTER = "BOTTOM_CENTER"
|
||||
|
||||
@classmethod
|
||||
def list(cls):
|
||||
return list(map(lambda c: c.value, cls))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Point:
|
||||
x: float
|
||||
@@ -0,0 +1,28 @@
|
||||
import numpy as np
|
||||
|
||||
from supervision.geometry.core import Point
|
||||
|
||||
|
||||
def get_polygon_center(polygon: np.ndarray) -> Point:
|
||||
"""
|
||||
Calculate the center of a polygon.
|
||||
|
||||
This function takes in a polygon as a 2-dimensional numpy ndarray and returns the center of the polygon as a Point object. The center is calculated as the mean of the polygon's vertices along each axis, and is rounded down to the nearest integer.
|
||||
|
||||
Parameters:
|
||||
polygon (np.ndarray): A 2-dimensional numpy ndarray representing the vertices of the polygon.
|
||||
|
||||
Returns:
|
||||
Point: The center of the polygon, represented as a Point object with x and y attributes.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> from supervision.geometry.utils import get_polygon_center
|
||||
|
||||
>>> vertices = np.array([[0, 0], [0, 1], [1, 1], [1, 0]])
|
||||
>>> get_center(vertices)
|
||||
Point(x=0.5, y=0.5)
|
||||
```
|
||||
"""
|
||||
center = np.mean(polygon, axis=0).astype(int)
|
||||
return Point(x=center[0], y=center[1])
|
||||
@@ -1,24 +0,0 @@
|
||||
from typing import Tuple
|
||||
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
def show_frame_in_notebook(
|
||||
frame: np.ndarray, size: Tuple[int, int] = (10, 10), cmap: str = "gray"
|
||||
):
|
||||
"""
|
||||
Display a frame in Jupyter Notebook using Matplotlib
|
||||
|
||||
:param frame: np.ndarray : The frame to be displayed.
|
||||
:param size: Tuple[int, int] : The size of the plot. default:(10,10)
|
||||
:param cmap: str : the colormap to use for single channel images. default:gray
|
||||
"""
|
||||
if frame.ndim == 2:
|
||||
plt.figure(figsize=size)
|
||||
plt.imshow(frame, cmap=cmap)
|
||||
else:
|
||||
plt.figure(figsize=size)
|
||||
plt.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
|
||||
plt.show()
|
||||
@@ -1,205 +0,0 @@
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
|
||||
|
||||
class Detections:
|
||||
def __init__(
|
||||
self,
|
||||
xyxy: np.ndarray,
|
||||
confidence: np.ndarray,
|
||||
class_id: np.ndarray,
|
||||
tracker_id: Optional[np.ndarray] = None,
|
||||
):
|
||||
"""
|
||||
Data class containing information about the detections in a video frame.
|
||||
|
||||
:param xyxy: np.ndarray : An array of shape (n, 4) containing the bounding boxes coordinates in format [x1, y1, x2, y2]
|
||||
:param confidence: np.ndarray : An array of shape (n,) containing the confidence scores of the detections.
|
||||
:param class_id: np.ndarray : An array of shape (n,) containing the class ids of the detections.
|
||||
:param tracker_id: Optional[np.ndarray] : An array of shape (n,) containing the tracker ids of the detections.
|
||||
"""
|
||||
self.xyxy: np.ndarray = xyxy
|
||||
self.confidence: np.ndarray = confidence
|
||||
self.class_id: np.ndarray = class_id
|
||||
self.tracker_id: Optional[np.ndarray] = tracker_id
|
||||
|
||||
n = len(self.xyxy)
|
||||
validators = [
|
||||
(isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)),
|
||||
(isinstance(self.confidence, np.ndarray) and self.confidence.shape == (n,)),
|
||||
(isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)),
|
||||
self.tracker_id is None
|
||||
or (
|
||||
isinstance(self.tracker_id, np.ndarray)
|
||||
and self.tracker_id.shape == (n,)
|
||||
),
|
||||
]
|
||||
if not all(validators):
|
||||
raise ValueError(
|
||||
"xyxy must be 2d np.ndarray with (n, 4) shape, "
|
||||
"confidence must be 1d np.ndarray with (n,) shape, "
|
||||
"class_id must be 1d np.ndarray with (n,) shape, "
|
||||
"tracker_id must be None or 1d np.ndarray with (n,) shape"
|
||||
)
|
||||
|
||||
def __len__(self):
|
||||
"""
|
||||
Returns the number of detections in the Detections object.
|
||||
"""
|
||||
return len(self.xyxy)
|
||||
|
||||
def __iter__(self):
|
||||
"""
|
||||
Iterates over the Detections object and yield a tuple of (xyxy, confidence, class_id, tracker_id) for each detection.
|
||||
"""
|
||||
for i in range(len(self.xyxy)):
|
||||
yield (
|
||||
self.xyxy[i],
|
||||
self.confidence[i],
|
||||
self.class_id[i],
|
||||
self.tracker_id[i] if self.tracker_id is not None else None,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_yolov5(cls, yolov5_output: np.ndarray):
|
||||
"""
|
||||
Creates a Detections instance from a YOLOv5 output tensor
|
||||
|
||||
:param yolov5_output: np.ndarray : The output tensor from YOLOv5
|
||||
:return: Detections : A Detections instance representing the detections in the frame
|
||||
|
||||
Example:
|
||||
detections = Detections.from_yolov5(yolov5_output)
|
||||
"""
|
||||
xyxy = yolov5_output[:, :4]
|
||||
confidence = yolov5_output[:, 4]
|
||||
class_id = yolov5_output[:, 5].astype(int)
|
||||
return cls(xyxy, confidence, class_id)
|
||||
|
||||
def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[np.ndarray]:
|
||||
"""
|
||||
Filter the detections by applying a mask
|
||||
|
||||
:param mask: np.ndarray : A mask of shape (n,) containing a boolean value for each detection indicating if it should be included in the filtered detections
|
||||
:param inplace: bool : If True, the original data will be modified and self will be returned.
|
||||
:return: Optional[np.ndarray] : A new instance of Detections with the filtered detections, if inplace is set to False. None otherwise.
|
||||
"""
|
||||
if inplace:
|
||||
self.xyxy = self.xyxy[mask]
|
||||
self.confidence = self.confidence[mask]
|
||||
self.class_id = self.class_id[mask]
|
||||
self.tracker_id = (
|
||||
self.tracker_id[mask] if self.tracker_id is not None else None
|
||||
)
|
||||
return self
|
||||
else:
|
||||
return Detections(
|
||||
xyxy=self.xyxy[mask],
|
||||
confidence=self.confidence[mask],
|
||||
class_id=self.class_id[mask],
|
||||
tracker_id=self.tracker_id[mask]
|
||||
if self.tracker_id is not None
|
||||
else None,
|
||||
)
|
||||
|
||||
|
||||
class BoxAnnotator:
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette],
|
||||
thickness: int = 2,
|
||||
text_color: Color = Color.black(),
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
):
|
||||
"""
|
||||
A class for drawing bounding boxes on an image using detections provided.
|
||||
|
||||
:param color: Union[Color, ColorPalette] : The color to draw the bounding box, can be a single color or a color palette
|
||||
:param thickness: int : The thickness of the bounding box lines, default is 2
|
||||
:param text_color: Color : The color of the text on the bounding box, default is white
|
||||
:param text_scale: float : The scale of the text on the bounding box, default is 0.5
|
||||
:param text_thickness: int : The thickness of the text on the bounding box, default is 1
|
||||
:param text_padding: int : The padding around the text on the bounding box, default is 5
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.text_color: Color = text_color
|
||||
self.text_scale: float = text_scale
|
||||
self.text_thickness: int = text_thickness
|
||||
self.text_padding: int = text_padding
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
frame: np.ndarray,
|
||||
detections: Detections,
|
||||
labels: Optional[List[str]] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draws bounding boxes on the frame using the detections provided.
|
||||
|
||||
:param frame: np.ndarray : The image on which the bounding boxes will be drawn
|
||||
:param detections: Detections : The detections for which the bounding boxes will be drawn
|
||||
:param labels: Optional[List[str]] : An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.
|
||||
:return: np.ndarray : The image with the bounding boxes drawn on it
|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections):
|
||||
color = (
|
||||
self.color.by_idx(class_id)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
text = (
|
||||
f"{confidence:0.2f}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[i]
|
||||
)
|
||||
|
||||
x1, y1, x2, y2 = xyxy.astype(int)
|
||||
text_width, text_height = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
thickness=self.text_thickness,
|
||||
)[0]
|
||||
|
||||
text_x = x1 + self.text_padding
|
||||
text_y = y1 - self.text_padding
|
||||
|
||||
text_background_x1 = x1
|
||||
text_background_y1 = y1 - 2 * self.text_padding - text_height
|
||||
|
||||
text_background_x2 = x1 + 2 * self.text_padding + text_width
|
||||
text_background_y2 = y1
|
||||
|
||||
cv2.rectangle(
|
||||
img=frame,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
cv2.rectangle(
|
||||
img=frame,
|
||||
pt1=(text_background_x1, text_background_y1),
|
||||
pt2=(text_background_x2, text_background_y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=cv2.FILLED,
|
||||
)
|
||||
cv2.putText(
|
||||
img=frame,
|
||||
text=text,
|
||||
org=(text_x, text_y),
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
color=self.text_color.as_rgb(),
|
||||
thickness=self.text_thickness,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return frame
|
||||
@@ -0,0 +1,140 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import List, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import yaml
|
||||
|
||||
|
||||
class NumpyJsonEncoder(json.JSONEncoder):
|
||||
def default(self, obj):
|
||||
if isinstance(obj, np.integer):
|
||||
return int(obj)
|
||||
if isinstance(obj, np.floating):
|
||||
return float(obj)
|
||||
if isinstance(obj, np.ndarray):
|
||||
return obj.tolist()
|
||||
return super(NumpyJsonEncoder, self).default(obj)
|
||||
|
||||
|
||||
def list_files_with_extensions(
|
||||
directory: Union[str, Path], extensions: Optional[List[str]] = None
|
||||
) -> List[Path]:
|
||||
"""
|
||||
List files in a directory with specified extensions or all files if no extensions are provided.
|
||||
|
||||
Args:
|
||||
directory (Union[str, Path]): The directory path as a string or Path object.
|
||||
extensions (Optional[List[str]]): A list of file extensions to filter. Default is None, which lists all files.
|
||||
|
||||
Returns:
|
||||
(List[Path]): A list of Path objects for the matching files.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> # List all files in the directory
|
||||
>>> files = sv.list_files_with_extensions(directory='my_directory')
|
||||
|
||||
>>> # List only files with '.txt' and '.md' extensions
|
||||
>>> files = sv.list_files_with_extensions(directory='my_directory', extensions=['txt', 'md'])
|
||||
```
|
||||
"""
|
||||
directory = Path(directory)
|
||||
files_with_extensions = []
|
||||
|
||||
if extensions is not None:
|
||||
for ext in extensions:
|
||||
files_with_extensions.extend(directory.glob(f"*.{ext}"))
|
||||
else:
|
||||
files_with_extensions.extend(directory.glob("*"))
|
||||
|
||||
return files_with_extensions
|
||||
|
||||
|
||||
def read_txt_file(file_path: str) -> List[str]:
|
||||
"""
|
||||
Read a text file and return a list of strings without newline characters.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the text file.
|
||||
|
||||
Returns:
|
||||
List[str]: A list of strings representing the lines in the text file.
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
lines = file.readlines()
|
||||
lines = [line.rstrip("\n") for line in lines]
|
||||
|
||||
return lines
|
||||
|
||||
|
||||
def save_text_file(lines: List[str], file_path: str):
|
||||
"""
|
||||
Write a list of strings to a text file, each string on a new line.
|
||||
|
||||
Args:
|
||||
lines (List[str]): The list of strings to be written to the file.
|
||||
file_path (str): The path to the text file.
|
||||
"""
|
||||
with open(file_path, "w") as file:
|
||||
for line in lines:
|
||||
file.write(line + "\n")
|
||||
|
||||
|
||||
def read_json_file(file_path: str) -> dict:
|
||||
"""
|
||||
Read a json file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the json file.
|
||||
|
||||
Returns:
|
||||
dict: A dict of annotations information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = json.load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_json_file(data: dict, file_path: str, indent: int = 3) -> None:
|
||||
"""
|
||||
Write a dict to a json file.
|
||||
|
||||
Args:
|
||||
indent:
|
||||
data (dict): dict with unique keys and value as pair.
|
||||
file_path (str): The path to the json file.
|
||||
"""
|
||||
with open(file_path, "w") as fp:
|
||||
json.dump(data, fp, cls=NumpyJsonEncoder, indent=indent)
|
||||
|
||||
|
||||
def read_yaml_file(file_path: str) -> dict:
|
||||
"""
|
||||
Read a yaml file and return a dict.
|
||||
|
||||
Args:
|
||||
file_path (str): The path to the yaml file.
|
||||
|
||||
Returns:
|
||||
dict: A dict of content information
|
||||
"""
|
||||
with open(file_path, "r") as file:
|
||||
data = yaml.safe_load(file)
|
||||
return data
|
||||
|
||||
|
||||
def save_yaml_file(data: dict, file_path: str) -> None:
|
||||
"""
|
||||
Save a dict to a json file.
|
||||
|
||||
Args:
|
||||
indent:
|
||||
data (dict): dict with unique keys and value as pair.
|
||||
file_path (str): The path to the json file.
|
||||
"""
|
||||
|
||||
with open(file_path, "w") as outfile:
|
||||
yaml.dump(data, outfile, sort_keys=False, default_flow_style=None)
|
||||
@@ -0,0 +1,93 @@
|
||||
import os
|
||||
import shutil
|
||||
from typing import Optional
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
def crop(image: np.ndarray, xyxy: np.ndarray) -> np.ndarray:
|
||||
"""
|
||||
Crops the given image based on the given bounding box.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The image to be cropped, represented as a numpy array.
|
||||
xyxy (np.ndarray): A numpy array containing the bounding box coordinates in the format (x1, y1, x2, y2).
|
||||
|
||||
Returns:
|
||||
(np.ndarray): The cropped image as a numpy array.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> detection = sv.Detections(...)
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path') as sink:
|
||||
... for xyxy in detection.xyxy:
|
||||
... cropped_image = sv.crop(image=image, xyxy=xyxy)
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
"""
|
||||
|
||||
xyxy = np.round(xyxy).astype(int)
|
||||
x1, y1, x2, y2 = xyxy
|
||||
cropped_img = image[y1:y2, x1:x2]
|
||||
return cropped_img
|
||||
|
||||
|
||||
class ImageSink:
|
||||
def __init__(
|
||||
self,
|
||||
target_dir_path: str,
|
||||
overwrite: bool = False,
|
||||
image_name_pattern: str = "image_{:05d}.png",
|
||||
):
|
||||
"""
|
||||
Initialize a context manager for saving images.
|
||||
|
||||
Args:
|
||||
target_dir_path (str): The target directory where images will be saved.
|
||||
overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
|
||||
image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
|
||||
... for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
|
||||
... sink.save_image(image=image)
|
||||
```
|
||||
"""
|
||||
self.target_dir_path = target_dir_path
|
||||
self.overwrite = overwrite
|
||||
self.image_name_pattern = image_name_pattern
|
||||
self.image_count = 0
|
||||
|
||||
def __enter__(self):
|
||||
if os.path.exists(self.target_dir_path):
|
||||
if self.overwrite:
|
||||
shutil.rmtree(self.target_dir_path)
|
||||
os.makedirs(self.target_dir_path)
|
||||
else:
|
||||
os.makedirs(self.target_dir_path)
|
||||
|
||||
return self
|
||||
|
||||
def save_image(self, image: np.ndarray, image_name: Optional[str] = None):
|
||||
"""
|
||||
Save a given image in the target directory.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The image to be saved.
|
||||
image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
|
||||
"""
|
||||
if image_name is None:
|
||||
image_name = self.image_name_pattern.format(self.image_count)
|
||||
|
||||
image_path = os.path.join(self.target_dir_path, image_name)
|
||||
cv2.imwrite(image_path, image)
|
||||
self.image_count += 1
|
||||
|
||||
def __exit__(self, exc_type, exc_value, exc_traceback):
|
||||
pass
|
||||
@@ -0,0 +1,18 @@
|
||||
import functools
|
||||
import warnings
|
||||
|
||||
|
||||
def deprecated(reason: str):
|
||||
def decorator(func):
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
warnings.warn(
|
||||
f"{func.__name__} is deprecated: {reason}",
|
||||
category=DeprecationWarning,
|
||||
stacklevel=2,
|
||||
)
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,97 @@
|
||||
from typing import List, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
def plot_image(
|
||||
image: np.ndarray, size: Tuple[int, int] = (12, 12), cmap: Optional[str] = "gray"
|
||||
) -> None:
|
||||
"""
|
||||
Plots image using matplotlib.
|
||||
|
||||
Args:
|
||||
image (np.ndarray): The frame to be displayed.
|
||||
size (Tuple[int, int]): The size of the plot.
|
||||
cmap (str): the colormap to use for single channel images.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = cv2.imread("path/to/image.jpg")
|
||||
|
||||
%matplotlib inline
|
||||
>>> sv.plot_image(image=image, size=(16, 16))
|
||||
```
|
||||
"""
|
||||
plt.figure(figsize=size)
|
||||
|
||||
if image.ndim == 2:
|
||||
plt.imshow(image, cmap=cmap)
|
||||
else:
|
||||
plt.imshow(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
|
||||
|
||||
plt.axis("off")
|
||||
plt.show()
|
||||
|
||||
|
||||
def plot_images_grid(
|
||||
images: List[np.ndarray],
|
||||
grid_size: Tuple[int, int],
|
||||
titles: Optional[List[str]] = None,
|
||||
size: Tuple[int, int] = (12, 12),
|
||||
cmap: Optional[str] = "gray",
|
||||
) -> None:
|
||||
"""
|
||||
Plots images in a grid using matplotlib.
|
||||
|
||||
Args:
|
||||
images (List[np.ndarray]): A list of images as numpy arrays.
|
||||
grid_size (Tuple[int, int]): A tuple specifying the number of rows and columns for the grid.
|
||||
titles (Optional[List[str]]): A list of titles for each image. Defaults to None.
|
||||
size (Tuple[int, int]): A tuple specifying the width and height of the entire plot in inches.
|
||||
cmap (str): the colormap to use for single channel images.
|
||||
|
||||
Raises:
|
||||
ValueError: If the number of images exceeds the grid size.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import cv2
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image1 = cv2.imread("path/to/image1.jpg")
|
||||
>>> image2 = cv2.imread("path/to/image2.jpg")
|
||||
>>> image3 = cv2.imread("path/to/image3.jpg")
|
||||
|
||||
>>> images = [image1, image2, image3]
|
||||
>>> titles = ["Image 1", "Image 2", "Image 3"]
|
||||
|
||||
%matplotlib inline
|
||||
>>> plot_images_grid(images, grid_size=(2, 2), titles=titles, size=(16, 16))
|
||||
```
|
||||
"""
|
||||
nrows, ncols = grid_size
|
||||
|
||||
if len(images) > nrows * ncols:
|
||||
raise ValueError(
|
||||
"The number of images exceeds the grid size. Please increase the grid size or reduce the number of images."
|
||||
)
|
||||
|
||||
fig, axes = plt.subplots(nrows=nrows, ncols=ncols, figsize=size)
|
||||
|
||||
for idx, ax in enumerate(axes.flat):
|
||||
if idx < len(images):
|
||||
if images[idx].ndim == 2:
|
||||
ax.imshow(images[idx], cmap=cmap)
|
||||
else:
|
||||
ax.imshow(cv2.cvtColor(images[idx], cv2.COLOR_BGR2RGB))
|
||||
|
||||
if titles is not None and idx < len(titles):
|
||||
ax.set_title(titles[idx])
|
||||
|
||||
ax.axis("off")
|
||||
plt.show()
|
||||
@@ -0,0 +1,184 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Generator, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoInfo:
|
||||
"""
|
||||
A class to store video information, including width, height, fps and total number of frames.
|
||||
|
||||
Attributes:
|
||||
width (int): width of the video in pixels
|
||||
height (int): height of the video in pixels
|
||||
fps (int): frames per second of the video
|
||||
total_frames (int, optional): total number of frames in the video, default is None
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='video.mp4')
|
||||
|
||||
>>> video_info
|
||||
VideoInfo(width=3840, height=2160, fps=25, total_frames=538)
|
||||
|
||||
>>> video_info.resolution_wh
|
||||
(3840, 2160)
|
||||
```
|
||||
"""
|
||||
|
||||
width: int
|
||||
height: int
|
||||
fps: int
|
||||
total_frames: Optional[int] = None
|
||||
|
||||
@classmethod
|
||||
def from_video_path(cls, video_path: str) -> VideoInfo:
|
||||
video = cv2.VideoCapture(video_path)
|
||||
if not video.isOpened():
|
||||
raise Exception(f"Could not open video at {video_path}")
|
||||
|
||||
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
fps = int(video.get(cv2.CAP_PROP_FPS))
|
||||
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
video.release()
|
||||
return VideoInfo(width, height, fps, total_frames)
|
||||
|
||||
@property
|
||||
def resolution_wh(self) -> Tuple[int, int]:
|
||||
return self.width, self.height
|
||||
|
||||
|
||||
class VideoSink:
|
||||
"""
|
||||
Context manager that saves video frames to a file using OpenCV.
|
||||
|
||||
Attributes:
|
||||
target_path (str): The path to the output file where the video will be saved.
|
||||
video_info (VideoInfo): Information about the video resolution, fps, and total frame count.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4')
|
||||
|
||||
>>> with sv.VideoSink(target_path='target_video.mp4', video_info=video_info) as sink:
|
||||
... for frame in get_video_frames_generator(source_path='source_video.mp4', stride=2):
|
||||
... sink.write_frame(frame=frame)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, target_path: str, video_info: VideoInfo):
|
||||
self.target_path = target_path
|
||||
self.video_info = video_info
|
||||
self.__fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
self.__writer = None
|
||||
|
||||
def __enter__(self):
|
||||
self.__writer = cv2.VideoWriter(
|
||||
self.target_path,
|
||||
self.__fourcc,
|
||||
self.video_info.fps,
|
||||
self.video_info.resolution_wh,
|
||||
)
|
||||
return self
|
||||
|
||||
def write_frame(self, frame: np.ndarray):
|
||||
self.__writer.write(frame)
|
||||
|
||||
def __exit__(self, exc_type, exc_value, exc_traceback):
|
||||
self.__writer.release()
|
||||
|
||||
|
||||
def _validate_and_setup_video(source_path: str, start: int, end: Optional[int]):
|
||||
video = cv2.VideoCapture(source_path)
|
||||
if not video.isOpened():
|
||||
raise Exception(f"Could not open video at {source_path}")
|
||||
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
if end is not None and end > total_frames:
|
||||
raise Exception(f"Requested frames are outbound")
|
||||
start = max(start, 0)
|
||||
end = min(end, total_frames) if end is not None else total_frames
|
||||
video.set(cv2.CAP_PROP_POS_FRAMES, start)
|
||||
return video, start, end
|
||||
|
||||
|
||||
def get_video_frames_generator(
|
||||
source_path: str, stride: int = 1, start: int = 0, end: Optional[int] = None
|
||||
) -> Generator[np.ndarray, None, None]:
|
||||
"""
|
||||
Get a generator that yields the frames of the video.
|
||||
|
||||
Args:
|
||||
source_path (str): The path of the video file.
|
||||
stride (int): Indicates the interval at which frames are returned, skipping stride - 1 frames between each.
|
||||
start (int): Indicates the starting position from which video should generate frames
|
||||
end (Optional[int]): Indicates the ending position at which video should stop generating frames. If None, video will be read to the end.
|
||||
|
||||
Returns:
|
||||
(Generator[np.ndarray, None, None]): A generator that yields the frames of the video.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'):
|
||||
... ...
|
||||
```
|
||||
"""
|
||||
video, start, end = _validate_and_setup_video(source_path, start, end)
|
||||
frame_position = start
|
||||
while True:
|
||||
success, frame = video.read()
|
||||
if not success or frame_position >= end:
|
||||
break
|
||||
yield frame
|
||||
for _ in range(stride - 1):
|
||||
success = video.grab()
|
||||
if not success:
|
||||
break
|
||||
frame_position += stride
|
||||
video.release()
|
||||
|
||||
|
||||
def process_video(
|
||||
source_path: str,
|
||||
target_path: str,
|
||||
callback: Callable[[np.ndarray, int], np.ndarray],
|
||||
) -> None:
|
||||
"""
|
||||
Process a video file by applying a callback function on each frame and saving the result to a target video file.
|
||||
|
||||
Args:
|
||||
source_path (str): The path to the source video file.
|
||||
target_path (str): The path to the target video file.
|
||||
callback (Callable[[np.ndarray, int], np.ndarray]): A function that takes in a numpy ndarray representation of a video frame and an int index of the frame and returns a processed numpy ndarray representation of the frame.
|
||||
|
||||
Examples:
|
||||
```python
|
||||
>>> from supervision import process_video
|
||||
|
||||
>>> def process_frame(scene: np.ndarray) -> np.ndarray:
|
||||
... ...
|
||||
|
||||
>>> process_video(
|
||||
... source_path='source_video.mp4',
|
||||
... target_path='target_video.mp4',
|
||||
... callback=process_frame
|
||||
... )
|
||||
```
|
||||
"""
|
||||
source_video_info = VideoInfo.from_video_path(video_path=source_path)
|
||||
with VideoSink(target_path=target_path, video_info=source_video_info) as sink:
|
||||
for index, frame in enumerate(
|
||||
get_video_frames_generator(source_path=source_path)
|
||||
):
|
||||
result_frame = callback(frame, index)
|
||||
sink.write_frame(frame=result_frame)
|
||||
@@ -1,46 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Tuple
|
||||
|
||||
import cv2
|
||||
|
||||
|
||||
@dataclass
|
||||
class VideoInfo:
|
||||
"""
|
||||
Data class containing information about the video resolution, fps, and total frame count.
|
||||
|
||||
:param width: int : The width of the video frames in pixels.
|
||||
:param height: int : The height of the video frames in pixels.
|
||||
:param fps: int : The frames per second of the video.
|
||||
:param total_frames: int : The total number of frames in the video.
|
||||
"""
|
||||
|
||||
width: int
|
||||
height: int
|
||||
fps: int
|
||||
total_frames: Optional[int] = None
|
||||
|
||||
@classmethod
|
||||
def from_video_path(cls, video_path: str) -> VideoInfo:
|
||||
"""
|
||||
Returns a VideoInfo data class containing information about the video resolution, fps, and total frame count.
|
||||
|
||||
:param video_path: str : The path of the video file.
|
||||
:return: VideoInfo : A data class containing information about the video resolution, fps, and total frame count.
|
||||
"""
|
||||
video = cv2.VideoCapture(video_path)
|
||||
if not video.isOpened():
|
||||
raise Exception(f"Could not open video at {video_path}")
|
||||
|
||||
width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
|
||||
height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
||||
fps = int(video.get(cv2.CAP_PROP_FPS))
|
||||
total_frames = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
|
||||
video.release()
|
||||
return VideoInfo(width, height, fps, total_frames)
|
||||
|
||||
@property
|
||||
def resolution(self) -> Tuple[int, int]:
|
||||
return self.width, self.height
|
||||
@@ -1,48 +0,0 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.video.dataclasses import VideoInfo
|
||||
|
||||
|
||||
class VideoSink:
|
||||
"""
|
||||
A context manager that uses OpenCV to save video frames to a file.
|
||||
|
||||
:param output_path: str : The path to the output file where the video will be saved.
|
||||
:param video_info: VideoInfo : An instance of VideoInfo containing information about the video resolution, fps, and total frame count.
|
||||
"""
|
||||
|
||||
def __init__(self, output_path: str, video_info: VideoInfo):
|
||||
"""
|
||||
Initializes the VideoSink with the specified output path and video information.
|
||||
"""
|
||||
self.output_path = output_path
|
||||
self.video_info = video_info
|
||||
self.fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
self.writer = None
|
||||
|
||||
def __enter__(self):
|
||||
"""
|
||||
Opens the output file and returns the VideoSink instance.
|
||||
"""
|
||||
self.writer = cv2.VideoWriter(
|
||||
self.output_path,
|
||||
self.fourcc,
|
||||
self.video_info.fps,
|
||||
self.video_info.resolution,
|
||||
)
|
||||
return self
|
||||
|
||||
def write_frame(self, frame: np.ndarray):
|
||||
"""
|
||||
Writes a frame to the output video file.
|
||||
|
||||
:param frame: np.ndarray : The frame to be written.
|
||||
"""
|
||||
self.writer.write(frame)
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
"""
|
||||
Closes the output file.
|
||||
"""
|
||||
self.writer.release()
|
||||
@@ -1,20 +0,0 @@
|
||||
from typing import Generator
|
||||
|
||||
import cv2
|
||||
|
||||
|
||||
def get_video_frames_generator(video_path: str) -> Generator[int, None, None]:
|
||||
"""
|
||||
Returns a generator that yields the frames of the video.
|
||||
|
||||
:param video_path: str : The path of the video file.
|
||||
:return: Generator[int, None, None] : Generator that yields the frames of the video.
|
||||
"""
|
||||
video = cv2.VideoCapture(video_path)
|
||||
if not video.isOpened():
|
||||
raise Exception(f"Could not open video at {video_path}")
|
||||
success, frame = video.read()
|
||||
while success:
|
||||
yield frame
|
||||
success, frame = video.read()
|
||||
video.release()
|
||||
@@ -0,0 +1,63 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
|
||||
import pytest
|
||||
import numpy as np
|
||||
|
||||
from supervision.classification.core import Classifications
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"class_id, confidence, k, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.array([0, 1, 2, 3, 4]),
|
||||
np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
|
||||
5,
|
||||
(np.array([2, 4, 3, 1, 0]), np.array([0.9, 0.5, 0.4, 0.2, 0.1])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences
|
||||
(
|
||||
np.array([5, 1, 2, 3, 4]),
|
||||
np.array([0.1, 0.2, 0.9, 0.4, 0.5]),
|
||||
1,
|
||||
(np.array([2]), np.array([0.9])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences, retrieve where k = 1
|
||||
(
|
||||
np.array([4, 1, 2, 3, 6, 5]),
|
||||
np.array([0.8, 0.2, 0.9, 0.4, 0.5, 0.1]),
|
||||
2,
|
||||
(np.array([2, 4]), np.array([0.9, 0.8])),
|
||||
DoesNotRaise(),
|
||||
), # class_id with 5 numbers and 5 confidences, retrieve where k = 3
|
||||
(
|
||||
np.array([0, 1, 2, 3, 4]),
|
||||
np.array([]),
|
||||
5,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id with 5 numbers and 0 confidences
|
||||
(
|
||||
[0, 1, 2, 3, 4],
|
||||
[0.1, 0.2, 0.3, 0.4],
|
||||
5,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # class_id with 5 numbers and 4 confidences
|
||||
],
|
||||
)
|
||||
def test_top_k(
|
||||
class_id: np.ndarray,
|
||||
confidence: Optional[np.ndarray],
|
||||
k: int,
|
||||
expected_result: Optional[Tuple[np.ndarray, np.ndarray]],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = Classifications(
|
||||
class_id=np.array(class_id), confidence=np.array(confidence)
|
||||
).get_top_k(k)
|
||||
|
||||
assert np.array_equal(result[0], expected_result[0])
|
||||
assert np.array_equal(result[1], expected_result[1])
|
||||
@@ -0,0 +1,292 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, Tuple
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision import Detections
|
||||
from supervision.dataset.formats.coco import classes_to_coco_categories, coco_categories_to_classes, \
|
||||
group_coco_annotations_by_image_id, coco_annotations_to_detections
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
def mock_cock_coco_annotation(
|
||||
annotation_id: int = 0,
|
||||
image_id: int = 0,
|
||||
category_id: int = 0,
|
||||
bbox: Tuple[float, float, float, float] = (0.0, 0.0, 0.0, 0.0),
|
||||
area: float = 0.0
|
||||
) -> dict:
|
||||
return {
|
||||
"id": annotation_id,
|
||||
"image_id": image_id,
|
||||
"category_id": category_id,
|
||||
"bbox": list(bbox),
|
||||
"area": area,
|
||||
"iscrowd": 0
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_categories, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # empty coco categories
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
}
|
||||
],
|
||||
["fashion-assistant"],
|
||||
DoesNotRaise()
|
||||
), # single coco category with supercategory == "none"
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"name": "baseball cap",
|
||||
"supercategory": "fashion-assistant"
|
||||
}
|
||||
],
|
||||
[
|
||||
"fashion-assistant", "baseball cap"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # two coco categories; one with supercategory == "none" and one with supercategory != "none"
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"name": "baseball cap",
|
||||
"supercategory": "fashion-assistant"
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"name": "hoodie",
|
||||
"supercategory": "fashion-assistant"
|
||||
}
|
||||
],
|
||||
[
|
||||
"fashion-assistant",
|
||||
"baseball cap",
|
||||
"hoodie"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # three coco categories; one with supercategory == "none" and two with supercategory != "none"
|
||||
(
|
||||
[
|
||||
{
|
||||
"id": 0,
|
||||
"name": "fashion-assistant",
|
||||
"supercategory": "none"
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"name": "hoodie",
|
||||
"supercategory": "fashion-assistant"
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"name": "baseball cap",
|
||||
"supercategory": "fashion-assistant"
|
||||
}
|
||||
],
|
||||
[
|
||||
"fashion-assistant",
|
||||
"baseball cap",
|
||||
"hoodie"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # three coco categories; one with supercategory == "none" and two with supercategory != "none" (different order)
|
||||
]
|
||||
)
|
||||
def test_coco_categories_to_classes(
|
||||
coco_categories: List[dict],
|
||||
expected_result: List[str],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = coco_categories_to_classes(coco_categories=coco_categories)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"classes, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # empty classes
|
||||
(
|
||||
[
|
||||
"baseball cap"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # single class
|
||||
(
|
||||
[
|
||||
"baseball cap",
|
||||
"hoodie"
|
||||
],
|
||||
DoesNotRaise()
|
||||
), # two classes
|
||||
]
|
||||
)
|
||||
def test_classes_to_coco_categories_and_back_to_classes(classes: List[str], exception: Exception) -> None:
|
||||
with exception:
|
||||
coco_categories = classes_to_coco_categories(classes=classes)
|
||||
result = coco_categories_to_classes(coco_categories=coco_categories)
|
||||
assert result == classes
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"coco_annotations, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty coco annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
|
||||
],
|
||||
{
|
||||
0: [
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # single coco annotation
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
|
||||
],
|
||||
{
|
||||
0: [
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0)
|
||||
],
|
||||
1: [
|
||||
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=0)
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # two coco annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
|
||||
mock_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
|
||||
mock_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
|
||||
mock_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
|
||||
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
|
||||
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
|
||||
],
|
||||
{
|
||||
0: [
|
||||
mock_cock_coco_annotation(annotation_id=0, image_id=0, category_id=0),
|
||||
],
|
||||
1: [
|
||||
mock_cock_coco_annotation(annotation_id=1, image_id=1, category_id=1),
|
||||
mock_cock_coco_annotation(annotation_id=2, image_id=1, category_id=2),
|
||||
],
|
||||
2: [
|
||||
mock_cock_coco_annotation(annotation_id=3, image_id=2, category_id=3),
|
||||
],
|
||||
3: [
|
||||
mock_cock_coco_annotation(annotation_id=4, image_id=3, category_id=1),
|
||||
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=2),
|
||||
mock_cock_coco_annotation(annotation_id=5, image_id=3, category_id=3),
|
||||
]
|
||||
},
|
||||
DoesNotRaise()
|
||||
), # two coco annotations
|
||||
]
|
||||
)
|
||||
def test_group_coco_annotations_by_image_id(
|
||||
coco_annotations: List[dict],
|
||||
expected_result: dict,
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = group_coco_annotations_by_image_id(coco_annotations=coco_annotations)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"image_annotations, resolution_wh, with_masks, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
), # empty image annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100)
|
||||
],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 0, 0, 100, 100]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0
|
||||
], dtype=int)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # single image annotations
|
||||
(
|
||||
[
|
||||
mock_cock_coco_annotation(category_id=0, bbox=(0, 0, 100, 100), area=100 * 100),
|
||||
mock_cock_coco_annotation(category_id=0, bbox=(100, 100, 100, 100), area=100 * 100),
|
||||
],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[ 0, 0, 100, 100],
|
||||
[ 100, 100, 200, 200]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([
|
||||
0, 0
|
||||
], dtype=int)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two image annotations
|
||||
]
|
||||
)
|
||||
def test_coco_annotations_to_detections(
|
||||
image_annotations: List[dict],
|
||||
resolution_wh: Tuple[int, int],
|
||||
with_masks: bool,
|
||||
expected_result: Detections,
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = coco_annotations_to_detections(
|
||||
image_annotations=image_annotations,
|
||||
resolution_wh=resolution_wh,
|
||||
with_masks=with_masks
|
||||
)
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,304 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, Tuple, Optional
|
||||
|
||||
import pytest
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.dataset.formats.yolo import yolo_annotations_to_detections, _with_mask, _image_name_to_annotation_name, \
|
||||
object_to_yolo
|
||||
|
||||
|
||||
def _mock_simple_mask(resolution_wh: Tuple[int, int], box: List[int]) -> np.array:
|
||||
x_min, y_min, x_max, y_max = box
|
||||
mask = np.full(resolution_wh, False, dtype=bool)
|
||||
mask[y_min:y_max, x_min:x_max] = True
|
||||
return mask
|
||||
|
||||
|
||||
# The result of _mock_simple_mask is a little different from the result produced by cv2.
|
||||
def _arrays_almost_equal(arr1: np.ndarray, arr2: np.ndarray, threshold: float = 0.99) -> bool:
|
||||
equal_elements = np.equal(arr1, arr2)
|
||||
proportion_equal = np.mean(equal_elements)
|
||||
return proportion_equal >= threshold
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'lines, expected_result, exception',
|
||||
[
|
||||
(
|
||||
[],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
), # empty yolo annotation file
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with single line with box
|
||||
(
|
||||
[
|
||||
'0 0.50 0.50 0.20 0.20',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with two lines with box
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
False,
|
||||
DoesNotRaise()
|
||||
),
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
|
||||
],
|
||||
True,
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with single line with polygon
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
True,
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with two lines - one box and one polygon
|
||||
]
|
||||
)
|
||||
def test_with_mask(lines: List[str], expected_result: Optional[bool], exception: Exception) -> None:
|
||||
with exception:
|
||||
result = _with_mask(lines=lines)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'lines, resolution_wh, with_masks, expected_result, exception',
|
||||
[
|
||||
(
|
||||
[],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
), # empty yolo annotation file
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with single line with box
|
||||
(
|
||||
[
|
||||
'0 0.50 0.50 0.20 0.20',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600],
|
||||
[ 0, 320, 220, 620]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0, 1], dtype=int)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with two lines with box
|
||||
(
|
||||
[
|
||||
'0 0.5 0.5 0.2 0.2'
|
||||
],
|
||||
(1000, 1000),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array([
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
|
||||
], dtype=bool)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with single line with box in with_masks mode
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6'
|
||||
],
|
||||
(1000, 1000),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0], dtype=int),
|
||||
mask=np.array([
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600])
|
||||
], dtype=bool)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with single line with polygon
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
(1000, 1000),
|
||||
True,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600],
|
||||
[ 0, 320, 220, 620]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0, 1], dtype=int),
|
||||
mask=np.array([
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[400, 400, 600, 600]),
|
||||
_mock_simple_mask(resolution_wh=(1000, 1000), box=[ 0, 320, 220, 620])
|
||||
], dtype=bool)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with two lines - one box and one polygon in with_masks mode
|
||||
(
|
||||
[
|
||||
'0 0.4 0.4 0.6 0.4 0.6 0.6 0.4 0.6',
|
||||
'1 0.11 0.47 0.22 0.30'
|
||||
],
|
||||
(1000, 1000),
|
||||
False,
|
||||
Detections(
|
||||
xyxy=np.array([
|
||||
[400, 400, 600, 600],
|
||||
[ 0, 320, 220, 620]
|
||||
], dtype=np.float32),
|
||||
class_id=np.array([0, 1], dtype=int)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # yolo annotation file with two lines - one box and one polygon
|
||||
]
|
||||
)
|
||||
def test_yolo_annotations_to_detections(
|
||||
lines: List[str],
|
||||
resolution_wh: Tuple[int, int],
|
||||
with_masks: bool,
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = yolo_annotations_to_detections(
|
||||
lines=lines,
|
||||
resolution_wh=resolution_wh,
|
||||
with_masks=with_masks)
|
||||
assert np.array_equal(result.xyxy, expected_result.xyxy)
|
||||
assert np.array_equal(result.class_id, expected_result.class_id)
|
||||
assert (result.mask is None and expected_result.mask is None) or _arrays_almost_equal(result.mask, expected_result.mask)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'image_name, expected_result, exception',
|
||||
[
|
||||
(
|
||||
'image.png',
|
||||
'image.txt',
|
||||
DoesNotRaise()
|
||||
), # simple png image
|
||||
(
|
||||
'image.jpeg',
|
||||
'image.txt',
|
||||
DoesNotRaise()
|
||||
), # simple jpeg image
|
||||
(
|
||||
'image.jpg',
|
||||
'image.txt',
|
||||
DoesNotRaise()
|
||||
), # simple jpg image
|
||||
(
|
||||
'image.000.jpg',
|
||||
'image.000.txt',
|
||||
DoesNotRaise()
|
||||
), # jpg image with multiple dots in name
|
||||
]
|
||||
)
|
||||
def test_image_name_to_annotation_name(
|
||||
image_name: str,
|
||||
expected_result: Optional[str],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = _image_name_to_annotation_name(image_name=image_name)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'xyxy, class_id, image_shape, polygon, expected_result, exception',
|
||||
[
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
None,
|
||||
'1 0.15000 0.15000 0.10000 0.10000',
|
||||
DoesNotRaise()
|
||||
), # square bounding box on square image
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(800, 1000, 3),
|
||||
None,
|
||||
'1 0.15000 0.18750 0.10000 0.12500',
|
||||
DoesNotRaise()
|
||||
), # square bounding box on horizontal image
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 800, 3),
|
||||
None,
|
||||
'1 0.18750 0.15000 0.12500 0.10000',
|
||||
DoesNotRaise()
|
||||
), # square bounding box on vertical image
|
||||
(
|
||||
np.array([100, 200, 200, 400], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
None,
|
||||
'1 0.15000 0.30000 0.10000 0.20000',
|
||||
DoesNotRaise()
|
||||
), # horizontal bounding box on square image
|
||||
(
|
||||
np.array([200, 100, 400, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
None,
|
||||
'1 0.30000 0.15000 0.20000 0.10000',
|
||||
DoesNotRaise()
|
||||
), # vertical bounding box on square image
|
||||
(
|
||||
np.array([100, 100, 200, 200], dtype=np.float32),
|
||||
1,
|
||||
(1000, 1000, 3),
|
||||
np.array([
|
||||
[100, 100],
|
||||
[200, 100],
|
||||
[200, 200],
|
||||
[100, 100]
|
||||
], dtype=np.float32),
|
||||
'1 0.10000 0.10000 0.20000 0.10000 0.20000 0.20000 0.10000 0.10000',
|
||||
DoesNotRaise()
|
||||
), # square mask on square image
|
||||
]
|
||||
)
|
||||
def test_object_to_yolo(xyxy: np.ndarray, class_id: int, image_shape: Tuple[int, int, int], polygon: Optional[np.ndarray], expected_result: Optional[str], exception: Exception) -> None:
|
||||
with exception:
|
||||
result = object_to_yolo(xyxy=xyxy, class_id=class_id, image_shape=image_shape, polygon=polygon)
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,176 @@
|
||||
from typing import List, Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision import DetectionDataset
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
|
||||
import numpy as np
|
||||
|
||||
from test.utils import mock_detections
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dataset_list, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[],
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
DoesNotRaise()
|
||||
), # empty dataset list
|
||||
(
|
||||
[
|
||||
DetectionDataset(classes=[], images={}, annotations={})
|
||||
],
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
DoesNotRaise()
|
||||
), # single empty dataset
|
||||
(
|
||||
[
|
||||
DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
|
||||
DetectionDataset(classes=['dog', 'person'], images={}, annotations={})
|
||||
],
|
||||
DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
|
||||
DoesNotRaise()
|
||||
), # two datasets; no images and annotations, the same classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(classes=['dog', 'person'], images={}, annotations={}),
|
||||
DetectionDataset(classes=['cat'], images={}, annotations={})
|
||||
],
|
||||
DetectionDataset(classes=['cat', 'dog', 'person'], images={}, annotations={}),
|
||||
DoesNotRaise()
|
||||
), # two datasets; no images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=['dog', 'person'],
|
||||
images={
|
||||
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(classes=[], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(
|
||||
classes=['dog', 'person'],
|
||||
images={
|
||||
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
}
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two datasets; images and annotations, the same classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=['dog', 'person'],
|
||||
images={
|
||||
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(classes=['cat'], images={}, annotations={}),
|
||||
],
|
||||
DetectionDataset(
|
||||
classes=['cat', 'dog', 'person'],
|
||||
images={
|
||||
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
|
||||
}
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two datasets; images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=['dog', 'person'],
|
||||
images={
|
||||
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(
|
||||
classes=['cat'],
|
||||
images={
|
||||
'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
}
|
||||
),
|
||||
],
|
||||
DetectionDataset(
|
||||
classes=['cat', 'dog', 'person'],
|
||||
images={
|
||||
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
|
||||
'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
}
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two datasets; images and annotations, different classes
|
||||
(
|
||||
[
|
||||
DetectionDataset(
|
||||
classes=['dog', 'person'],
|
||||
images={
|
||||
'image-1.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-1.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
}
|
||||
),
|
||||
DetectionDataset(
|
||||
classes=['dog', 'person'],
|
||||
images={
|
||||
'image-2.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
'image-3.png': np.zeros((100, 100, 3), dtype=np.uint8),
|
||||
},
|
||||
annotations={
|
||||
'image-2.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
'image-3.png': mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
}
|
||||
),
|
||||
],
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
)
|
||||
]
|
||||
)
|
||||
def test_dataset_merge(
|
||||
dataset_list: List[DetectionDataset],
|
||||
expected_result: Optional[DetectionDataset],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = DetectionDataset.merge(dataset_list=dataset_list)
|
||||
print(result.images.keys())
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,241 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import List, TypeVar, Optional, Tuple, Dict
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision import Detections
|
||||
from supervision.dataset.utils import train_test_split, merge_class_lists, build_class_index_mapping, \
|
||||
map_detections_class_id
|
||||
from test.utils import mock_detections
|
||||
|
||||
T = TypeVar("T")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'data, train_ratio, random_state, shuffle, expected_result, exception',
|
||||
[
|
||||
(
|
||||
[],
|
||||
0.5,
|
||||
None,
|
||||
False,
|
||||
([], []),
|
||||
DoesNotRaise()
|
||||
), # empty data
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
0.5,
|
||||
None,
|
||||
False,
|
||||
([0, 1, 2, 3, 4], [5, 6, 7, 8, 9]),
|
||||
DoesNotRaise()
|
||||
), # data with 10 numbers and 50% train split
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
1.0,
|
||||
None,
|
||||
False,
|
||||
([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], []),
|
||||
DoesNotRaise()
|
||||
), # data with 10 numbers and 100% train split
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
0.0,
|
||||
None,
|
||||
False,
|
||||
([], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]),
|
||||
DoesNotRaise()
|
||||
), # data with 10 numbers and 0% train split
|
||||
(
|
||||
['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'],
|
||||
0.5,
|
||||
None,
|
||||
False,
|
||||
(['a', 'b', 'c', 'd', 'e'], ['f', 'g', 'h', 'i', 'j']),
|
||||
DoesNotRaise()
|
||||
), # data with 10 chars and 50% train split
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
0.5,
|
||||
23,
|
||||
True,
|
||||
([7, 8, 5, 6, 3], [2, 9, 0, 1, 4]),
|
||||
DoesNotRaise()
|
||||
), # data with 10 numbers and 50% train split with 23 random seed
|
||||
(
|
||||
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
|
||||
0.5,
|
||||
32,
|
||||
True,
|
||||
([4, 6, 0, 8, 9], [5, 7, 2, 3, 1]),
|
||||
DoesNotRaise()
|
||||
), # data with 10 numbers and 50% train split with 23 random seed
|
||||
]
|
||||
)
|
||||
def test_train_test_split(
|
||||
data: List[T],
|
||||
train_ratio: float,
|
||||
random_state: int,
|
||||
shuffle: bool,
|
||||
expected_result: Optional[Tuple[List[T], List[T]]],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = train_test_split(data=data, train_ratio=train_ratio, random_state=random_state, shuffle=shuffle)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'class_lists, expected_result, exception',
|
||||
[
|
||||
(
|
||||
[],
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # empty class lists
|
||||
(
|
||||
[
|
||||
['dog', 'person']
|
||||
],
|
||||
['dog', 'person'],
|
||||
DoesNotRaise()
|
||||
), # single class list; already alphabetically sorted
|
||||
(
|
||||
[
|
||||
['person', 'dog']
|
||||
],
|
||||
['dog', 'person'],
|
||||
DoesNotRaise()
|
||||
), # single class list; not alphabetically sorted
|
||||
(
|
||||
[
|
||||
['dog', 'person'],
|
||||
['dog', 'person']
|
||||
],
|
||||
['dog', 'person'],
|
||||
DoesNotRaise()
|
||||
), # two class lists; the same classes; already alphabetically sorted
|
||||
(
|
||||
[
|
||||
['dog', 'person'],
|
||||
['cat']
|
||||
],
|
||||
['cat', 'dog', 'person'],
|
||||
DoesNotRaise()
|
||||
), # two class lists; different classes; already alphabetically sorted
|
||||
]
|
||||
)
|
||||
def test_merge_class_maps(class_lists: List[List[str]], expected_result: List[str], exception: Exception) -> None:
|
||||
with exception:
|
||||
result = merge_class_lists(class_lists=class_lists)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'source_classes, target_classes, expected_result, exception',
|
||||
[
|
||||
(
|
||||
[],
|
||||
[],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty class lists
|
||||
(
|
||||
[],
|
||||
['dog', 'person'],
|
||||
{},
|
||||
DoesNotRaise()
|
||||
), # empty source class list
|
||||
(
|
||||
['dog', 'person'],
|
||||
[],
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
), # empty target class list
|
||||
(
|
||||
['dog', 'person'],
|
||||
['dog', 'person'],
|
||||
{0: 0, 1: 1},
|
||||
DoesNotRaise()
|
||||
), # same class lists
|
||||
(
|
||||
['dog', 'person'],
|
||||
['person', 'dog'],
|
||||
{0: 1, 1: 0},
|
||||
DoesNotRaise()
|
||||
), # same class lists but not alphabetically sorted
|
||||
(
|
||||
['dog', 'person'],
|
||||
['cat', 'dog', 'person'],
|
||||
{0: 1, 1: 2},
|
||||
DoesNotRaise()
|
||||
), # source class list is a subset of target class list
|
||||
(
|
||||
['dog', 'person'],
|
||||
['cat', 'dog'],
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
), # source class list is not a subset of target class list
|
||||
]
|
||||
)
|
||||
def test_build_class_index_mapping(
|
||||
source_classes: List[str],
|
||||
target_classes: List[str],
|
||||
expected_result: Optional[Dict[int, int]],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = build_class_index_mapping(source_classes=source_classes, target_classes=target_classes)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'source_to_target_mapping, detections, expected_result, exception',
|
||||
[
|
||||
(
|
||||
{},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
), # empty mapping
|
||||
(
|
||||
{0: 1},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
DoesNotRaise()
|
||||
), # single mapping
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1]),
|
||||
DoesNotRaise()
|
||||
), # multiple mappings
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[0, 1]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10], [0, 0, 10, 10]], class_id=[1, 2]),
|
||||
DoesNotRaise()
|
||||
), # multiple mappings
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[2]),
|
||||
None,
|
||||
pytest.raises(ValueError)
|
||||
), # class_id not in mapping
|
||||
(
|
||||
{0: 1, 1: 2},
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[0], confidence=[0.5]),
|
||||
mock_detections(xyxy=[[0, 0, 10, 10]], class_id=[1], confidence=[0.5]),
|
||||
DoesNotRaise(),
|
||||
), # confidence is not None
|
||||
]
|
||||
)
|
||||
def test_map_detections_class_id(
|
||||
source_to_target_mapping: Dict[int, int],
|
||||
detections: Detections,
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = map_detections_class_id(source_to_target_mapping=source_to_target_mapping, detections=detections)
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,245 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision import Detections
|
||||
|
||||
from typing import Optional, Union, List
|
||||
|
||||
import numpy as np
|
||||
|
||||
from test.utils import mock_detections
|
||||
|
||||
PREDICTIONS = np.array([
|
||||
[ 2254, 906, 2447, 1353, 0.90538, 0],
|
||||
[ 2049, 1133, 2226, 1371, 0.59002, 56],
|
||||
[ 727, 1224, 838, 1601, 0.51119, 39],
|
||||
[ 808, 1214, 910, 1564, 0.45287, 39],
|
||||
[ 6, 52, 1131, 2133, 0.45057, 72],
|
||||
[ 299, 1225, 512, 1663, 0.45029, 39],
|
||||
[ 529, 874, 645, 945, 0.31101, 39],
|
||||
[ 8, 47, 1935, 2135, 0.28192, 72],
|
||||
[ 2265, 813, 2328, 901, 0.2714, 62]
|
||||
], dtype=np.float32)
|
||||
|
||||
DETECTIONS = Detections(
|
||||
xyxy=PREDICTIONS[:, :4],
|
||||
confidence=PREDICTIONS[:, 4],
|
||||
class_id=PREDICTIONS[:, 5].astype(int)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'detections, index, expected_result, exception',
|
||||
[
|
||||
(
|
||||
DETECTIONS,
|
||||
DETECTIONS.class_id == 0,
|
||||
mock_detections(
|
||||
xyxy=[[2254, 906, 2447, 1353]],
|
||||
confidence=[0.90538],
|
||||
class_id=[0]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only detections with class_id = 0
|
||||
(
|
||||
DETECTIONS,
|
||||
DETECTIONS.confidence > 0.5,
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2254, 906, 2447, 1353],
|
||||
[2049, 1133, 2226, 1371],
|
||||
[727, 1224, 838, 1601]
|
||||
],
|
||||
confidence=[0.90538, 0.59002, 0.51119],
|
||||
class_id=[0, 56, 39]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only detections with confidence > 0.5
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([True, True, True, True, True, True, True, True, True], dtype=bool),
|
||||
DETECTIONS,
|
||||
DoesNotRaise()
|
||||
), # take all detections
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([False, False, False, False, False, False, False, False, False], dtype=bool),
|
||||
Detections(
|
||||
xyxy=np.empty((0, 4), dtype=np.float32),
|
||||
confidence=np.array([], dtype=np.float32),
|
||||
class_id=np.array([], dtype=int)
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take no detections
|
||||
(
|
||||
DETECTIONS,
|
||||
[0, 2],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2254, 906, 2447, 1353],
|
||||
[727, 1224, 838, 1601]
|
||||
],
|
||||
confidence=[0.90538, 0.51119],
|
||||
class_id=[0, 39]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only first and third detection using List[int] index
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([0, 2]),
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2254, 906, 2447, 1353],
|
||||
[727, 1224, 838, 1601]
|
||||
],
|
||||
confidence=[0.90538, 0.51119],
|
||||
class_id=[0, 39]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only first and third detection using np.ndarray index
|
||||
(
|
||||
DETECTIONS,
|
||||
0,
|
||||
mock_detections(
|
||||
xyxy=[[2254, 906, 2447, 1353]],
|
||||
confidence=[0.90538],
|
||||
class_id=[0]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only first detection by index
|
||||
(
|
||||
DETECTIONS,
|
||||
slice(1, 3),
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[2049, 1133, 2226, 1371],
|
||||
[727, 1224, 838, 1601]
|
||||
],
|
||||
confidence=[0.59002, 0.51119],
|
||||
class_id=[56, 39]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # take only first detection by index slice (1, 3)
|
||||
(
|
||||
DETECTIONS,
|
||||
10,
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
), # index out of range
|
||||
(
|
||||
DETECTIONS,
|
||||
[0, 2, 10],
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
), # index out of range
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([0, 2, 10]),
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
),
|
||||
(
|
||||
DETECTIONS,
|
||||
np.array([True, True, True, True, True, True, True, True, True, True, True]),
|
||||
None,
|
||||
pytest.raises(IndexError)
|
||||
)
|
||||
]
|
||||
)
|
||||
def test_getitem(
|
||||
detections: Detections,
|
||||
index: Union[int, slice, List[int], np.ndarray],
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = detections[index]
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'detections_list, expected_result, exception',
|
||||
[
|
||||
(
|
||||
[],
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
), # empty detections list
|
||||
(
|
||||
[
|
||||
Detections.empty()
|
||||
],
|
||||
Detections.empty(),
|
||||
DoesNotRaise()
|
||||
), # single empty detections
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]])
|
||||
],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
DoesNotRaise()
|
||||
), # single detection with xyxy field
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
Detections.empty()
|
||||
],
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
DoesNotRaise()
|
||||
), # single detection with xyxy field + empty detection
|
||||
(
|
||||
[
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
mock_detections(xyxy=[[20, 20, 30, 30]])
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[10, 10, 20, 20],
|
||||
[20, 20, 30, 30]
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # two detections with xyxy field
|
||||
(
|
||||
[
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20]],
|
||||
class_id=[0]),
|
||||
mock_detections(
|
||||
xyxy=[[20, 20, 30, 30]])
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[10, 10, 20, 20],
|
||||
[20, 20, 30, 30]
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # detection with xyxy, class_id fields + detection with xyxy field
|
||||
(
|
||||
[
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20]],
|
||||
class_id=[0]),
|
||||
mock_detections(
|
||||
xyxy=[[20, 20, 30, 30]],
|
||||
class_id=[1]),
|
||||
],
|
||||
mock_detections(
|
||||
xyxy=[
|
||||
[10, 10, 20, 20],
|
||||
[20, 20, 30, 30]
|
||||
],
|
||||
class_id=[0, 1]
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # two detections with xyxy, class_id fields
|
||||
]
|
||||
)
|
||||
def test_merge(
|
||||
detections_list: List[Detections],
|
||||
expected_result: Optional[Detections],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = Detections.merge(detections_list=detections_list)
|
||||
assert result == expected_result
|
||||
@@ -0,0 +1,337 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from typing import Optional, Tuple, List
|
||||
|
||||
import pytest
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.utils import non_max_suppression, clip_boxes, filter_polygons_by_area, \
|
||||
process_roboflow_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"predictions, iou_threshold, expected_result, exception",
|
||||
[
|
||||
(
|
||||
np.empty(shape=(0, 5)),
|
||||
0.5,
|
||||
np.array([]),
|
||||
DoesNotRaise()
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8]
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # single box with no category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0]
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # single box with category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9],
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # two boxes with no category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 1],
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
True,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # two boxes with different category
|
||||
(
|
||||
np.array([
|
||||
[10.0, 10.0, 40.0, 40.0, 0.8, 0],
|
||||
[15.0, 15.0, 40.0, 40.0, 0.9, 0],
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # two boxes with same category
|
||||
(
|
||||
np.array([
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85],
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True,
|
||||
False
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # three boxes with no category
|
||||
(
|
||||
np.array([
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 1],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 2],
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
True,
|
||||
True,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # three boxes with same category
|
||||
(
|
||||
np.array([
|
||||
[0.0, 0.0, 30.0, 40.0, 0.8, 0],
|
||||
[5.0, 5.0, 35.0, 45.0, 0.9, 0],
|
||||
[10.0, 10.0, 40.0, 50.0, 0.85, 1],
|
||||
]),
|
||||
0.5,
|
||||
np.array([
|
||||
False,
|
||||
True,
|
||||
True
|
||||
]),
|
||||
DoesNotRaise()
|
||||
), # three boxes with different category
|
||||
]
|
||||
)
|
||||
def test_non_max_suppression(
|
||||
predictions: np.ndarray,
|
||||
iou_threshold: float,
|
||||
expected_result: Optional[np.ndarray],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = non_max_suppression(predictions=predictions, iou_threshold=iou_threshold)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"boxes_xyxy, frame_resolution_wh, expected_result",
|
||||
[
|
||||
(
|
||||
np.empty(shape=(0, 4)),
|
||||
(1280, 720),
|
||||
np.empty(shape=(0, 4)),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[-1.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[0.0, 1.0, 1279.0, 719.0]
|
||||
]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, -1.0, 1279.0, 719.0]
|
||||
]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 0.0, 1279.0, 719.0]
|
||||
]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, 1.0, 1281.0, 719.0]
|
||||
]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 1.0, 1280.0, 719.0]
|
||||
]),
|
||||
),
|
||||
(
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 721.0]
|
||||
]),
|
||||
(1280, 720),
|
||||
np.array([
|
||||
[1.0, 1.0, 1279.0, 720.0]
|
||||
]),
|
||||
),
|
||||
]
|
||||
)
|
||||
def test_clip_boxes(boxes_xyxy: np.ndarray, frame_resolution_wh: Tuple[int, int], expected_result: np.ndarray) -> None:
|
||||
result = clip_boxes(boxes_xyxy=boxes_xyxy, frame_resolution_wh=frame_resolution_wh)
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"polygons, min_area, max_area, expected_result, exception",
|
||||
[
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
None,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # single polygon without area constraints
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
50,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # single polygon with min_area constraint
|
||||
(
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
None,
|
||||
50,
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # single polygon with max_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
200,
|
||||
None,
|
||||
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
None,
|
||||
200,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with max_area constraint
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
200,
|
||||
200,
|
||||
[],
|
||||
DoesNotRaise()
|
||||
), # two polygons with both area constraints
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
100,
|
||||
100,
|
||||
[np.array([[0, 0], [0, 10], [10, 10], [10, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area and max_area equal to the area of the first polygon
|
||||
(
|
||||
[
|
||||
np.array([[0, 0], [0, 10], [10, 10], [10, 0]]),
|
||||
np.array([[0, 0], [0, 20], [20, 20], [20, 0]])
|
||||
],
|
||||
400,
|
||||
400,
|
||||
[np.array([[0, 0], [0, 20], [20, 20], [20, 0]])],
|
||||
DoesNotRaise()
|
||||
), # two polygons with min_area and max_area equal to the area of the second polygon
|
||||
]
|
||||
)
|
||||
def test_filter_polygons_by_area(
|
||||
polygons: List[np.ndarray],
|
||||
min_area: Optional[float],
|
||||
max_area: Optional[float],
|
||||
expected_result: List[np.ndarray],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = filter_polygons_by_area(polygons=polygons, min_area=min_area, max_area=max_area)
|
||||
assert len(result) == len(expected_result)
|
||||
for result_polygon, expected_result_polygon in zip(result, expected_result):
|
||||
assert np.array_equal(result_polygon, expected_result_polygon)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"roboflow_result, class_list, expected_result, exception",
|
||||
[
|
||||
(
|
||||
{
|
||||
"predictions": [],
|
||||
"image": {"width": 1000, "height": 1000}
|
||||
},
|
||||
["person", "car", "truck"],
|
||||
(
|
||||
np.empty((0, 4)),
|
||||
np.empty(0),
|
||||
np.empty(0),
|
||||
None
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # empty result
|
||||
(
|
||||
{
|
||||
"predictions": [
|
||||
{
|
||||
"x": 200.0,
|
||||
"y": 300.0,
|
||||
"width": 50.0,
|
||||
"height": 50.0,
|
||||
"confidence": 0.9,
|
||||
"class": "person"
|
||||
}
|
||||
],
|
||||
"image": {"width": 1000, "height": 1000}
|
||||
},
|
||||
["person", "car", "truck"],
|
||||
(
|
||||
np.array([
|
||||
[175.0, 275.0, 225.0, 325.0]
|
||||
]),
|
||||
np.array([0.9]),
|
||||
np.array([0]),
|
||||
None
|
||||
),
|
||||
DoesNotRaise()
|
||||
), # single bounding box
|
||||
]
|
||||
)
|
||||
def test_process_roboflow_result(
|
||||
roboflow_result: dict,
|
||||
class_list: List[str],
|
||||
expected_result: Tuple[np.ndarray, np.ndarray, np.ndarray, Optional[np.ndarray]],
|
||||
exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = process_roboflow_result(roboflow_result=roboflow_result, class_list=class_list)
|
||||
assert np.array_equal(result[0], expected_result[0])
|
||||
assert np.array_equal(result[1], expected_result[1])
|
||||
assert np.array_equal(result[2], expected_result[2])
|
||||
assert (result[3] is None and expected_result[3] is None) or (np.array_equal(result[3], expected_result[3]))
|
||||
+16
-18
@@ -7,27 +7,25 @@ from supervision.draw.color import Color
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'color_hex, expected_result, exception',
|
||||
"color_hex, expected_result, exception",
|
||||
[
|
||||
('fff', Color.white(), DoesNotRaise()),
|
||||
('#fff', Color.white(), DoesNotRaise()),
|
||||
('ffffff', Color.white(), DoesNotRaise()),
|
||||
('#ffffff', Color.white(), DoesNotRaise()),
|
||||
('f00', Color.red(), DoesNotRaise()),
|
||||
('0f0', Color.green(), DoesNotRaise()),
|
||||
('00f', Color.blue(), DoesNotRaise()),
|
||||
('#808000', Color(r=128, g=128, b=0), DoesNotRaise()),
|
||||
('', None, pytest.raises(ValueError)),
|
||||
('00', None, pytest.raises(ValueError)),
|
||||
('0000', None, pytest.raises(ValueError)),
|
||||
('0000000', None, pytest.raises(ValueError)),
|
||||
('ffg', None, pytest.raises(ValueError)),
|
||||
]
|
||||
("fff", Color.white(), DoesNotRaise()),
|
||||
("#fff", Color.white(), DoesNotRaise()),
|
||||
("ffffff", Color.white(), DoesNotRaise()),
|
||||
("#ffffff", Color.white(), DoesNotRaise()),
|
||||
("f00", Color.red(), DoesNotRaise()),
|
||||
("0f0", Color.green(), DoesNotRaise()),
|
||||
("00f", Color.blue(), DoesNotRaise()),
|
||||
("#808000", Color(r=128, g=128, b=0), DoesNotRaise()),
|
||||
("", None, pytest.raises(ValueError)),
|
||||
("00", None, pytest.raises(ValueError)),
|
||||
("0000", None, pytest.raises(ValueError)),
|
||||
("0000000", None, pytest.raises(ValueError)),
|
||||
("ffg", None, pytest.raises(ValueError)),
|
||||
],
|
||||
)
|
||||
def test_color_from_hex(
|
||||
color_hex,
|
||||
expected_result: Optional[Color],
|
||||
exception: Exception
|
||||
color_hex, expected_result: Optional[Color], exception: Exception
|
||||
) -> None:
|
||||
with exception:
|
||||
result = Color.from_hex(color_hex=color_hex)
|
||||
|
||||
@@ -1,40 +1,31 @@
|
||||
import pytest
|
||||
|
||||
from supervision.geometry.dataclasses import Vector, Point
|
||||
from supervision.geometry.core import Vector, Point
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
'vector, point, expected_result',
|
||||
"vector, point, expected_result",
|
||||
[
|
||||
(Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=-1, y=1), False),
|
||||
(Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=6, y=6), False),
|
||||
(Vector(start=Point(x=0, y=0), end=Point(x=5, y=5)), Point(x=3, y=6), False),
|
||||
|
||||
(Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=-1, y=1), True),
|
||||
(Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=6, y=6), False),
|
||||
(Vector(start=Point(x=5, y=5), end=Point(x=0, y=0)), Point(x=3, y=6), True),
|
||||
|
||||
(Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=0), False),
|
||||
(Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=-1), True),
|
||||
(Vector(start=Point(x=0, y=0), end=Point(x=1, y=0)), Point(x=0, y=1), False),
|
||||
|
||||
(Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=0), False),
|
||||
(Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=-1), False),
|
||||
(Vector(start=Point(x=1, y=0), end=Point(x=0, y=0)), Point(x=0, y=1), True),
|
||||
|
||||
(Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=0, y=0), False),
|
||||
(Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=1, y=4), False),
|
||||
(Vector(start=Point(x=1, y=1), end=Point(x=1, y=3)), Point(x=2, y=4), True),
|
||||
|
||||
(Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=0, y=0), True),
|
||||
(Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=1, y=4), False),
|
||||
(Vector(start=Point(x=1, y=3), end=Point(x=1, y=1)), Point(x=2, y=4), False),
|
||||
]
|
||||
],
|
||||
)
|
||||
def test_vector_is_in(
|
||||
vector: Vector,
|
||||
point: Point,
|
||||
expected_result: bool
|
||||
) -> None:
|
||||
def test_vector_is_in(vector: Vector, point: Point, expected_result: bool) -> None:
|
||||
result = vector.is_in(point=point)
|
||||
assert result == expected_result
|
||||
|
||||
@@ -0,0 +1,19 @@
|
||||
from typing import List
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision import Detections
|
||||
|
||||
|
||||
def mock_detections(
|
||||
xyxy: List[List[float]],
|
||||
confidence: List[float] = None,
|
||||
class_id: List[int] = None,
|
||||
tracker_id: List[int] = None
|
||||
) -> Detections:
|
||||
return Detections(
|
||||
xyxy=np.array(xyxy, dtype=np.float32),
|
||||
confidence=confidence if confidence is None else np.array(confidence, dtype=np.float32),
|
||||
class_id=class_id if class_id is None else np.array(class_id, dtype=int),
|
||||
tracker_id=tracker_id if tracker_id is None else np.array(tracker_id, dtype=int)
|
||||
)
|
||||
Reference in New Issue
Block a user