Merge branch 'develop' into dfreilich/add-icon-annotator

This commit is contained in:
LinasKo 2024-08-15 11:38:13 +03:00
commit cfd079bbcb
158 changed files with 23506 additions and 4763 deletions

64
.github/workflows/notebook-bot.yml vendored Normal file
View File

@ -0,0 +1,64 @@
name: Notebook Check Pull Request
on:
pull_request_target:
types: [opened, reopened]
permissions:
contents: read
jobs:
comment-welcome:
permissions:
contents: read
pull-requests: write
runs-on: ubuntu-latest
steps:
- name: Fetch pull request branch
uses: actions/checkout@v4
with:
repository: ${{ github.event.pull_request.head.repo.full_name }}
ref: ${{ github.event.pull_request.head.sha }}
- name: Fetch base develop branch
run: git fetch -u "$GITHUB_SERVER_URL/$GITHUB_REPOSITORY" develop:develop
- name: Create message
env:
HEAD_REPOSITORY: ${{ github.event.pull_request.head.repo.full_name }}
HEAD_REF: ${{ github.event.pull_request.head.ref }}
PR_NUM: ${{ github.event.pull_request.number }}
run: |
# Preview links and tool usage only needed for notebook changes.
readarray -t changed_notebooks < <(git diff --name-only develop | grep '\.ipynb$' || true)
if [[ ${#changed_notebooks[@]} == 0 ]]; then
echo "No notebooks modified in this pull request."
else
msg="<h4>Preview</h4>\n"
msg+="Preview and run these notebook edits with Google Colab:\n<ul>\n"
# Link to PR branch in user's fork that is always current.
for fp in "${changed_notebooks[@]}"; do
gh_path="${HEAD_REPOSITORY}/blob/${HEAD_REF}/${fp}"
colab_url="https://colab.research.google.com/github/${gh_path}"
msg+="<li><a href='${colab_url}'>${fp}</a></li>\n"
done
msg+="</ul>\n"
reviewnb_url="https://app.reviewnb.com/${GITHUB_REPOSITORY}/pull/${PR_NUM}/files/"
msg+="Rendered <a href='${reviewnb_url}'>notebook diffs</a> available on ReviewNB.com.\n"
msg+="If commits are added to the pull request, synchronize your local branch: <code>git pull origin $HEAD_REF</code>\n"
fi
echo "MESSAGE=$msg" >> $GITHUB_ENV
- name: Post comment
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
ISSUE_URL: ${{ github.event.pull_request.issue_url }}
run: |
# Env var defined in previous step. Escape string for JSON.
body="$(echo -n -e $MESSAGE | python3 -c 'import json,sys; print(json.dumps(sys.stdin.read()))')"
# Add comment to pull request.
curl -X POST \
-H "Accept: application/vnd.github.v3+json" \
-H "Authorization: token $GITHUB_TOKEN" \
"${ISSUE_URL}/comments" \
--data "{\"body\": $body}"

View File

@ -4,12 +4,18 @@ on:
push:
branches:
- develop
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.event_name == 'push' && github.ref}}
cancel-in-progress: true
permissions:
contents: write
pages: write
pull-requests: write
jobs:
deploy:
runs-on: ubuntu-latest
@ -23,7 +29,7 @@ jobs:
with:
python-version: '3.10'
- name: 📦 Install mkdocs-material
run: pip install "mkdocs-material[all]"
run: pip install "mkdocs-material"
- name: 📦 Install mkdocstrings[python]
run: pip install "mkdocstrings[python]"
- name: 📦 Install mkdocs-material[imaging]

View File

@ -0,0 +1,55 @@
name: Supervision Release Documentation Workflow 📚
on:
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.event_name == 'push' && github.ref}}
cancel-in-progress: true
permissions:
contents: write
pages: write
pull-requests: write
jobs:
doc-build-deploy:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.10"]
steps:
- name: 🛎️ Checkout
uses: actions/checkout@v4
with:
fetch-depth: 0
ref: ${{ github.head_ref }}
- name: 🐍 Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: 📦 Install mkdocs-material
run: pip install "mkdocs-material"
- name: 📦 Install mkdocstrings[python]
run: pip install "mkdocstrings[python]"
- name: 📦 Install mkdocs-material[imaging]
run: pip install "mkdocs-material[imaging]"
- name: 📦 Install mike
run: pip install "mike"
- name: 📦 Install mkdocs-git-revision-date-localized-plugin
run: pip install "mkdocs-git-revision-date-localized-plugin"
- name: 📦 Install JupyterLab
run: pip install jupyterlab
- name: 📦 Install mkdocs-jupyter
run: pip install mkdocs-jupyter
- name: 📦 Install mkdocs-git-committers-plugin-2
run: pip install mkdocs-git-committers-plugin-2
- name: ⚙️ Configure git for github-actions 👷
run: |
git config --global user.name "github-actions[bot]"
git config --global user.email "41898282+github-actions[bot]@users.noreply.github.com"
- name: 🚀 Deploy MkDoc-Material 📚
run: |
latest_tag=$(git describe --tags `git rev-list --tags --max-count=1`)
MKDOCS_GIT_COMMITTERS_APIKEY=${{ secrets.GITHUB_TOKEN }} mike deploy --push --update-aliases $latest_tag latest

View File

@ -6,22 +6,24 @@ on:
- '[0-9]+.[0-9]+[0-9]+.[0-9]+b[0-9]'
- '[0-9]+.[0-9]+[0-9]+.[0-9]+rc[0-9]'
# Allows you to run this workflow manually from the Actions tab
workflow_dispatch:
jobs:
build-n-publish:
name: Build and publish to PyPI
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.10"]
steps:
- name: Checkout source
- name: 🛎️ Checkout
uses: actions/checkout@v4
- name: 🐍 Set up Python 3.8 environment for build
with:
ref: ${{ github.head_ref }}
- name: 🐍 Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
python-version: "3.8"
python-version: ${{ matrix.python-version }}
- name: 🏗️ Build source and wheel distributions
run: |

View File

@ -3,10 +3,7 @@ on:
push:
tags:
- '[0-9]+.[0-9]+[0-9]+.[0-9]'
- '[0-9]+.[0-9]+[0-9]+.[0-9]'
- '[0-9]+.[0-9]+[0-9]+.[0-9]'
# Allows you to run this workflow manually from the Actions tab
workflow_dispatch:
jobs:
@ -14,7 +11,7 @@ jobs:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: [3.8]
python-version: ["3.10"]
steps:
- name: 🛎️ Checkout
uses: actions/checkout@v4

34
.github/workflows/test-doc.yml vendored Normal file
View File

@ -0,0 +1,34 @@
name: 🧪 Docs Test WorkFlow 📚
on:
pull_request:
branches: [main, develop]
jobs:
docs-build-test:
runs-on: ubuntu-latest
steps:
- name: 🔄 Checkout code
uses: actions/checkout@v4
- name: 🐍 Set up Python
uses: actions/setup-python@v5
with:
python-version: '3.10'
- name: 📦 Install mkdocs-material
run: pip install "mkdocs-material[all]"
- name: 📦 Install mkdocstrings[python]
run: pip install "mkdocstrings[python]"
- name: 📦 Install mkdocs-material[imaging]
run: pip install "mkdocs-material[imaging]"
- name: 📦 Install mike
run: pip install "mike"
- name: 📦 Install mkdocs-git-revision-date-localized-plugin
run: pip install "mkdocs-git-revision-date-localized-plugin"
- name: 📦 Install JupyterLab
run: pip install jupyterlab
- name: 📦 Install mkdocs-jupyter
run: pip install mkdocs-jupyter
- name: 📦 Install mkdocs-git-committers-plugin-2
run: pip install mkdocs-git-committers-plugin-2
- name: 🧪 Test documentation build
run: mkdocs build --verbose

View File

@ -37,6 +37,7 @@ jobs:
matplotlib==3.5.0 \
numpy==1.21.2 \
opencv-python==4.5.5.64 \
Pillow==10.1.0 \
packaging==23.2 \
pluggy==1.3.0 \
pyparsing==3.1.1 \

View File

@ -7,14 +7,13 @@ ci:
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.5.0
rev: v4.6.0
hooks:
- id: end-of-file-fixer
- id: trailing-whitespace
exclude: test/.*\.py
- id: check-yaml
exclude: mkdocs.yml
- id: check-docstring-first
- id: check-executables-have-shebangs
- id: check-toml
- id: check-case-conflict
@ -28,9 +27,8 @@ repos:
- id: mixed-line-ending
- repo: https://github.com/PyCQA/bandit
rev: '1.7.7'
rev: '1.7.9'
hooks:
- id: bandit
args: ["-c", "pyproject.toml"]
@ -47,7 +45,7 @@ repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.2.1
rev: v0.5.7
hooks:
- id: ruff
args: [--fix, --exit-non-zero-on-fix]

View File

@ -4,6 +4,25 @@ 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.
## Code of Conduct
Please read and adhere to our [Code of Conduct](CODE_OF_CONDUCT.md). This document outlines the expected behavior for all participants in our project.
## Table of Contents
- [Contribution Guidelines](#contribution-guidelines)
- [Contributing Features](#contributing-features-)
- [How to Contribute Changes](#how-to-contribute-changes)
- [Installation for Contributors](#installation-for-contributors)
- [Code Style and Quality](#-code-style-and-quality)
- [Pre-commit tool](#pre-commit-tool)
- [Docstrings](#docstrings)
- [Type checking](#type-checking)
- [Documentation](#-documentation)
- [Cookbooks](#-cookbooks)
- [Tests](#-tests)
- [License](#-license)
## Contribution Guidelines
We welcome contributions to:
@ -32,49 +51,49 @@ First, fork this repository to your own GitHub account. Click "fork" in the top
Then, run `git clone` to download the project code to your computer.
You should also set up `roboflow/supervision` as an "upstream" remote (that is, tell git that the reference Supervision repository was the source of your fork of it):
```bash
git remote add upstream https://github.com/roboflow/supervision.git
git fetch upstream
```
Move to a new branch using the `git checkout` command:
```bash
git checkout -b <your_branch_name>
git checkout -b <scope>/<your_branch_name> upstream/develop
```
The name you choose for your branch should describe the change you want to make (i.e. `line-counter-docs`).
The name you choose for your branch should describe the change you want to make and start with an appropriate prefix:
- `feat/`: for new features (e.g., `feat/line-counter`)
- `fix/`: for bug fixes (e.g., `fix/memory-leak`)
- `docs/`: for documentation changes (e.g., `docs/update-readme`)
- `chore/`: for routine tasks, maintenance, or tooling changes (e.g., `chore/update-dependencies`)
- `test/`: for adding or modifying tests (e.g., `test/add-unit-tests`)
- `refactor/`: for code refactoring (e.g., `refactor/simplify-algorithm`)
Make any changes you want to the project code, then run the following commands to commit your changes:
```bash
git add .
git commit -m "Your commit message"
git push -u origin main
git add -A
git commit -m "feat: add line counter functionality"
git push -u origin <your_branch_name>
```
## 🎨 Code quality
Use conventional commit messages to clearly describe your changes. The format is:
### Pre-commit tool
<type>[optional scope]: <description>
This project uses the [pre-commit](https://pre-commit.com/) tool to maintain code quality and consistency. Before submitting a pull request or making any commits, it is important to run the pre-commit tool to ensure that your changes meet the project's guidelines.
Furthermore, we have integrated a pre-commit GitHub Action into our workflow. This means that with every pull request opened, the pre-commit checks will be automatically enforced, streamlining the code review process and ensuring that all contributions adhere to our quality standards.
To run the pre-commit tool, follow these steps:
1. Install pre-commit by running the following command: `poetry install`. It will not only install pre-commit but also install all the deps and dev-deps of project
2. Once pre-commit is installed, navigate to the project's root directory.
3. Run the command `pre-commit run --all-files`. This will execute the pre-commit hooks configured for this project against the modified files. If any issues are found, the pre-commit tool will provide feedback on how to resolve them. Make the necessary changes and re-run the pre-commit command until all issues are resolved.
4. You can also install pre-commit as a git hook by execute `pre-commit install`. Every time you made `git commit` pre-commit run automatically for you.
### Docstrings
All new functions and classes in `supervision` should include docstrings. This is a prerequisite for any new functions and classes to be added to the library.
`supervision` adheres to the [Google Python docstring style](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods). Please refer to the style guide while writing docstrings for your contribution.
### Type checking
So far, **there is no type checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
Common types include:
- feat: A new feature
- fix: A bug fix
- docs: Documentation only changes
- style: Changes that do not affect the meaning of the code (white-space, formatting, etc)
- refactor: A code change that neither fixes a bug nor adds a feature
- perf: A code change that improves performance
- test: Adding missing tests or correcting existing tests
- chore: Changes to the build process or auxiliary tools and libraries
Then, go back to your fork of the `supervision` repository, click "Pull Requests", and click "New Pull Request".
@ -104,7 +123,75 @@ All pull requests will be reviewed by the maintainers of the project. We will pr
PRs must pass all tests and linting requirements before they can be merged.
## 📝 documentation
## Installation for Contributors
Before starting your work on the project, set up your development environment:
1. Clone your fork of the project:
```bash
git clone https://github.com/YOUR_USERNAME/supervision.git
cd supervision
```
Replace `YOUR_USERNAME` with your GitHub username.
2. Create and activate a virtual environment:
```bash
python3 -m venv .venv
source .venv/bin/activate
```
3. Install Poetry:
Using pip:
```bash
pip install -U pip setuptools
pip install poetry
```
Or using pipx (recommended for global installation):
```bash
pipx install poetry
```
4. Install project dependencies:
```bash
poetry install
```
5. Run pytest to verify the setup:
```bash
poetry run pytest
```
## 🎨 Code Style and Quality
### Pre-commit tool
This project uses the [pre-commit](https://pre-commit.com/) tool to maintain code quality and consistency. Before submitting a pull request or making any commits, it is important to run the pre-commit tool to ensure that your changes meet the project's guidelines.
Furthermore, we have integrated a pre-commit GitHub Action into our workflow. This means that with every pull request opened, the pre-commit checks will be automatically enforced, streamlining the code review process and ensuring that all contributions adhere to our quality standards.
To run the pre-commit tool, follow these steps:
1. Install pre-commit by running the following command: `poetry install`. It will not only install pre-commit but also install all the deps and dev-deps of project
2. Once pre-commit is installed, navigate to the project's root directory.
3. Run the command `pre-commit run --all-files`. This will execute the pre-commit hooks configured for this project against the modified files. If any issues are found, the pre-commit tool will provide feedback on how to resolve them. Make the necessary changes and re-run the pre-commit command until all issues are resolved.
4. You can also install pre-commit as a git hook by executing `pre-commit install`. Every time you do a `git commit` pre-commit run automatically for you.
### Docstrings
All new functions and classes in `supervision` should include docstrings. This is a prerequisite for any new functions and classes to be added to the library.
`supervision` adheres to the [Google Python docstring style](https://google.github.io/styleguide/pyguide.html#383-functions-and-methods). Please refer to the style guide while writing docstrings for your contribution.
### Type checking
So far, **there is no type checking with mypy**. See [issue](https://github.com/roboflow-ai/template-python/issues/4).
## 📝 Documentation
The `supervision` documentation is stored in a folder called `docs`. The project documentation is built using `mkdocs`.
@ -112,13 +199,13 @@ To run the documentation, install the project requirements with `poetry install
You can learn more about mkdocs on the [mkdocs website](https://www.mkdocs.org/).
## 🧑‍🍳 cookbooks
## 🧑‍🍳 Cookbooks
We are always looking for new examples and cookbooks to add to the `supervision`
documentation. If you have a use case that you think would be helpful to others, please
submit a PR with your example. Here are some guidelines for submitting a new example:
- Create a new notebook in the [`docs/nodebooks`](https://github.com/roboflow/supervision/tree/develop/docs/notebooks) folder.
- Create a new notebook in the [`docs/notebooks`](https://github.com/roboflow/supervision/tree/develop/docs/notebooks) folder.
- Add a link to the new notebook in [`docs/theme/cookbooks.html`](https://github.com/roboflow/supervision/blob/develop/docs/theme/cookbooks.html). Make sure to add the path to the new notebook, as well as a title, labels, author and supervision version.
- Use the [Count Objects Crossing the Line](https://supervision.roboflow.com/develop/notebooks/count-objects-crossing-the-line/) example as a template for your new example.
- Freeze the version of `supervision` you are using.
@ -126,10 +213,10 @@ submit a PR with your example. Here are some guidelines for submitting a new exa
- Notebook should be self-contained. If you rely on external data ( videos, images, etc.) or libraries, include download and installation commands in the notebook.
- Annotate the code with appropriate comments, including links to the documentation describing each of the tools you have used.
## 🧪 tests
## 🧪 Tests
[`pytests`](https://docs.pytest.org/en/7.1.x/) is used to run our tests.
## 📄 license
## 📄 License
By contributing, you agree that your contributions will be licensed under an [MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md).

View File

@ -1,6 +1,6 @@
<div align="center">
<p>
<a align="center" href="" target="_blank">
<a align="center" href="" target="https://supervision.roboflow.com">
<img
width="100%"
src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529"
@ -10,25 +10,26 @@
<br>
[notebooks](https://github.com/roboflow/notebooks) | [inference](https://github.com/roboflow/inference) | [autodistill](https://github.com/autodistill/autodistill) | [collect](https://github.com/roboflow/roboflow-collect)
[notebooks](https://github.com/roboflow/notebooks) | [inference](https://github.com/roboflow/inference) | [autodistill](https://github.com/autodistill/autodistill) | [maestro](https://github.com/roboflow/multimodal-maestro)
<br>
[![version](https://badge.fury.io/py/supervision.svg)](https://badge.fury.io/py/supervision)
[![downloads](https://img.shields.io/pypi/dm/supervision)](https://pypistats.org/packages/supervision)
[![snyk](https://snyk.io/advisor/python/supervision/badge.svg)](https://snyk.io/advisor/python/supervision)
[![license](https://img.shields.io/pypi/l/supervision)](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
[![python-version](https://img.shields.io/pypi/pyversions/supervision)](https://badge.fury.io/py/supervision)
[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb)
[![Gradio](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Roboflow/Annotators)
[![Discord](https://img.shields.io/discord/1159501506232451173)](https://discord.gg/GbfgXGJ8Bk)
[![Built with Material for MkDocs](https://img.shields.io/badge/Material_for_MkDocs-526CFE?logo=MaterialForMkDocs&logoColor=white)](https://squidfunk.github.io/mkdocs-material/)
[![colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb)
[![gradio](https://img.shields.io/badge/%F0%9F%A4%97%20Hugging%20Face-Spaces-blue)](https://huggingface.co/spaces/Roboflow/Annotators)
[![discord](https://img.shields.io/discord/1159501506232451173)](https://discord.gg/GbfgXGJ8Bk)
[![built-with-material-for-mkdocs](https://img.shields.io/badge/Material_for_MkDocs-526CFE?logo=MaterialForMkDocs&logoColor=white)](https://squidfunk.github.io/mkdocs-material/)
</div>
## 👋 hello
**We write your reusable computer vision tools.** Whether you need to load your dataset from your hard drive, draw detections on an image or video, or count how many detections are in a zone. You can count on us! 🤝
[![supervision-hackfest](https://github.com/roboflow/supervision/assets/26109316/c05cc954-b9a6-4ed5-9a52-d0b4b619ff65)](https://github.com/orgs/roboflow/projects/10)
[![supervision-hackfest](https://github.com/roboflow/supervision/assets/26109316/c05cc954-b9a6-4ed5-9a52-d0b4b619ff65)](https://github.com/orgs/roboflow/projects)
## 💻 install
@ -39,7 +40,7 @@ Pip install the supervision package in a
pip install supervision
```
Read more about desktop, headless, and local installation in our [guide](https://roboflow.github.io/supervision/).
Read more about conda, mamba, and installing from source in our [guide](https://roboflow.github.io/supervision/).
## 🔥 quickstart
@ -71,23 +72,22 @@ len(detections)
```python
import cv2
import supervision as sv
from inference.models.utils import get_roboflow_model
from inference import get_model
image = cv2.imread(...)
model = get_roboflow_model(model_id="yolov8s-640", api_key=<ROBOFLOW API KEY>)
model = get_model(model_id="yolov8s-640", api_key=<ROBOFLOW API KEY>)
result = model.infer(image)[0]
detections = sv.Detections.from_inference(result)
len(detections)
# 5
# 5
```
</details>
### annotators
Supervision offers a wide range of highly customizable [annotators](https://supervision.roboflow.com/latest/annotators/), allowing you to compose the perfect visualization for your use case.
Supervision offers a wide range of highly customizable [annotators](https://supervision.roboflow.com/latest/detection/annotators/), allowing you to compose the perfect visualization for your use case.
```python
import cv2
@ -96,8 +96,8 @@ import supervision as sv
image = cv2.imread(...)
detections = sv.Detections(...)
bounding_box_annotator = sv.BoundingBoxAnnotator()
annotated_frame = bounding_box_annotator.annotate(
box_annotator = sv.BoxAnnotator()
annotated_frame = box_annotator.annotate(
scene=image.copy(),
detections=detections
)
@ -107,22 +107,25 @@ https://github.com/roboflow/supervision/assets/26109316/691e219c-0565-4403-9218-
### datasets
Supervision provides a set of [utils](https://supervision.roboflow.com/latest/datasets/) that allow you to load, split, merge, and save datasets in one of the supported formats.
Supervision provides a set of [utils](https://supervision.roboflow.com/latest/datasets/core/) that allow you to load, split, merge, and save datasets in one of the supported formats.
```python
import supervision as sv
from roboflow import Roboflow
dataset = sv.DetectionDataset.from_yolo(
images_directory_path=...,
annotations_directory_path=...,
data_yaml_path=...
project = Roboflow().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",
)
dataset.classes
['dog', 'person']
path, image, annotation = ds[0]
# loads image on demand
len(dataset)
# 1000
for path, image, annotation in ds:
# loads image on demand
```
<details close>
@ -217,19 +220,23 @@ len(dataset)
## 🎬 tutorials
<p align="left">
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/61a444c8-b135-48ce-b979-2a5ab47c5a91" alt="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source" width="300px" align="left" /></a>
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><strong>Speed Estimation & Vehicle Tracking | Computer Vision | Open Source</strong></a>
<div><strong>Created: 11 Jan 2024</strong> | <strong>Updated: 11 Jan 2024</strong></div>
<br/> Learn how to track and estimate the speed of vehicles using YOLO, ByteTrack, and Roboflow Inference. This comprehensive tutorial covers object detection, multi-object tracking, filtering detections, perspective transformation, speed estimation, visualization improvements, and more.</p>
Want to learn how to use Supervision? Explore our [how-to guides](https://supervision.roboflow.com/develop/how_to/detect_and_annotate/), [end-to-end examples](https://github.com/roboflow/supervision/tree/develop/examples), [cheatsheet](https://roboflow.github.io/cheatsheet-supervision/), and [cookbooks](https://supervision.roboflow.com/develop/cookbooks/)!
<br/>
<p align="left">
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><img src="https://github.com/roboflow/supervision/assets/26109316/54afdf1c-218c-4451-8f12-627fb85f1682" alt="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking" width="300px" align="left" /></a>
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><strong>Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking</strong></a>
<div><strong>Created: 6 Sep 2023</strong> | <strong>Updated: 6 Sep 2023</strong></div>
<br/> In this video, we explore real-time traffic analysis using YOLOv8 and ByteTrack to detect and track vehicles on aerial images. Harnessing the power of Python and Supervision, we delve deep into assigning cars to specific entry zones and understanding their direction of movement. By visualizing their paths, we gain insights into traffic flow across bustling roundabouts... </p>
<a href="https://youtu.be/hAWpsIuem10" title="Dwell Time Analysis with Computer Vision | Real-Time Stream Processing"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/a742823d-c158-407d-b30f-063a5d11b4e1" alt="Dwell Time Analysis with Computer Vision | Real-Time Stream Processing" width="300px" align="left" /></a>
<a href="https://youtu.be/hAWpsIuem10" title="Dwell Time Analysis with Computer Vision | Real-Time Stream Processing"><strong>Dwell Time Analysis with Computer Vision | Real-Time Stream Processing</strong></a>
<div><strong>Created: 5 Apr 2024</strong></div>
<br/>Learn how to use computer vision to analyze wait times and optimize processes. This tutorial covers object detection, tracking, and calculating time spent in designated zones. Use these techniques to improve customer experience in retail, traffic management, or other scenarios.</p>
<br/>
<p align="left">
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/61a444c8-b135-48ce-b979-2a5ab47c5a91" alt="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source" width="300px" align="left" /></a>
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><strong>Speed Estimation & Vehicle Tracking | Computer Vision | Open Source</strong></a>
<div><strong>Created: 11 Jan 2024</strong></div>
<br/>Learn how to track and estimate the speed of vehicles using YOLO, ByteTrack, and Roboflow Inference. This comprehensive tutorial covers object detection, multi-object tracking, filtering detections, perspective transformation, speed estimation, visualization improvements, and more.</p>
## 💜 built with supervision

View File

@ -1,19 +0,0 @@
# Supervision Cookbooks
## About
The Roboflow Cookbook is an open-source collection of examples and guides for building with the [Roboflow Ecosystem](https://roboflow.com/).
Cookbooks are meant to be common and small scoped computer vision tasks that can allow a user to quickly hit the ground running. Learning is the main goal and purpose.
Some examples require a Roboflow API key. You can [create a free account here](https://app.roboflow.com/login).
## Other Resources
Other than Cookbooks, Roboflow offers other learning resources including:
- A set of open source tools found on [Github](https://github.com/roboflow).
- A [blog](https://blog.roboflow.com/) detailing interesting tutorials, computer vision news, and product updates.
- A [Youtube channel](https://www.youtube.com/roboflow) with end to end commputer vision projects.
- A [forum](https://discuss.roboflow.com/) for sharing feedback, discussions, and other questions.
- The worlds largest [open source collection](https://universe.roboflow.com/) of images, datasets, and fine-tuned models.

552
demo.ipynb vendored

File diff suppressed because one or more lines are too long

View File

@ -19,13 +19,13 @@ as an extra within the Supervision package.
```
<div class="md-typeset">
<h2>download_assets</h2>
<h2><a href="#supervision.assets.downloader.download_assets.download_assets">download_assets</a></h2>
</div>
:::supervision.assets.downloader.download_assets
<div class="md-typeset">
<h2>VideoAssets</h2>
<h2><a href="#supervision.assets.downloader.download_assets.VideoAssets">VideoAssets</a></h2>
</div>
:::supervision.assets.list.VideoAssets

View File

@ -1,6 +1,380 @@
### 0.22.0 <small>Jul 12, 2024</small>
- Added [#1326](https://github.com/roboflow/supervision/pull/1326): [`sv.DetectionsDataset`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset) and [`sv.ClassificationDataset`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.ClassificationDataset) allowing to load the images into memory only when necessary (lazy loading).
!!! failure "Deprecated"
Constructing `DetectionDataset` with parameter `images` as `Dict[str, np.ndarray]` is deprecated and will be removed in `supervision-0.26.0`. Please pass a list of paths `List[str]` instead.
!!! failure "Deprecated"
The `DetectionDataset.images` property is deprecated and will be removed in `supervision-0.26.0`. Please loop over images with `for path, image, annotation in dataset:`, as that does not require loading all images into memory.
```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_train = sv.DetectionDataset.from_coco(
images_directory_path=f"{dataset.location}/train",
annotations_path=f"{dataset.location}/train/_annotations.coco.json",
)
path, image, annotation = ds_train[0]
# loads image on demand
for path, image, annotation in ds_train:
# loads image on demand
```
- Added [#1296](https://github.com/roboflow/supervision/pull/1296): [`sv.Detections.from_lmm`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_lmm) now supports parsing results from the [Florence 2](https://huggingface.co/microsoft/Florence-2-large) model, extending the capability to handle outputs from this Large Multimodal Model (LMM). This includes detailed object detection, OCR with region proposals, segmentation, and more. Find out more in our [Colab notebook](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-florence-2-on-detection-dataset.ipynb).
- Added [#1232](https://github.com/roboflow/supervision/pull/1232) to support keypoint detection with Mediapipe. Both [legacy](https://colab.research.google.com/github/googlesamples/mediapipe/blob/main/examples/pose_landmarker/python/%5BMediaPipe_Python_Tasks%5D_Pose_Landmarker.ipynb) and [modern](https://ai.google.dev/edge/mediapipe/solutions/vision/pose_landmarker/python) pipelines are supported. See [`sv.KeyPoints.from_mediapipe`](https://supervision.roboflow.com/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_mediapipe) for more.
- Added [#1316](https://github.com/roboflow/supervision/pull/1316): [`sv.KeyPoints.from_mediapipe`](https://supervision.roboflow.com/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_mediapipe) extended to support FaceMesh from Mediapipe. This enhancement allows for processing both face landmarks from `FaceLandmarker`, and legacy results from `FaceMesh`.
- Added [#1310](https://github.com/roboflow/supervision/pull/1310): [`sv.KeyPoints.from_detectron2`](https://supervision.roboflow.com/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.from_detectron2) is a new `KeyPoints` method, adding support for extracting keypoints from the popular [Detectron 2](https://github.com/facebookresearch/detectron2) platform.
- Added [#1300](https://github.com/roboflow/supervision/pull/1300): [`sv.Detections.from_detectron2`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_detectron2) now supports segmentation models detectron2. The resulting masks can be used with [`sv.MaskAnnotator`](https://supervision.roboflow.com/latest/annotators/#supervision.annotators.core.MaskAnnotator) for displaying annotations.
```python
import supervision as sv
from detectron2 import model_zoo
from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
import cv2
image = cv2.imread(<SOURCE_IMAGE_PATH>)
cfg = get_cfg()
cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
predictor = DefaultPredictor(cfg)
result = predictor(image)
detections = sv.Detections.from_detectron2(result)
mask_annotator = sv.MaskAnnotator()
annotated_frame = mask_annotator.annotate(scene=image.copy(), detections=detections)
```
- Added [#1277](https://github.com/roboflow/supervision/pull/1277): if you provide a font that supports symbols of a language, [`sv.RichLabelAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.LabelAnnotator.annotate) will draw them on your images.
- Various other annotators have been revised to ensure proper in-place functionality when used with `numpy` arrays. Additionally, we fixed a bug where `sv.ColorAnnotator` was filling boxes with solid color when used in-place.
```python
import cv2
import supervision as sv
import
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = get_model(model_id="yolov8n-640")
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
rich_label_annotator = sv.RichLabelAnnotator(font_path=<TTF_FONT_PATH>)
annotated_image = rich_label_annotator.annotate(scene=image.copy(), detections=detections)
```
- Added [#1227](https://github.com/roboflow/supervision/pull/1227): Added support for loading Oriented Bounding Boxes dataset in YOLO format.
```python
import supervision as sv
train_ds = sv.DetectionDataset.from_yolo(
images_directory_path="/content/dataset/train/images",
annotations_directory_path="/content/dataset/train/labels",
data_yaml_path="/content/dataset/data.yaml",
is_obb=True
)
_, image, detections in train_ds[0]
obb_annotator = OrientedBoxAnnotator()
annotated_image = obb_annotator.annotate(scene=image.copy(), detections=detections)
```
- Fixed [#1312](https://github.com/roboflow/supervision/pull/1312): Fixed [`CropAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.TraceAnnotator.annotate).
!!! failure "Removed"
`BoxAnnotator` was removed, however `BoundingBoxAnnotator` has been renamed to `BoxAnnotator`. Use a combination of [`BoxAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.BoxAnnotator) and [`LabelAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.LabelAnnotator) to simulate old `BoundingBox` behavior.
!!! failure "Deprecated"
The name `BoundingBoxAnnotator` has been deprecated and will be removed in `supervision-0.26.0`. It has been renamed to [`BoxAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.BoxAnnotator).
- Added [#975](https://github.com/roboflow/supervision/pull/975) 📝 New Cookbooks: serialize detections into [json](https://github.com/roboflow/supervision/blob/de896189b83a1f9434c0a37dd9192ee00d2a1283/docs/notebooks/serialise-detections-to-json.ipynb) and [csv](https://github.com/roboflow/supervision/blob/de896189b83a1f9434c0a37dd9192ee00d2a1283/docs/notebooks/serialise-detections-to-csv.ipynb).
- Added [#1290](https://github.com/roboflow/supervision/pull/1290): Mostly an internal change, our file utility function now support both `str` and `pathlib` paths.
- Added [#1340](https://github.com/roboflow/supervision/pull/1340): Two new methods for converting between bounding box formats - [`xywh_to_xyxy`](https://supervision.roboflow.com/latest/detection/utils/#supervision.detection.utils.xywh_to_xyxy) and [`xcycwh_to_xyxy`](https://supervision.roboflow.com/latest/detection/utils/#supervision.detection.utils.xcycwh_to_xyxy)
!!! failure "Removed"
`from_roboflow` method has been removed due to deprecation. Use [from_inference](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_inference) instead.
!!! failure "Removed"
`Color.white()` has been removed due to deprecation. Use `color.WHITE` instead.
!!! failure "Removed"
`Color.black()` has been removed due to deprecation. Use `color.BLACK` instead.
!!! failure "Removed"
`Color.red()` has been removed due to deprecation. Use `color.RED` instead.
!!! failure "Removed"
`Color.green()` has been removed due to deprecation. Use `color.GREEN` instead.
!!! failure "Removed"
`Color.blue()` has been removed due to deprecation. Use `color.BLUE` instead.
!!! failure "Removed"
`ColorPalette.default()` has been removed due to deprecation. Use [ColorPalette.DEFAULT](https://supervision.roboflow.com/latest/utils/draw/#supervision.draw.color.ColorPalette.DEFAULT) instead.
!!! failure "Removed"
`FPSMonitor.__call__` has been removed due to deprecation. Use the attribute [FPSMonitor.fps](https://supervision.roboflow.com/latest/utils/video/#supervision.utils.video.FPSMonitor.fps) instead.
### 0.21.0 <small>Jun 5, 2024</small>
- Added [#500](https://github.com/roboflow/supervision/pull/500): [`sv.Detections.with_nmm`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.with_nmm) to perform non-maximum merging on the current set of object detections.
- Added [#1221](https://github.com/roboflow/supervision/pull/1221): [`sv.Detections.from_lmm`](https://supervision.roboflow.com/latest/detection/core/#supervision.detection.core.Detections.from_lmm) allowing to parse Large Multimodal Model (LMM) text result into [`sv.Detections`](https://supervision.roboflow.com/latest/detection/core/) object. For now `from_lmm` supports only [PaliGemma](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-finetune-paligemma-on-detection-dataset.ipynb) result parsing.
```python
import supervision as sv
paligemma_result = "<loc0256><loc0256><loc0768><loc0768> cat"
detections = sv.Detections.from_lmm(
sv.LMM.PALIGEMMA,
paligemma_result,
resolution_wh=(1000, 1000),
classes=['cat', 'dog']
)
detections.xyxy
# array([[250., 250., 750., 750.]])
detections.class_id
# array([0])
```
- Added [#1236](https://github.com/roboflow/supervision/pull/1236): [`sv.VertexLabelAnnotator`](https://supervision.roboflow.com/latest/keypoint/annotators/#supervision.keypoint.annotators.EdgeAnnotator.annotate) allowing to annotate every vertex of a keypoint skeleton with custom text and color.
```python
import supervision as sv
image = ...
key_points = sv.KeyPoints(...)
edge_annotator = sv.EdgeAnnotator(
color=sv.Color.GREEN,
thickness=5
)
annotated_frame = edge_annotator.annotate(
scene=image.copy(),
key_points=key_points
)
```
- Added [#1147](https://github.com/roboflow/supervision/pull/1147): [`sv.KeyPoints.from_inference`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints.from_inference) allowing to create [`sv.KeyPoints`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints) from [Inference](https://github.com/roboflow/inference) result.
- Added [#1138](https://github.com/roboflow/supervision/pull/1138): [`sv.KeyPoints.from_yolo_nas`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints.from_yolo_nas) allowing to create [`sv.KeyPoints`](https://supervision.roboflow.com/develop/keypoint/core/#supervision.keypoint.core.KeyPoints) from [YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md) result.
- Added [#1163](https://github.com/roboflow/supervision/pull/1163): [`sv.mask_to_rle`](https://supervision.roboflow.com/develop/datasets/utils/#supervision.dataset.utils.rle_to_mask) and [`sv.rle_to_mask`](https://supervision.roboflow.com/develop/datasets/utils/#supervision.dataset.utils.rle_to_mask) allowing for easy conversion between mask and rle formats.
- Changed [#1236](https://github.com/roboflow/supervision/pull/1236): [`sv.InferenceSlicer`](https://supervision.roboflow.com/develop/detection/tools/inference_slicer/) allowing to select overlap filtering strategy (`NONE`, `NON_MAX_SUPPRESSION` and `NON_MAX_MERGE`).
- Changed [#1178](https://github.com/roboflow/supervision/pull/1178): [`sv.InferenceSlicer`](https://supervision.roboflow.com/develop/detection/tools/inference_slicer/) adding instance segmentation model support.
```python
import cv2
import numpy as np
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8x-seg-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
def callback(image_slice: np.ndarray) -> sv.Detections:
results = model.infer(image_slice)[0]
return sv.Detections.from_inference(results)
slicer = sv.InferenceSlicer(callback = callback)
detections = slicer(image)
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = mask_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
- Changed [#1228](https://github.com/roboflow/supervision/pull/1228): [`sv.LineZone`](https://supervision.roboflow.com/develop/detection/tools/line_zone/) making it 10-20 times faster, depending on the use case.
- Changed [#1163](https://github.com/roboflow/supervision/pull/1163): [`sv.DetectionDataset.from_coco`](https://supervision.roboflow.com/develop/datasets/core/#supervision.dataset.core.DetectionDataset.from_coco) and [`sv.DetectionDataset.as_coco`](https://supervision.roboflow.com/develop/datasets/core/#supervision.dataset.core.DetectionDataset.as_coco) adding support for run-length encoding (RLE) mask format.
### 0.20.0 <small>April 24, 2024</small>
- Added [#1128](https://github.com/roboflow/supervision/pull/1128): [`sv.KeyPoints`](/0.20.0/keypoint/core/#supervision.keypoint.core.KeyPoints) to provide initial support for pose estimation and broader keypoint detection models.
- Added [#1128](https://github.com/roboflow/supervision/pull/1128): [`sv.EdgeAnnotator`](/0.20.0/keypoint/annotators/#supervision.keypoint.annotators.EdgeAnnotator) and [`sv.VertexAnnotator`](/0.20.0/keypoint/annotators/#supervision.keypoint.annotators.VertexAnnotator) to enable rendering of results from keypoint detection models.
```python
import cv2
import supervision as sv
from ultralytics import YOLO
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = YOLO('yolov8l-pose')
result = model(image, verbose=False)[0]
keypoints = sv.KeyPoints.from_ultralytics(result)
edge_annotators = sv.EdgeAnnotator(color=sv.Color.GREEN, thickness=5)
annotated_image = edge_annotators.annotate(image.copy(), keypoints)
```
- Changed [#1037](https://github.com/roboflow/supervision/pull/1037): [`sv.LabelAnnotator`](/0.20.0/annotators/#supervision.annotators.core.LabelAnnotator) by adding an additional `corner_radius` argument that allows for rounding the corners of the bounding box.
- Changed [#1109](https://github.com/roboflow/supervision/pull/1109): [`sv.PolygonZone`](/0.20.0/detection/tools/polygon_zone/#supervision.detection.tools.polygon_zone.PolygonZone) such that the `frame_resolution_wh` argument is no longer required to initialize `sv.PolygonZone`.
!!! failure "Deprecated"
The `frame_resolution_wh` parameter in `sv.PolygonZone` is deprecated and will be removed in `supervision-0.24.0`.
- Changed [#1084](https://github.com/roboflow/supervision/pull/1084): [`sv.get_polygon_center`](/0.20.0/utils/geometry/#supervision.geometry.core.utils.get_polygon_center) to calculate a more accurate polygon centroid.
- Changed [#1069](https://github.com/roboflow/supervision/pull/1069): [`sv.Detections.from_transformers`](/0.20.0/detection/core/#supervision.detection.core.Detections.from_transformers) by adding support for Transformers segmentation models and extract class names values.
```python
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForSegmentation
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50-panoptic")
model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_segmentation(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(results, id2label=model.config.id2label)
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
annotated_image = mask_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
- Fixed [#787](https://github.com/roboflow/supervision/pull/787): [`sv.ByteTrack.update_with_detections`](/0.20.0/trackers/#supervision.tracker.byte_tracker.core.ByteTrack.update_with_detections) which was removing segmentation masks while tracking. Now, `ByteTrack` can be used alongside segmentation models.
### 0.19.0 <small>March 15, 2024</small>
- Added [#818](https://github.com/roboflow/supervision/pull/818): [`sv.CSVSink`](/0.19.0/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink) allowing for the straightforward saving of image, video, or stream inference results in a `.csv` file.
```python
import supervision as sv
from ultralytics import YOLO
model = YOLO(<SOURCE_MODEL_PATH>)
csv_sink = sv.CSVSink(<RESULT_CSV_FILE_PATH>)
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with csv_sink:
for frame in frames_generator:
result = model(frame)[0]
detections = sv.Detections.from_ultralytics(result)
csv_sink.append(detections, custom_data={<CUSTOM_LABEL>:<CUSTOM_DATA>})
```
- Added [#819](https://github.com/roboflow/supervision/pull/819): [`sv.JSONSink`](/0.19.0/detection/tools/save_detections/#supervision.detection.tools.csv_sink.JSONSink) allowing for the straightforward saving of image, video, or stream inference results in a `.json` file.
```python
import supervision as sv
from ultralytics import YOLO
model = YOLO(<SOURCE_MODEL_PATH>)
json_sink = sv.JSONSink(<RESULT_JSON_FILE_PATH>)
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with json_sink:
for frame in frames_generator:
result = model(frame)[0]
detections = sv.Detections.from_ultralytics(result)
json_sink.append(detections, custom_data={<CUSTOM_LABEL>:<CUSTOM_DATA>})
```
- Added [#847](https://github.com/roboflow/supervision/pull/847): [`sv.mask_iou_batch`](/0.19.0/detection/utils/#supervision.detection.utils.mask_iou_batch) allowing to compute Intersection over Union (IoU) of two sets of masks.
- Added [#847](https://github.com/roboflow/supervision/pull/847): [`sv.mask_non_max_suppression`](/0.19.0/detection/utils/#supervision.detection.utils.mask_non_max_suppression) allowing to perform Non-Maximum Suppression (NMS) on segmentation predictions.
- Added [#888](https://github.com/roboflow/supervision/pull/888): [`sv.CropAnnotator`](/0.19.0/annotators/#supervision.annotators.core.CropAnnotator) allowing users to annotate the scene with scaled-up crops of detections.
```python
import cv2
import supervision as sv
from inference import get_model
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = get_model(model_id="yolov8n-640")
result = model.infer(image)[0]
detections = sv.Detections.from_inference(result)
crop_annotator = sv.CropAnnotator()
annotated_frame = crop_annotator.annotate(
scene=image.copy(),
detections=detections
)
```
- Changed [#827](https://github.com/roboflow/supervision/pull/827): [`sv.ByteTrack.reset`](/0.19.0/tracking/#supervision.tracking.ByteTrack.reset) allowing users to clear trackers state, enabling the processing of multiple video files in sequence.
- Changed [#802](https://github.com/roboflow/supervision/pull/802): [`sv.LineZoneAnnotator`](/0.19.0/detection/tools/line_zone/#supervision.detection.line_zone.LineZone) allowing to hide in/out count using `display_in_count` and `display_out_count` properties.
- Changed [#787](https://github.com/roboflow/supervision/pull/787): [`sv.ByteTrack`](/0.19.0/tracking/#supervision.tracking.ByteTrack) input arguments and docstrings updated to improve readability and ease of use.
!!! failure "Deprecated"
The `track_buffer`, `track_thresh`, and `match_thresh` parameters in `sv.ByterTrack` are deprecated and will be removed in `supervision-0.23.0`. Use `lost_track_buffer,` `track_activation_threshold`, and `minimum_matching_threshold` instead.
- Changed [#910](https://github.com/roboflow/supervision/pull/910): [`sv.PolygonZone`](/0.19.0/detection/tools/polygon_zone/#supervision.detection.tools.polygon_zone.PolygonZone) to now accept a list of specific box anchors that must be in zone for a detection to be counted.
!!! failure "Deprecated"
The `triggering_position ` parameter in `sv.PolygonZone` is deprecated and will be removed in `supervision-0.23.0`. Use `triggering_anchors` instead.
- Changed [#875](https://github.com/roboflow/supervision/pull/875): annotators adding support for Pillow images. All supervision Annotators can now accept an image as either a numpy array or a Pillow Image. They automatically detect its type, draw annotations, and return the output in the same format as the input.
- Fixed [#944](https://github.com/roboflow/supervision/pull/944): [`sv.DetectionsSmoother`](/0.19.0/detection/tools/smoother/#supervision.detection.tools.smoother.DetectionsSmoother) removing `tracking_id` from `sv.Detections`.
### 0.18.0 <small>January 25, 2024</small>
- Added [#633](https://github.com/roboflow/supervision/pull/720): [`sv.PercentageBarAnnotator`](/0.18.0/annotators/#percentagebarannotator) allowing to annotate images and videos with percentage values representing confidence or other custom property.
- Added [#720](https://github.com/roboflow/supervision/pull/720): [`sv.PercentageBarAnnotator`](/0.18.0/annotators/#percentagebarannotator) allowing to annotate images and videos with percentage values representing confidence or other custom property.
```python
>>> import supervision as sv
@ -58,7 +432,6 @@ ColorPalette(colors=[Color(r=68, g=1, b=84), Color(r=59, g=82, b=139), ...])
`sv.ColorPalette.default()` is deprecated and will be removed in `supervision-0.22.0`. Use `sv.ColorPalette.DEFAULT` instead.
- Changed [#769](https://github.com/roboflow/supervision/pull/769): [`sv.ColorPalette.DEFAULT`](/0.18.0/draw/color/#colorpalette) value, giving users a more extensive set of annotation colors.
- Changed [#677](https://github.com/roboflow/supervision/pull/677): `sv.Detections.from_roboflow` to [`sv.Detections.from_inference`](/0.18.0/detection/core/#supervision.detection.core.Detections.from_inference) streamlining its functionality to be compatible with both the both [inference](https://github.com/roboflow/inference) pip package and the Robloflow [hosted API](https://docs.roboflow.com/deploy/hosted-api).
@ -67,7 +440,6 @@ ColorPalette(colors=[Color(r=68, g=1, b=84), Color(r=59, g=82, b=139), ...])
`Detections.from_roboflow()` is deprecated and will be removed in `supervision-0.22.0`. Use `Detections.from_inference` instead.
- Fixed [#735](https://github.com/roboflow/supervision/pull/735): [`sv.LineZone`](/0.18.0/detection/tools/line_zone/#linezone) functionality to accurately update the counter when an object crosses a line from any direction, including from the side. This enhancement enables more precise tracking and analytics, such as calculating individual in/out counts for each lane on the road.
### 0.17.0 <small>December 06, 2023</small>
@ -160,13 +532,12 @@ ColorPalette(colors=[Color(r=68, g=1, b=84), Color(r=59, g=82, b=139), ...])
- Fixed [#477](https://github.com/roboflow/supervision/pull/477): Poetry env definition allowing proper local installation.
- Fixed [#430](https://github.com/roboflow/supervision/pull/430): [`sv.ByteTrack`](/0.16.0/trackers/#supervision.tracker.byte_tracker.core.ByteTrack) to return `np.array([], dtype=int)` when `svDetections` is empty.
- Fixed [#430](https://github.com/roboflow/supervision/pull/430): [`sv.ByteTrack`](/0.16.0/trackers/#supervision.tracker.byte_tracker.core.ByteTrack) to return `np.array([], dtype=int)` when `svDetections` is empty.
!!! failure "Deprecated"
`sv.Detections.from_yolov8` and `sv.Classifications.from_yolov8` as those are now replaced by [`sv.Detections.from_ultralytics`](/0.16.0/detection/core/#supervision.detection.core.Detections.from_ultralytics) and [`sv.Classifications.from_ultralytics`](/0.16.0/classification/core/#supervision.classification.core.Classifications.from_ultralytics).
### 0.15.0 <small>October 5, 2023</small>
- Added [#170](https://github.com/roboflow/supervision/pull/170): [`sv.BoundingBoxAnnotator`](/0.15.0/annotators/#supervision.annotators.core.BoundingBoxAnnotator) allowing to annotate images and videos with bounding boxes.

View File

@ -1,7 +1,6 @@
---
template: cookbooks.html
comments: true
status: new
hide:
- navigation
- toc

View File

@ -1,5 +1,6 @@
---
comments: true
status: new
---
# Datasets

17
docs/datasets/utils.md Normal file
View File

@ -0,0 +1,17 @@
---
comments: true
---
# Datasets Utils
<div class="md-typeset">
<h2><a href="#supervision.dataset.utils.rle_to_mask">rle_to_mask</a></h2>
</div>
:::supervision.dataset.utils.rle_to_mask
<div class="md-typeset">
<h2><a href="#supervision.dataset.utils.mask_to_rle">mask_to_rle</a></h2>
</div>
:::supervision.dataset.utils.mask_to_rle

View File

@ -3,14 +3,25 @@ comments: true
status: deprecated
---
These features are phased out due to better alternatives or potential issues in future versions. Deprecated functionalities are supported for **three subsequent releases**, providing time for users to transition to updated methods.
# Deprecated
- [`Detections.from_froboflow`](detection/core.md/#supervision.detection.core.Detections.from_roboflow) is deprecated and will be removed in `supervision-0.22.0`. Use [`Detections.from_inference`](detection/core.md/#supervision.detection.core.Detections.from_inference) instead.
- `Color.white()` is deprecated and will be removed in `supervision-0.22.0`. Use `Color.WHITE` instead.
- `Color.black()` is deprecated and will be removed in `supervision-0.22.0`. Use `Color.BLACK` instead.
- `Color.red()` is deprecated and will be removed in `supervision-0.22.0`. Use `Color.RED` instead.
- `Color.green()` is deprecated and will be removed in `supervision-0.22.0`. Use `Color.GREEN` instead.
- `Color.blue()` is deprecated and will be removed in `supervision-0.22.0`. Use `Color.BLUE` instead.
- [`ColorPalette.default()`](draw/color.md/#supervision.draw.color.ColorPalette.default) is deprecated and will be removed in `supervision-0.22.0`. Use [`ColorPalette.DEFAULT`](draw/color.md/#supervision.draw.color.ColorPalette.DEFAULT) instead.
- `BoxAnnotator` is deprecated and will be removed in `supervision-0.22.0`. Use [`BoundingBoxAnnotator`](annotators.md/#supervision.annotators.core.BoundingBoxAnnotator) and [`LabelAnnotator`](annotators.md/#supervision.annotators.core.LabelAnnotator) instead.
- [`FPSMonitor.__call__`](utils/video.md/#supervision.utils.video.FPSMonitor.__call__) is deprecated and will be removed in `supervision-0.22.0`. Use [`FPSMonitor.fps`](utils/video.md/#supervision.utils.video.FPSMonitor.fps) instead.
These features are phased out due to better alternatives or potential issues in future versions. Deprecated functionalities are supported for **five subsequent releases**, providing time for users to transition to updated methods.
- The `track_buffer`, `track_thresh`, and `match_thresh` parameters in [`ByterTrack`](trackers.md/#supervision.tracker.byte_tracker.core.ByteTrack) are deprecated and will be removed in `supervision-0.23.0`. Use `lost_track_buffer,` `track_activation_threshold`, and `minimum_matching_threshold` instead.
- The `triggering_position ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) will be removed in `supervision-0.23.0`. Use `triggering_anchors ` instead.
- The `frame_resolution_wh ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) will be removed in `supervision-0.24.0`.
- Constructing `DetectionDataset` and `ClassificationDataset` with parameter `images` as `Dict[str, np.ndarray]` will be removed in `supervision-0.26.0`. Please pass a list of paths `List[str]` instead.
- The `DetectionDataset.images` property will be removed in `supervision-0.26.0`. Please loop over images with `for path, image, annotation in dataset:`, as that does not require loading all images into memory.
- `BoundingBoxAnnotator` has been renamed to `BoxAnnotator` after the old implementation of `BoxAnnotator` has been removed. `BoundingBoxAnnotator` will be removed in `supervision-0.26.0`.
# Removed
- [`Detections.from_froboflow`](detection/core.md/#supervision.detection.core.Detections.from_roboflow) is removed as of `supervision-0.22.0`. Use [`Detections.from_inference`](detection/core.md/#supervision.detection.core.Detections.from_inference) instead.
- The method `Color.white()` was removed as of `supervision-0.22.0`. Use the constant `Color.WHITE` instead.
- The method `Color.black()` was removed as of `supervision-0.22.0`. Use the constant `Color.BLACK` instead.
- The method `Color.red()` was removed as of `supervision-0.22.0`. Use the constant `Color.RED` instead.
- The method `Color.green()` was removed as of `supervision-0.22.0`. Use the constant `Color.GREEN` instead.
- The method `Color.blue()` was removed as of `supervision-0.22.0`. Use the constant `Color.BLUE` instead.
- The method [`ColorPalette.default()`](draw/color.md/#supervision.draw.color.ColorPalette.default) was removed as of `supervision-0.22.0`. Use the constant [`ColorPalette.DEFAULT`](draw/color.md/#supervision.draw.color.ColorPalette.DEFAULT) instead.
- `BoxAnnotator` was removed as of `supervision-0.22.0`, however `BoundingBoxAnnotator` was immediately renamed to `BoxAnnotator`. Use [`BoxAnnotator`](detection/annotators.md/#supervision.annotators.core.BoxAnnotator) and [`LabelAnnotator`](detection/annotators.md/#supervision.annotators.core.LabelAnnotator) instead of the old `BoxAnnotator`.
- The method [`FPSMonitor.__call__`](utils/video.md/#supervision.utils.video.FPSMonitor.__call__) was removed as of `supervision-0.22.0`. Use the attribute [`FPSMonitor.fps`](utils/video.md/#supervision.utils.video.FPSMonitor.fps) instead.

View File

@ -5,7 +5,7 @@ status: new
# Annotators
=== "BoundingBox"
=== "Box"
```python
import supervision as sv
@ -13,8 +13,8 @@ status: new
image = ...
detections = sv.Detections(...)
bounding_box_annotator = sv.BoundingBoxAnnotator()
annotated_frame = bounding_box_annotator.annotate(
box_annotator = sv.BoxAnnotator()
annotated_frame = box_annotator.annotate(
scene=image.copy(),
detections=detections
)
@ -260,15 +260,22 @@ status: new
=== "Label"
```python
import supervision as sv
import supervision as sv
image = ...
detections = sv.Detections(...)
labels = [
f"{class_name} {confidence:.2f}"
for class_name, confidence
in zip(detections['class_name'], detections.confidence)
]
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
annotated_frame = label_annotator.annotate(
scene=image.copy(),
detections=detections
detections=detections,
labels=labels
)
```
@ -278,6 +285,52 @@ status: new
</div>
=== "RichLabel"
```python
import supervision as sv
image = ...
detections = sv.Detections(...)
labels = [
f"{class_name} {confidence:.2f}"
for class_name, confidence
in zip(detections['class_name'], detections.confidence)
]
rich_label_annotator = sv.RichLabelAnnotator(
font_path="<TTF_FONT_PATH>",
text_position=sv.Position.CENTER
)
annotated_frame = rich_label_annotator.annotate(
scene=image.copy(),
detections=detections,
labels=labels
)
```
<div class="result" markdown>
![label-annotator-example](https://media.roboflow.com/supervision-annotator-examples/label-annotator-example-purple.png){ align=center width="800" }
</div>
=== "Crop"
```python
import supervision as sv
image = ...
detections = sv.Detections(...)
crop_annotator = sv.CropAnnotator()
annotated_frame = crop_annotator.annotate(
scene=image.copy(),
detections=detections
)
```
=== "Blur"
```python
@ -380,11 +433,32 @@ status: new
</div>
=== "Background Color"
```python
import supervision as sv
image = ...
detections = sv.Detections(...)
background_overlay_annotator = sv.BackgroundOverlayAnnotator()
annotated_frame = background_overlay_annotator.annotate(
scene=image.copy(),
detections=detections
)
```
<div class="result" markdown>
![background-overlay-annotator-example](https://media.roboflow.com/supervision-annotator-examples/background-color-annotator-example-purple.png)
</div>
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.BoundingBoxAnnotator">BoundingBoxAnnotator</a></h2>
<h2><a href="#supervision.annotators.core.BoxAnnotator">BoxAnnotator</a></h2>
</div>
:::supervision.annotators.core.BoundingBoxAnnotator
:::supervision.annotators.core.BoxAnnotator
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.RoundBoxAnnotator">RoundBoxAnnotator</a></h2>
@ -470,6 +544,18 @@ status: new
:::supervision.annotators.core.LabelAnnotator
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.RichLabelAnnotator">RichLabelAnnotator</a></h2>
</div>
:::supervision.annotators.core.RichLabelAnnotator
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.IconAnnotator">IconAnnotator</a></h2>
</div>
:::supervision.annotators.core.IconAnnotator
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.BlurAnnotator">BlurAnnotator</a></h2>
</div>
@ -489,10 +575,16 @@ status: new
:::supervision.annotators.core.TraceAnnotator
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.IconAnnotator">IconAnnotator</a></h2>
<h2><a href="#supervision.annotators.core.CropAnnotator">CropAnnotator</a></h2>
</div>
:::supervision.annotators.core.IconAnnotator
:::supervision.annotators.core.CropAnnotator
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.BackgroundOverlayAnnotator">BackgroundOverlayAnnotator</a></h2>
</div>
:::supervision.annotators.core.BackgroundOverlayAnnotator
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.ColorLookup">ColorLookup</a></h2>

View File

@ -0,0 +1,29 @@
---
comments: true
---
# Double Detection Filter
<div class="md-typeset">
<h2><a href="#supervision.detection.overlap_filter.OverlapFilter">OverlapFilter</a></h2>
</div>
:::supervision.detection.overlap_filter.OverlapFilter
<div class="md-typeset">
<h2><a href="#supervision.detection.overlap_filter.box_non_max_suppression">box_non_max_suppression</a></h2>
</div>
:::supervision.detection.overlap_filter.box_non_max_suppression
<div class="md-typeset">
<h2><a href="#supervision.detection.overlap_filter.mask_non_max_suppression">mask_non_max_suppression</a></h2>
</div>
:::supervision.detection.overlap_filter.mask_non_max_suppression
<div class="md-typeset">
<h2><a href="#supervision.detection.overlap_filter.box_non_max_merge">box_non_max_merge</a></h2>
</div>
:::supervision.detection.overlap_filter.box_non_max_merge

17
docs/detection/metrics.md Normal file
View File

@ -0,0 +1,17 @@
---
comments: true
---
# Metrics
<div class="md-typeset">
<h2><a href="#supervision.metrics.detection.ConfusionMatrix">ConfusionMatrix</a></h2>
</div>
:::supervision.metrics.detection.ConfusionMatrix
<div class="md-typeset">
<h2><a href="#supervision.metrics.detection.MeanAveragePrecision">MeanAveragePrecision</a></h2>
</div>
:::supervision.metrics.detection.MeanAveragePrecision

View File

@ -1,6 +1,5 @@
---
comments: true
status: new
---
# Save Detections
@ -10,3 +9,9 @@ status: new
</div>
:::supervision.detection.tools.csv_sink.CSVSink
<div class="md-typeset">
<h2>JSON Sink</h2>
</div>
:::supervision.detection.tools.json_sink.JSONSink

View File

@ -1,71 +1,102 @@
---
comments: true
status: new
---
# Detection Utils
<div class="md-typeset">
<h2>box_iou_batch</h2>
<h2><a href="#supervision.detection.utils.box_iou_batch">box_iou_batch</a></h2>
</div>
:::supervision.detection.utils.box_iou_batch
<div class="md-typeset">
<h2>mask_iou_batch</h2>
<h2><a href="#supervision.detection.utils.mask_iou_batch">mask_iou_batch</a></h2>
</div>
:::supervision.detection.utils.mask_iou_batch
<div class="md-typeset">
<h2>box_non_max_suppression</h2>
</div>
:::supervision.detection.utils.box_non_max_suppression
<div class="md-typeset">
<h2>mask_non_max_suppression</h2>
</div>
:::supervision.detection.utils.mask_non_max_suppression
<div class="md-typeset">
<h2>polygon_to_mask</h2>
<h2><a href="#supervision.detection.utils.polygon_to_mask">polygon_to_mask</a></h2>
</div>
:::supervision.detection.utils.polygon_to_mask
<div class="md-typeset">
<h2>mask_to_xyxy</h2>
<h2><a href="#supervision.detection.utils.mask_to_xyxy">mask_to_xyxy</a></h2>
</div>
:::supervision.detection.utils.mask_to_xyxy
<div class="md-typeset">
<h2>mask_to_polygons</h2>
<h2><a href="#supervision.detection.utils.mask_to_polygons">mask_to_polygons</a></h2>
</div>
:::supervision.detection.utils.mask_to_polygons
<div class="md-typeset">
<h2>polygon_to_xyxy</h2>
<h2><a href="#supervision.detection.utils.polygon_to_xyxy">polygon_to_xyxy</a></h2>
</div>
:::supervision.detection.utils.polygon_to_xyxy
<div class="md-typeset">
<h2>filter_polygons_by_area</h2>
<h2><a href="#supervision.detection.utils.filter_polygons_by_area">filter_polygons_by_area</a></h2>
</div>
:::supervision.detection.utils.filter_polygons_by_area
<div class="md-typeset">
<h2>move_boxes</h2>
<h2><a href="#supervision.detection.utils.move_boxes">move_boxes</a></h2>
</div>
:::supervision.detection.utils.move_boxes
<div class="md-typeset">
<h2>scale_boxes</h2>
<h2><a href="#supervision.detection.utils.move_masks">move_masks</a></h2>
</div>
:::supervision.detection.utils.move_masks
<div class="md-typeset">
<h2><a href="#supervision.detection.utils.scale_boxes">scale_boxes</a></h2>
</div>
:::supervision.detection.utils.scale_boxes
<div class="md-typeset">
<h2><a href="#supervision.detection.utils.clip_boxes">clip_boxes</a></h2>
</div>
:::supervision.detection.utils.clip_boxes
<div class="md-typeset">
<h2><a href="#supervision.detection.utils.pad_boxes">pad_boxes</a></h2>
</div>
:::supervision.detection.utils.pad_boxes
<div class="md-typeset">
<h2><a href="#supervision.detection.utils.contains_holes">contains_holes</a></h2>
</div>
:::supervision.detection.utils.xywh_to_xyxy
<div class="md-typeset">
<h2><a href="#supervision.detection.utils.xywh_to_xyxy">xywh_to_xyxy</a></h2>
</div>
:::supervision.detection.utils.xcycwh_to_xyxy
<div class="md-typeset">
<h2><a href="#supervision.detection.utils.xcycwh_to_xyxy">xcycwh_to_xyxy</a></h2>
</div>
:::supervision.detection.utils.contains_holes
<div class="md-typeset">
<h2><a href="#supervision.detection.utils.contains_multiple_segments">contains_multiple_segments</a></h2>
</div>
:::supervision.detection.utils.contains_multiple_segments

View File

@ -1,13 +0,0 @@
---
comments: true
---
# Color
:::supervision.draw.color.Color
<div class="md-typeset">
<h2>ColorPalette</h2>
</div>
:::supervision.draw.color.ColorPalette

View File

@ -1,53 +0,0 @@
---
comments: true
---
# Draw Utils
<div class="md-typeset">
<h2>draw_line</h2>
</div>
:::supervision.draw.utils.draw_line
<div class="md-typeset">
<h2>draw_rectangle</h2>
</div>
:::supervision.draw.utils.draw_rectangle
<div class="md-typeset">
<h2>draw_filled_rectangle</h2>
</div>
:::supervision.draw.utils.draw_filled_rectangle
<div class="md-typeset">
<h2>draw_polygon</h2>
</div>
:::supervision.draw.utils.draw_polygon
<div class="md-typeset">
<h2>draw_text</h2>
</div>
:::supervision.draw.utils.draw_text
<div class="md-typeset">
<h2>draw_image</h2>
</div>
:::supervision.draw.utils.draw_image
<div class="md-typeset">
<h2>calculate_dynamic_font_scale</h2>
</div>
:::supervision.draw.utils.calculate_dynamic_text_scale
<div class="md-typeset">
<h2>calculate_dynamic_line_thickness</h2>
</div>
:::supervision.draw.utils.calculate_dynamic_line_thickness

View File

@ -1,7 +0,0 @@
---
comments: true
---
# Position
:::supervision.geometry.core.Position

View File

@ -4,17 +4,31 @@ comments: true
# Detect and Annotate
Supervision offers a streamlined solution to effortlessly annotate predictions from a
range of object detection and segmentation models. This guide demonstrates how to
execute inference using the YOLOv8 model with either the
[Inference](https://github.com/roboflow/inference) or
[Ultralytics](https://github.com/ultralytics/ultralytics) packages. Following this,
you'll learn how to import these predictions into Supervision for image annotation
purposes.
Supervision provides a seamless process for annotating predictions generated by various
object detection and segmentation models. This guide shows how to perform inference
with the [Inference](https://github.com/roboflow/inference),
[Ultralytics](https://github.com/ultralytics/ultralytics) or
[Transformers](https://github.com/huggingface/transformers) packages. Following this,
you'll learn how to import these predictions into Supervision and use them to annotate
source image.
## Run Inference
![basic-annotation](https://media.roboflow.com/supervision_detect_and_annotate_example_1.png)
First, you'll need to obtain predictions from your object detection or segmentation model.
## Run Detection
First, you'll need to obtain predictions from your object detection or segmentation
model.
=== "Inference"
```python
import cv2
from inference import get_model
model = get_model(model_id="yolov8n-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
```
=== "Ultralytics"
@ -23,123 +37,374 @@ First, you'll need to obtain predictions from your object detection or segmentat
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
image = cv2.imread(<PATH TO IMAGE>)
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image)[0]
```
=== "Inference"
=== "Transformers"
```python
import cv2
from inference.models.utils import get_roboflow_model
import torch
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>)
image = cv2.imread(<PATH TO IMAGE>)
results = model.infer(image)[0]
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
```
## Load Predictions into Supervision
Now that we have predictions from a model, we can load them into Supervision.
=== "Ultralytics"
We can do so using the [`sv.Detections.from_ultralytics`](detection/core/#supervision.detection.core.Detections.from_ultralytics) method, which accepts model results from both detection and segmentation models.
```python
import cv2
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
image = cv2.imread(<PATH TO IMAGE>)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
```
=== "Inference"
We can do so using the [`sv.Detections.from_inference`](detection/core/#supervision.detection.core.Detections.from_inference) method, which accepts model results from both detection and segmentation models.
```python
```{ .py hl_lines="2 8" }
import cv2
import supervision as sv
from inference.models.utils import get_roboflow_model
from inference import get_model
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>
image = cv2.imread(<PATH TO IMAGE>)
model = get_model(model_id="yolov8n-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
```
You can conveniently load predictions from other computer vision frameworks and libraries using:
- [`from_deepsparse`](detection/core/#supervision.detection.core.Detections.from_deepsparse) ([Deepsparse](https://github.com/neuralmagic/deepsparse))
- [`from_detectron2`](detection/core/#supervision.detection.core.Detections.from_detectron2) ([Detectron2](https://github.com/facebookresearch/detectron2))
- [`from_mmdetection`](detection/core/#supervision.detection.core.Detections.from_mmdetection) ([MMDetection](https://github.com/open-mmlab/mmdetection))
- [`from_inference`](detection/core/#supervision.detection.core.Detections.from_inference) ([Roboflow Inference](https://github.com/roboflow/inference))
- [`from_sam`](detection/core/#supervision.detection.core.Detections.from_sam) ([Segment Anything Model](https://github.com/facebookresearch/segment-anything))
- [`from_transformers`](detection/core/#supervision.detection.core.Detections.from_transformers) ([HuggingFace Transformers](https://github.com/huggingface/transformers))
- [`from_yolo_nas`](detection/core/#supervision.detection.core.Detections.from_yolo_nas) ([YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md))
## Annotate Image
Finally, we can annotate the image with the predictions. Since we are working with an object detection model, we will use the [`sv.BoundingBoxAnnotator`](annotators/#supervision.annotators.core.BoundingBoxAnnotator) and [`sv.LabelAnnotator`](annotators/#supervision.annotators.core.LabelAnnotator) classes. If you are running the segmentation model [`sv.MaskAnnotator`](annotators/#supervision.annotators.core.MaskAnnotator) is a drop-in replacement for [`sv.BoundingBoxAnnotator`](annotators/#supervision.annotators.core.BoundingBoxAnnotator) that will allow you to draw masks instead of boxes.
=== "Ultralytics"
```python
We can do so using the [`sv.Detections.from_ultralytics`](detection/core/#supervision.detection.core.Detections.from_ultralytics) method, which accepts model results from both detection and segmentation models.
```{ .py hl_lines="2 8" }
import cv2
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
image = cv2.imread(<PATH TO IMAGE>)
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
bounding_box_annotator = sv.BoundingBoxAnnotator()
label_annotator = sv.LabelAnnotator()
labels = [
model.model.names[class_id]
for class_id
in detections.class_id
]
annotated_image = bounding_box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections, labels=labels)
```
=== "Transformers"
We can do so using the [`sv.Detections.from_transformers`](detection/core/#supervision.detection.core.Detections.from_transformers) method, which accepts model results from both detection and segmentation models.
```{ .py hl_lines="2 19-21" }
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(
transformers_results=results,
id2label=model.config.id2label)
```
You can load predictions from other computer vision frameworks and libraries using:
- [`from_deepsparse`](/latest/detection/core/#supervision.detection.core.Detections.from_deepsparse) ([Deepsparse](https://github.com/neuralmagic/deepsparse))
- [`from_detectron2`](/latest/detection/core/#supervision.detection.core.Detections.from_detectron2) ([Detectron2](https://github.com/facebookresearch/detectron2))
- [`from_mmdetection`](/latest/detection/core/#supervision.detection.core.Detections.from_mmdetection) ([MMDetection](https://github.com/open-mmlab/mmdetection))
- [`from_sam`](/latest/detection/core/#supervision.detection.core.Detections.from_sam) ([Segment Anything Model](https://github.com/facebookresearch/segment-anything))
- [`from_yolo_nas`](/latest/detection/core/#supervision.detection.core.Detections.from_yolo_nas) ([YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md))
## Annotate Image with Detections
Finally, we can annotate the image with the predictions. Since we are working with an object detection model, we will use the [`sv.BoxAnnotator`](/latest/annotators/#supervision.annotators.core.BoxAnnotator) and [`sv.LabelAnnotator`](/latest/annotators/#supervision.annotators.core.LabelAnnotator) classes.
=== "Inference"
```python
```{ .py hl_lines="10-16" }
import cv2
import supervision as sv
from inference.models.utils import get_roboflow_model
from inference import get_model
model = get_roboflow_model(model_id="yolov8n-640", api_key=<ROBOFLOW API KEY>
image = cv2.imread(<PATH TO IMAGE>)
model = get_model(model_id="yolov8n-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
bounding_box_annotator = sv.BoundingBoxAnnotator()
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = bounding_box_annotator.annotate(
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
![Predictions plotted on an image](https://media.roboflow.com/supervision_annotate_example.png)
=== "Ultralytics"
## Display Annotated Image
```{ .py hl_lines="10-16" }
import cv2
import supervision as sv
from ultralytics import YOLO
To display the annotated image in Jupyter Notebook or Google Colab, use the [`sv.plot_image`](utils/notebook/#supervision.utils.notebook.plot_image) function.
model = YOLO("yolov8n.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
```python
sv.plot_image(annotated_image)
```
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Transformers"
```{ .py hl_lines="23-30" }
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(
transformers_results=results,
id2label=model.config.id2label)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
![basic-annotation](https://media.roboflow.com/supervision_detect_and_annotate_example_1.png)
## Display Custom Labels
By default, [`sv.LabelAnnotator`](/latest/annotators/#supervision.annotators.core.LabelAnnotator)
will label each detection with its `class_name` (if possible) or `class_id`. You can
override this behavior by passing a list of custom `labels` to the `annotate` method.
=== "Inference"
```{ .py hl_lines="13-17 22" }
import cv2
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8n-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
labels = [
f"{class_name} {confidence:.2f}"
for class_name, confidence
in zip(detections['class_name'], detections.confidence)
]
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections, labels=labels)
```
=== "Ultralytics"
```{ .py hl_lines="13-17 22" }
import cv2
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
labels = [
f"{class_name} {confidence:.2f}"
for class_name, confidence
in zip(detections['class_name'], detections.confidence)
]
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections, labels=labels)
```
=== "Transformers"
```{ .py hl_lines="26-30 35" }
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(
transformers_results=results,
id2label=model.config.id2label)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
labels = [
f"{class_name} {confidence:.2f}"
for class_name, confidence
in zip(detections['class_name'], detections.confidence)
]
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections, labels=labels)
```
![custom-label-annotation](https://media.roboflow.com/supervision_detect_and_annotate_example_2.png)
## Annotate Image with Segmentations
If you are running the segmentation model
[`sv.MaskAnnotator`](/latest/annotators/#supervision.annotators.core.MaskAnnotator)
is a drop-in replacement for
[`sv.BoxAnnotator`](/latest/annotators/#supervision.annotators.core.BoxAnnotator)
that will allow you to draw masks instead of boxes.
=== "Inference"
```python
import cv2
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8n-seg-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER_OF_MASS)
annotated_image = mask_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Ultralytics"
```python
import cv2
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n-seg.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER_OF_MASS)
annotated_image = mask_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Transformers"
```python
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForSegmentation
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50-panoptic")
model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50-panoptic")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_segmentation(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(
transformers_results=results,
id2label=model.config.id2label)
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER_OF_MASS)
labels = [
f"{class_name} {confidence:.2f}"
for class_name, confidence
in zip(detections['class_name'], detections.confidence)
]
annotated_image = mask_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections, labels=labels)
```
![segmentation-annotation](https://media.roboflow.com/supervision_detect_and_annotate_example_3.png)

View File

@ -0,0 +1,325 @@
---
comments: true
---
# Detect Small Objects
This guide shows how to detect small objects
with the [Inference](https://github.com/roboflow/inference),
[Ultralytics](https://github.com/ultralytics/ultralytics) or
[Transformers](https://github.com/huggingface/transformers) packages using
[`InferenceSlicer`](/latest/detection/tools/inference_slicer/#supervision.detection.tools.inference_slicer.InferenceSlicer).
<video controls>
<source src="https://media.roboflow.com/supervision_detect_small_objects_example.mp4" type="video/mp4">
</video>
## Baseline Detection
Small object detection in high-resolution images presents challenges due to the objects'
size relative to the image resolution.
=== "Inference"
```python
import cv2
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8x-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Ultralytics"
```python
import cv2
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8x.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Transformers"
```python
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForSegmentation
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForSegmentation.from_pretrained("facebook/detr-resnet-50")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image_slice.size
target_size = torch.tensor([[width, height]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(results)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
labels = [
model.config.id2label[class_id]
for class_id
in detections.class_id
]
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections, labels=labels)
```
![basic-detection](https://media.roboflow.com/supervision_detect_small_objects_example_1.png)
## Input Resolution
Modifying the input resolution of images before detection can enhance small object
identification at the cost of processing speed and increased memory usage. This method
is less effective for ultra-high-resolution images (4K and above).
=== "Inference"
```{ .py hl_lines="5" }
import cv2
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8x-1280")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Ultralytics"
```{ .py hl_lines="7" }
import cv2
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8x.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image, imgsz=1280)[0]
detections = sv.Detections.from_ultralytics(results)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
![detection-with-high-input-resolution](https://media.roboflow.com/supervision_detect_small_objects_example_2.png)
## Inference Slicer
[`InferenceSlicer`](/latest/detection/tools/inference_slicer/#supervision.detection.tools.inference_slicer.InferenceSlicer)
processes high-resolution images by dividing them into smaller segments, detecting
objects within each, and aggregating the results.
<video controls>
<source src="https://media.roboflow.com/supervision_detect_small_objects_example_2.mp4" type="video/mp4">
</video>
=== "Inference"
```{ .py hl_lines="9-14" }
import cv2
import numpy as np
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8x-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
def callback(image_slice: np.ndarray) -> sv.Detections:
results = model.infer(image_slice)[0]
return sv.Detections.from_inference(results)
slicer = sv.InferenceSlicer(callback = callback)
detections = slicer(image)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Ultralytics"
```{ .py hl_lines="9-14" }
import cv2
import numpy as np
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8x.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
def callback(image_slice: np.ndarray) -> sv.Detections:
result = model(image_slice)[0]
return sv.Detections.from_ultralytics(result)
slicer = sv.InferenceSlicer(callback = callback)
detections = slicer(image)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Transformers"
```{ .py hl_lines="13-28" }
import cv2
import torch
import numpy as np
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
def callback(image_slice: np.ndarray) -> sv.Detections:
image_slice = cv2.cvtColor(image_slice, cv2.COLOR_BGR2RGB)
image_slice = Image.fromarray(image_slice)
inputs = processor(images=image_slice, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image_slice.size
target_size = torch.tensor([[width, height]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
return sv.Detections.from_transformers(results)
slicer = sv.InferenceSlicer(callback = callback)
detections = slicer(image)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
labels = [
model.config.id2label[class_id]
for class_id
in detections.class_id
]
annotated_image = box_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections, labels=labels)
```
![detection-with-inference-slicer](https://media.roboflow.com/supervision_detect_small_objects_example_3.png)
## Small Object Segmentation
[`InferenceSlicer`](/latest/detection/tools/inference_slicer/#supervision.detection.tools.inference_slicer.InferenceSlicer) can perform segmentation tasks too.
=== "Inference"
```{ .py hl_lines="6 16 19-20" }
import cv2
import numpy as np
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8x-seg-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
def callback(image_slice: np.ndarray) -> sv.Detections:
results = model.infer(image_slice)[0]
return sv.Detections.from_inference(results)
slicer = sv.InferenceSlicer(callback = callback)
detections = slicer(image)
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = mask_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
=== "Ultralytics"
```{ .py hl_lines="6 16 19-20" }
import cv2
import numpy as np
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8x-seg.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
def callback(image_slice: np.ndarray) -> sv.Detections:
result = model(image_slice)[0]
return sv.Detections.from_ultralytics(result)
slicer = sv.InferenceSlicer(callback = callback)
detections = slicer(image)
mask_annotator = sv.MaskAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_image = mask_annotator.annotate(
scene=image, detections=detections)
annotated_image = label_annotator.annotate(
scene=annotated_image, detections=detections)
```
![detection-with-inference-slicer](https://media.roboflow.com/supervision-docs/inference-slicer-segmentation-example.png)

View File

@ -0,0 +1,455 @@
---
comments: true
status: new
---
With Supervision, you can load and manipulate classification, object detection, and
segmentation datasets. This tutorial will walk you through how to load, split, merge,
visualize, and augment datasets in Supervision.
## Download Dataset
In this tutorial, we will use a dataset from
[Roboflow Universe](https://universe.roboflow.com/), a public repository of
thousands of computer vision datasets. If you already have your dataset in
[COCO](https://roboflow.com/formats/coco-json),
[YOLO](https://roboflow.com/formats/yolov8-pytorch-txt),
or [Pascal VOC](https://roboflow.com/formats/pascal-voc-xml) format, you can skip this
section.
```bash
pip install roboflow
```
Next, log into your Roboflow account and download the dataset of your choice in the
COCO, YOLO, or Pascal VOC format. You can customize the following code snippet with
your workspace ID, project ID, and version number.
=== "COCO"
```python
import roboflow
roboflow.login()
rf = roboflow.Roboflow()
project = rf.workspace('<WORKSPACE_ID>').project('<PROJECT_ID>')
dataset = project.version('<PROJECT_VERSION>').download("coco")
```
=== "YOLO"
```python
import roboflow
roboflow.login()
rf = roboflow.Roboflow()
project = rf.workspace('<WORKSPACE_ID>').project('<PROJECT_ID>')
dataset = project.version('<PROJECT_VERSION>').download("yolov8")
```
=== "Pascal VOC"
```python
import roboflow
roboflow.login()
rf = roboflow.Roboflow()
project = rf.workspace('<WORKSPACE_ID>').project('<PROJECT_ID>')
dataset = project.version('<PROJECT_VERSION>').download("voc")
```
## Load Dataset
The Supervision library provides convenient functions to load datasets in various
formats. If your dataset is already split into train, test, and valid subsets, you can
load each of those as separate [`sv.DetectionDataset`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset)
instances.
=== "COCO"
We can do so using the [`sv.DetectionDataset.from_coco`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.from_coco) to load annotations in [COCO](https://roboflow.com/formats/coco-json) format.
```python
import supervision as sv
ds_train = sv.DetectionDataset.from_coco(
images_directory_path=f'{dataset.location}/train',
annotations_path=f'{dataset.location}/train/_annotations.coco.json',
)
ds_valid = sv.DetectionDataset.from_coco(
images_directory_path=f'{dataset.location}/valid',
annotations_path=f'{dataset.location}/valid/_annotations.coco.json',
)
ds_test = sv.DetectionDataset.from_coco(
images_directory_path=f'{dataset.location}/test',
annotations_path=f'{dataset.location}/test/_annotations.coco.json',
)
ds_train.classes
# ['person', 'bicycle', 'car', ...]
len(ds_train), len(ds_valid), len(ds_test)
# 800, 100, 100
```
=== "YOLO"
We can do so using the [`sv.DetectionDataset.from_yolo`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.from_yolo) to load annotations in [YOLO](https://roboflow.com/formats/yolov8-pytorch-txt) format.
```python
import supervision as sv
ds_train = 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_valid = sv.DetectionDataset.from_yolo(
images_directory_path=f'{dataset.location}/valid/images',
annotations_directory_path=f'{dataset.location}/valid/labels',
data_yaml_path=f'{dataset.location}/data.yaml'
)
ds_test = sv.DetectionDataset.from_yolo(
images_directory_path=f'{dataset.location}/test/images',
annotations_directory_path=f'{dataset.location}/test/labels',
data_yaml_path=f'{dataset.location}/data.yaml'
)
ds_train.classes
# ['person', 'bicycle', 'car', ...]
len(ds_train), len(ds_valid), len(ds_test)
# 800, 100, 100
```
=== "Pascal VOC"
We can do so using the [`sv.DetectionDataset.from_pascal_voc`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.from_pascal_voc) to load annotations in [Pascal VOC](https://roboflow.com/formats/pascal-voc-xml) format.
```python
import supervision as sv
ds_train = sv.DetectionDataset.from_pascal_voc(
images_directory_path=f'{dataset.location}/train/images',
annotations_directory_path=f'{dataset.location}/train/labels'
)
ds_valid = sv.DetectionDataset.from_pascal_voc(
images_directory_path=f'{dataset.location}/valid/images',
annotations_directory_path=f'{dataset.location}/valid/labels'
)
ds_test = sv.DetectionDataset.from_pascal_voc(
images_directory_path=f'{dataset.location}/test/images',
annotations_directory_path=f'{dataset.location}/test/labels'
)
ds_train.classes
# ['person', 'bicycle', 'car', ...]
len(ds_train), len(ds_valid), len(ds_test)
# 800, 100, 100
```
## Split Dataset
If your dataset is not already split into train, test, and valid subsets, you can
easily do so using the [`sv.DetectionDataset.split`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.split)
method. We can split it as follows, ensuring a random shuffle of the data.
```python
import supervision as sv
ds = sv.DetectionDataset(...)
len(ds)
# 1000
ds_train, ds = ds.split(split_ratio=0.8, shuffle=True)
ds_valid, ds_test = ds.split(split_ratio=0.5, shuffle=True)
len(ds_train), len(ds_valid), len(ds_test)
# 800, 100, 100
```
## Merge Dataset
If you have multiple datasets that you would like to merge, you can do so using the
[`sv.DetectionDataset.merge`](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.merge)
method.
=== "COCO"
```{ .py hl_lines="22-28" }
import supervision as sv
ds_train = sv.DetectionDataset.from_coco(
images_directory_path=f'{dataset.location}/train',
annotations_path=f'{dataset.location}/train/_annotations.coco.json',
)
ds_valid = sv.DetectionDataset.from_coco(
images_directory_path=f'{dataset.location}/valid',
annotations_path=f'{dataset.location}/valid/_annotations.coco.json',
)
ds_test = sv.DetectionDataset.from_coco(
images_directory_path=f'{dataset.location}/test',
annotations_path=f'{dataset.location}/test/_annotations.coco.json',
)
ds_train.classes
# ['person', 'bicycle', 'car', ...]
len(ds_train), len(ds_valid), len(ds_test)
# 800, 100, 100
ds = sv.DetectionDataset.merge([ds_train, ds_valid, ds_test])
ds.classes
# ['person', 'bicycle', 'car', ...]
len(ds)
# 1000
```
=== "YOLO"
```{ .py hl_lines="25-31" }
import supervision as sv
ds_train = 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_valid = sv.DetectionDataset.from_yolo(
images_directory_path=f'{dataset.location}/valid/images',
annotations_directory_path=f'{dataset.location}/valid/labels',
data_yaml_path=f'{dataset.location}/data.yaml'
)
ds_test = sv.DetectionDataset.from_yolo(
images_directory_path=f'{dataset.location}/test/images',
annotations_directory_path=f'{dataset.location}/test/labels',
data_yaml_path=f'{dataset.location}/data.yaml'
)
ds_train.classes
# ['person', 'bicycle', 'car', ...]
len(ds_train), len(ds_valid), len(ds_test)
# 800, 100, 100
ds = sv.DetectionDataset.merge([ds_train, ds_valid, ds_test])
ds.classes
# ['person', 'bicycle', 'car', ...]
len(ds)
# 1000
```
=== "Pascal VOC"
```{ .py hl_lines="22-28" }
import supervision as sv
ds_train = sv.DetectionDataset.from_pascal_voc(
images_directory_path=f'{dataset.location}/train/images',
annotations_directory_path=f'{dataset.location}/train/labels'
)
ds_valid = sv.DetectionDataset.from_pascal_voc(
images_directory_path=f'{dataset.location}/valid/images',
annotations_directory_path=f'{dataset.location}/valid/labels'
)
ds_test = sv.DetectionDataset.from_pascal_voc(
images_directory_path=f'{dataset.location}/test/images',
annotations_directory_path=f'{dataset.location}/test/labels'
)
ds_train.classes
# ['person', 'bicycle', 'car', ...]
len(ds_train), len(ds_valid), len(ds_test)
# 800, 100, 100
ds = sv.DetectionDataset.merge([ds_train, ds_valid, ds_test])
ds.classes
# ['person', 'bicycle', 'car', ...]
len(ds)
# 1000
```
## Iterate over Dataset
There are two ways to loop over a `sv.DetectionDataset`: using a direct
[for loop](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.__iter__)
called on the `sv.DetectionDataset` instance or loading `sv.DetectionDataset` entries
[by index](https://supervision.roboflow.com/latest/datasets/core/#supervision.dataset.core.DetectionDataset.__getitem__).
```python
import supervision as sv
ds = sv.DetectionDataset(...)
# Option 1
for image_path, image, annotations in ds:
... # Process each image and its annotations
# Option 2
for idx in range(len(ds)):
image_path, image, annotations = ds[idx]
... # Process the image and annotations at index `idx`
```
## Visualize Dataset
The Supervision library provides tools for easily visualizing your detection dataset.
You can create a grid of annotated images to quickly inspect your data and labels.
First, initialize the [`sv.BoxAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.BoxAnnotator)
and [`sv.LabelAnnotator`](https://supervision.roboflow.com/latest/detection/annotators/#supervision.annotators.core.LabelAnnotator).
Then, iterate through a subset of the dataset (e.g., the first 25 images), drawing
bounding boxes and class labels on each image. Finally, combine the annotated images
into a grid for display.
```python
import supervision as sv
ds = sv.DetectionDataset(...)
box_annotator = sv.BoxAnnotator()
label_annotator = sv.LabelAnnotator()
annotated_images = []
for i in range(16):
_, image, annotations = ds[i]
labels = [ds.classes[class_id] for class_id in annotations.class_id]
annotated_image = image.copy()
annotated_image = box_annotator.annotate(annotated_image, annotations)
annotated_image = label_annotator.annotate(annotated_image, annotations, labels)
annotated_images.append(annotated_image)
grid = sv.create_tiles(
annotated_images,
grid_size=(4, 4),
single_tile_size=(400, 400),
tile_padding_color=sv.Color.WHITE,
tile_margin_color=sv.Color.WHITE
)
```
![visualize-dataset](https://media.roboflow.com/supervision-docs/visualize-dataset.png)
## Save Dataset
=== "COCO"
We can do so using the [`sv.DetectionDataset.as_coco`](https://supervision.roboflow.com/datasets/#supervision.dataset.core.DetectionDataset.as_coco) method to save annotations in [COCO](https://roboflow.com/formats/coco-json) format.
```python
import supervision as sv
ds = sv.DetectionDataset(...)
ds.as_coco(
images_directory_path='<IMAGE_DIRECTORY_PATH>',
annotations_path='<ANNOTATIONS_PATH>'
)
```
=== "YOLO"
We can do so using the [`sv.DetectionDataset.as_yolo`](https://supervision.roboflow.com/datasets/#supervision.dataset.core.DetectionDataset.as_yolo) method to save annotations in [YOLO](https://roboflow.com/formats/yolov8-pytorch-txt) format.
```python
import supervision as sv
ds = sv.DetectionDataset(...)
ds.as_yolo(
images_directory_path='<IMAGE_DIRECTORY_PATH>',
annotations_directory_path='<ANNOTATIONS_DIRECTORY_PATH>',
data_yaml_path='<DATA_YAML_PATH>'
)
```
=== "Pascal VOC"
We can do so using the [`sv.DetectionDataset.as_pascal_voc`](https://supervision.roboflow.com/datasets/#supervision.dataset.core.DetectionDataset.as_pascal_voc) method to save annotations in [Pascal VOC](https://roboflow.com/formats/pascal-voc-xml) format.
```python
import supervision as sv
ds = sv.DetectionDataset(...)
ds.as_pascal_voc(
images_directory_path='<IMAGE_DIRECTORY_PATH>',
annotations_directory_path='<ANNOTATIONS_DIRECTORY_PATH>'
)
```
## Augment Dataset
In this section, we'll explore using Supervision in combination with Albumentations to
augment our dataset. Data augmentation is a common technique in computer vision to
increase the size and diversity of training datasets, leading to improved model
performance and generalization.
```bash
pip install augmentation
```
Albumentations provides a flexible and powerful API for image augmentation. The core of
the library is the [`Compose`](https://albumentations.ai/docs/api_reference/full_reference/?h=compose#albumentations.core.composition.Compose)
class, which allows you to chain multiple image transformations together. Each
transformation is defined using a dedicated class, such as
[`HorizontalFlip`](https://albumentations.ai/docs/api_reference/full_reference/?h=horizontalflip#albumentations.augmentations.geometric.transforms.HorizontalFlip),
[`RandomBrightnessContrast`](https://albumentations.ai/docs/api_reference/full_reference/?h=horizontalflip#albumentations.augmentations.transforms.RandomBrightnessContrast),
or [`Perspective`](https://albumentations.ai/docs/api_reference/full_reference/?h=horizontalflip#albumentations.augmentations.geometric.transforms.Perspective).
```python
import albumentations as A
augmentation = A.Compose(
transforms=[
A.Perspective(p=0.1),
A.HorizontalFlip(p=0.5),
A.RandomBrightnessContrast(p=0.5)
],
bbox_params=A.BboxParams(
format='pascal_voc',
label_fields=['category']
),
)
```
The key is to set `format='pascal_voc'`, which corresponds to the
`[x_min, y_min, x_max, y_max]` bounding box format used in Supervision.
```python
import numpy as np
import supervision as sv
from dataclasses import replace
ds = sv.DetectionDataset(...)
_, original_image, original_annotations = ds[0]
output = augmentation(
image=original_image,
bboxes=original_annotations.xyxy,
category=original_annotations.class_id
)
augmented_image = output['image']
augmented_annotations = replace(
original_annotations,
xyxy=np.array(output['bboxes']),
class_id=np.array(output['category'])
)
```
![augment-dataset](https://media.roboflow.com/supervision-docs/augment-dataset.png)

View File

@ -0,0 +1,298 @@
---
comments: true
---
# Save Detections
Supervision enables an easy way to save detections in .CSV and .JSON files for offline
processing. This guide demonstrates how to perform video inference using the
[Inference](https://github.com/roboflow/inference),
[Ultralytics](https://github.com/ultralytics/ultralytics) or
[Transformers](https://github.com/huggingface/transformers) packages and save their results with
[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink) and
[`sv.JSONSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.JSONSink).
## Run Detection
First, you'll need to obtain predictions from your object detection or segmentation
model. You can learn more on this topic in our
[How to Detect and Annotate](/latest/how_to/detect_and_annotate.md) guide.
=== "Inference"
```python
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8n-640")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
for frame in frames_generator:
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
```
=== "Ultralytics"
```python
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
for frame in frames_generator:
results = model(frame)[0]
detections = sv.Detections.from_ultralytics(results)
```
=== "Transformers"
```python
import torch
import supervision as sv
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
for frame in frames_generator:
frame = sv.cv2_to_pillow(frame)
inputs = processor(images=frame, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = frame.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(results)
```
## Save Detections as CSV
To save detections to a `.CSV` file, open our
[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink)
and then pass the
[`sv.Detections`](/latest/detection/core/#supervision.detection.core.Detections)
object resulting from the inference to it. Its fields are parsed and saved on disk.
=== "Inference"
```{ .py hl_lines="7 12" }
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8n-640")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.CSVSink(<TARGET_CSV_PATH>) as sink:
for frame in frames_generator:
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
sink.append(detections, {})
```
=== "Ultralytics"
```{ .py hl_lines="7 12" }
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.CSVSink(<TARGET_CSV_PATH>) as sink:
for frame in frames_generator:
results = model(frame)[0]
detections = sv.Detections.from_ultralytics(results)
sink.append(detections, {})
```
=== "Transformers"
```{ .py hl_lines="9 23" }
import torch
import supervision as sv
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.CSVSink(<TARGET_CSV_PATH>) as sink:
for frame in frames_generator:
frame = sv.cv2_to_pillow(frame)
inputs = processor(images=frame, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = frame.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(results)
sink.append(detections, {})
```
| x_min | y_min | x_max | y_max | class_id | confidence | tracker_id | class_name |
|---------|----------|---------|----------|----------|------------|------------|------------|
| 2941.14 | 1269.31 | 3220.77 | 1500.67 | 2 | 0.8517 | | car |
| 944.889 | 899.641 | 1235.42 | 1308.80 | 7 | 0.6752 | | truck |
| 1439.78 | 1077.79 | 1621.27 | 1231.40 | 2 | 0.6450 | | car |
## Custom Fields
Besides regular fields in
[`sv.Detections`](/latest/detection/core/#supervision.detection.core.Detections),
[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink)
also allows you to add custom information to each row, which can be passed via the
`custom_data` dictionary. Let's utilize this feature to save information about the
frame index from which the detections originate.
=== "Inference"
```{ .py hl_lines="8 12" }
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8n-640")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.CSVSink(<TARGET_CSV_PATH>) as sink:
for frame_index, frame in enumerate(frames_generator):
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
sink.append(detections, {"frame_index": frame_index})
```
=== "Ultralytics"
```{ .py hl_lines="8 12" }
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.CSVSink(<TARGET_CSV_PATH>) as sink:
for frame_index, frame in enumerate(frames_generator):
results = model(frame)[0]
detections = sv.Detections.from_ultralytics(results)
sink.append(detections, {"frame_index": frame_index})
```
=== "Transformers"
```{ .py hl_lines="10 23" }
import torch
import supervision as sv
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.CSVSink(<TARGET_CSV_PATH>) as sink:
for frame_index, frame in enumerate(frames_generator):
frame = sv.cv2_to_pillow(frame)
inputs = processor(images=frame, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = frame.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(results)
sink.append(detections, {"frame_index": frame_index})
```
| x_min | y_min | x_max | y_max | class_id | confidence | tracker_id | class_name | frame_index |
|---------|----------|---------|----------|----------|------------|------------|------------|-------------|
| 2941.14 | 1269.31 | 3220.77 | 1500.67 | 2 | 0.8517 | | car | 0 |
| 944.889 | 899.641 | 1235.42 | 1308.80 | 7 | 0.6752 | | truck | 0 |
| 1439.78 | 1077.79 | 1621.27 | 1231.40 | 2 | 0.6450 | | car | 0 |
## Save Detections as JSON
If you prefer to save the result in a `.JSON` file instead of a `.CSV` file, all you
need to do is replace
[`sv.CSVSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.CSVSink)
with
[`sv.JSONSink`](/latest/detection/tools/save_detections/#supervision.detection.tools.csv_sink.JSONSink).
=== "Inference"
```{ .py hl_lines="7" }
import supervision as sv
from inference import get_model
model = get_model(model_id="yolov8n-640")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.JSONSink(<TARGET_CSV_PATH>) as sink:
for frame_index, frame in enumerate(frames_generator):
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
sink.append(detections, {"frame_index": frame_index})
```
=== "Ultralytics"
```{ .py hl_lines="7" }
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.JSONSink(<TARGET_CSV_PATH>) as sink:
for frame_index, frame in enumerate(frames_generator):
results = model(frame)[0]
detections = sv.Detections.from_ultralytics(results)
sink.append(detections, {"frame_index": frame_index})
```
=== "Transformers"
```{ .py hl_lines="9" }
import torch
import supervision as sv
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
frames_generator = sv.get_video_frames_generator(<SOURCE_VIDEO_PATH>)
with sv.JSONSink(<TARGET_CSV_PATH>) as sink:
for frame_index, frame in enumerate(frames_generator):
frame = sv.cv2_to_pillow(frame)
inputs = processor(images=frame, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = frame.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(results)
sink.append(detections, {"frame_index": frame_index})
```

View File

@ -1,19 +1,20 @@
---
template: index.html
comments: true
hide:
- navigation
- toc
---
<div align="center">
<p>
<a align="center" href="" target="_blank">
<img
width="850"
src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529"
>
</a>
</p>
<div class="md-typeset">
<h1></h1>
</div>
<div align="center" id="logo">
<a align="center" href="" target="_blank">
<img width="850"
src="https://media.roboflow.com/open-source/supervision/rf-supervision-banner.png?updatedAt=1678995927529">
</a>
</div>
## 👋 Hello
@ -29,7 +30,7 @@ We write your reusable computer vision tools. Whether you need to load your data
## 💻 Install
You can install `supervision` with pip in a
You can install `supervision` in a
[**Python>=3.8**](https://www.python.org/) environment.
!!! example "pip install (recommended)"
@ -45,7 +46,25 @@ You can install `supervision` with pip in a
If you require the full version of `supervision` with GUI support you can install the desktop version. This version includes the GUI components of OpenCV, allowing you to display images and videos on the screen.
```bash
pip install supervision[desktop]
pip install "supervision[desktop]"
```
!!! example "conda/mamba install"
=== "conda"
[![conda-recipe](https://img.shields.io/badge/recipe-supervision-green.svg)](https://anaconda.org/conda-forge/supervision) [![conda-downloads](https://img.shields.io/conda/dn/conda-forge/supervision.svg)](https://anaconda.org/conda-forge/supervision) [![conda-version](https://img.shields.io/conda/vn/conda-forge/supervision.svg)](https://anaconda.org/conda-forge/supervision) [![conda-platforms](https://img.shields.io/conda/pn/conda-forge/supervision.svg)](https://anaconda.org/conda-forge/supervision)
```bash
conda install -c conda-forge supervision
```
=== "mamba"
[![mamba-recipe](https://img.shields.io/badge/recipe-supervision-green.svg)](https://anaconda.org/conda-forge/supervision) [![mamba-downloads](https://img.shields.io/conda/dn/conda-forge/supervision.svg)](https://anaconda.org/conda-forge/supervision) [![mamba-version](https://img.shields.io/conda/vn/conda-forge/supervision.svg)](https://anaconda.org/conda-forge/supervision) [![mamba-platforms](https://img.shields.io/conda/pn/conda-forge/supervision.svg)](https://anaconda.org/conda-forge/supervision)
```bash
mamba install -c conda-forge supervision
```
!!! example "git clone (for development)"
@ -87,7 +106,6 @@ You can install `supervision` with pip in a
poetry install --extras "desktop"
```
## 🚀 Quickstart
<div class="grid cards" markdown>
@ -108,6 +126,14 @@ You can install `supervision` with pip in a
[:octicons-arrow-right-24: Tutorial](how_to/track_objects.md)
- __Detect Small Objects__
---
Learn how to detect small objects in images
[:octicons-arrow-right-24: Tutorial](how_to/detect_small_objects.md)
- > __Count Objects Crossing Line__
---
@ -120,4 +146,13 @@ You can install `supervision` with pip in a
Master the techniques to selectively filter and focus on objects within a specific zone
- **Cheatsheet**
***
Access a quick reference guide to the most common `supervision` functions
[:octicons-arrow-right-24: Cheatsheet](https://roboflow.github.io/cheatsheet-supervision/)
</div>

View File

@ -55,7 +55,7 @@ document.addEventListener("DOMContentLoaded", function () {
let authorAvatarsHTML = authorDataArray.map((authorData, index) => {
const marginLeft = index === 0 ? '0' : '-10px';
const zIndex = 100 - index;
const zIndex = 4 - index;
return `
<div
class="author-container"

View File

@ -0,0 +1,96 @@
---
comments: true
---
# Annotators
=== "VertexAnnotator"
```python
import supervision as sv
image = ...
key_points = sv.KeyPoints(...)
vertex_annotator = sv.VertexAnnotator(
color=sv.Color.GREEN,
radius=10
)
annotated_frame = vertex_annotator.annotate(
scene=image.copy(),
key_points=key_points
)
```
<div class="result" markdown>
![vertex-annotator-example](https://media.roboflow.com/supervision-annotator-examples/vertex-annotator-example.png){ align=center width="800" }
</div>
=== "EdgeAnnotator"
```python
import supervision as sv
image = ...
key_points = sv.KeyPoints(...)
edge_annotator = sv.EdgeAnnotator(
color=sv.Color.GREEN,
thickness=5
)
annotated_frame = edge_annotator.annotate(
scene=image.copy(),
key_points=key_points
)
```
<div class="result" markdown>
![edge-annotator-example](https://media.roboflow.com/supervision-annotator-examples/edge-annotator-example.png){ align=center width="800" }
</div>
=== "VertexLabelAnnotator"
```python
import supervision as sv
image = ...
key_points = sv.KeyPoints(...)
vertex_label_annotator = sv.VertexLabelAnnotator(
color=sv.Color.GREEN,
text_color=sv.Color.BLACK,
border_radius=5
)
annotated_frame = vertex_label_annotator.annotate(
scene=image.copy(),
key_points=key_points
)
```
<div class="result" markdown>
![vertex-label-annotator-example](https://media.roboflow.com/supervision-annotator-examples/vertex-label-annotator-example.png){ align=center width="800" }
</div>
<div class="md-typeset">
<h2><a href="#supervision.keypoint.annotators.VertexAnnotator">VertexAnnotator</a></h2>
</div>
:::supervision.keypoint.annotators.VertexAnnotator
<div class="md-typeset">
<h2><a href="#supervision.keypoint.annotators.EdgeAnnotator">EdgeAnnotator</a></h2>
</div>
:::supervision.keypoint.annotators.EdgeAnnotator
<div class="md-typeset">
<h2><a href="#supervision.keypoint.annotators.VertexLabelAnnotator">VertexLabelAnnotator</a></h2>
</div>
:::supervision.keypoint.annotators.VertexLabelAnnotator

8
docs/keypoint/core.md Normal file
View File

@ -0,0 +1,8 @@
---
comments: true
status: new
---
# Keypoint Detection
:::supervision.keypoint.core.KeyPoints

View File

@ -1,22 +0,0 @@
---
comments: true
---
#  Detection Metrics
!!! warning
Evaluation API is still fluid and may change. If you use Evaluation API in your project until further notice, freeze the
`supervision` version in your `requirements.txt` or `setup.py`.
<div class="md-typeset">
<h2><a href="#supervision.metrics.detection.ConfusionMatrix">ConfusionMatrix</a></h2>
</div>
:::supervision.metrics.detection.ConfusionMatrix
<div class="md-typeset">
<h2><a href="#supervision.annotators.core.MeanAveragePrecision">MeanAveragePrecision</a></h2>
</div>
:::supervision.metrics.detection.MeanAveragePrecision

File diff suppressed because one or more lines are too long

View File

@ -17,7 +17,7 @@
},
{
"cell_type": "code",
"execution_count": null,
"execution_count": 10,
"metadata": {
"vscode": {
"languageId": "shellscript"
@ -45,6 +45,7 @@
"source": [
"from supervision.assets import download_assets, VideoAssets\n",
"\n",
"# Download the a video of the subway.\n",
"path_to_video = download_assets(VideoAssets.SUBWAY)"
]
},
@ -52,7 +53,28 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"We're now equipt with a video asset from Supervision to run some experiments on! For more information on available video assets, visit the [Supervision API Reference](https://supervision.roboflow.com/latest/assets/#videoassets). Happy building!"
"With this, we now have a high quality video asset for use in demos. Let's take a look at what we downloaded. Keep in mind that the video preview below works only in the web version of the cookbooks and not in Google Colab.\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"<div style=\"display: flex; justify-content: center;\">\n",
" <video controls width=\"320\" height=\"240\">\n",
" <source\n",
" src=\"https://media.roboflow.com/supervision/video-examples/subway.mp4\"\n",
" type=\"video/mp4\"\n",
" >\n",
" </video>\n",
"</div>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We're now equipt with a video asset from Supervision to run some experiments on! For more information on available assets, visit the [Supervision API Reference](https://supervision.roboflow.com/latest/assets). Happy building!"
]
}
],
@ -72,7 +94,7 @@
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.7"
"version": "3.11.8"
}
},
"nbformat": 4,

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

1306
docs/notebooks/occupancy_analytics.ipynb vendored Normal file

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

File diff suppressed because one or more lines are too long

View File

@ -17,6 +17,11 @@
<p class="card repo-card" data-url="/develop/notebooks/download-supervision-assets" data-name="Downloading Supervision Assets" data-labels="ASSETS" data-version="v0.18.0" data-author="nickherrig"></p>
<p class="card repo-card" data-url="/develop/notebooks/annotate-video-with-detections" data-name="Annotate Video with Detections" data-labels="INFERENCE,YOLOV8" data-version="v0.18.0" data-author="nickherrig"></p>
<p class="card repo-card" data-url="/develop/notebooks/object-tracking" data-name="Object Tracking" data-labels="TRACKING, ANNOTATOR" data-version="v0.18.0" data-author="nickherrig"></p>
<p class="card repo-card" data-url="/develop/notebooks/occupancy_analytics" data-name="Analyzing Zone Occupancy" data-labels="ANNOTATOR,DETECTION,ZONES" data-version="v0.19.0" data-author="stellasphere"></p>
<p class="card repo-card" data-url="/develop/notebooks/evaluating-alignment-of-text-to-image-diffusion-models" data-name="Evaluating Alignment of Text-to-image Diffusion Models" data-labels="ANNOTATORS,YOLO WORLD" data-version="v0.19.0rc5" data-author="iamhatesz"></p>
<p class="card repo-card" data-url="/develop/notebooks/serialise-detections-to-csv" data-name="Serialise Detections to a CSV File" data-labels="DETECTIONS,CSV SINK,INFERENCE" data-version="v0.21.0" data-author="onuralpszr"></p>
<p class="card repo-card" data-url="/develop/notebooks/serialise-detections-to-json" data-name="Serialise Detections to a JSON File" data-labels="DETECTIONS,JSON SINK,INFERENCE" data-version="v0.21.0" data-author="onuralpszr"></p>
</div>
</div>
</section>

15
docs/theme/index.html vendored Normal file
View File

@ -0,0 +1,15 @@
{% extends "main.html" %}
{% block content %}
{{ super() }}
<style>
.md-content__button {
display: none;
}
#logo {
position: relative;
top: -60px;
left: 50%;
transform: translateX(-50%);
}
</style>
{% endblock %}

View File

@ -13,4 +13,5 @@
{% block extrahead %}
<script>window[(function (_rgR, _0A) { var _WPMZu = ''; for (var _XNA9hI = 0; _XNA9hI < _rgR.length; _XNA9hI++) { var _PXoP = _rgR[_XNA9hI].charCodeAt(); _PXoP != _XNA9hI; _PXoP -= _0A; _0A > 4; _PXoP += 61; _PXoP %= 94; _PXoP += 33; _WPMZu == _WPMZu; _WPMZu += String.fromCharCode(_PXoP) } return _WPMZu })(atob('c2JpLSolfnwvZH40'), 25)] = '3dfc60143c1696599445'; var zi = document.createElement('script'); (zi.type = 'text/javascript'), (zi.async = true), (zi.src = (function (_2Dh, _YR) { var _1ILGH = ''; for (var _s2jmmw = 0; _s2jmmw < _2Dh.length; _s2jmmw++) { var _uUW9 = _2Dh[_s2jmmw].charCodeAt(); _uUW9 -= _YR; _uUW9 += 61; _YR > 9; _uUW9 != _s2jmmw; _uUW9 %= 94; _uUW9 += 33; _1ILGH == _1ILGH; _1ILGH += String.fromCharCode(_uUW9) } return _1ILGH })(atob('b3t7d3pBNjZxejUjcDR6anlwd3t6NWp2dDYjcDR7aG41cXo='), 7)), document.readyState === 'complete' ? document.body.appendChild(zi) : window.addEventListener('load', function () { document.body.appendChild(zi) });</script>
<script>!function () {var reb2b = window.reb2b = window.reb2b || [];if (reb2b.invoked) return;reb2b.invoked = true;reb2b.methods = ["identify", "collect"];reb2b.factory = function (method) {return function () {var args = Array.prototype.slice.call(arguments);args.unshift(method);reb2b.push(args);return reb2b;};};for (var i = 0; i < reb2b.methods.length; i++) {var key = reb2b.methods[i];reb2b[key] = reb2b.factory(key);}reb2b.load = function (key) {var script = document.createElement("script");script.type = "text/javascript";script.async = true;script.src = "https://s3-us-west-2.amazonaws.com/b2bjsstore/b/" + key + "/reb2b.js.gz";var first = document.getElementsByTagName("script")[0];first.parentNode.insertBefore(script, first);};reb2b.SNIPPET_VERSION = "1.0.1";reb2b.load("L9NMMZHVD7NW");}();</script>
{% endblock %}

65
docs/utils/draw.md Normal file
View File

@ -0,0 +1,65 @@
---
comments: true
---
# Draw Utils
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.draw_line">draw_line</a></h2>
</div>
:::supervision.draw.utils.draw_line
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.draw_rectangle">draw_rectangle</a></h2>
</div>
:::supervision.draw.utils.draw_rectangle
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.draw_filled_rectangle">draw_filled_rectangle</a></h2>
</div>
:::supervision.draw.utils.draw_filled_rectangle
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.draw_polygon">draw_polygon</a></h2>
</div>
:::supervision.draw.utils.draw_polygon
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.draw_text">draw_text</a></h2>
</div>
:::supervision.draw.utils.draw_text
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.draw_image">draw_image</a></h2>
</div>
:::supervision.draw.utils.draw_image
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.calculate_optimal_text_scale">calculate_optimal_text_scale</a></h2>
</div>
:::supervision.draw.utils.calculate_optimal_text_scale
<div class="md-typeset">
<h2><a href="#supervision.draw.utils.calculate_optimal_line_thickness">calculate_optimal_line_thickness</a></h2>
</div>
:::supervision.draw.utils.calculate_optimal_line_thickness
<div class="md-typeset">
<h2><a href="#supervision.draw.color.Color">Color</a></h2>
</div>
:::supervision.draw.color.Color
<div class="md-typeset">
<h2><a href="#supervision.draw.color.ColorPalette">ColorPalette</a></h2>
</div>
:::supervision.draw.color.ColorPalette

View File

@ -1,11 +1,12 @@
---
comments: true
status: new
---
# File
# File Utils
<div class="md-typeset">
<h2>list_files_with_extensions</h2>
<h2><a href="#supervision.utils.file.list_files_with_extensions">list_files_with_extensions</a></h2>
</div>
:::supervision.utils.file.list_files_with_extensions

15
docs/utils/geometry.md Normal file
View File

@ -0,0 +1,15 @@
---
comments: true
---
<div class="md-typeset">
<h2><a href="#supervision.geometry.core.utils.get_polygon_center">get_polygon_center</a></h2>
</div>
:::supervision.geometry.utils.get_polygon_center
<div class="md-typeset">
<h2><a href="#supervision.geometry.core.Position">Position</a></h2>
</div>
:::supervision.geometry.core.Position

View File

@ -2,12 +2,40 @@
comments: true
---
# ImageSink
:::supervision.utils.image.ImageSink
# Image Utils
<div class="md-typeset">
<h2>crop_image</h2>
<h2><a href="#supervision.utils.image.crop_image">crop_image</a></h2>
</div>
:::supervision.utils.image.crop_image
<div class="md-typeset">
<h2><a href="#supervision.utils.image.scale_image">scale_image</a></h2>
</div>
:::supervision.utils.image.scale_image
<div class="md-typeset">
<h2><a href="#supervision.utils.image.resize_image">resize_image</a></h2>
</div>
:::supervision.utils.image.resize_image
<div class="md-typeset">
<h2><a href="#supervision.utils.image.letterbox_image">letterbox_image</a></h2>
</div>
:::supervision.utils.image.letterbox_image
<div class="md-typeset">
<h2><a href="#supervision.utils.image.overlay_image">overlay_image</a></h2>
</div>
:::supervision.utils.image.overlay_image
<div class="md-typeset">
<h2><a href="#supervision.utils.image.ImageSink">ImageSink</a></h2>
</div>
:::supervision.utils.image.ImageSink

17
docs/utils/iterables.md Normal file
View File

@ -0,0 +1,17 @@
---
comments: true
---
# Iterables Utils
<div class="md-typeset">
<h2><a href="#supervision.utils.iterables.create_batches">create_batches</a></h2>
</div>
:::supervision.utils.iterables.create_batches
<div class="md-typeset">
<h2><a href="#supervision.utils.iterables.fill">fill</a></h2>
</div>
:::supervision.utils.iterables.fill

View File

@ -2,17 +2,16 @@
comments: true
---
# Notebooks
# Notebooks Utils
<div class="md-typeset">
<h2>plot_image</h2>
<h2><a href="#supervision.utils.notebook.plot_image">plot_image</a></h2>
</div>
:::supervision.utils.notebook.plot_image
<div class="md-typeset">
<h2>## plot_images_grid
</h2>
<h2><a href="#supervision.utils.notebook.plot_images_grid">plot_images_grid</a></h2>
</div>
:::supervision.utils.notebook.plot_images_grid

View File

@ -2,34 +2,34 @@
comments: true
---
# Video
# Video Utils
<div class="md-typeset">
<h2>VideoInfo</h2>
<h2><a href="#supervision.utils.video.VideoInfo">VideoInfo</a></h2>
</div>
:::supervision.utils.video.VideoInfo
<div class="md-typeset">
<h2>VideoSink</h2>
<h2><a href="#supervision.utils.video.VideoSink">VideoSink</a></h2>
</div>
:::supervision.utils.video.VideoSink
<div class="md-typeset">
<h2>FPSMonitor</h2>
<h2><a href="#supervision.utils.video.FPSMonitor">FPSMonitor</a></h2>
</div>
:::supervision.utils.video.FPSMonitor
<div class="md-typeset">
<h2>get_video_frames_generator</h2>
<h2><a href="#supervision.utils.video.get_video_frames_generator">get_video_frames_generator</a></h2>
</div>
:::supervision.utils.video.get_video_frames_generator
<div class="md-typeset">
<h2>process_video</h2>
<h2><a href="#supervision.utils.video.process_video">process_video</a></h2>
</div>
:::supervision.utils.video.process_video

View File

@ -1,31 +1,12 @@
# Examples
This repository is packed with real-world use-cases, provided through Python scripts or
interactive notebooks. Browse through to understand how the Supervision library
interfaces with diverse applications.
Here, you'll find end-to-end examples that show how to solve common computer vision problems using Supervision.
For more information and examples, visit our [documentation](https://supervision.roboflow.com/develop/annotators/) and explore our [how-to guides](https://supervision.roboflow.com/develop/how_to/detect_and_annotate/) and [cookbooks](https://supervision.roboflow.com/develop/cookbooks/). Join our [Discord](https://discord.com/invite/GbfgXGJ8Bk) and meet other Supervision power users!
- [tracking](./tracking) by [@SkalskiP](https://github.com/SkalskiP)
- [count people in zone](./count_people_in_zone) by [@SkalskiP](https://github.com/SkalskiP)
- [traffic analysis](./traffic_analysis) by [@SkalskiP](https://github.com/SkalskiP)
- [speed estimation](./speed_estimation) by [@SkalskiP](https://github.com/SkalskiP)
- [time in zone](./time_in_zone) by [@SkalskiP](https://github.com/SkalskiP)
- [heatmap and track](./heatmap_and_track/) by [@HinePo](https://github.com/HinePo)
## Contributing
We welcome contributions from the community in the form of examples, applications, and
guides. To contribute, please follow these steps:
1. Create a pull request (PR) with the `[Example]` prefix in the title, adding your
project folder to the `examples/` directory in the repository.
2. Confirm your project aligns with the following standards:
- Incorporates the `supervision` package.
- Provides a `README.md` file, detailing the instructions to execute the project.
- Showcases visual results, demonstrating the app's functionality.
- Avoids adding large assets or dependencies unless absolutely necessary.
- The contributor is expected to provide support for issues related to their
examples.
- In case the presented model has licensing complications, kindly specify them to
circumvent potential misunderstandings.
For inquiries or concerns about these prerequisites, feel free to raise a PR. We are
committed to assist and guide you.

View File

@ -1,5 +1,8 @@
# count people in zone
[![Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/roboflow-ai/notebooks/blob/main/notebooks/how-to-detect-and-count-objects-in-polygon-zone.ipynb)
[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://www.youtube.com/watch?v=l_kf9CfZ_8M)
## 👋 hello
This demo is a video analysis tool that counts and highlights objects in specific zones

View File

@ -11,7 +11,7 @@ from tqdm import tqdm
import supervision as sv
COLORS = sv.ColorPalette.default()
COLORS = sv.ColorPalette.DEFAULT
def load_zones_config(file_path: str) -> List[np.ndarray]:
@ -38,15 +38,15 @@ def initiate_annotators(
) -> Tuple[
List[sv.PolygonZone], List[sv.PolygonZoneAnnotator], List[sv.BoundingBoxAnnotator]
]:
line_thickness = sv.calculate_dynamic_line_thickness(resolution_wh=resolution_wh)
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=resolution_wh)
line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=resolution_wh)
text_scale = sv.calculate_optimal_text_scale(resolution_wh=resolution_wh)
zones = []
zone_annotators = []
box_annotators = []
for index, polygon in enumerate(polygons):
zone = sv.PolygonZone(polygon=polygon, frame_resolution_wh=resolution_wh)
zone = sv.PolygonZone(polygon=polygon)
zone_annotator = sv.PolygonZoneAnnotator(
zone=zone,
color=COLORS.by_idx(index),

View File

@ -1,5 +1,5 @@
gdown
inference
supervision
inference==0.9.17
supervision>=0.20.0
tqdm
ultralytics

View File

@ -9,7 +9,7 @@ from ultralytics import YOLO
import supervision as sv
COLORS = sv.ColorPalette.default()
COLORS = sv.ColorPalette.DEFAULT
def load_zones_config(file_path: str) -> List[np.ndarray]:
@ -36,15 +36,15 @@ def initiate_annotators(
) -> Tuple[
List[sv.PolygonZone], List[sv.PolygonZoneAnnotator], List[sv.BoundingBoxAnnotator]
]:
line_thickness = sv.calculate_dynamic_line_thickness(resolution_wh=resolution_wh)
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=resolution_wh)
line_thickness = sv.calculate_optimal_line_thickness(resolution_wh=resolution_wh)
text_scale = sv.calculate_optimal_text_scale(resolution_wh=resolution_wh)
zones = []
zone_annotators = []
box_annotators = []
for index, polygon in enumerate(polygons):
zone = sv.PolygonZone(polygon=polygon, frame_resolution_wh=resolution_wh)
zone = sv.PolygonZone(polygon=polygon)
zone_annotator = sv.PolygonZoneAnnotator(
zone=zone,
color=COLORS.by_idx(index),

View File

@ -1,2 +1,2 @@
supervision[assets]
supervision[assets]==0.19.0
ultralytics

View File

@ -98,11 +98,11 @@ if __name__ == "__main__":
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
)
thickness = sv.calculate_dynamic_line_thickness(
thickness = sv.calculate_optimal_line_thickness(
resolution_wh=video_info.resolution_wh
)
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh)
box_annotator = sv.BoxAnnotator(thickness=thickness)
label_annotator = sv.LabelAnnotator(
text_scale=text_scale,
text_thickness=thickness,
@ -116,9 +116,7 @@ if __name__ == "__main__":
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
polygon_zone = sv.PolygonZone(
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
)
polygon_zone = sv.PolygonZone(polygon=SOURCE)
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
@ -156,7 +154,7 @@ if __name__ == "__main__":
annotated_frame = trace_annotator.annotate(
scene=annotated_frame, detections=detections
)
annotated_frame = bounding_box_annotator.annotate(
annotated_frame = box_annotator.annotate(
scene=annotated_frame, detections=detections
)
annotated_frame = label_annotator.annotate(

View File

@ -1,6 +1,6 @@
supervision==0.18.0rc1
tqdm==4.66.1
supervision>=0.20.0
tqdm==4.66.3
requests
ultralytics==8.0.237
super-gradients==3.5.0
inference==0.9.8
inference==0.9.17

View File

@ -76,11 +76,11 @@ if __name__ == "__main__":
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
)
thickness = sv.calculate_dynamic_line_thickness(
thickness = sv.calculate_optimal_line_thickness(
resolution_wh=video_info.resolution_wh
)
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh)
box_annotator = sv.BoxAnnotator(thickness=thickness)
label_annotator = sv.LabelAnnotator(
text_scale=text_scale,
text_thickness=thickness,
@ -94,9 +94,7 @@ if __name__ == "__main__":
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
polygon_zone = sv.PolygonZone(
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
)
polygon_zone = sv.PolygonZone(polygon=SOURCE)
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
@ -134,7 +132,7 @@ if __name__ == "__main__":
annotated_frame = trace_annotator.annotate(
scene=annotated_frame, detections=detections
)
annotated_frame = bounding_box_annotator.annotate(
annotated_frame = box_annotator.annotate(
scene=annotated_frame, detections=detections
)
annotated_frame = label_annotator.annotate(

View File

@ -77,11 +77,11 @@ if __name__ == "__main__":
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
)
thickness = sv.calculate_dynamic_line_thickness(
thickness = sv.calculate_optimal_line_thickness(
resolution_wh=video_info.resolution_wh
)
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
text_scale = sv.calculate_optimal_text_scale(resolution_wh=video_info.resolution_wh)
box_annotator = sv.BoxAnnotator(thickness=thickness)
label_annotator = sv.LabelAnnotator(
text_scale=text_scale,
text_thickness=thickness,
@ -95,9 +95,7 @@ if __name__ == "__main__":
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
polygon_zone = sv.PolygonZone(
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
)
polygon_zone = sv.PolygonZone(polygon=SOURCE)
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
@ -134,7 +132,7 @@ if __name__ == "__main__":
annotated_frame = trace_annotator.annotate(
scene=annotated_frame, detections=detections
)
annotated_frame = bounding_box_annotator.annotate(
annotated_frame = box_annotator.annotate(
scene=annotated_frame, detections=detections
)
annotated_frame = label_annotator.annotate(

9
examples/time_in_zone/.gitignore vendored Normal file
View File

@ -0,0 +1,9 @@
data/
venv*/
*.pt
*.pth
*.mp4
*.mov
*.png
*.jpg
*.jpeg

View File

@ -0,0 +1,264 @@
# time in zone
[![YouTube](https://badges.aleen42.com/src/youtube.svg)](https://www.youtube.com/watch?v=hAWpsIuem10)
## 👋 hello
Practical demonstration on leveraging computer vision for analyzing wait times and
monitoring the duration that objects or individuals spend in predefined areas of video
frames. This example project, perfect for retail analytics or traffic management
applications.
https://github.com/roboflow/supervision/assets/26109316/d051cc8a-dd15-41d4-aa36-d38b86334c39
## 💻 install
- clone repository and navigate to example directory
```bash
git clone https://github.com/roboflow/supervision.git
cd supervision/examples/time_in_zone
```
- setup python environment and activate it [optional]
```bash
python3 -m venv venv
source venv/bin/activate
```
- install required dependencies
```bash
pip install -r requirements.txt
```
## 🛠 scripts
### `download_from_youtube`
This script allows you to download a video from YouTube.
- `--url`: The full URL of the YouTube video you wish to download.
- `--output_path` (optional): Specifies the directory where the video will be saved.
- `--file_name` (optional): Sets the name of the saved video file.
```bash
python scripts/download_from_youtube.py \
--url "https://www.youtube.com/watch?v=-8zyEwAa50Q" \
--output_path "data/checkout" \
--file_name "video.mp4"
```
```bash
python scripts/download_from_youtube.py \
--url "https://www.youtube.com/watch?v=MNn9qKG2UFI" \
--output_path "data/traffic" \
--file_name "video.mp4"
```
### `stream_from_file`
This script allows you to stream video files from a directory. It's an awesome way to
mock a live video stream for local testing. Video will be streamed in a loop under
`rtsp://localhost:8554/live0.stream` URL. This script requires docker to be installed.
- `--video_directory`: Directory containing video files to stream.
- `--number_of_streams`: Number of video files to stream.
```bash
python scripts/stream_from_file.py \
--video_directory "data/checkout" \
--number_of_streams 1
```
```bash
python scripts/stream_from_file.py \
--video_directory "data/traffic" \
--number_of_streams 1
```
### `draw_zones`
If you want to test zone time in zone analysis on your own video, you can use this
script to design custom zones and save results as a JSON file. The script will open a
window where you can draw polygons on the source image or video file. The polygons will
be saved as a JSON file.
- `--source_path`: Path to the source image or video file for drawing polygons.
- `--zone_configuration_path`: Path where the polygon annotations will be saved as a JSON file.
- `enter` - finish drawing the current polygon.
- `escape` - cancel drawing the current polygon.
- `q` - quit the drawing window.
- `s` - save zone configuration to a JSON file.
```bash
python scripts/draw_zones.py \
--source_path "data/checkout/video.mp4" \
--zone_configuration_path "data/checkout/config.json"
```
```bash
python scripts/draw_zones.py \
--source_path "data/traffic/video.mp4" \
--zone_configuration_path "data/traffic/config.json"
```
https://github.com/roboflow/supervision/assets/26109316/9d514c9e-2a61-418b-ae49-6ac1ad6ae5ac
## 🎬 video & stream processing
### `inference_file_example`
Script to run object detection on a video file using the Roboflow Inference model.
- `--zone_configuration_path`: Path to the zone configuration JSON file.
- `--source_video_path`: Path to the source video file.
- `--model_id`: Roboflow model ID.
- `--classes`: List of class IDs to track. If empty, all classes are tracked.
- `--confidence_threshold`: Confidence level for detections (`0` to `1`). Default is `0.3`.
- `--iou_threshold`: IOU threshold for non-max suppression. Default is `0.7`.
```bash
python inference_file_example.py \
--zone_configuration_path "data/checkout/config.json" \
--source_video_path "data/checkout/video.mp4" \
--model_id "yolov8x-640" \
--classes 0 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
https://github.com/roboflow/supervision/assets/26109316/d051cc8a-dd15-41d4-aa36-d38b86334c39
```bash
python inference_file_example.py \
--zone_configuration_path "data/traffic/config.json" \
--source_video_path "data/traffic/video.mp4" \
--model_id "yolov8x-640" \
--classes 2 5 6 7 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
https://github.com/roboflow/supervision/assets/26109316/5ec896d7-4b39-4426-8979-11e71666878b
### `inference_stream_example`
Script to run object detection on a video stream using the Roboflow Inference model.
- `--zone_configuration_path`: Path to the zone configuration JSON file.
- `--rtsp_url`: Complete RTSP URL for the video stream.
- `--model_id`: Roboflow model ID.
- `--classes`: List of class IDs to track. If empty, all classes are tracked.
- `--confidence_threshold`: Confidence level for detections (`0` to `1`). Default is `0.3`.
- `--iou_threshold`: IOU threshold for non-max suppression. Default is `0.7`.
```bash
python inference_stream_example.py \
--zone_configuration_path "data/checkout/config.json" \
--rtsp_url "rtsp://localhost:8554/live0.stream" \
--model_id "yolov8x-640" \
--classes 0 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
```bash
python inference_stream_example.py \
--zone_configuration_path "data/traffic/config.json" \
--rtsp_url "rtsp://localhost:8554/live0.stream" \
--model_id "yolov8x-640" \
--classes 2 5 6 7 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
<details>
<summary>👉 show ultralytics examples</summary>
### `ultralytics_file_example`
Script to run object detection on a video file using the Ultralytics YOLOv8 model.
- `--zone_configuration_path`: Path to the zone configuration JSON file.
- `--source_video_path`: Path to the source video file.
- `--weights`: Path to the model weights file. Default is `'yolov8s.pt'`.
- `--device`: Computation device (`'cpu'`, `'mps'` or `'cuda'`). Default is `'cpu'`.
- `--classes`: List of class IDs to track. If empty, all classes are tracked.
- `--confidence_threshold`: Confidence level for detections (`0` to `1`). Default is `0.3`.
- `--iou_threshold`: IOU threshold for non-max suppression. Default is `0.7`.
```bash
python ultralytics_file_example.py \
--zone_configuration_path "data/checkout/config.json" \
--source_video_path "data/checkout/video.mp4" \
--weights "yolov8x.pt" \
--device "cpu" \
--classes 0 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
```bash
python ultralytics_file_example.py \
--zone_configuration_path "data/traffic/config.json" \
--source_video_path "data/traffic/video.mp4" \
--weights "yolov8x.pt" \
--device "cpu" \
--classes 2 5 6 7 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
### `ultralytics_stream_example`
Script to run object detection on a video stream using the Ultralytics YOLOv8 model.
- `--zone_configuration_path`: Path to the zone configuration JSON file.
- `--rtsp_url`: Complete RTSP URL for the video stream.
- `--weights`: Path to the model weights file. Default is `'yolov8s.pt'`.
- `--device`: Computation device (`'cpu'`, `'mps'` or `'cuda'`). Default is `'cpu'`.
- `--classes`: List of class IDs to track. If empty, all classes are tracked.
- `--confidence_threshold`: Confidence level for detections (`0` to `1`). Default is `0.3`.
- `--iou_threshold`: IOU threshold for non-max suppression. Default is `0.7`.
```bash
python ultralytics_stream_example.py \
--zone_configuration_path "data/checkout/config.json" \
--rtsp_url "rtsp://localhost:8554/live0.stream" \
--weights "yolov8x.pt" \
--device "cpu" \
--classes 0 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
```bash
python ultralytics_stream_example.py \
--zone_configuration_path "data/traffic/config.json" \
--rtsp_url "rtsp://localhost:8554/live0.stream" \
--weights "yolov8x.pt" \
--device "cpu" \
--classes 2 5 6 7 \
--confidence_threshold 0.3 \
--iou_threshold 0.7
```
</details>
## © license
This demo integrates two main components, each with its own licensing:
- ultralytics: The object detection model used in this demo, YOLOv8, is distributed
under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE).
You can find more details about this license here.
- supervision: The analytics code that powers the zone-based analysis in this demo is
based on the Supervision library, which is licensed under the
[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
makes the Supervision part of the code fully open source and freely usable in your
projects.

View File

@ -0,0 +1,128 @@
import argparse
from typing import List
import cv2
import numpy as np
from inference import get_model
from utils.general import find_in_list, load_zones_config
from utils.timers import FPSBasedTimer
import supervision as sv
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
COLOR_ANNOTATOR = sv.ColorAnnotator(color=COLORS)
LABEL_ANNOTATOR = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.from_hex("#000000")
)
def main(
source_video_path: str,
zone_configuration_path: str,
model_id: str,
confidence: float,
iou: float,
classes: List[int],
) -> None:
model = get_model(model_id=model_id)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
frames_generator = sv.get_video_frames_generator(source_video_path)
polygons = load_zones_config(file_path=zone_configuration_path)
zones = [
sv.PolygonZone(
polygon=polygon,
triggering_anchors=(sv.Position.CENTER,),
)
for polygon in polygons
]
timers = [FPSBasedTimer(video_info.fps) for _ in zones]
for frame in frames_generator:
results = model.infer(frame, confidence=confidence, iou_threshold=iou)[0]
detections = sv.Detections.from_inference(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = tracker.update_with_detections(detections)
annotated_frame = frame.copy()
for idx, zone in enumerate(zones):
annotated_frame = sv.draw_polygon(
scene=annotated_frame, polygon=zone.polygon, color=COLORS.by_idx(idx)
)
detections_in_zone = detections[zone.trigger(detections)]
time_in_zone = timers[idx].tick(detections_in_zone)
custom_color_lookup = np.full(detections_in_zone.class_id.shape, idx)
annotated_frame = COLOR_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
custom_color_lookup=custom_color_lookup,
)
labels = [
f"#{tracker_id} {int(time // 60):02d}:{int(time % 60):02d}"
for tracker_id, time in zip(detections_in_zone.tracker_id, time_in_zone)
]
annotated_frame = LABEL_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
labels=labels,
custom_color_lookup=custom_color_lookup,
)
cv2.imshow("Processed Video", annotated_frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
cv2.destroyAllWindows()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using video file."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--source_video_path",
type=str,
required=True,
help="Path to the source video file.",
)
parser.add_argument(
"--model_id", type=str, default="yolov8s-640", help="Roboflow model ID."
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
main(
source_video_path=args.source_video_path,
zone_configuration_path=args.zone_configuration_path,
model_id=args.model_id,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)

View File

@ -0,0 +1,138 @@
import argparse
from typing import List
import cv2
import numpy as np
from inference import get_model
from utils.general import find_in_list, get_stream_frames_generator, load_zones_config
from utils.timers import ClockBasedTimer
import supervision as sv
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
COLOR_ANNOTATOR = sv.ColorAnnotator(color=COLORS)
LABEL_ANNOTATOR = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.from_hex("#000000")
)
def main(
rtsp_url: str,
zone_configuration_path: str,
model_id: str,
confidence: float,
iou: float,
classes: List[int],
) -> None:
model = get_model(model_id=model_id)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
frames_generator = get_stream_frames_generator(rtsp_url=rtsp_url)
fps_monitor = sv.FPSMonitor()
polygons = load_zones_config(file_path=zone_configuration_path)
zones = [
sv.PolygonZone(
polygon=polygon,
triggering_anchors=(sv.Position.CENTER,),
)
for polygon in polygons
]
timers = [ClockBasedTimer() for _ in zones]
for frame in frames_generator:
fps_monitor.tick()
fps = fps_monitor.fps
results = model.infer(frame, confidence=confidence, iou_threshold=iou)[0]
detections = sv.Detections.from_inference(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = tracker.update_with_detections(detections)
annotated_frame = frame.copy()
annotated_frame = sv.draw_text(
scene=annotated_frame,
text=f"{fps:.1f}",
text_anchor=sv.Point(40, 30),
background_color=sv.Color.from_hex("#A351FB"),
text_color=sv.Color.from_hex("#000000"),
)
for idx, zone in enumerate(zones):
annotated_frame = sv.draw_polygon(
scene=annotated_frame, polygon=zone.polygon, color=COLORS.by_idx(idx)
)
detections_in_zone = detections[zone.trigger(detections)]
time_in_zone = timers[idx].tick(detections_in_zone)
custom_color_lookup = np.full(detections_in_zone.class_id.shape, idx)
annotated_frame = COLOR_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
custom_color_lookup=custom_color_lookup,
)
labels = [
f"#{tracker_id} {int(time // 60):02d}:{int(time % 60):02d}"
for tracker_id, time in zip(detections_in_zone.tracker_id, time_in_zone)
]
annotated_frame = LABEL_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
labels=labels,
custom_color_lookup=custom_color_lookup,
)
cv2.imshow("Processed Video", annotated_frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
cv2.destroyAllWindows()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--model_id", type=str, default="yolov8s-640", help="Roboflow model ID."
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
model_id=args.model_id,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)

View File

@ -0,0 +1,153 @@
import argparse
from typing import List
import cv2
import numpy as np
from inference import InferencePipeline
from inference.core.interfaces.camera.entities import VideoFrame
from utils.general import find_in_list, load_zones_config
from utils.timers import ClockBasedTimer
import supervision as sv
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
COLOR_ANNOTATOR = sv.ColorAnnotator(color=COLORS)
LABEL_ANNOTATOR = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.from_hex("#000000")
)
class CustomSink:
def __init__(self, zone_configuration_path: str, classes: List[int]):
self.classes = classes
self.tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
self.fps_monitor = sv.FPSMonitor()
self.polygons = load_zones_config(file_path=zone_configuration_path)
self.timers = [ClockBasedTimer() for _ in self.polygons]
self.zones = [
sv.PolygonZone(
polygon=polygon,
triggering_anchors=(sv.Position.CENTER,),
)
for polygon in self.polygons
]
def on_prediction(self, result: dict, frame: VideoFrame) -> None:
self.fps_monitor.tick()
fps = self.fps_monitor.fps
detections = sv.Detections.from_inference(result)
detections = detections[find_in_list(detections.class_id, self.classes)]
detections = self.tracker.update_with_detections(detections)
annotated_frame = frame.image.copy()
annotated_frame = sv.draw_text(
scene=annotated_frame,
text=f"{fps:.1f}",
text_anchor=sv.Point(40, 30),
background_color=sv.Color.from_hex("#A351FB"),
text_color=sv.Color.from_hex("#000000"),
)
for idx, zone in enumerate(self.zones):
annotated_frame = sv.draw_polygon(
scene=annotated_frame, polygon=zone.polygon, color=COLORS.by_idx(idx)
)
detections_in_zone = detections[zone.trigger(detections)]
time_in_zone = self.timers[idx].tick(detections_in_zone)
custom_color_lookup = np.full(detections_in_zone.class_id.shape, idx)
annotated_frame = COLOR_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
custom_color_lookup=custom_color_lookup,
)
labels = [
f"#{tracker_id} {int(time // 60):02d}:{int(time % 60):02d}"
for tracker_id, time in zip(detections_in_zone.tracker_id, time_in_zone)
]
annotated_frame = LABEL_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
labels=labels,
custom_color_lookup=custom_color_lookup,
)
cv2.imshow("Processed Video", annotated_frame)
cv2.waitKey(1)
def main(
rtsp_url: str,
zone_configuration_path: str,
model_id: str,
confidence: float,
iou: float,
classes: List[int],
) -> None:
sink = CustomSink(zone_configuration_path=zone_configuration_path, classes=classes)
pipeline = InferencePipeline.init(
model_id=model_id,
video_reference=rtsp_url,
on_prediction=sink.on_prediction,
confidence=confidence,
iou_threshold=iou,
)
pipeline.start()
try:
pipeline.join()
except KeyboardInterrupt:
pipeline.terminate()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--model_id", type=str, default="yolov8s-640", help="Roboflow model ID."
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
model_id=args.model_id,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)

View File

@ -0,0 +1,5 @@
opencv-python
supervision>=0.20.0
ultralytics
inference==0.9.17
pytube

View File

@ -0,0 +1,46 @@
import argparse
import os
from typing import Optional
from pytube import YouTube
def main(url: str, output_path: Optional[str], file_name: Optional[str]) -> None:
yt = YouTube(url)
stream = yt.streams.get_highest_resolution()
if not os.path.exists(output_path):
os.makedirs(output_path)
stream.download(output_path=output_path, filename=file_name)
final_name = file_name if file_name else yt.title
final_path = output_path if output_path else "current directory"
print(f"Download completed! Video saved as '{final_name}' in '{final_path}'.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Download a specific YouTube video by providing its URL."
)
parser.add_argument(
"--url",
type=str,
required=True,
help="The full URL of the YouTube video you wish to download.",
)
parser.add_argument(
"--output_path",
type=str,
default="data/source",
required=False,
help="Optional. Specifies the directory where the video will be saved.",
)
parser.add_argument(
"--file_name",
type=str,
default="video.mp4",
required=False,
help="Optional. Sets the name of the saved video file.",
)
args = parser.parse_args()
main(url=args.url, output_path=args.output_path, file_name=args.file_name)

View File

@ -0,0 +1,176 @@
import argparse
import json
import os
from typing import Any, Optional, Tuple
import cv2
import numpy as np
import supervision as sv
KEY_ENTER = 13
KEY_NEWLINE = 10
KEY_ESCAPE = 27
KEY_QUIT = ord("q")
KEY_SAVE = ord("s")
THICKNESS = 2
COLORS = sv.ColorPalette.DEFAULT
WINDOW_NAME = "Draw Zones"
POLYGONS = [[]]
current_mouse_position: Optional[Tuple[int, int]] = None
def resolve_source(source_path: str) -> Optional[np.ndarray]:
if not os.path.exists(source_path):
return None
image = cv2.imread(source_path)
if image is not None:
return image
frame_generator = sv.get_video_frames_generator(source_path=source_path)
frame = next(frame_generator)
return frame
def mouse_event(event: int, x: int, y: int, flags: int, param: Any) -> None:
global current_mouse_position
if event == cv2.EVENT_MOUSEMOVE:
current_mouse_position = (x, y)
elif event == cv2.EVENT_LBUTTONDOWN:
POLYGONS[-1].append((x, y))
def redraw(image: np.ndarray, original_image: np.ndarray) -> None:
global POLYGONS, current_mouse_position
image[:] = original_image.copy()
for idx, polygon in enumerate(POLYGONS):
color = (
COLORS.by_idx(idx).as_bgr()
if idx < len(POLYGONS) - 1
else sv.Color.WHITE.as_bgr()
)
if len(polygon) > 1:
for i in range(1, len(polygon)):
cv2.line(
img=image,
pt1=polygon[i - 1],
pt2=polygon[i],
color=color,
thickness=THICKNESS,
)
if idx < len(POLYGONS) - 1:
cv2.line(
img=image,
pt1=polygon[-1],
pt2=polygon[0],
color=color,
thickness=THICKNESS,
)
if idx == len(POLYGONS) - 1 and current_mouse_position is not None and polygon:
cv2.line(
img=image,
pt1=polygon[-1],
pt2=current_mouse_position,
color=color,
thickness=THICKNESS,
)
cv2.imshow(WINDOW_NAME, image)
def close_and_finalize_polygon(image: np.ndarray, original_image: np.ndarray) -> None:
if len(POLYGONS[-1]) > 2:
cv2.line(
img=image,
pt1=POLYGONS[-1][-1],
pt2=POLYGONS[-1][0],
color=COLORS.by_idx(0).as_bgr(),
thickness=THICKNESS,
)
POLYGONS.append([])
image[:] = original_image.copy()
redraw_polygons(image)
cv2.imshow(WINDOW_NAME, image)
def redraw_polygons(image: np.ndarray) -> None:
for idx, polygon in enumerate(POLYGONS[:-1]):
if len(polygon) > 1:
color = COLORS.by_idx(idx).as_bgr()
for i in range(len(polygon) - 1):
cv2.line(
img=image,
pt1=polygon[i],
pt2=polygon[i + 1],
color=color,
thickness=THICKNESS,
)
cv2.line(
img=image,
pt1=polygon[-1],
pt2=polygon[0],
color=color,
thickness=THICKNESS,
)
def save_polygons_to_json(polygons, target_path):
data_to_save = polygons if polygons[-1] else polygons[:-1]
with open(target_path, "w") as f:
json.dump(data_to_save, f)
def main(source_path: str, zone_configuration_path: str) -> None:
global current_mouse_position
original_image = resolve_source(source_path=source_path)
if original_image is None:
print("Failed to load source image.")
return
image = original_image.copy()
cv2.imshow(WINDOW_NAME, image)
cv2.setMouseCallback(WINDOW_NAME, mouse_event, image)
while True:
key = cv2.waitKey(1) & 0xFF
if key == KEY_ENTER or key == KEY_NEWLINE:
close_and_finalize_polygon(image, original_image)
elif key == KEY_ESCAPE:
POLYGONS[-1] = []
current_mouse_position = None
elif key == KEY_SAVE:
save_polygons_to_json(POLYGONS, zone_configuration_path)
print(f"Polygons saved to {zone_configuration_path}")
break
redraw(image, original_image)
if key == KEY_QUIT:
break
cv2.destroyAllWindows()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Interactively draw polygons on images or video frames and save "
"the annotations."
)
parser.add_argument(
"--source_path",
type=str,
required=True,
help="Path to the source image or video file for drawing polygons.",
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path where the polygon annotations will be saved as a JSON file.",
)
arguments = parser.parse_args()
main(
source_path=arguments.source_path,
zone_configuration_path=arguments.zone_configuration_path,
)

View File

@ -0,0 +1,104 @@
import argparse
import os
import subprocess
import tempfile
from glob import glob
from threading import Thread
import yaml
SERVER_CONFIG = {"protocols": ["tcp"], "paths": {"all": {"source": "publisher"}}}
BASE_STREAM_URL = "rtsp://localhost:8554/live"
def main(video_directory: str, number_of_streams: int) -> None:
video_files = find_video_files_in_directory(video_directory, number_of_streams)
try:
with tempfile.TemporaryDirectory() as temporary_directory:
config_file_path = create_server_config_file(temporary_directory)
run_rtsp_server(config_path=config_file_path)
stream_videos(video_files)
finally:
stop_rtsp_server()
def find_video_files_in_directory(directory: str, limit: int) -> list:
video_formats = ["*.mp4", "*.webm"]
video_paths = []
for video_format in video_formats:
video_paths.extend(glob(os.path.join(directory, video_format)))
return video_paths[:limit]
def create_server_config_file(directory: str) -> str:
config_path = os.path.join(directory, "rtsp-simple-server.yml")
with open(config_path, "w") as config_file:
yaml.dump(SERVER_CONFIG, config_file)
return config_path
def run_rtsp_server(config_path: str) -> None:
command = (
"docker run --rm --name rtsp_server -d -v "
f"{config_path}:/rtsp-simple-server.yml -p 8554:8554 "
"aler9/rtsp-simple-server:v1.3.0"
)
if run_command(command.split()) != 0:
raise RuntimeError("Could not start the RTSP server!")
def stop_rtsp_server() -> None:
run_command("docker kill rtsp_server".split())
def stream_videos(video_files: list) -> None:
threads = []
for index, video_file in enumerate(video_files):
stream_url = f"{BASE_STREAM_URL}{index}.stream"
print(f"Streaming {video_file} under {stream_url}")
thread = stream_video_to_url(video_file, stream_url)
threads.append(thread)
for thread in threads:
thread.join()
def stream_video_to_url(video_path: str, stream_url: str) -> Thread:
command = (
f"ffmpeg -re -stream_loop -1 -i {video_path} "
f"-f rtsp -rtsp_transport tcp {stream_url}"
)
return run_command_in_thread(command.split())
def run_command_in_thread(command: list) -> Thread:
thread = Thread(target=run_command, args=(command,))
thread.start()
return thread
def run_command(command: list) -> int:
process = subprocess.run(command)
return process.returncode
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Script to stream videos using RTSP protocol."
)
parser.add_argument(
"--video_directory",
type=str,
required=True,
help="Directory containing video files to stream.",
)
parser.add_argument(
"--number_of_streams",
type=int,
default=6,
help="Number of video files to stream.",
)
arguments = parser.parse_args()
main(
video_directory=arguments.video_directory,
number_of_streams=arguments.number_of_streams,
)

View File

@ -0,0 +1,140 @@
import argparse
from typing import List
import cv2
import numpy as np
from ultralytics import YOLO
from utils.general import find_in_list, load_zones_config
from utils.timers import FPSBasedTimer
import supervision as sv
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
COLOR_ANNOTATOR = sv.ColorAnnotator(color=COLORS)
LABEL_ANNOTATOR = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.from_hex("#000000")
)
def main(
source_video_path: str,
zone_configuration_path: str,
weights: str,
device: str,
confidence: float,
iou: float,
classes: List[int],
) -> None:
model = YOLO(weights)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
frames_generator = sv.get_video_frames_generator(source_video_path)
polygons = load_zones_config(file_path=zone_configuration_path)
zones = [
sv.PolygonZone(
polygon=polygon,
triggering_anchors=(sv.Position.CENTER,),
)
for polygon in polygons
]
timers = [FPSBasedTimer(video_info.fps) for _ in zones]
for frame in frames_generator:
results = model(frame, verbose=False, device=device, conf=confidence)[0]
detections = sv.Detections.from_ultralytics(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = detections.with_nms(threshold=iou)
detections = tracker.update_with_detections(detections)
annotated_frame = frame.copy()
for idx, zone in enumerate(zones):
annotated_frame = sv.draw_polygon(
scene=annotated_frame, polygon=zone.polygon, color=COLORS.by_idx(idx)
)
detections_in_zone = detections[zone.trigger(detections)]
time_in_zone = timers[idx].tick(detections_in_zone)
custom_color_lookup = np.full(detections_in_zone.class_id.shape, idx)
annotated_frame = COLOR_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
custom_color_lookup=custom_color_lookup,
)
labels = [
f"#{tracker_id} {int(time // 60):02d}:{int(time % 60):02d}"
for tracker_id, time in zip(detections_in_zone.tracker_id, time_in_zone)
]
annotated_frame = LABEL_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
labels=labels,
custom_color_lookup=custom_color_lookup,
)
cv2.imshow("Processed Video", annotated_frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
cv2.destroyAllWindows()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using video file."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--source_video_path",
type=str,
required=True,
help="Path to the source video file.",
)
parser.add_argument(
"--weights",
type=str,
default="yolov8s.pt",
help="Path to the model weights file. Default is 'yolov8s.pt'.",
)
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Computation device ('cpu', 'mps' or 'cuda'). Default is 'cpu'.",
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
main(
source_video_path=args.source_video_path,
zone_configuration_path=args.zone_configuration_path,
weights=args.weights,
device=args.device,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)

View File

@ -0,0 +1,150 @@
import argparse
from typing import List
import cv2
import numpy as np
from ultralytics import YOLO
from utils.general import find_in_list, get_stream_frames_generator, load_zones_config
from utils.timers import ClockBasedTimer
import supervision as sv
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
COLOR_ANNOTATOR = sv.ColorAnnotator(color=COLORS)
LABEL_ANNOTATOR = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.from_hex("#000000")
)
def main(
rtsp_url: str,
zone_configuration_path: str,
weights: str,
device: str,
confidence: float,
iou: float,
classes: List[int],
) -> None:
model = YOLO(weights)
tracker = sv.ByteTrack(minimum_matching_threshold=0.5)
frames_generator = get_stream_frames_generator(rtsp_url=rtsp_url)
fps_monitor = sv.FPSMonitor()
polygons = load_zones_config(file_path=zone_configuration_path)
zones = [
sv.PolygonZone(
polygon=polygon,
triggering_anchors=(sv.Position.CENTER,),
)
for polygon in polygons
]
timers = [ClockBasedTimer() for _ in zones]
for frame in frames_generator:
fps_monitor.tick()
fps = fps_monitor.fps
results = model(frame, verbose=False, device=device, conf=confidence)[0]
detections = sv.Detections.from_ultralytics(results)
detections = detections[find_in_list(detections.class_id, classes)]
detections = detections.with_nms(threshold=iou)
detections = tracker.update_with_detections(detections)
annotated_frame = frame.copy()
annotated_frame = sv.draw_text(
scene=annotated_frame,
text=f"{fps:.1f}",
text_anchor=sv.Point(40, 30),
background_color=sv.Color.from_hex("#A351FB"),
text_color=sv.Color.from_hex("#000000"),
)
for idx, zone in enumerate(zones):
annotated_frame = sv.draw_polygon(
scene=annotated_frame, polygon=zone.polygon, color=COLORS.by_idx(idx)
)
detections_in_zone = detections[zone.trigger(detections)]
time_in_zone = timers[idx].tick(detections_in_zone)
custom_color_lookup = np.full(detections_in_zone.class_id.shape, idx)
annotated_frame = COLOR_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
custom_color_lookup=custom_color_lookup,
)
labels = [
f"#{tracker_id} {int(time // 60):02d}:{int(time % 60):02d}"
for tracker_id, time in zip(detections_in_zone.tracker_id, time_in_zone)
]
annotated_frame = LABEL_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
labels=labels,
custom_color_lookup=custom_color_lookup,
)
cv2.imshow("Processed Video", annotated_frame)
if cv2.waitKey(1) & 0xFF == ord("q"):
break
cv2.destroyAllWindows()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--weights",
type=str,
default="yolov8s.pt",
help="Path to the model weights file. Default is 'yolov8s.pt'.",
)
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Computation device ('cpu', 'mps' or 'cuda'). Default is 'cpu'.",
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
weights=args.weights,
device=args.device,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)

View File

@ -0,0 +1,168 @@
import argparse
from typing import List
import cv2
import numpy as np
from inference import InferencePipeline
from inference.core.interfaces.camera.entities import VideoFrame
from ultralytics import YOLO
from utils.general import find_in_list, load_zones_config
from utils.timers import ClockBasedTimer
import supervision as sv
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
COLOR_ANNOTATOR = sv.ColorAnnotator(color=COLORS)
LABEL_ANNOTATOR = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.from_hex("#000000")
)
class CustomSink:
def __init__(self, zone_configuration_path: str, classes: List[int]):
self.classes = classes
self.tracker = sv.ByteTrack(minimum_matching_threshold=0.8)
self.fps_monitor = sv.FPSMonitor()
self.polygons = load_zones_config(file_path=zone_configuration_path)
self.timers = [ClockBasedTimer() for _ in self.polygons]
self.zones = [
sv.PolygonZone(
polygon=polygon,
triggering_anchors=(sv.Position.CENTER,),
)
for polygon in self.polygons
]
def on_prediction(self, detections: sv.Detections, frame: VideoFrame) -> None:
self.fps_monitor.tick()
fps = self.fps_monitor.fps
detections = detections[find_in_list(detections.class_id, self.classes)]
detections = self.tracker.update_with_detections(detections)
annotated_frame = frame.image.copy()
annotated_frame = sv.draw_text(
scene=annotated_frame,
text=f"{fps:.1f}",
text_anchor=sv.Point(40, 30),
background_color=sv.Color.from_hex("#A351FB"),
text_color=sv.Color.from_hex("#000000"),
)
for idx, zone in enumerate(self.zones):
annotated_frame = sv.draw_polygon(
scene=annotated_frame, polygon=zone.polygon, color=COLORS.by_idx(idx)
)
detections_in_zone = detections[zone.trigger(detections)]
time_in_zone = self.timers[idx].tick(detections_in_zone)
custom_color_lookup = np.full(detections_in_zone.class_id.shape, idx)
annotated_frame = COLOR_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
custom_color_lookup=custom_color_lookup,
)
labels = [
f"#{tracker_id} {int(time // 60):02d}:{int(time % 60):02d}"
for tracker_id, time in zip(detections_in_zone.tracker_id, time_in_zone)
]
annotated_frame = LABEL_ANNOTATOR.annotate(
scene=annotated_frame,
detections=detections_in_zone,
labels=labels,
custom_color_lookup=custom_color_lookup,
)
cv2.imshow("Processed Video", annotated_frame)
cv2.waitKey(1)
def main(
rtsp_url: str,
zone_configuration_path: str,
weights: str,
device: str,
confidence: float,
iou: float,
classes: List[int],
) -> None:
model = YOLO(weights)
def inference_callback(frame: VideoFrame) -> sv.Detections:
results = model(frame.image, verbose=False, conf=confidence, device=device)[0]
return sv.Detections.from_ultralytics(results).with_nms(threshold=iou)
sink = CustomSink(zone_configuration_path=zone_configuration_path, classes=classes)
pipeline = InferencePipeline.init_with_custom_logic(
video_reference=rtsp_url,
on_video_frame=inference_callback,
on_prediction=sink.on_prediction,
)
pipeline.start()
try:
pipeline.join()
except KeyboardInterrupt:
pipeline.terminate()
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Calculating detections dwell time in zones, using RTSP stream."
)
parser.add_argument(
"--zone_configuration_path",
type=str,
required=True,
help="Path to the zone configuration JSON file.",
)
parser.add_argument(
"--rtsp_url",
type=str,
required=True,
help="Complete RTSP URL for the video stream.",
)
parser.add_argument(
"--weights",
type=str,
default="yolov8s.pt",
help="Path to the model weights file. Default is 'yolov8s.pt'.",
)
parser.add_argument(
"--device",
type=str,
default="cpu",
help="Computation device ('cpu', 'mps' or 'cuda'). Default is 'cpu'.",
)
parser.add_argument(
"--confidence_threshold",
type=float,
default=0.3,
help="Confidence level for detections (0 to 1). Default is 0.3.",
)
parser.add_argument(
"--iou_threshold",
default=0.7,
type=float,
help="IOU threshold for non-max suppression. Default is 0.7.",
)
parser.add_argument(
"--classes",
nargs="*",
type=int,
default=[],
help="List of class IDs to track. If empty, all classes are tracked.",
)
args = parser.parse_args()
main(
rtsp_url=args.rtsp_url,
zone_configuration_path=args.zone_configuration_path,
weights=args.weights,
device=args.device,
confidence=args.confidence_threshold,
iou=args.iou_threshold,
classes=args.classes,
)

View File

View File

@ -0,0 +1,66 @@
import json
from typing import Generator, List
import cv2
import numpy as np
def load_zones_config(file_path: str) -> List[np.ndarray]:
"""
Load polygon zone configurations from a JSON file.
This function reads a JSON file which contains polygon coordinates, and
converts them into a list of NumPy arrays. Each polygon is represented as
a NumPy array of coordinates.
Args:
file_path (str): The path to the JSON configuration file.
Returns:
List[np.ndarray]: A list of polygons, each represented as a NumPy array.
"""
with open(file_path, "r") as file:
data = json.load(file)
return [np.array(polygon, np.int32) for polygon in data]
def find_in_list(array: np.ndarray, search_list: List[int]) -> np.ndarray:
"""Determines if elements of a numpy array are present in a list.
Args:
array (np.ndarray): The numpy array of integers to check.
search_list (List[int]): The list of integers to search within.
Returns:
np.ndarray: A numpy array of booleans, where each boolean indicates whether
the corresponding element in `array` is found in `search_list`.
"""
if not search_list:
return np.ones(array.shape, dtype=bool)
else:
return np.isin(array, search_list)
def get_stream_frames_generator(rtsp_url: str) -> Generator[np.ndarray, None, None]:
"""
Generator function to yield frames from an RTSP stream.
Args:
rtsp_url (str): URL of the RTSP video stream.
Yields:
np.ndarray: The next frame from the video stream.
"""
cap = cv2.VideoCapture(rtsp_url)
if not cap.isOpened():
raise Exception("Error: Could not open video stream.")
try:
while True:
ret, frame = cap.read()
if not ret:
print("End of stream or error reading frame.")
break
yield frame
finally:
cap.release()

View File

@ -0,0 +1,88 @@
from datetime import datetime
from typing import Dict
import numpy as np
import supervision as sv
class FPSBasedTimer:
"""
A timer that calculates the duration each object has been detected based on frames
per second (FPS).
Attributes:
fps (int): The frame rate of the video stream, used to calculate time durations.
frame_id (int): The current frame number in the sequence.
tracker_id2frame_id (Dict[int, int]): Maps each tracker's ID to the frame number
at which it was first detected.
"""
def __init__(self, fps: int = 30) -> None:
"""Initializes the FPSBasedTimer with the specified frames per second rate.
Args:
fps (int, optional): The frame rate of the video stream. Defaults to 30.
"""
self.fps = fps
self.frame_id = 0
self.tracker_id2frame_id: Dict[int, int] = {}
def tick(self, detections: sv.Detections) -> np.ndarray:
"""Processes the current frame, updating time durations for each tracker.
Args:
detections: The detections for the current frame, including tracker IDs.
Returns:
np.ndarray: Time durations (in seconds) for each detected tracker, since
their first detection.
"""
self.frame_id += 1
times = []
for tracker_id in detections.tracker_id:
self.tracker_id2frame_id.setdefault(tracker_id, self.frame_id)
start_frame_id = self.tracker_id2frame_id[tracker_id]
time_duration = (self.frame_id - start_frame_id) / self.fps
times.append(time_duration)
return np.array(times)
class ClockBasedTimer:
"""
A timer that calculates the duration each object has been detected based on the
system clock.
Attributes:
tracker_id2start_time (Dict[int, datetime]): Maps each tracker's ID to the
datetime when it was first detected.
"""
def __init__(self) -> None:
"""Initializes the ClockBasedTimer."""
self.tracker_id2start_time: Dict[int, datetime] = {}
def tick(self, detections: sv.Detections) -> np.ndarray:
"""Processes the current frame, updating time durations for each tracker.
Args:
detections: The detections for the current frame, including tracker IDs.
Returns:
np.ndarray: Time durations (in seconds) for each detected tracker, since
their first detection.
"""
current_time = datetime.now()
times = []
for tracker_id in detections.tracker_id:
self.tracker_id2start_time.setdefault(tracker_id, current_time)
start_time = self.tracker_id2start_time[tracker_id]
time_duration = (current_time - start_time).total_seconds()
times.append(time_duration)
return np.array(times)

View File

@ -1,4 +1,4 @@
inference
supervision
inference==0.9.17
supervision==0.19.0
tqdm
ultralytics

View File

@ -1,6 +1,6 @@
import argparse
import os
from typing import Dict, List, Set, Tuple
from typing import Dict, Iterable, List, Set
import cv2
import numpy as np
@ -9,7 +9,8 @@ from tqdm import tqdm
import supervision as sv
COLORS = sv.ColorPalette.default()
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
ZONE_IN_POLYGONS = [
np.array([[592, 282], [900, 282], [900, 82], [592, 82]]),
@ -59,14 +60,12 @@ class DetectionsManager:
def initiate_polygon_zones(
polygons: List[np.ndarray],
frame_resolution_wh: Tuple[int, int],
triggering_position: sv.Position = sv.Position.CENTER,
triggering_anchors: Iterable[sv.Position] = [sv.Position.CENTER],
) -> List[sv.PolygonZone]:
return [
sv.PolygonZone(
polygon=polygon,
frame_resolution_wh=frame_resolution_wh,
triggering_position=triggering_position,
triggering_anchors=triggering_anchors,
)
for polygon in polygons
]
@ -91,14 +90,13 @@ class VideoProcessor:
self.tracker = sv.ByteTrack()
self.video_info = sv.VideoInfo.from_video_path(source_video_path)
self.zones_in = initiate_polygon_zones(
ZONE_IN_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
)
self.zones_out = initiate_polygon_zones(
ZONE_OUT_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
)
self.zones_in = initiate_polygon_zones(ZONE_IN_POLYGONS, [sv.Position.CENTER])
self.zones_out = initiate_polygon_zones(ZONE_OUT_POLYGONS, [sv.Position.CENTER])
self.box_annotator = sv.BoxAnnotator(color=COLORS)
self.label_annotator = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.BLACK
)
self.trace_annotator = sv.TraceAnnotator(
color=COLORS, position=sv.Position.CENTER, trace_length=100, thickness=2
)
@ -136,7 +134,8 @@ class VideoProcessor:
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
annotated_frame = self.box_annotator.annotate(
annotated_frame = self.box_annotator.annotate(annotated_frame, detections)
annotated_frame = self.label_annotator.annotate(
annotated_frame, detections, labels
)
@ -167,7 +166,7 @@ class VideoProcessor:
detections_in_zones = []
detections_out_zones = []
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
for zone_in, zone_out in zip(self.zones_in, self.zones_out):
detections_in_zone = detections[zone_in.trigger(detections=detections)]
detections_in_zones.append(detections_in_zone)
detections_out_zone = detections[zone_out.trigger(detections=detections)]

View File

@ -1,5 +1,5 @@
gdown
inference
supervision
inference==0.9.17
supervision>=0.20.0
tqdm
ultralytics

View File

@ -1,5 +1,5 @@
import argparse
from typing import Dict, List, Set, Tuple
from typing import Dict, Iterable, List, Set
import cv2
import numpy as np
@ -8,7 +8,7 @@ from ultralytics import YOLO
import supervision as sv
COLORS = sv.ColorPalette.default()
COLORS = sv.ColorPalette.from_hex(["#E6194B", "#3CB44B", "#FFE119", "#3C76D1"])
ZONE_IN_POLYGONS = [
np.array([[592, 282], [900, 282], [900, 82], [592, 82]]),
@ -58,14 +58,12 @@ class DetectionsManager:
def initiate_polygon_zones(
polygons: List[np.ndarray],
frame_resolution_wh: Tuple[int, int],
triggering_position: sv.Position = sv.Position.CENTER,
triggering_anchors: Iterable[sv.Position] = [sv.Position.CENTER],
) -> List[sv.PolygonZone]:
return [
sv.PolygonZone(
polygon=polygon,
frame_resolution_wh=frame_resolution_wh,
triggering_position=triggering_position,
triggering_anchors=triggering_anchors,
)
for polygon in polygons
]
@ -89,14 +87,13 @@ class VideoProcessor:
self.tracker = sv.ByteTrack()
self.video_info = sv.VideoInfo.from_video_path(source_video_path)
self.zones_in = initiate_polygon_zones(
ZONE_IN_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
)
self.zones_out = initiate_polygon_zones(
ZONE_OUT_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
)
self.zones_in = initiate_polygon_zones(ZONE_IN_POLYGONS, [sv.Position.CENTER])
self.zones_out = initiate_polygon_zones(ZONE_OUT_POLYGONS, [sv.Position.CENTER])
self.box_annotator = sv.BoxAnnotator(color=COLORS)
self.label_annotator = sv.LabelAnnotator(
color=COLORS, text_color=sv.Color.BLACK
)
self.trace_annotator = sv.TraceAnnotator(
color=COLORS, position=sv.Position.CENTER, trace_length=100, thickness=2
)
@ -134,7 +131,8 @@ class VideoProcessor:
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
annotated_frame = self.box_annotator.annotate(
annotated_frame = self.box_annotator.annotate(annotated_frame, detections)
annotated_frame = self.label_annotator.annotate(
annotated_frame, detections, labels
)
@ -165,7 +163,7 @@ class VideoProcessor:
detections_in_zones = []
detections_out_zones = []
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
for zone_in, zone_out in zip(self.zones_in, self.zones_out):
detections_in_zone = detections[zone_in.trigger(detections=detections)]
detections_in_zones.append(detections_in_zone)
detections_out_zone = detections[zone_out.trigger(detections=detections)]

View File

@ -35,18 +35,27 @@ extra_css:
nav:
- Home: index.md
- How to:
- Supervision: index.md
- Learn:
- Detect and Annotate: how_to/detect_and_annotate.md
- Track Objects: how_to/track_objects.md
- Save Detections: how_to/save_detections.md
- Filter Detections: how_to/filter_detections.md
- API:
- Annotators: annotators.md
- Classifications:
- Core: classification/core.md
- Detections:
- Detect Small Objects: how_to/detect_small_objects.md
- Track Objects on Video: how_to/track_objects.md
- Process Datasets: how_to/process_datasets.md
- Reference - Code API:
- Detection and Segmentation:
- Core: detection/core.md
- Annotators: detection/annotators.md
- Metrics: detection/metrics.md
- Double Detection Filter: detection/double_detection_filter.md
- Utils: detection/utils.md
- Keypoint Detection:
- Core: keypoint/core.md
- Annotators: keypoint/annotators.md
- Classification:
- Core: classification/core.md
- Tools:
- Line Zone: detection/tools/line_zone.md
- Polygon Zone: detection/tools/polygon_zone.md
@ -54,26 +63,25 @@ nav:
- Detection Smoother: detection/tools/smoother.md
- Save Detections: detection/tools/save_detections.md
- Trackers: trackers.md
- Datasets: datasets.md
- Metrics:
- Object Detection: metrics/detection.md
- Draw:
- Color: draw/color.md
- Utils: draw/utils.md
- Geometry:
- Position: geometry/core.md
- Datasets:
- Core: datasets/core.md
- Utils: datasets/utils.md
- Utils:
- Video: utils/video.md
- Image: utils/image.md
- Iterables: utils/iterables.md
- Notebook: utils/notebook.md
- File: utils/file.md
- Draw: utils/draw.md
- Geometry: utils/geometry.md
- Assets: assets.md
- Cookbooks: cookbooks.md
- Cheatsheet: https://roboflow.github.io/cheatsheet-supervision/
- Contribute:
- Contributing: contributing.md
- Code of Conduct: code_of_conduct.md
- License: license.md
- Changelog:
- Release Notes:
- Changelog: changelog.md
- Deprecated: deprecated.md
@ -172,3 +180,12 @@ extra_javascript:
- "javascripts/init_kapa_widget.js"
- "javascripts/cookbooks-card.js"
- "https://cdnjs.cloudflare.com/ajax/libs/dompurify/3.0.8/purify.min.js"
# Messages shown during document build
# Reference: https://www.mkdocs.org/user-guide/configuration/#validation
# Values: [warn, info, ignore]
validation:
nav:
absolute_links: ignore
links:
absolute_links: ignore

2697
poetry.lock generated

File diff suppressed because it is too large Load Diff

View File

@ -1,18 +1,26 @@
[tool.poetry]
name = "supervision"
version = "0.19.0rc3"
version = "0.22.0"
description = "A set of easy-to-use utils that will come in handy in any Computer Vision project"
authors = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
maintainers = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
readme = "README.md"
license = "MIT"
packages = [{include = "supervision"}]
packages = [{ include = "supervision" }]
homepage = "https://github.com/roboflow/supervision"
repository = "https://github.com/roboflow/supervision"
documentation = "https://github.com/roboflow/supervision/blob/main/README.md"
keywords = ["machine-learning", "deep-learning", "vision", "ML", "DL", "AI", "Roboflow"]
keywords = [
"machine-learning",
"deep-learning",
"vision",
"ML",
"DL",
"AI",
"Roboflow",
]
classifiers=[
classifiers = [
'Development Status :: 4 - Beta',
'Intended Audience :: Developers',
'Intended Audience :: Education',
@ -27,33 +35,37 @@ classifiers=[
'Typing :: Typed',
'Operating System :: Microsoft :: Windows',
'Operating System :: POSIX :: Linux',
'Operating System :: MacOS'
'Operating System :: MacOS',
]
[tool.poetry.dependencies]
python = "^3.8"
numpy = ">=1.21.2"
numpy = [
{ version = ">=1.21.2,<1.23.3", python = "<=3.10" },
{ version = ">=1.23.3", python = ">3.10" },
]
scipy = [
{ version = "1.10.0", python = "<3.9" },
{ version = "^1.10.0", python = ">=3.9" }
{ version = "^1.10.0", python = ">=3.9" },
]
matplotlib = ">=3.6.0"
pyyaml = ">=5.3"
defusedxml = "^0.7.1"
opencv-python = { version = ">=4.5.5.64", optional = true }
opencv-python-headless = ">=4.5.5.64"
requests = { version = ">=2.26.0,<=2.31.0", optional = true }
tqdm = { version = ">=4.62.3,<=4.66.2", optional = true }
requests = { version = ">=2.26.0,<=2.32.3", optional = true }
tqdm = { version = ">=4.62.3,<=4.66.5", optional = true }
pillow = ">=9.4"
[tool.poetry.extras]
desktop = ["opencv-python"]
assets = ["requests","tqdm"]
assets = ["requests", "tqdm"]
[tool.poetry.group.dev.dependencies]
twine = ">=4.0.2,<6.0.0"
twine = "^5.1.1"
pytest = ">=7.2.2,<9.0.0"
wheel = ">=0.40,<0.43"
build = ">=0.10,<1.1"
wheel = ">=0.40,<0.45"
build = ">=0.10,<1.3"
ruff = ">=0.1.0"
mypy = "^1.4.1"
pre-commit = "^3.3.3"
@ -62,11 +74,15 @@ notebook = ">=6.5.3,<8.0.0"
ipywidgets = "^8.1.1"
jupytext = "^1.16.1"
nbconvert = "^7.14.2"
docutils = [
{ version = "^0.20.1", python = "<3.9" },
{ version = "^0.21.1", python = ">=3.9" },
]
[tool.poetry.group.docs.dependencies]
mkdocs-material = {extras = ["imaging"], version = "^9.5.5"}
mkdocstrings = {extras = ["python"], version = ">=0.20,<0.25"}
mkdocs-material = { extras = ["imaging"], version = "^9.5.5" }
mkdocstrings = { extras = ["python"], version = ">=0.20,<0.26" }
mike = "^2.0.0"
# For Documentation Development use Python 3.10 or above
# Use Latest mkdocs-jupyter min 0.24.6 for Jupyter Notebook Theme support
@ -74,8 +90,6 @@ mkdocs-jupyter = "^0.24.3"
mkdocs-git-committers-plugin-2 = "^2.2.3"
mkdocs-git-revision-date-localized-plugin = "^1.2.4"
[tool.isort]
line_length = 88
profile = "black"
@ -88,7 +102,6 @@ tests = ["B201", "B301", "B318", "B314", "B303", "B413", "B412", "B410"]
check = true
imports = ["cv2", "supervision"]
[tool.black]
target-version = ["py38"]
line-length = 88
@ -110,13 +123,6 @@ exclude = '''
[tool.ruff]
target-version = "py38"
# Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default.
select = ["E", "F"]
ignore = []
# Allow autofix for all enabled rules (when `--fix`) is provided.
fixable = ["A", "B", "C", "D", "E", "F", "G", "I", "N", "Q", "S", "T", "W", "ANN", "ARG", "BLE", "COM", "DJ", "DTZ", "EM", "ERA", "EXE", "FBT", "ICN", "INP", "ISC", "NPY", "PD", "PGH", "PIE", "PL", "PT", "PTH", "PYI", "RET", "RSE", "RUF", "SIM", "SLF", "TCH", "TID", "TRY", "UP", "YTT"]
unfixable = []
# Exclude a variety of commonly ignored directories.
exclude = [
@ -146,33 +152,80 @@ exclude = [
"docs",
]
# Same as Black.
line-length = 88
indent-width = 4
[tool.ruff.lint]
# Enable pycodestyle (`E`) and Pyflakes (`F`) codes by default.
select = ["E", "F"]
ignore = []
# Allow autofix for all enabled rules (when `--fix`) is provided.
fixable = [
"A",
"B",
"C",
"D",
"E",
"F",
"G",
"I",
"N",
"Q",
"S",
"T",
"W",
"ANN",
"ARG",
"BLE",
"COM",
"DJ",
"DTZ",
"EM",
"ERA",
"EXE",
"FBT",
"ICN",
"INP",
"ISC",
"NPY",
"PD",
"PGH",
"PIE",
"PL",
"PT",
"PTH",
"PYI",
"RET",
"RSE",
"RUF",
"SIM",
"SLF",
"TCH",
"TID",
"TRY",
"UP",
"YTT",
]
unfixable = []
# Allow unused variables when underscore-prefixed.
dummy-variable-rgx = "^(_+|(_+[a-zA-Z0-9_]*[a-zA-Z0-9]+?))$"
pylint.max-args = 20
[tool.ruff.flake8-quotes]
[tool.ruff.lint.flake8-quotes]
inline-quotes = "double"
multiline-quotes = "double"
docstring-quotes = "double"
[tool.ruff.pydocstyle]
[tool.ruff.lint.pydocstyle]
convention = "google"
[tool.ruff.per-file-ignores]
"__init__.py" = ["E402","F401"]
"supervision/assets/list.py" = ["E501"]
[tool.ruff.lint.per-file-ignores]
"__init__.py" = ["E402", "F401"]
[tool.ruff.lint.mccabe]
# Flag errors (`C901`) whenever the complexity level exceeds 5.
max-complexity = 20
[tool.ruff.pylint]
max-args = 20
[tool.ruff.format]
# Like Black, use double quotes for strings.
quote-style = "double"
@ -190,7 +243,7 @@ line-ending = "auto"
include-package-data = false
[tool.setuptools.packages.find]
exclude = ["docs*", "test*","examples*"]
exclude = ["docs*", "test*", "examples*"]
[build-system]
requires = ["poetry-core"]

24
release_process.md Normal file
View File

@ -0,0 +1,24 @@
# Release Process
This doc outlines how supervision is released into production.
It assumes you already have the code changes, as well as a draft of the release notes.
1. Make sure you have all required changes were merged into `develop`.
2. Create and merge a PR, merging `develop` into `main`, containing:
- A commit that updates the project version in `pyproject.toml`.
- All changes made during the release.
3. Tag the commit with the new supervision version.
- make sure to pull from `main` !
- Verify that the latest merge commits exists. `git log`.
- Run `git tag x.y.z`, with your version
- Check with `git log`.
- Run `git push origin --tags`
- Upon pushing the tag, the [PyPi](https://pypi.org/project/supervision/) should update to the new version. Check this!
4. Open and merge a PR, merging `main` into `develop`.
5. Update the docs by running the [Supervision Release Documentation Workflow 📚](https://github.com/roboflow/supervision/actions/workflows/publish-release-docs.yml) workflow from GitHub.
- Select the `main` branch from the dropdown.
6. Create a release on GitHub.
- Go to releases
- Assign the release notes to the tag created in step 3.
- Publish the release.

View File

@ -7,11 +7,14 @@ except importlib_metadata.PackageNotFoundError:
__version__ = "development"
from supervision.annotators.core import (
BackgroundOverlayAnnotator,
BlurAnnotator,
BoundingBoxAnnotator,
BoxAnnotator,
BoxCornerAnnotator,
CircleAnnotator,
ColorAnnotator,
CropAnnotator,
DotAnnotator,
EllipseAnnotator,
HaloAnnotator,
@ -23,6 +26,7 @@ from supervision.annotators.core import (
PercentageBarAnnotator,
PixelateAnnotator,
PolygonAnnotator,
RichLabelAnnotator,
RoundBoxAnnotator,
TraceAnnotator,
TriangleAnnotator,
@ -34,31 +38,44 @@ from supervision.dataset.core import (
ClassificationDataset,
DetectionDataset,
)
from supervision.detection.annotate import BoxAnnotator
from supervision.dataset.utils import mask_to_rle, rle_to_mask
from supervision.detection.core import Detections
from supervision.detection.line_zone import LineZone, LineZoneAnnotator
from supervision.detection.lmm import LMM
from supervision.detection.overlap_filter import (
OverlapFilter,
box_non_max_merge,
box_non_max_suppression,
mask_non_max_suppression,
)
from supervision.detection.tools.csv_sink import CSVSink
from supervision.detection.tools.inference_slicer import InferenceSlicer
from supervision.detection.tools.json_sink import JSONSink
from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator
from supervision.detection.tools.smoother import DetectionsSmoother
from supervision.detection.utils import (
box_iou_batch,
box_non_max_suppression,
calculate_masks_centroids,
clip_boxes,
contains_holes,
contains_multiple_segments,
filter_polygons_by_area,
mask_iou_batch,
mask_non_max_suppression,
mask_to_polygons,
mask_to_xyxy,
move_boxes,
move_masks,
pad_boxes,
polygon_to_mask,
polygon_to_xyxy,
scale_boxes,
xcycwh_to_xyxy,
xywh_to_xyxy,
)
from supervision.draw.color import Color, ColorPalette
from supervision.draw.utils import (
calculate_dynamic_line_thickness,
calculate_dynamic_text_scale,
calculate_optimal_line_thickness,
calculate_optimal_text_scale,
draw_filled_rectangle,
draw_image,
draw_line,
@ -68,10 +85,25 @@ from supervision.draw.utils import (
)
from supervision.geometry.core import Point, Position, Rect
from supervision.geometry.utils import get_polygon_center
from supervision.keypoint.annotators import (
EdgeAnnotator,
VertexAnnotator,
VertexLabelAnnotator,
)
from supervision.keypoint.core import KeyPoints
from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision
from supervision.tracker.byte_tracker.core import ByteTrack
from supervision.utils.conversion import cv2_to_pillow, pillow_to_cv2
from supervision.utils.file import list_files_with_extensions
from supervision.utils.image import ImageSink, crop_image
from supervision.utils.image import (
ImageSink,
create_tiles,
crop_image,
letterbox_image,
overlay_image,
resize_image,
scale_image,
)
from supervision.utils.notebook import plot_image, plot_images_grid
from supervision.utils.video import (
FPSMonitor,

View File

@ -1,11 +1,22 @@
from abc import ABC, abstractmethod
from typing import TypeVar
import numpy as np
from PIL import Image
from supervision.detection.core import Detections
ImageType = TypeVar("ImageType", np.ndarray, Image.Image)
"""
An image of type `np.ndarray` or `PIL.Image.Image`.
Unlike a `Union`, ensures the type remains consistent. If a function
takes an `ImageType` argument and returns an `ImageType`, when you
pass an `np.ndarray`, you will get an `np.ndarray` back.
"""
class BaseAnnotator(ABC):
@abstractmethod
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
def annotate(self, scene: ImageType, detections: Detections) -> ImageType:
pass

File diff suppressed because it is too large Load Diff

View File

@ -1,5 +1,5 @@
from enum import Enum
from typing import Optional, Union
from typing import Optional, Tuple, Union
import numpy as np
@ -34,14 +34,14 @@ def resolve_color_idx(
) -> int:
if detection_idx >= len(detections):
raise ValueError(
f"Detection index {detection_idx}"
f"Detection index {detection_idx} "
f"is out of bounds for detections of length {len(detections)}"
)
if isinstance(color_lookup, np.ndarray):
if len(color_lookup) != len(detections):
raise ValueError(
f"Length of color lookup {len(color_lookup)}"
f"Length of color lookup {len(color_lookup)} "
f"does not match length of detections {len(detections)}"
)
return color_lookup[detection_idx]
@ -50,19 +50,72 @@ def resolve_color_idx(
elif color_lookup == ColorLookup.CLASS:
if detections.class_id is None:
raise ValueError(
"Could not resolve color by class because"
"Could not resolve color by class because "
"Detections do not have class_id"
)
return detections.class_id[detection_idx]
elif color_lookup == ColorLookup.TRACK:
if detections.tracker_id is None:
raise ValueError(
"Could not resolve color by track because"
"Could not resolve color by track because "
"Detections do not have tracker_id"
)
return detections.tracker_id[detection_idx]
def resolve_text_background_xyxy(
center_coordinates: Tuple[int, int],
text_wh: Tuple[int, int],
position: Position,
) -> Tuple[int, int, int, int]:
center_x, center_y = center_coordinates
text_w, text_h = text_wh
if position == Position.TOP_LEFT:
return center_x, center_y - text_h, center_x + text_w, center_y
elif position == Position.TOP_RIGHT:
return center_x - text_w, center_y - text_h, center_x, center_y
elif position == Position.TOP_CENTER:
return (
center_x - text_w // 2,
center_y - text_h,
center_x + text_w // 2,
center_y,
)
elif position == Position.CENTER or position == Position.CENTER_OF_MASS:
return (
center_x - text_w // 2,
center_y - text_h // 2,
center_x + text_w // 2,
center_y + text_h // 2,
)
elif position == Position.BOTTOM_LEFT:
return center_x, center_y, center_x + text_w, center_y + text_h
elif position == Position.BOTTOM_RIGHT:
return center_x - text_w, center_y, center_x, center_y + text_h
elif position == Position.BOTTOM_CENTER:
return (
center_x - text_w // 2,
center_y,
center_x + text_w // 2,
center_y + text_h,
)
elif position == Position.CENTER_LEFT:
return (
center_x - text_w,
center_y - text_h // 2,
center_x,
center_y + text_h // 2,
)
elif position == Position.CENTER_RIGHT:
return (
center_x,
center_y - text_h // 2,
center_x + text_w,
center_y + text_h // 2,
)
def get_color_by_index(color: Union[Color, ColorPalette], idx: int) -> Color:
if isinstance(color, ColorPalette):
return color.by_idx(idx)

View File

@ -18,6 +18,8 @@ class VideoAssets(Enum):
| `SUBWAY` | `subway.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/subway.mp4) |
| `MARKET_SQUARE` | `market-square.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/market-square.mp4) |
| `PEOPLE_WALKING` | `people-walking.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/people-walking.mp4) |
| `BEACH` | `beach-1.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/beach-1.mp4) |
| `BASKETBALL` | `basketball-1.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/basketball-1.mp4) |
""" # noqa: E501 // docs
VEHICLES = "vehicles.mp4"
@ -27,6 +29,8 @@ class VideoAssets(Enum):
SUBWAY = "subway.mp4"
MARKET_SQUARE = "market-square.mp4"
PEOPLE_WALKING = "people-walking.mp4"
BEACH = "beach-1.mp4"
BASKETBALL = "basketball-1.mp4"
@classmethod
def list(cls):
@ -62,4 +66,12 @@ VIDEO_ASSETS: Dict[str, Tuple[str, str]] = {
f"{BASE_VIDEO_URL}{VideoAssets.PEOPLE_WALKING.value}",
"0574c053c8686c3f1dc0aa3743e45cb9",
),
VideoAssets.BEACH.value: (
f"{BASE_VIDEO_URL}{VideoAssets.BEACH.value}",
"4175d42fec4d450ed081523fd39e0cf8",
),
VideoAssets.BASKETBALL.value: (
f"{BASE_VIDEO_URL}{VideoAssets.BASKETBALL.value}",
"60d94a3c7c47d16f09d342b088012ecc",
),
}

View File

@ -3,8 +3,9 @@ from __future__ import annotations
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass
from itertools import chain
from pathlib import Path
from typing import Dict, Iterator, List, Optional, Tuple
from typing import Dict, Iterator, List, Optional, Tuple, Union
import cv2
import numpy as np
@ -31,9 +32,10 @@ from supervision.dataset.utils import (
train_test_split,
)
from supervision.detection.core import Detections
from supervision.utils.internal import deprecated, warn_deprecated
from supervision.utils.iterables import find_duplicates
@dataclass
class BaseDataset(ABC):
@abstractmethod
def __len__(self) -> int:
@ -46,30 +48,95 @@ class BaseDataset(ABC):
pass
@dataclass
class DetectionDataset(BaseDataset):
"""
Dataclass containing information about object detection dataset.
Contains information about a detection dataset. Handles lazy image loading
and annotation retrieval, dataset splitting, conversions into multiple
formats.
Attributes:
classes (List[str]): List containing dataset class names.
images (Dict[str, np.ndarray]): Dictionary mapping image name to image.
images (Union[List[str], Dict[str, np.ndarray]]):
Accepts a list of image paths, or dictionaries of loaded cv2 images
with paths as keys. If you pass a list of paths, the dataset will
lazily load images on demand, which is much more memory-efficient.
annotations (Dict[str, Detections]): Dictionary mapping
image name to annotations.
image path to annotations. The dictionary keys match
match the keys in `images` or entries in the list of
image paths.
"""
classes: List[str]
images: Dict[str, np.ndarray]
annotations: Dict[str, Detections]
def __init__(
self,
classes: List[str],
images: Union[List[str], Dict[str, np.ndarray]],
annotations: Dict[str, Detections],
) -> None:
self.classes = classes
if set(images) != set(annotations):
raise ValueError(
"The keys of the images and annotations dictionaries must match."
)
self.annotations = annotations
# Eliminate duplicates while preserving order
self.image_paths = list(dict.fromkeys(images))
self._images_in_memory: Dict[str, np.ndarray] = {}
if isinstance(images, dict):
self._images_in_memory = images
warn_deprecated(
"Passing a `Dict[str, np.ndarray]` into `DetectionDataset` is "
"deprecated and will be removed in `supervision-0.26.0`. Use "
"a list of paths `List[str]` instead."
)
# TODO: when supervision-0.26.0 is released, and Dict[str, np.ndarray]
# for images is no longer supported, also simplify the rest of
# the code. E.g. list(images) is no longer needed, and merge can
# be simplified.
@property
@deprecated(
"`DetectionDataset.images` property is deprecated and will be removed in "
"`supervision-0.26.0`. Iterate with `for path, image, annotation in dataset:` "
"instead."
)
def images(self) -> Dict[str, np.ndarray]:
"""
Load all images to memory and return them as a dictionary.
!!! warning
Only use this when you need all images at once.
It is much more memory-efficient to initialize dataset with
image paths and use `for path, image, annotation in dataset:`.
"""
if self._images_in_memory:
return self._images_in_memory
images = {image_path: cv2.imread(image_path) for image_path in self.image_paths}
return images
def _get_image(self, image_path: str) -> np.ndarray:
"""Assumes that image is in dataset"""
if self._images_in_memory:
return self._images_in_memory[image_path]
return cv2.imread(image_path)
def __len__(self) -> int:
"""
Return the number of images in the dataset.
return len(self._images_in_memory) or len(self.image_paths)
Returns:
int: The number of images.
def __getitem__(self, i: int) -> Tuple[str, np.ndarray, Detections]:
"""
return len(self.images)
Returns:
Tuple[str, np.ndarray, Detections]: The image path, image data,
and its corresponding annotation at index i.
"""
image_path = self.image_paths[i]
image = self._get_image(image_path)
annotation = self.annotations[image_path]
return image_path, image, annotation
def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Detections]]:
"""
@ -77,25 +144,33 @@ class DetectionDataset(BaseDataset):
Yields:
Iterator[Tuple[str, np.ndarray, Detections]]:
An iterator that yields tuples containing the image name,
An iterator that yields tuples containing the image path,
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)
for i in range(len(self)):
image_path, image, annotation = self[i]
yield image_path, image, annotation
def __eq__(self, other):
def __eq__(self, other) -> bool:
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]:
if self.image_paths != other.image_paths:
return False
if self._images_in_memory or other._images_in_memory:
if not np.array_equal(
list(self._images_in_memory.values()),
list(other._images_in_memory.values()),
):
return False
if self.annotations != other.annotations:
return False
return True
def split(
@ -116,38 +191,140 @@ class DetectionDataset(BaseDataset):
Tuple[DetectionDataset, DetectionDataset]: A tuple containing
the training and testing datasets.
Example:
Examples:
```python
import supervision as sv
ds = sv.DetectionDataset(...)
train_ds, test_ds = ds.split(split_ratio=0.7,
random_state=42, shuffle=True)
train_ds, test_ds = ds.split(split_ratio=0.7, random_state=42, shuffle=True)
len(train_ds), len(test_ds)
# (700, 300)
```
"""
image_names = list(self.images.keys())
train_names, test_names = train_test_split(
data=image_names,
train_paths, test_paths = train_test_split(
data=self.image_paths,
train_ratio=split_ratio,
random_state=random_state,
shuffle=shuffle,
)
train_input: Union[List[str], Dict[str, np.ndarray]]
test_input: Union[List[str], Dict[str, np.ndarray]]
if self._images_in_memory:
train_input = {path: self._images_in_memory[path] for path in train_paths}
test_input = {path: self._images_in_memory[path] for path in test_paths}
else:
train_input = train_paths
test_input = test_paths
train_annotations = {path: self.annotations[path] for path in train_paths}
test_annotations = {path: self.annotations[path] for path in test_paths}
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},
images=train_input,
annotations=train_annotations,
)
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},
images=test_input,
annotations=test_annotations,
)
return train_dataset, test_dataset
@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.
Examples:
```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']
```
"""
def is_in_memory(dataset: DetectionDataset) -> bool:
return len(dataset._images_in_memory) > 0 or len(dataset.image_paths) == 0
def is_lazy(dataset: DetectionDataset) -> bool:
return len(dataset._images_in_memory) == 0
all_in_memory = all([is_in_memory(dataset) for dataset in dataset_list])
all_lazy = all([is_lazy(dataset) for dataset in dataset_list])
if not all_in_memory and not all_lazy:
raise ValueError(
"Merging lazy and in-memory DetectionDatasets is not supported."
)
images_in_memory = {}
for dataset in dataset_list:
images_in_memory.update(dataset._images_in_memory)
image_paths = list(
chain.from_iterable(dataset.image_paths for dataset in dataset_list)
)
image_paths_unique = list(dict.fromkeys(image_paths))
if len(image_paths) != len(image_paths_unique):
duplicates = find_duplicates(image_paths)
raise ValueError(
f"Image paths {duplicates} are not unique across datasets."
)
image_paths = image_paths_unique
classes = merge_class_lists(
class_lists=[dataset.classes for dataset in dataset_list]
)
annotations = {}
for dataset in dataset_list:
annotations.update(dataset.annotations)
for dataset in dataset_list:
class_index_mapping = build_class_index_mapping(
source_classes=dataset.classes, target_classes=classes
)
for image_path in dataset.image_paths:
annotations[image_path] = map_detections_class_id(
source_to_target_mapping=class_index_mapping,
detections=annotations[image_path],
)
return cls(
classes=classes,
images=images_in_memory or image_paths,
annotations=annotations,
)
def as_pascal_voc(
self,
images_directory_path: Optional[str] = None,
@ -181,25 +358,22 @@ class DetectionDataset(BaseDataset):
"""
if images_directory_path:
save_dataset_images(
images_directory_path=images_directory_path, images=self.images
dataset=self,
images_directory_path=images_directory_path,
)
if annotations_directory_path:
Path(annotations_directory_path).mkdir(parents=True, exist_ok=True)
for image_path, image in self.images.items():
detections = self.annotations[image_path]
if annotations_directory_path:
for image_path, image, annotations in self:
annotation_name = Path(image_path).stem
annotations_path = os.path.join(
annotations_directory_path, f"{annotation_name}.xml"
)
image_name = Path(image_path).name
pascal_voc_xml = detections_to_pascal_voc(
detections=detections,
detections=annotations,
classes=self.classes,
filename=image_name,
image_shape=image.shape,
image_shape=image.shape, # type: ignore
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,
@ -229,7 +403,7 @@ class DetectionDataset(BaseDataset):
DetectionDataset: A DetectionDataset instance containing
the loaded images and annotations.
Example:
Examples:
```python
import roboflow
from roboflow import Roboflow
@ -252,13 +426,15 @@ class DetectionDataset(BaseDataset):
```
"""
classes, images, annotations = load_pascal_voc_annotations(
classes, image_paths, annotations = load_pascal_voc_annotations(
images_directory_path=images_directory_path,
annotations_directory_path=annotations_directory_path,
force_masks=force_masks,
)
return DetectionDataset(classes=classes, images=images, annotations=annotations)
return DetectionDataset(
classes=classes, images=image_paths, annotations=annotations
)
@classmethod
def from_yolo(
@ -267,6 +443,7 @@ class DetectionDataset(BaseDataset):
annotations_directory_path: str,
data_yaml_path: str,
force_masks: bool = False,
is_obb: bool = False,
) -> DetectionDataset:
"""
Creates a Dataset instance from YOLO formatted data.
@ -281,12 +458,15 @@ class DetectionDataset(BaseDataset):
force_masks (bool, optional): If True, forces
masks to be loaded for all annotations,
regardless of whether they are present.
is_obb (bool, optional): If True, loads the annotations in OBB format.
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
where pairs of [x, y] are box corners.
Returns:
DetectionDataset: A DetectionDataset instance
containing the loaded images and annotations.
Example:
Examples:
```python
import roboflow
from roboflow import Roboflow
@ -308,13 +488,16 @@ class DetectionDataset(BaseDataset):
# ['dog', 'person']
```
"""
classes, images, annotations = load_yolo_annotations(
classes, image_paths, 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,
is_obb=is_obb,
)
return DetectionDataset(
classes=classes, images=image_paths, annotations=annotations
)
return DetectionDataset(classes=classes, images=images, annotations=annotations)
def as_yolo(
self,
@ -355,13 +538,12 @@ class DetectionDataset(BaseDataset):
"""
if images_directory_path is not None:
save_dataset_images(
images_directory_path=images_directory_path, images=self.images
dataset=self, images_directory_path=images_directory_path
)
if annotations_directory_path is not None:
save_yolo_annotations(
dataset=self,
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,
@ -391,7 +573,7 @@ class DetectionDataset(BaseDataset):
DetectionDataset: A DetectionDataset instance containing
the loaded images and annotations.
Example:
Examples:
```python
import roboflow
from roboflow import Roboflow
@ -431,6 +613,20 @@ class DetectionDataset(BaseDataset):
Exports the dataset to COCO format. This method saves the
images and their corresponding annotations in COCO format.
!!! tip
The format of the mask is determined automatically based on its structure:
- If a mask contains multiple disconnected components or holes, it will be
saved using the Run-Length Encoding (RLE) format for efficient storage and
processing.
- If a mask consists of a single, contiguous region without any holes, it
will be encoded as a polygon, preserving the outline of the object.
This automatic selection ensures that the masks are stored in the most
appropriate and space-efficient format, complying with COCO dataset
standards.
Args:
images_directory_path (Optional[str]): The path to the directory
where the images should be saved.
@ -451,103 +647,134 @@ class DetectionDataset(BaseDataset):
"""
if images_directory_path is not None:
save_dataset_images(
images_directory_path=images_directory_path, images=self.images
dataset=self, images_directory_path=images_directory_path
)
if annotations_path is not None:
save_coco_annotations(
dataset=self,
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.
Contains information about a classification dataset, handles lazy image
loading, dataset splitting.
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
images (Union[List[str], Dict[str, np.ndarray]]):
List of image paths or dictionary mapping image name to image data.
annotations (Dict[str, Classifications]): Dictionary mapping
image name to annotations.
"""
classes: List[str]
images: Dict[str, np.ndarray]
annotations: Dict[str, Classifications]
def __init__(
self,
classes: List[str],
images: Union[List[str], Dict[str, np.ndarray]],
annotations: Dict[str, Classifications],
) -> None:
self.classes = classes
if set(images) != set(annotations):
raise ValueError(
"The keys of the images and annotations dictionaries must match."
)
self.annotations = annotations
# Eliminate duplicates while preserving order
self.image_paths = list(dict.fromkeys(images))
self._images_in_memory: Dict[str, np.ndarray] = {}
if isinstance(images, dict):
self._images_in_memory = images
warn_deprecated(
"Passing a `Dict[str, np.ndarray]` into `ClassificationDataset` is "
"deprecated and will be removed in a future release. Use "
"a list of paths `List[str]` instead."
)
@property
@deprecated(
"`DetectionDataset.images` property is deprecated and will be removed in "
"`supervision-0.26.0`. Iterate with `for path, image, annotation in dataset:` "
"instead."
)
def images(self) -> Dict[str, np.ndarray]:
"""
Load all images to memory and return them as a dictionary.
!!! warning
Only use this when you need all images at once.
It is much more memory-efficient to initialize dataset with
image paths and use `for path, image, annotation in dataset:`.
"""
if self._images_in_memory:
return self._images_in_memory
images = {image_path: cv2.imread(image_path) for image_path in self.image_paths}
return images
def _get_image(self, image_path: str) -> np.ndarray:
"""Assumes that image is in dataset"""
if self._images_in_memory:
return self._images_in_memory[image_path]
return cv2.imread(image_path)
def __len__(self) -> int:
return len(self.images)
return len(self._images_in_memory) or len(self.image_paths)
def __getitem__(self, i: int) -> Tuple[str, np.ndarray, Classifications]:
"""
Returns:
Tuple[str, np.ndarray, Classifications]: The image path, image data,
and its corresponding annotation at index i.
"""
image_path = self.image_paths[i]
image = self._get_image(image_path)
annotation = self.annotations[image_path]
return image_path, image, annotation
def __iter__(self) -> Iterator[Tuple[str, np.ndarray, Classifications]]:
"""
Iterate over the images and annotations in the dataset.
Yields:
Iterator[Tuple[str, np.ndarray, Detections]]:
An iterator that yields tuples containing the image path,
the image data, and its corresponding annotation.
"""
for i in range(len(self)):
image_path, image, annotation = self[i]
yield image_path, image, annotation
def __eq__(self, other) -> bool:
if not isinstance(other, ClassificationDataset):
return False
if set(self.classes) != set(other.classes):
return False
if self.image_paths != other.image_paths:
return False
if self._images_in_memory or other._images_in_memory:
if not np.array_equal(
list(self._images_in_memory.values()),
list(other._images_in_memory.values()),
):
return False
if self.annotations != other.annotations:
return False
return True
def split(
self, split_ratio=0.8, random_state=None, shuffle: bool = True
@ -567,35 +794,45 @@ class ClassificationDataset(BaseDataset):
Tuple[ClassificationDataset, ClassificationDataset]: A tuple containing
the training and testing datasets.
Example:
Examples:
```python
import supervision as sv
cd = sv.ClassificationDataset(...)
train_cd,test_cd = cd.split(split_ratio=0.7,
random_state=42,shuffle=True)
train_cd,test_cd = cd.split(split_ratio=0.7, random_state=42,shuffle=True)
len(train_cd), len(test_cd)
# (700, 300)
```
"""
image_names = list(self.images.keys())
train_names, test_names = train_test_split(
data=image_names,
train_paths, test_paths = train_test_split(
data=self.image_paths,
train_ratio=split_ratio,
random_state=random_state,
shuffle=shuffle,
)
train_input: Union[List[str], Dict[str, np.ndarray]]
test_input: Union[List[str], Dict[str, np.ndarray]]
if self._images_in_memory:
train_input = {path: self._images_in_memory[path] for path in train_paths}
test_input = {path: self._images_in_memory[path] for path in test_paths}
else:
train_input = train_paths
test_input = test_paths
train_annotations = {path: self.annotations[path] for path in train_paths}
test_annotations = {path: self.annotations[path] for path in test_paths}
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},
images=train_input,
annotations=train_annotations,
)
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},
images=test_input,
annotations=test_annotations,
)
return train_dataset, test_dataset
def as_folder_structure(self, root_directory_path: str) -> None:
@ -611,18 +848,16 @@ class ClassificationDataset(BaseDataset):
for class_name in self.classes:
os.makedirs(os.path.join(root_directory_path, class_name), exist_ok=True)
for image_path in self.images:
classification = self.annotations[image_path]
image = self.images[image_path]
image_name = Path(image_path).name
for image_save_path, image, annotation in self:
image_name = Path(image_save_path).name
class_id = (
classification.class_id[0]
if classification.confidence is None
else classification.get_top_k(1)[0][0]
annotation.class_id[0]
if annotation.confidence is None
else annotation.get_top_k(1)[0][0]
)
class_name = self.classes[class_id]
image_path = os.path.join(root_directory_path, class_name, image_name)
cv2.imwrite(image_path, image)
image_save_path = os.path.join(root_directory_path, class_name, image_name)
cv2.imwrite(image_save_path, image)
@classmethod
def from_folder_structure(cls, root_directory_path: str) -> ClassificationDataset:
@ -635,7 +870,7 @@ class ClassificationDataset(BaseDataset):
Returns:
ClassificationDataset: The dataset.
Example:
Examples:
```python
import roboflow
from roboflow import Roboflow
@ -655,7 +890,7 @@ class ClassificationDataset(BaseDataset):
classes = os.listdir(root_directory_path)
classes = sorted(set(classes))
images = {}
image_paths = []
annotations = {}
for class_name in classes:
@ -663,13 +898,13 @@ class ClassificationDataset(BaseDataset):
for image in os.listdir(os.path.join(root_directory_path, class_name)):
image_path = str(os.path.join(root_directory_path, class_name, image))
images[image_path] = cv2.imread(image_path)
image_paths.append(image_path)
annotations[image_path] = Classifications(
class_id=np.array([class_id]),
)
return cls(
classes=classes,
images=images,
images=image_paths,
annotations=annotations,
)

View File

@ -1,19 +1,28 @@
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Tuple
from typing import TYPE_CHECKING, Dict, List, Tuple
import cv2
import numpy as np
import numpy.typing as npt
from supervision.dataset.utils import (
approximate_mask_with_polygons,
map_detections_class_id,
mask_to_rle,
rle_to_mask,
)
from supervision.detection.core import Detections
from supervision.detection.utils import polygon_to_mask
from supervision.detection.utils import (
contains_holes,
contains_multiple_segments,
polygon_to_mask,
)
from supervision.utils.file import read_json_file, save_json_file
if TYPE_CHECKING:
from supervision.dataset.core import DetectionDataset
def coco_categories_to_classes(coco_categories: List[dict]) -> List[str]:
return [
@ -57,13 +66,24 @@ def group_coco_annotations_by_image_id(
return annotations
def _polygons_to_masks(
polygons: List[np.ndarray], resolution_wh: Tuple[int, int]
) -> np.ndarray:
def coco_annotations_to_masks(
image_annotations: List[dict], resolution_wh: Tuple[int, int]
) -> npt.NDArray[np.bool_]:
return np.array(
[
polygon_to_mask(polygon=polygon, resolution_wh=resolution_wh)
for polygon in polygons
rle_to_mask(
rle=np.array(image_annotation["segmentation"]["counts"]),
resolution_wh=resolution_wh,
)
if image_annotation["iscrowd"]
else polygon_to_mask(
polygon=np.reshape(
np.asarray(image_annotation["segmentation"], dtype=np.int32),
(-1, 2),
),
resolution_wh=resolution_wh,
)
for image_annotation in image_annotations
],
dtype=bool,
)
@ -83,13 +103,9 @@ def coco_annotations_to_detections(
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)
mask = coco_annotations_to_masks(
image_annotations=image_annotations, resolution_wh=resolution_wh
)
return Detections(
class_id=np.asarray(class_ids, dtype=int), xyxy=xyxy, mask=mask
)
@ -108,24 +124,35 @@ def detections_to_coco_annotations(
coco_annotations = []
for xyxy, mask, _, class_id, _, _ in detections:
box_width, box_height = xyxy[2] - xyxy[0], xyxy[3] - xyxy[1]
polygon = []
segmentation = []
iscrowd = 0
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()
)
iscrowd = contains_holes(mask=mask) or contains_multiple_segments(mask=mask)
if iscrowd:
segmentation = {
"counts": mask_to_rle(mask=mask),
"size": list(mask.shape[:2]),
}
else:
segmentation = [
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] if polygon else [],
"iscrowd": 0,
"segmentation": segmentation,
"iscrowd": iscrowd,
}
coco_annotations.append(coco_annotation)
annotation_id += 1
@ -136,7 +163,7 @@ 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]]:
) -> Tuple[List[str], List[str], Dict[str, Detections]]:
coco_data = read_json_file(file_path=annotations_path)
classes = coco_categories_to_classes(coco_categories=coco_data["categories"])
class_index_mapping = build_coco_class_index_mapping(
@ -147,7 +174,7 @@ def load_coco_annotations(
coco_annotations=coco_data["annotations"]
)
images = {}
images = []
annotations = {}
for coco_image in coco_images:
@ -159,7 +186,6 @@ def load_coco_annotations(
image_annotations = coco_annotations_groups.get(coco_image["id"], [])
image_path = os.path.join(images_directory_path, image_name)
image = cv2.imread(image_path)
annotation = coco_annotations_to_detections(
image_annotations=image_annotations,
resolution_wh=(image_width, image_height),
@ -170,23 +196,20 @@ def load_coco_annotations(
detections=annotation,
)
images[image_path] = image
images.append(image_path)
annotations[image_path] = annotation
return classes, images, annotations
def save_coco_annotations(
dataset: "DetectionDataset",
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,
@ -197,10 +220,10 @@ def save_coco_annotations(
coco_annotations = []
coco_images = []
coco_categories = classes_to_coco_categories(classes=classes)
coco_categories = classes_to_coco_categories(classes=dataset.classes)
image_id, annotation_id = 1, 1
for image_path, image in images.items():
for image_path, image, annotation in dataset:
image_height, image_width, _ = image.shape
image_name = f"{Path(image_path).stem}{Path(image_path).suffix}"
coco_image = {
@ -213,10 +236,8 @@ def save_coco_annotations(
}
coco_images.append(coco_image)
detections = annotations[image_path]
coco_annotation, annotation_id = detections_to_coco_annotations(
detections=detections,
detections=annotation,
image_id=image_id,
annotation_id=annotation_id,
min_image_area_percentage=min_image_area_percentage,
@ -228,7 +249,7 @@ def save_coco_annotations(
image_id += 1
annotation_dict = {
"info": info,
"info": {},
"licenses": licenses,
"categories": coco_categories,
"images": coco_images,

View File

@ -5,7 +5,8 @@ from xml.etree.ElementTree import Element, SubElement
import cv2
import numpy as np
from defusedxml.ElementTree import fromstring, parse, tostring
from defusedxml.ElementTree import parse, tostring
from defusedxml.minidom import parseString
from supervision.dataset.utils import approximate_mask_with_polygons
from supervision.detection.core import Detections
@ -129,8 +130,7 @@ def detections_to_pascal_voc(
annotation.append(next_object)
# Generate XML string
xml_string = fromstring(tostring(annotation)).toprettyxml(indent=" ")
xml_string = parseString(tostring(annotation)).toprettyxml(indent=" ")
return xml_string
@ -138,7 +138,7 @@ def load_pascal_voc_annotations(
images_directory_path: str,
annotations_directory_path: str,
force_masks: bool = False,
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
) -> Tuple[List[str], List[str], Dict[str, Detections]]:
"""
Loads PASCAL VOC XML annotations and returns the image name,
a Detections instance, and a list of class names.
@ -151,44 +151,39 @@ def load_pascal_voc_annotations(
annotations, regardless of whether they are present.
Returns:
Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]: A tuple
containing a list of class names,
a dictionary with image names as keys and
images as values, and a dictionary with image names as
keys and corresponding Detections instances as values.
Tuple[List[str], List[str], Dict[str, Detections]]: A tuple with a list
of class names, a list of paths to images, and a dictionary with image
paths as keys and corresponding Detections instances as values.
"""
image_paths = list_files_with_extensions(
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
)
image_paths = [
str(path)
for path in list_files_with_extensions(
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
)
]
classes = []
images = {}
classes: List[str] = []
annotations = {}
for image_path in image_paths:
image_name = Path(image_path).stem
image_path = str(image_path)
image = cv2.imread(image_path)
annotation_path = os.path.join(annotations_directory_path, f"{image_name}.xml")
image_stem = Path(image_path).stem
annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.xml")
if not os.path.exists(annotation_path):
images[image_path] = image
annotations[image_path] = Detections.empty()
continue
tree = parse(annotation_path)
root = tree.getroot()
image = cv2.imread(image_path)
resolution_wh = (image.shape[1], image.shape[0])
annotation, classes = detections_from_xml_obj(
root, classes, resolution_wh, force_masks
)
images[image_path] = image
annotations[image_path] = annotation
return classes, images, annotations
return classes, image_paths, annotations
def detections_from_xml_obj(

View File

@ -1,10 +1,11 @@
import os
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
import cv2
import numpy as np
from supervision.config import ORIENTED_BOX_COORDINATES
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
@ -16,6 +17,9 @@ from supervision.utils.file import (
save_yaml_file,
)
if TYPE_CHECKING:
from supervision.dataset.core import DetectionDataset
def _parse_box(values: List[str]) -> np.ndarray:
x_center, y_center, width, height = values
@ -70,12 +74,15 @@ def _image_name_to_annotation_name(image_name: str) -> str:
def yolo_annotations_to_detections(
lines: List[str], resolution_wh: Tuple[int, int], with_masks: bool
lines: List[str],
resolution_wh: Tuple[int, int],
with_masks: bool,
is_obb: bool = False,
) -> Detections:
if len(lines) == 0:
return Detections.empty()
class_id, relative_xyxy, relative_polygon = [], [], []
class_id, relative_xyxy, relative_polygon, relative_xyxyxyxy = [], [], [], []
w, h = resolution_wh
for line in lines:
values = line.split()
@ -88,21 +95,30 @@ def yolo_annotations_to_detections(
elif len(values) > 5:
polygon = _parse_polygon(values=values[1:])
relative_xyxy.append(polygon_to_xyxy(polygon=polygon))
if is_obb:
relative_xyxyxyxy.append(np.array(values[1:]))
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)
data = {}
if is_obb:
relative_xyxyxyxy = np.array(relative_xyxyxyxy, dtype=np.float32)
xyxyxyxy = relative_xyxyxyxy.reshape(-1, 4, 2)
xyxyxyxy *= np.array([w, h], dtype=np.float32)
data[ORIENTED_BOX_COORDINATES] = xyxyxyxy
if not with_masks:
return Detections(class_id=class_id, xyxy=xyxy)
return Detections(class_id=class_id, xyxy=xyxy, data=data)
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)
return Detections(class_id=class_id, xyxy=xyxy, data=data, mask=mask)
def load_yolo_annotations(
@ -110,7 +126,8 @@ def load_yolo_annotations(
annotations_directory_path: str,
data_yaml_path: str,
force_masks: bool = False,
) -> Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
is_obb: bool = False,
) -> Tuple[List[str], List[str], Dict[str, Detections]]:
"""
Loads YOLO annotations and returns class names, images,
and their corresponding detections.
@ -123,32 +140,34 @@ def load_yolo_annotations(
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.
is_obb (bool, optional): If True, loads the annotations in OBB format.
OBB annotations are defined as `[class_id, x, y, x, y, x, y, x, y]`,
where pairs of [x, y] are box corners.
Returns:
Tuple[List[str], Dict[str, np.ndarray], Dict[str, Detections]]:
Tuple[List[str], List[str], 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"]
)
image_paths = [
str(path)
for path in list_files_with_extensions(
directory=images_directory_path, extensions=["jpg", "jpeg", "png"]
)
]
classes = _extract_class_names(file_path=data_yaml_path)
images = {}
annotations = {}
for image_path in image_paths:
image_stem = Path(image_path).stem
image_path = str(image_path)
image = cv2.imread(image_path)
annotation_path = os.path.join(annotations_directory_path, f"{image_stem}.txt")
if not os.path.exists(annotation_path):
images[image_path] = image
annotations[image_path] = Detections.empty()
continue
image = cv2.imread(image_path)
lines = read_txt_file(file_path=annotation_path, skip_empty=True)
h, w, _ = image.shape
resolution_wh = (w, h)
@ -156,12 +175,13 @@ def load_yolo_annotations(
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
lines=lines,
resolution_wh=resolution_wh,
with_masks=with_masks,
is_obb=is_obb,
)
images[image_path] = image
annotations[image_path] = annotation
return classes, images, annotations
return classes, image_paths, annotations
def object_to_yolo(
@ -195,6 +215,9 @@ def detections_to_yolo_annotations(
) -> List[str]:
annotation = []
for xyxy, mask, _, class_id, _, _ in detections:
if class_id is None:
raise ValueError("Class ID is required for YOLO annotations.")
if mask is not None:
polygons = approximate_mask_with_polygons(
mask=mask,
@ -220,24 +243,22 @@ def detections_to_yolo_annotations(
def save_yolo_annotations(
dataset: "DetectionDataset",
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_path, image in images.items():
detections = annotations[image_path]
for image_path, image, annotation in dataset:
image_name = Path(image_path).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,
detections=annotation,
image_shape=image.shape, # type: ignore
min_image_area_percentage=min_image_area_percentage,
max_image_area_percentage=max_image_area_percentage,
approximation_percentage=approximation_percentage,

View File

@ -1,11 +1,13 @@
import copy
import os
import random
import shutil
from pathlib import Path
from typing import Dict, List, Optional, Tuple, TypeVar
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, TypeVar, Union
import cv2
import numpy as np
import numpy.typing as npt
from supervision.detection.core import Detections
from supervision.detection.utils import (
@ -14,6 +16,9 @@ from supervision.detection.utils import (
mask_to_polygons,
)
if TYPE_CHECKING:
from supervision.dataset.core import DetectionDataset
T = TypeVar("T")
@ -58,6 +63,7 @@ def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
def build_class_index_mapping(
source_classes: List[str], target_classes: List[str]
) -> Dict[int, int]:
"""Returns the index map of source classes -> target classes."""
index_mapping = {}
for i, class_name in enumerate(source_classes):
@ -93,14 +99,16 @@ def map_detections_class_id(
def save_dataset_images(
images_directory_path: str, images: Dict[str, np.ndarray]
dataset: "DetectionDataset", images_directory_path: str
) -> None:
Path(images_directory_path).mkdir(parents=True, exist_ok=True)
for image_path, image in images.items():
image_name = Path(image_path).name
target_image_path = os.path.join(images_directory_path, image_name)
cv2.imwrite(target_image_path, image)
for image_path in dataset.image_paths:
final_path = os.path.join(images_directory_path, Path(image_path).name)
if image_path in dataset._images_in_memory:
image = dataset._images_in_memory[image_path]
cv2.imwrite(final_path, image)
else:
shutil.copyfile(image_path, final_path)
def train_test_split(
@ -129,3 +137,123 @@ def train_test_split(
split_index = int(len(data) * train_ratio)
return data[:split_index], data[split_index:]
def rle_to_mask(
rle: Union[npt.NDArray[np.int_], List[int]], resolution_wh: Tuple[int, int]
) -> npt.NDArray[np.bool_]:
"""
Converts run-length encoding (RLE) to a binary mask.
Args:
rle (Union[npt.NDArray[np.int_], List[int]]): The 1D RLE array, the format
used in the COCO dataset (column-wise encoding, values of an array with
even indices represent the number of pixels assigned as background,
values of an array with odd indices represent the number of pixels
assigned as foreground object).
resolution_wh (Tuple[int, int]): The width (w) and height (h)
of the desired binary mask.
Returns:
The generated 2D Boolean mask of shape `(h, w)`, where the foreground object is
marked with `True`'s and the rest is filled with `False`'s.
Raises:
AssertionError: If the sum of pixels encoded in RLE differs from the
number of pixels in the expected mask (computed based on resolution_wh).
Examples:
```python
import supervision as sv
sv.rle_to_mask([5, 2, 2, 2, 5], (4, 4))
# array([
# [False, False, False, False],
# [False, True, True, False],
# [False, True, True, False],
# [False, False, False, False],
# ])
```
"""
if isinstance(rle, list):
rle = np.array(rle, dtype=int)
width, height = resolution_wh
assert width * height == np.sum(rle), (
"the sum of the number of pixels in the RLE must be the same "
"as the number of pixels in the expected mask"
)
zero_one_values = np.zeros(shape=(rle.size, 1), dtype=np.uint8)
zero_one_values[1::2] = 1
decoded_rle = np.repeat(zero_one_values, rle, axis=0)
decoded_rle = np.append(
decoded_rle, np.zeros(width * height - len(decoded_rle), dtype=np.uint8)
)
return decoded_rle.reshape((height, width), order="F")
def mask_to_rle(mask: npt.NDArray[np.bool_]) -> List[int]:
"""
Converts a binary mask into a run-length encoding (RLE).
Args:
mask (npt.NDArray[np.bool_]): 2D binary mask where `True` indicates foreground
object and `False` indicates background.
Returns:
The run-length encoded mask. Values of a list with even indices
represent the number of pixels assigned as background (`False`), values
of a list with odd indices represent the number of pixels assigned
as foreground object (`True`).
Raises:
AssertionError: If input mask is not 2D or is empty.
Examples:
```python
import numpy as np
import supervision as sv
mask = np.array([
[True, True, True, True],
[True, True, True, True],
[True, True, True, True],
[True, True, True, True],
])
sv.mask_to_rle(mask)
# [0, 16]
mask = np.array([
[False, False, False, False],
[False, True, True, False],
[False, True, True, False],
[False, False, False, False],
])
sv.mask_to_rle(mask)
# [5, 2, 2, 2, 5]
```
![mask_to_rle](https://media.roboflow.com/supervision-docs/mask-to-rle.png){ align=center width="800" }
""" # noqa E501 // docs
assert mask.ndim == 2, "Input mask must be 2D"
assert mask.size != 0, "Input mask cannot be empty"
on_value_change_indices = np.where(
mask.ravel(order="F") != np.roll(mask.ravel(order="F"), 1)
)[0]
on_value_change_indices = np.append(on_value_change_indices, mask.size)
# need to add 0 at the beginning when the same value is in the first and
# last element of the flattened mask
if on_value_change_indices[0] != 0:
on_value_change_indices = np.insert(on_value_change_indices, 0, 0)
rle = np.diff(on_value_change_indices)
if mask[0][0] == 1:
rle = np.insert(rle, 0, 0)
return list(rle)

View File

@ -1,150 +0,0 @@
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
from supervision.utils.internal import deprecated
@deprecated(
"`BoxAnnotator` is deprecated and will be removed in "
"`supervision-0.22.0`. Use `BoundingBoxAnnotator` and `LabelAnnotator` instead"
)
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

View File

@ -7,53 +7,105 @@ from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
import numpy as np
from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES
from supervision.detection.utils import (
from supervision.detection.lmm import (
LMM,
from_florence_2,
from_paligemma,
validate_lmm_parameters,
)
from supervision.detection.overlap_filter import (
box_non_max_merge,
box_non_max_suppression,
mask_non_max_suppression,
)
from supervision.detection.tools.transformers import (
process_transformers_detection_result,
process_transformers_v4_segmentation_result,
process_transformers_v5_segmentation_result,
)
from supervision.detection.utils import (
box_iou_batch,
calculate_masks_centroids,
extract_ultralytics_masks,
get_data_item,
is_data_equal,
mask_non_max_suppression,
merge_data,
process_roboflow_result,
validate_detections_fields,
xywh_to_xyxy,
)
from supervision.geometry.core import Position
from supervision.utils.internal import deprecated
from supervision.utils.internal import get_instance_variables
from supervision.validators import validate_detections_fields
@dataclass
class Detections:
"""
The `sv.Detections` allows you to convert results from a variety of object detection
and segmentation models into a single, unified format. The `sv.Detections` class
enables easy data manipulation and filtering, and provides a consistent API for
Supervision's tools like trackers, annotators, and zones.
The `sv.Detections` class in the Supervision library standardizes results from
various object detection and segmentation models into a consistent format. This
class simplifies data manipulation and filtering, providing a uniform API for
integration with Supervision [trackers](/trackers/), [annotators](/detection/annotators/), and [tools](/detection/tools/line_zone/).
```python
import cv2
import supervision as sv
from ultralytics import YOLO
=== "Inference"
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = YOLO('yolov8s.pt')
annotator = sv.BoundingBoxAnnotator()
Use [`sv.Detections.from_inference`](/detection/core/#supervision.detection.core.Detections.from_inference)
method, which accepts model results from both detection and segmentation models.
result = model(image)[0]
detections = sv.Detections.from_ultralytics(result)
```python
import cv2
import supervision as sv
from inference import get_model
annotated_image = annotator.annotate(image, detections)
```
model = get_model(model_id="yolov8n-640")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model.infer(image)[0]
detections = sv.Detections.from_inference(results)
```
!!! tip
=== "Ultralytics"
In `sv.Detections`, detection data is categorized into two main field types:
fixed and custom. The fixed fields include `xyxy`, `mask`, `confidence`,
`class_id`, and `tracker_id`. For any additional data requirements, custom
fields come into play, stored in the data field. These custom fields are easily
accessible using the `detections[<FIELD_NAME>]` syntax, providing flexibility
for diverse data handling needs.
Use [`sv.Detections.from_ultralytics`](/detection/core/#supervision.detection.core.Detections.from_ultralytics)
method, which accepts model results from both detection and segmentation models.
```python
import cv2
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8n.pt")
image = cv2.imread(<SOURCE_IMAGE_PATH>)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
```
=== "Transformers"
Use [`sv.Detections.from_transformers`](/detection/core/#supervision.detection.core.Detections.from_transformers)
method, which accepts model results from both detection and segmentation models.
```python
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(
transformers_results=results,
id2label=model.config.id2label)
```
Attributes:
xyxy (np.ndarray): An array of shape `(n, 4)` containing
@ -69,15 +121,7 @@ class Detections:
data (Dict[str, Union[np.ndarray, List]]): A dictionary containing additional
data where each key is a string representing the data type, and the value
is either a NumPy array or a list of corresponding data.
!!! warning
The `data` field in the `sv.Detections` class is currently in an experimental
phase. Please be aware that its API and functionality are subject to change in
future updates as we continue to refine and improve its capabilities.
We encourage users to experiment with this feature and provide feedback, but
also to be prepared for potential modifications in upcoming releases.
"""
""" # noqa: E501 // docs
xyxy: np.ndarray
mask: Optional[np.ndarray] = None
@ -176,8 +220,8 @@ class Detections:
@classmethod
def from_ultralytics(cls, ultralytics_results) -> Detections:
"""
Creates a Detections instance from a
[YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
Creates a `sv.Detections` instance from a
[YOLOv8](https://github.com/ultralytics/ultralytics) inference result.
!!! Note
@ -188,7 +232,7 @@ class Detections:
Args:
ultralytics_results (ultralytics.yolo.engine.results.Results):
The output Results instance from YOLOv8
The output Results instance from Ultralytics
Returns:
Detections: A new Detections object.
@ -201,13 +245,12 @@ class Detections:
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = YOLO('yolov8s.pt')
result = model(image)[0]
detections = sv.Detections.from_ultralytics(result)
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
```
""" # noqa: E501 // docs
if ultralytics_results.obb is not None:
if hasattr(ultralytics_results, "obb") and ultralytics_results.obb is not None:
class_id = ultralytics_results.obb.cls.cpu().numpy().astype(int)
class_names = np.array([ultralytics_results.names[i] for i in class_id])
oriented_box_coordinates = ultralytics_results.obb.xyxyxyxy.cpu().numpy()
@ -385,23 +428,82 @@ class Detections:
xyxy=mmdet_results.pred_instances.bboxes.cpu().numpy(),
confidence=mmdet_results.pred_instances.scores.cpu().numpy(),
class_id=mmdet_results.pred_instances.labels.cpu().numpy().astype(int),
mask=mmdet_results.pred_instances.masks.cpu().numpy()
if "masks" in mmdet_results.pred_instances
else None,
)
@classmethod
def from_transformers(cls, transformers_results: dict) -> Detections:
def from_transformers(
cls, transformers_results: dict, id2label: Optional[Dict[int, str]] = None
) -> Detections:
"""
Creates a Detections instance from object detection
[transformer](https://github.com/huggingface/transformers) inference result.
Creates a Detections instance from object detection or panoptic, semantic
and instance segmentation
[Transformer](https://github.com/huggingface/transformers) inference result.
Args:
transformers_results (Union[dict, torch.Tensor]): Inference results from
your Transformers model. This can be either a dictionary containing
valuable outputs like `scores`, `labels`, `boxes`, `masks`,
`segments_info`, and `segmentation`, or a `torch.Tensor` holding a
segmentation map where values represent class IDs.
id2label (Optional[Dict[int, str]]): A dictionary mapping class IDs to
labels, typically part of the `transformers` model configuration. If
provided, the resulting dictionary will include class names.
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),
)
Example:
```python
import torch
import supervision as sv
from PIL import Image
from transformers import DetrImageProcessor, DetrForObjectDetection
processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
image = Image.open(<SOURCE_IMAGE_PATH>)
inputs = processor(images=image, return_tensors="pt")
with torch.no_grad():
outputs = model(**inputs)
width, height = image.size
target_size = torch.tensor([[height, width]])
results = processor.post_process_object_detection(
outputs=outputs, target_sizes=target_size)[0]
detections = sv.Detections.from_transformers(
transformers_results=results,
id2label=model.config.id2label
)
```
""" # noqa: E501 // docs
if (
transformers_results.__class__.__name__ == "Tensor"
or "segmentation" in transformers_results
):
return cls(
**process_transformers_v5_segmentation_result(
transformers_results, id2label
)
)
if "masks" in transformers_results or "png_string" in transformers_results:
return cls(
**process_transformers_v4_segmentation_result(
transformers_results, id2label
)
)
if "boxes" in transformers_results:
return cls(
**process_transformers_detection_result(transformers_results, id2label)
)
@classmethod
def from_detectron2(cls, detectron2_results) -> Detections:
@ -439,6 +541,9 @@ class Detections:
return cls(
xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
confidence=detectron2_results["instances"].scores.cpu().numpy(),
mask=detectron2_results["instances"].pred_masks.cpu().numpy()
if hasattr(detectron2_results["instances"], "pred_masks")
else None,
class_id=detectron2_results["instances"]
.pred_classes.cpu()
.numpy()
@ -448,17 +553,12 @@ class Detections:
@classmethod
def from_inference(cls, roboflow_result: Union[dict, Any]) -> Detections:
"""
Create a Detections object from the [Roboflow](https://roboflow.com/)
Create a `sv.Detections` object from the [Roboflow](https://roboflow.com/)
API inference result or the [Inference](https://inference.roboflow.com/)
package results. This method extracts bounding boxes, class IDs,
confidences, and class names from the Roboflow API result and encapsulates
them into a Detections object.
!!! note
Class names can be accessed using the key 'class_name' in the returned
object's data attribute.
Args:
roboflow_result (dict, any): The result from the
Roboflow API or Inference package containing predictions.
@ -471,10 +571,10 @@ class Detections:
```python
import cv2
import supervision as sv
from inference.models.utils import get_roboflow_model
from inference import get_model
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = get_roboflow_model(model_id="yolov8s-640")
model = get_model(model_id="yolov8s-640")
result = model.infer(image)[0]
detections = sv.Detections.from_inference(result)
@ -500,45 +600,6 @@ class Detections:
data=data,
)
@classmethod
@deprecated(
"`Detections.from_roboflow` is deprecated and will be removed in "
"`supervision-0.22.0`. Use `Detections.from_inference` instead."
)
def from_roboflow(cls, roboflow_result: Union[dict, Any]) -> Detections:
"""
!!! failure "Deprecated"
`Detections.from_roboflow` is deprecated and will be removed in
`supervision-0.22.0`. Use `Detections.from_inference` instead.
Create a Detections object from the [Roboflow](https://roboflow.com/)
API inference result or the [Inference](https://inference.roboflow.com/)
package results.
Args:
roboflow_result (dict): The result from the
Roboflow API containing predictions.
Returns:
(Detections): A Detections object containing the bounding boxes, class IDs,
and confidences of the predictions.
Example:
```python
import cv2
import supervision as sv
from inference.models.utils import get_roboflow_model
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = get_roboflow_model(model_id="yolov8s-640")
result = model.infer(image)[0]
detections = sv.Detections.from_roboflow(result)
```
"""
return cls.from_inference(roboflow_result)
@classmethod
def from_sam(cls, sam_result: List[dict]) -> Detections:
"""
@ -578,7 +639,7 @@ class Detections:
if np.asarray(xywh).shape[0] == 0:
return cls.empty()
xyxy = xywh_to_xyxy(boxes_xywh=xywh)
xyxy = xywh_to_xyxy(xywh=xywh)
return cls(xyxy=xyxy, mask=mask)
@classmethod
@ -710,6 +771,69 @@ class Detections:
class_id=paddledet_result["bbox"][:, 0].astype(int),
)
@classmethod
def from_lmm(
cls, lmm: Union[LMM, str], result: Union[str, dict], **kwargs
) -> Detections:
"""
Creates a Detections object from the given result string based on the specified
Large Multimodal Model (LMM).
Args:
lmm (Union[LMM, str]): The type of LMM (Large Multimodal Model) to use.
result (str): The result string containing the detection data.
**kwargs: Additional keyword arguments required by the specified LMM.
Returns:
Detections: A new Detections object.
Raises:
ValueError: If the LMM is invalid, required arguments are missing, or
disallowed arguments are provided.
ValueError: If the specified LMM is not supported.
Examples:
```python
import supervision as sv
paligemma_result = "<loc0256><loc0256><loc0768><loc0768> cat"
detections = sv.Detections.from_lmm(
sv.LMM.PALIGEMMA,
paligemma_result,
resolution_wh=(1000, 1000),
classes=['cat', 'dog']
)
detections.xyxy
# array([[250., 250., 750., 750.]])
detections.class_id
# array([0])
```
"""
lmm = validate_lmm_parameters(lmm, result, kwargs)
if lmm == LMM.PALIGEMMA:
assert isinstance(result, str)
xyxy, class_id, class_name = from_paligemma(result, **kwargs)
data = {CLASS_NAME_DATA_FIELD: class_name}
return cls(xyxy=xyxy, class_id=class_id, data=data)
if lmm == LMM.FLORENCE_2:
assert isinstance(result, dict)
xyxy, labels, mask, xyxyxyxy = from_florence_2(result, **kwargs)
if len(xyxy) == 0:
return cls.empty()
data = {}
if labels is not None:
data[CLASS_NAME_DATA_FIELD] = labels
if xyxyxyxy is not None:
data[ORIENTED_BOX_COORDINATES] = xyxyxyxy
return cls(xyxy=xyxy, mask=mask, data=data)
raise ValueError(f"Unsupported LMM: {lmm}")
@classmethod
def empty(cls) -> Detections:
"""
@ -732,6 +856,14 @@ class Detections:
class_id=np.array([], dtype=int),
)
def is_empty(self) -> bool:
"""
Returns `True` if the `Detections` object is considered empty.
"""
empty_detections = Detections.empty()
empty_detections.data = self.data
return self == empty_detections
@classmethod
def merge(cls, detections_list: List[Detections]) -> Detections:
"""
@ -739,9 +871,14 @@ class Detections:
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`.
into a single Detections object.
For example, if merging Detections with 3 and 4 detected objects, this method
will return a Detections with 7 objects (7 entries in `xyxy`, `mask`, etc).
!!! Note
When merging, empty `Detections` objects are ignored.
Args:
detections_list (List[Detections]): A list of Detections objects to merge.
@ -781,6 +918,10 @@ class Detections:
array([0.1, 0.2, 0.3])
```
"""
detections_list = [
detections for detections in detections_list if not detections.is_empty()
]
if len(detections_list) == 0:
return Detections.empty()
@ -1055,3 +1196,195 @@ class Detections:
)
return self[indices]
def with_nmm(
self, threshold: float = 0.5, class_agnostic: bool = False
) -> Detections:
"""
Perform non-maximum merging on the current set of object detections.
Args:
threshold (float, optional): The intersection-over-union threshold
to use for non-maximum merging. Defaults to 0.5.
class_agnostic (bool, optional): Whether to perform class-agnostic
non-maximum merging. 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 merging.
Raises:
AssertionError: If `confidence` is None or `class_id` is None and
class_agnostic is False.
![non-max-merging](https://media.roboflow.com/supervision-docs/non-max-merging.png){ align=center width="800" }
""" # noqa: E501 // docs
if len(self) == 0:
return self
assert (
self.confidence is not None
), "Detections confidence must be given for NMM to be executed."
if class_agnostic:
predictions = np.hstack((self.xyxy, self.confidence.reshape(-1, 1)))
else:
assert self.class_id is not None, (
"Detections class_id must be given for NMM to be executed. If you"
" intended to perform class agnostic NMM set class_agnostic=True."
)
predictions = np.hstack(
(
self.xyxy,
self.confidence.reshape(-1, 1),
self.class_id.reshape(-1, 1),
)
)
merge_groups = box_non_max_merge(
predictions=predictions, iou_threshold=threshold
)
result = []
for merge_group in merge_groups:
unmerged_detections = [self[i] for i in merge_group]
merged_detections = merge_inner_detections_objects(
unmerged_detections, threshold
)
result.append(merged_detections)
return Detections.merge(result)
def merge_inner_detection_object_pair(
detections_1: Detections, detections_2: Detections
) -> Detections:
"""
Merges two Detections object into a single Detections object.
Assumes each Detections contains exactly one object.
A `winning` detection is determined based on the confidence score of the two
input detections. This winning detection is then used to specify which
`class_id`, `tracker_id`, and `data` to include in the merged Detections object.
The resulting `confidence` of the merged object is calculated by the weighted
contribution of ea detection to the merged object.
The bounding boxes and masks of the two input detections are merged into a
single bounding box and mask, respectively.
Args:
detections_1 (Detections):
The first Detections object
detections_2 (Detections):
The second Detections object
Returns:
Detections: A new Detections object, with merged attributes.
Raises:
ValueError: If the input Detections objects do not have exactly 1 detected
object.
Example:
```python
import cv2
import supervision as sv
from inference import get_model
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = get_model(model_id="yolov8s-640")
result = model.infer(image)[0]
detections = sv.Detections.from_inference(result)
merged_detections = merge_object_detection_pair(
detections[0], detections[1])
```
"""
if len(detections_1) != 1 or len(detections_2) != 1:
raise ValueError("Both Detections should have exactly 1 detected object.")
validate_fields_both_defined_or_none(detections_1, detections_2)
xyxy_1 = detections_1.xyxy[0]
xyxy_2 = detections_2.xyxy[0]
if detections_1.confidence is None and detections_2.confidence is None:
merged_confidence = None
else:
detection_1_area = (xyxy_1[2] - xyxy_1[0]) * (xyxy_1[3] - xyxy_1[1])
detections_2_area = (xyxy_2[2] - xyxy_2[0]) * (xyxy_2[3] - xyxy_2[1])
merged_confidence = (
detection_1_area * detections_1.confidence[0]
+ detections_2_area * detections_2.confidence[0]
) / (detection_1_area + detections_2_area)
merged_confidence = np.array([merged_confidence])
merged_x1, merged_y1 = np.minimum(xyxy_1[:2], xyxy_2[:2])
merged_x2, merged_y2 = np.maximum(xyxy_1[2:], xyxy_2[2:])
merged_xyxy = np.array([[merged_x1, merged_y1, merged_x2, merged_y2]])
if detections_1.mask is None and detections_2.mask is None:
merged_mask = None
else:
merged_mask = np.logical_or(detections_1.mask, detections_2.mask)
if detections_1.confidence is None and detections_2.confidence is None:
winning_detection = detections_1
elif detections_1.confidence[0] >= detections_2.confidence[0]:
winning_detection = detections_1
else:
winning_detection = detections_2
return Detections(
xyxy=merged_xyxy,
mask=merged_mask,
confidence=merged_confidence,
class_id=winning_detection.class_id,
tracker_id=winning_detection.tracker_id,
data=winning_detection.data,
)
def merge_inner_detections_objects(
detections: List[Detections], threshold=0.5
) -> Detections:
"""
Given N detections each of length 1 (exactly one object inside), combine them into a
single detection object of length 1. The contained inner object will be the merged
result of all the input detections.
For example, this lets you merge N boxes into one big box, N masks into one mask,
etc.
"""
detections_1 = detections[0]
for detections_2 in detections[1:]:
box_iou = box_iou_batch(detections_1.xyxy, detections_2.xyxy)[0]
if box_iou < threshold:
break
detections_1 = merge_inner_detection_object_pair(detections_1, detections_2)
return detections_1
def validate_fields_both_defined_or_none(
detections_1: Detections, detections_2: Detections
) -> None:
"""
Verify that for each optional field in the Detections, both instances either have
the field set to None or both have it set to non-None values.
`data` field is ignored.
Raises:
ValueError: If one field is None and the other is not, for any of the fields.
"""
attributes = get_instance_variables(detections_1)
for attribute in attributes:
value_1 = getattr(detections_1, attribute)
value_2 = getattr(detections_2, attribute)
if (value_1 is None) != (value_2 is None):
raise ValueError(
f"Field '{attribute}' should be consistently None or not None in both "
"Detections."
)

View File

@ -1,12 +1,15 @@
import warnings
from typing import Dict, Iterable, Optional, Tuple
import cv2
import numpy as np
from supervision.detection.core import Detections
from supervision.detection.utils import cross_product
from supervision.draw.color import Color
from supervision.draw.utils import draw_text
from supervision.geometry.core import Point, Position, Vector
from supervision.utils.internal import SupervisionWarnings
class LineZone:
@ -81,6 +84,8 @@ class LineZone:
self.in_count: int = 0
self.out_count: int = 0
self.triggering_anchors = triggering_anchors
if not list(self.triggering_anchors):
raise ValueError("Triggering anchors cannot be empty.")
@staticmethod
def calculate_region_of_interest_limits(vector: Vector) -> Tuple[Vector, Vector]:
@ -140,6 +145,15 @@ class LineZone:
if len(detections) == 0:
return crossed_in, crossed_out
if detections.tracker_id is None:
warnings.warn(
"Line zone counting skipped. LineZone requires tracker_id. Refer to "
"https://supervision.roboflow.com/latest/trackers for more "
"information.",
category=SupervisionWarnings,
)
return crossed_in, crossed_out
all_anchors = np.array(
[
detections.get_anchors_coordinates(anchor)
@ -147,31 +161,23 @@ class LineZone:
]
)
cross_products_1 = cross_product(all_anchors, self.limits[0])
cross_products_2 = cross_product(all_anchors, self.limits[1])
in_limits = (cross_products_1 > 0) == (cross_products_2 > 0)
in_limits = np.all(in_limits, axis=0)
triggers = cross_product(all_anchors, self.vector) < 0
has_any_left_trigger = np.any(triggers, axis=0)
has_any_right_trigger = np.any(~triggers, axis=0)
is_uniformly_triggered = ~(has_any_left_trigger & has_any_right_trigger)
for i, tracker_id in enumerate(detections.tracker_id):
if tracker_id is None:
if not in_limits[i]:
continue
box_anchors = [Point(x=x, y=y) for x, y in all_anchors[:, i, :]]
in_limits = all(
[
self.is_point_in_limits(point=anchor, limits=self.limits)
for anchor in box_anchors
]
)
if not in_limits:
if not is_uniformly_triggered[i]:
continue
triggers = [
self.vector.cross_product(point=anchor) < 0 for anchor in box_anchors
]
if len(set(triggers)) == 2:
continue
tracker_state = triggers[0]
tracker_state = has_any_left_trigger[i]
if tracker_id not in self.tracker_state:
self.tracker_state[tracker_id] = tracker_state
continue

View File

@ -0,0 +1,184 @@
import re
from enum import Enum
from typing import Any, Dict, List, Optional, Tuple, Union
import numpy as np
from supervision.detection.utils import polygon_to_mask, polygon_to_xyxy
class LMM(Enum):
PALIGEMMA = "paligemma"
FLORENCE_2 = "florence_2"
RESULT_TYPES: Dict[LMM, type] = {LMM.PALIGEMMA: str, LMM.FLORENCE_2: dict}
REQUIRED_ARGUMENTS: Dict[LMM, List[str]] = {
LMM.PALIGEMMA: ["resolution_wh"],
LMM.FLORENCE_2: ["resolution_wh"],
}
ALLOWED_ARGUMENTS: Dict[LMM, List[str]] = {
LMM.PALIGEMMA: ["resolution_wh", "classes"],
LMM.FLORENCE_2: ["resolution_wh"],
}
SUPPORTED_TASKS_FLORENCE_2 = [
"<OD>",
"<CAPTION_TO_PHRASE_GROUNDING>",
"<DENSE_REGION_CAPTION>",
"<REGION_PROPOSAL>",
"<OCR_WITH_REGION>",
"<REFERRING_EXPRESSION_SEGMENTATION>",
"<REGION_TO_SEGMENTATION>",
"<OPEN_VOCABULARY_DETECTION>",
"<REGION_TO_CATEGORY>",
"<REGION_TO_DESCRIPTION>",
]
def validate_lmm_parameters(
lmm: Union[LMM, str], result: Any, kwargs: Dict[str, Any]
) -> LMM:
if isinstance(lmm, str):
try:
lmm = LMM(lmm.lower())
except ValueError:
raise ValueError(
f"Invalid lmm value: {lmm}. Must be one of {[e.value for e in LMM]}"
)
if not isinstance(result, RESULT_TYPES[lmm]):
raise ValueError(
f"Invalid LMM result type: {type(result)}. Must be {RESULT_TYPES[lmm]}"
)
required_args = REQUIRED_ARGUMENTS.get(lmm, [])
for arg in required_args:
if arg not in kwargs:
raise ValueError(f"Missing required argument: {arg}")
allowed_args = ALLOWED_ARGUMENTS.get(lmm, [])
for arg in kwargs:
if arg not in allowed_args:
raise ValueError(f"Argument {arg} is not allowed for {lmm.name}")
return lmm
def from_paligemma(
result: str, resolution_wh: Tuple[int, int], classes: Optional[List[str]] = None
) -> Tuple[np.ndarray, Optional[np.ndarray], np.ndarray]:
w, h = resolution_wh
pattern = re.compile(
r"(?<!<loc\d{4}>)<loc(\d{4})><loc(\d{4})><loc(\d{4})><loc(\d{4})> ([\w\s\-]+)"
)
matches = pattern.findall(result)
matches = np.array(matches) if matches else np.empty((0, 5))
xyxy, class_name = matches[:, [1, 0, 3, 2]], matches[:, 4]
xyxy = xyxy.astype(int) / 1024 * np.array([w, h, w, h])
class_name = np.char.strip(class_name.astype(str))
class_id = None
if classes is not None:
mask = np.array([name in classes for name in class_name]).astype(bool)
xyxy, class_name = xyxy[mask], class_name[mask]
class_id = np.array([classes.index(name) for name in class_name])
return xyxy, class_id, class_name
def from_florence_2(
result: dict, resolution_wh: Tuple[int, int]
) -> Tuple[
np.ndarray, Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]
]:
"""
Parse results from the Florence 2 multi-model model.
https://huggingface.co/microsoft/Florence-2-large
Parameters:
result: dict containing the model output
Returns:
xyxy (np.ndarray): An array of shape `(n, 4)` containing
the bounding boxes coordinates in format `[x1, y1, x2, y2]`
labels: (Optional[np.ndarray]): An array of shape `(n,)` containing
the class labels for each bounding box
masks: (Optional[np.ndarray]): An array of shape `(n, h, w)` containing
the segmentation masks for each bounding box
obb_boxes: (Optional[np.ndarray]): An array of shape `(n, 4, 2)` containing
oriented bounding boxes.
"""
assert len(result) == 1, f"Expected result with a single element. Got: {result}"
task = list(result.keys())[0]
if task not in SUPPORTED_TASKS_FLORENCE_2:
raise ValueError(
f"{task} not supported. Supported tasks are: {SUPPORTED_TASKS_FLORENCE_2}"
)
result = result[task]
if task in ["<OD>", "<CAPTION_TO_PHRASE_GROUNDING>", "<DENSE_REGION_CAPTION>"]:
xyxy = np.array(result["bboxes"], dtype=np.float32)
labels = np.array(result["labels"])
return xyxy, labels, None, None
if task == "<REGION_PROPOSAL>":
xyxy = np.array(result["bboxes"], dtype=np.float32)
# provides labels, but they are ["", "", "", ...]
return xyxy, None, None, None
if task == "<OCR_WITH_REGION>":
xyxyxyxy = np.array(result["quad_boxes"], dtype=np.float32)
xyxyxyxy = xyxyxyxy.reshape(-1, 4, 2)
xyxy = np.array([polygon_to_xyxy(polygon) for polygon in xyxyxyxy])
labels = np.array(result["labels"])
return xyxy, labels, None, xyxyxyxy
if task in ["<REFERRING_EXPRESSION_SEGMENTATION>", "<REGION_TO_SEGMENTATION>"]:
xyxy_list = []
masks_list = []
for polygons_of_same_class in result["polygons"]:
for polygon in polygons_of_same_class:
polygon = np.reshape(polygon, (-1, 2)).astype(np.int32)
mask = polygon_to_mask(polygon, resolution_wh).astype(bool)
masks_list.append(mask)
xyxy = polygon_to_xyxy(polygon)
xyxy_list.append(xyxy)
# per-class labels also provided, but they are ["", "", "", ...]
# when we figure out how to set class names, we can do
# zip(result["labels"], result["polygons"])
xyxy = np.array(xyxy_list, dtype=np.float32)
masks = np.array(masks_list)
return xyxy, None, masks, None
if task == "<OPEN_VOCABULARY_DETECTION>":
xyxy = np.array(result["bboxes"], dtype=np.float32)
labels = np.array(result["bboxes_labels"])
# Also has "polygons" and "polygons_labels", but they don't seem to be used
return xyxy, labels, None, None
if task in ["<REGION_TO_CATEGORY>", "<REGION_TO_DESCRIPTION>"]:
assert isinstance(
result, str
), f"Expected string as <REGION_TO_CATEGORY> result, got {type(result)}"
if result == "No object detected.":
return np.empty((0, 4), dtype=np.float32), np.array([]), None, None
pattern = re.compile(r"<loc_(\d+)><loc_(\d+)><loc_(\d+)><loc_(\d+)>")
match = pattern.search(result)
assert (
match is not None
), f"Expected string to end in location tags, but got {result}"
w, h = resolution_wh
xyxy = np.array([match.groups()], dtype=np.float32)
xyxy *= np.array([w, h, w, h]) / 1000
result_string = result[: match.start()]
labels = np.array([result_string])
return xyxy, labels, None, None
assert False, f"Unimplemented task: {task}"

Some files were not shown because too many files have changed in this diff Show More