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

Merge pull request #1655 from roboflow/develop

Browse Source 
`supervision-0.25.0` release
This commit is contained in:
LinasKo authored and GitHub committed 2024-11-12 23:58:16 +02:00
2 parents 58e27f44f1 8d562c16a6
commit eaeb51a8db
55 files changed
+5226 -1021
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.github/workflows/poetry-test.yml
+37
View File
@@ -0,0 +1,37 @@
name: 🔧 Poetry Check and Installation Test Workflow
on:
push:
paths:
- 'poetry.lock'
- 'pyproject.toml'
pull_request:
paths:
- 'poetry.lock'
- 'pyproject.toml'
workflow_dispatch:
jobs:
poetry-tests:
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
runs-on: ${{ matrix.os }}
steps:
- name: 📥 Checkout the repository
uses: actions/checkout@v4
- name: 🐍 Set up Python ${{ matrix.python-version }}
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
- name: 📦 Install the base dependencies
run: python -m pip install --upgrade poetry
- name: 🔍 Check the correctness of the project config
run: poetry check
- name: 🚀 Do Install the package Test
run: poetry install
.github/workflows/publish-test.yml
+4 -3
View File
@@ -2,9 +2,9 @@ name: Publish Supervision Pre-Releases to PyPI and TestPyPI
on:
push:
tags:
- '[0-9]+.[0-9]+[0-9]+.[0-9]+a[0-9]'
- '[0-9]+.[0-9]+[0-9]+.[0-9]+b[0-9]'
- '[0-9]+.[0-9]+[0-9]+.[0-9]+rc[0-9]'
- "[0-9]+.[0-9]+[0-9]+.[0-9]+a[0-9]"
- "[0-9]+.[0-9]+[0-9]+.[0-9]+b[0-9]"
- "[0-9]+.[0-9]+[0-9]+.[0-9]+rc[0-9]"
workflow_dispatch:
@@ -12,6 +12,7 @@ jobs:
build-and-publish-pre-release-pypi:
name: Build and publish to PyPI
runs-on: ubuntu-latest
environment: test
permissions:
id-token: write
strategy:
.github/workflows/publish.yml
+3 -2
View File
@@ -2,13 +2,14 @@ name: Publish Supervision Releases to PyPI and TestPyPI
on:
push:
tags:
- '[0-9]+.[0-9]+[0-9]+.[0-9]'
- "[0-9]+.[0-9]+[0-9]+.[0-9]"
workflow_dispatch:
jobs:
build-and-publish-pre-release:
runs-on: ubuntu-latest
environment: release
permissions:
id-token: write
strategy:
@@ -24,7 +25,7 @@ jobs:
with:
python-version: ${{ matrix.python-version }}
- name: 🏗️ Build source and wheel distributions
- name: 🏗️ Build source and wheel distributions
run: |
python -m pip install --upgrade build twine
python -m build
.github/workflows/test.yml
+1 -1
View File
@@ -9,7 +9,7 @@ jobs:
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
python-version: ["3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
steps:
- name: 🛎️ Checkout
uses: actions/checkout@v4
.github/workflows/welcome.yml
-18
View File
@@ -1,18 +0,0 @@
name: Welcome WorkFlow
on:
issues:
types: [opened]
pull_request_target:
types: [opened]
jobs:
build:
name: 👋 Welcome
runs-on: ubuntu-latest
steps:
- uses: actions/first-interaction@v1.3.0
with:
repo-token: ${{ secrets.GITHUB_TOKEN }}
issue-message: "Hello there, thank you for opening an Issue ! 🙏🏻 The team was notified and they will get back to you asap."
pr-message: "Hello there, thank you for opening an PR ! 🙏🏻 The team was notified and they will get back to you asap."
.pre-commit-config.yaml
+2 -2
View File
@@ -7,7 +7,7 @@ ci:
repos:
- repo: https://github.com/pre-commit/pre-commit-hooks
rev: v4.6.0
rev: v5.0.0
hooks:
- id: trailing-whitespace
exclude: test/.*\.py
@@ -32,7 +32,7 @@ repos:
additional_dependencies: ["bandit[toml]"]
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.6.8
rev: v0.7.3
hooks:
- id: ruff
args: [--fix, --exit-non-zero-on-fix]
README.md
+64 -66
View File
@@ -21,7 +21,7 @@
[![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)
[![discord](https://img.shields.io/discord/1159501506232451173?logo=discord&label=discord&labelColor=fff&color=5865f2&link=https%3A%2F%2Fdiscord.gg%2FGbfgXGJ8Bk)](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 align="center">
@@ -34,8 +34,6 @@
**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://media.roboflow.com/supervision/supervision-hacktoberfest-banner-2024.png)](https://github.com/roboflow/supervision/issues?q=is%3Aissue+is%3Aopen+label%3Ahacktoberfest)
## 💻 install
Pip install the supervision package in a
@@ -137,88 +135,88 @@ for path, image, annotation in ds:
- load
```python
dataset = sv.DetectionDataset.from_yolo(
images_directory_path=...,
annotations_directory_path=...,
data_yaml_path=...
)
```python
dataset = sv.DetectionDataset.from_yolo(
images_directory_path=...,
annotations_directory_path=...,
data_yaml_path=...
)
dataset = sv.DetectionDataset.from_pascal_voc(
images_directory_path=...,
annotations_directory_path=...
)
dataset = sv.DetectionDataset.from_pascal_voc(
images_directory_path=...,
annotations_directory_path=...
)
dataset = sv.DetectionDataset.from_coco(
images_directory_path=...,
annotations_path=...
)
```
dataset = sv.DetectionDataset.from_coco(
images_directory_path=...,
annotations_path=...
)
```
- split
```python
train_dataset, test_dataset = dataset.split(split_ratio=0.7)
test_dataset, valid_dataset = test_dataset.split(split_ratio=0.5)
```python
train_dataset, test_dataset = dataset.split(split_ratio=0.7)
test_dataset, valid_dataset = test_dataset.split(split_ratio=0.5)
len(train_dataset), len(test_dataset), len(valid_dataset)
# (700, 150, 150)
```
len(train_dataset), len(test_dataset), len(valid_dataset)
# (700, 150, 150)
```
- merge
```python
ds_1 = sv.DetectionDataset(...)
len(ds_1)
# 100
ds_1.classes
# ['dog', 'person']
```python
ds_1 = sv.DetectionDataset(...)
len(ds_1)
# 100
ds_1.classes
# ['dog', 'person']
ds_2 = sv.DetectionDataset(...)
len(ds_2)
# 200
ds_2.classes
# ['cat']
ds_2 = sv.DetectionDataset(...)
len(ds_2)
# 200
ds_2.classes
# ['cat']
ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
len(ds_merged)
# 300
ds_merged.classes
# ['cat', 'dog', 'person']
```
ds_merged = sv.DetectionDataset.merge([ds_1, ds_2])
len(ds_merged)
# 300
ds_merged.classes
# ['cat', 'dog', 'person']
```
- save
```python
dataset.as_yolo(
images_directory_path=...,
annotations_directory_path=...,
data_yaml_path=...
)
```python
dataset.as_yolo(
images_directory_path=...,
annotations_directory_path=...,
data_yaml_path=...
)
dataset.as_pascal_voc(
images_directory_path=...,
annotations_directory_path=...
)
dataset.as_pascal_voc(
images_directory_path=...,
annotations_directory_path=...
)
dataset.as_coco(
images_directory_path=...,
annotations_path=...
)
```
dataset.as_coco(
images_directory_path=...,
annotations_path=...
)
```
- convert
```python
sv.DetectionDataset.from_yolo(
images_directory_path=...,
annotations_directory_path=...,
data_yaml_path=...
).as_pascal_voc(
images_directory_path=...,
annotations_directory_path=...
)
```
```python
sv.DetectionDataset.from_yolo(
images_directory_path=...,
annotations_directory_path=...,
data_yaml_path=...
).as_pascal_voc(
images_directory_path=...,
annotations_directory_path=...
)
```
</details>
docs/assets.md
+2
View File
@@ -1,5 +1,6 @@
---
comments: true
status: new
---
# Assets
@@ -13,6 +14,7 @@ To install the Supervision assets utility, you can use `pip`. This utility is av
as an extra within the Supervision package.
!!! example "pip install"
```bash
pip install "supervision[assets]"
```
docs/changelog.md
+144
View File
@@ -1,5 +1,149 @@
# CHANGELOG
### 0.25.0 <small>Nov 12, 2024</small>
- No removals or deprecations in this release!
- Essential update to the [`LineZone`](https://supervision.roboflow.com/0.25.0/detection/tools/line_zone/): when computing line crossings, detections that jitter might be counted twice (or more). This can now be solved with the `minimum_crossing_threshold` argument. If you set it to `2` or more, extra frames will be used to confirm the crossing, improving the accuracy significantly. ([#1540](https://github.com/roboflow/supervision/pull/1540))
- It is now possible to track objects detected as [`KeyPoints`](https://supervision.roboflow.com/0.25.0/keypoint/core/#supervision.keypoint.core.KeyPoints). See the complete step-by-step guide in the [Object Tracking Guide](https://supervision.roboflow.com/latest/how_to/track_objects/#keypoints). ([#1658](https://github.com/roboflow/supervision/pull/1658))
```python
import numpy as np
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8m-pose.pt")
tracker = sv.ByteTrack()
trace_annotator = sv.TraceAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model(frame)[0]
key_points = sv.KeyPoints.from_ultralytics(results)
detections = key_points.as_detections()
detections = tracker.update_with_detections(detections)
annotated_image = trace_annotator.annotate(frame.copy(), detections)
return annotated_image
sv.process_video(
source_path="input_video.mp4",
target_path="output_video.mp4",
callback=callback
)
```
- Added `is_empty` method to [`KeyPoints`](https://supervision.roboflow.com/0.25.0/keypoint/core/#supervision.keypoint.core.KeyPoints) to check if there are any keypoints in the object. ([#1658](https://github.com/roboflow/supervision/pull/1658))
- Added `as_detections` method to [`KeyPoints`](https://supervision.roboflow.com/0.25.0/keypoint/core/#supervision.keypoint.core.KeyPoints) that converts `KeyPoints` to `Detections`. ([#1658](https://github.com/roboflow/supervision/pull/1658))
- Added a new video to `supervision[assets]`. ([#1657](https://github.com/roboflow/supervision/pull/1657))
```python
from supervision.assets import download_assets, VideoAssets
path_to_video = download_assets(VideoAssets.SKIING)
```
- Supervision can now be used with [`Python 3.13`](https://docs.python.org/3/whatsnew/3.13.html). The most renowned update is the ability to run Python [without Global Interpreter Lock (GIL)](https://docs.python.org/3/whatsnew/3.13.html#whatsnew313-free-threaded-cpython). We expect support for this among our dependencies to be inconsistent, but if you do attempt it - let us know the results! ([#1595](https://github.com/roboflow/supervision/pull/1595))
- Added [`Mean Average Recall`](https://supervision.roboflow.com/latest/metrics/mean_average_recall/) mAR metric, which returns a recall score, averaged over IoU thresholds, detected object classes, and limits imposed on maximum considered detections. ([#1661](https://github.com/roboflow/supervision/pull/1661))
```python
import supervision as sv
from supervision.metrics import MeanAverageRecall
predictions = sv.Detections(...)
targets = sv.Detections(...)
map_metric = MeanAverageRecall()
map_result = map_metric.update(predictions, targets).compute()
map_result.plot()
```
- Added [`Precision`](https://supervision.roboflow.com/latest/metrics/precision/) and [`Recall`](https://supervision.roboflow.com/latest/metrics/recall/) metrics, providing a baseline for comparing model outputs to ground truth or another model ([#1609](https://github.com/roboflow/supervision/pull/1609))
```python
import supervision as sv
from supervision.metrics import Recall
predictions = sv.Detections(...)
targets = sv.Detections(...)
recall_metric = Recall()
recall_result = recall_metric.update(predictions, targets).compute()
recall_result.plot()
```
- All Metrics now support Oriented Bounding Boxes (OBB) ([#1593](https://github.com/roboflow/supervision/pull/1593))
```python
import supervision as sv
from supervision.metrics import F1_Score
predictions = sv.Detections(...)
targets = sv.Detections(...)
f1_metric = MeanAverageRecall(metric_target=sv.MetricTarget.ORIENTED_BOUNDING_BOXES)
f1_result = f1_metric.update(predictions, targets).compute()
```
- Introducing Smart Labels! When `smart_position` is set for [`LabelAnnotator`](https://supervision.roboflow.com/0.25.0/detection/annotators/#supervision.annotators.core.LabelAnnotator), [`RichLabelAnnotator`](https://supervision.roboflow.com/0.25.0/detection/annotators/#supervision.annotators.core.RichLabelAnnotator) or [`VertexLabelAnnotator`](https://supervision.roboflow.com/0.25.0/detection/annotators/#supervision.annotators.core.RichLabelAnnotator), the labels will move around to avoid overlapping others. ([#1625](https://github.com/roboflow/supervision/pull/1625))
```python
import supervision as sv
from ultralytics import YOLO
image = cv2.imread("image.jpg")
label_annotator = sv.LabelAnnotator(smart_position=True)
model = YOLO("yolo11m.pt")
results = model(image)[0]
detections = sv.Detections.from_ultralytics(results)
annotated_frame = label_annotator.annotate(first_frame.copy(), detections)
sv.plot_image(annotated_frame)
```
- Added the `metadata` variable to [`Detections`](https://supervision.roboflow.com/0.25.0/detection/core/#supervision.detection.core.Detections). It allows you to store custom data per-image, rather than per-detected-object as was possible with `data` variable. For example, `metadata` could be used to store the source video path, camera model or camera parameters. ([#1589](https://github.com/roboflow/supervision/pull/1589))
```python
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8m")
result = model("image.png")[0]
detections = sv.Detections.from_ultralytics(result)
# Items in `data` must match length of detections
object_ids = [num for num in range(len(detections))]
detections.data["object_number"] = object_ids
# Items in `metadata` can be of any length.
detections.metadata["camera_model"] = "Luxonis OAK-D"
```
- Added a `py.typed` type hints metafile. It should provide a stronger signal to type annotators and IDEs that type support is available. ([#1586](https://github.com/roboflow/supervision/pull/1586))
- `ByteTrack` no longer requires `detections` to have a `class_id` ([#1637](https://github.com/roboflow/supervision/pull/1637))
- `draw_line`, `draw_rectangle`, `draw_filled_rectangle`, `draw_polygon`, `draw_filled_polygon` and `PolygonZoneAnnotator` now comes with a default color ([#1591](https://github.com/roboflow/supervision/pull/1591))
- Dataset classes are treated as case-sensitive when merging multiple datasets. ([#1643](https://github.com/roboflow/supervision/pull/1643))
- Expanded [metrics documentation](https://supervision.roboflow.com/0.25.0/metrics/f1_score/) with example plots and printed results ([#1660](https://github.com/roboflow/supervision/pull/1660))
- Added usage example for polygon zone ([#1608](https://github.com/roboflow/supervision/pull/1608))
- Small improvements to error handling in polygons: ([#1602](https://github.com/roboflow/supervision/pull/1602))
- Updated [`ByteTrack`](https://supervision.roboflow.com/0.25.0/trackers/#supervision.tracker.byte_tracker.core.ByteTrack), removing shared variables. Previously, multiple instances of `ByteTrack` would share some date, requiring liberal use of `tracker.reset()`. ([#1603](https://github.com/roboflow/supervision/pull/1603)), ([#1528](https://github.com/roboflow/supervision/pull/1528))
- Fixed a bug where `class_agnostic` setting in `MeanAveragePrecision` would not work. ([#1577](https://github.com/roboflow/supervision/pull/1577)) hacktoberfest
- Removed welcome workflow from our CI system. ([#1596](https://github.com/roboflow/supervision/pull/1596))
- Large refactor of `ByteTrack`: STrack moved to separate class, removed superfluous `BaseTrack` class, removed unused variables ([#1603](https://github.com/roboflow/supervision/pull/1603))
- Large refactor of `RichLabelAnnotator`, matching its contents with `LabelAnnotator`. ([#1625](https://github.com/roboflow/supervision/pull/1625))
### 0.24.0 <small>Oct 4, 2024</small>
- Added [F1 score](https://supervision.roboflow.com/0.24.0/metrics/f1_score/#supervision.metrics.f1_score.F1Score) as a new metric for detection and segmentation. [#1521](https://github.com/roboflow/supervision/pull/1521)
docs/deprecated.md
+4
View File
@@ -19,6 +19,10 @@ These features are phased out due to better alternatives or potential issues in
# Removed
### 0.25.0
No removals in this version!
### 0.24.0
- The `frame_resolution_wh ` parameter in [`sv.PolygonZone`](detection/tools/polygon_zone.md/#supervision.detection.tools.polygon_zone.PolygonZone) has been removed.
docs/detection/tools/inference_slicer.md
-1
View File
@@ -1,6 +1,5 @@
---
comments: true
status: new
---
# InferenceSlicer
docs/detection/tools/polygon_zone.md
-1
View File
@@ -1,6 +1,5 @@
---
comments: true
status: new
---
<div class="md-typeset">
docs/how_to/track_objects.md
+353 -8
View File
@@ -1,18 +1,21 @@
---
comments: true
status: new
---
# Track Objects
Leverage Supervision's advanced capabilities for enhancing your video analysis by
seamlessly [tracking](/latest/trackers/) objects recognized by
a multitude of object detection and segmentation models. This comprehensive guide will
a multitude of object detection, segmentation and keypoint models. This comprehensive guide will
take you through the steps to perform inference using the YOLOv8 model via either the
[Inference](https://github.com/roboflow/inference) or
[Ultralytics](https://github.com/ultralytics/ultralytics) packages. Following this,
you'll discover how to track these objects efficiently and annotate your video content
for a deeper analysis.
## Object Detection & Segmentation
To make it easier for you to follow our tutorial download the video we will use as an
example. You can do this using
[`supervision[assets]`](/latest/assets/) extension.
@@ -27,7 +30,7 @@ download_assets(VideoAssets.PEOPLE_WALKING)
<source src="https://media.roboflow.com/supervision/video-examples/people-walking.mp4" type="video/mp4">
</video>
## Run Inference
### Run Inference
First, you'll need to obtain predictions from your object detection or segmentation
model. In this tutorial, we are using the YOLOv8 model as an example. However,
@@ -40,7 +43,12 @@ by obtaining model predictions and then annotating the frame based on these pred
This `callback` function will be essential in the subsequent steps of the tutorial, as
it will be modified to include tracking, labeling, and trace annotations.
!!! tip
Both object detection and segmentation models are supported. Try it with `yolov8n.pt` or `yolov8n-640-seg`!
=== "Ultralytics"
```{ .py }
import numpy as np
import supervision as sv
@@ -62,6 +70,7 @@ it will be modified to include tracking, labeling, and trace annotations.
```
=== "Inference"
```{ .py }
import numpy as np
import supervision as sv
@@ -86,7 +95,7 @@ it will be modified to include tracking, labeling, and trace annotations.
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/run-inference.mp4" type="video/mp4">
</video>
## Tracking
### Tracking
After running inference and obtaining predictions, the next step is to track the
detected objects throughout the video. Utilizing Supervision’s
@@ -95,6 +104,7 @@ functionality, each detected object is assigned a unique tracker ID,
enabling the continuous following of the object's motion path across different frames.
=== "Ultralytics"
```{ .py hl_lines="6 12" }
import numpy as np
import supervision as sv
@@ -118,6 +128,7 @@ enabling the continuous following of the object's motion path across different f
```
=== "Inference"
```{ .py hl_lines="6 12" }
import numpy as np
import supervision as sv
@@ -140,7 +151,7 @@ enabling the continuous following of the object's motion path across different f
)
```
## Annotate Video with Tracking IDs
### Annotate Video with Tracking IDs
Annotating the video with tracking IDs helps in distinguishing and following each object
distinctly. With the
@@ -149,6 +160,7 @@ in Supervision, we can overlay the tracker IDs and class labels on the detected
offering a clear visual representation of each object's class and unique identifier.
=== "Ultralytics"
```{ .py hl_lines="8 15-19 23-24" }
import numpy as np
import supervision as sv
@@ -183,6 +195,7 @@ offering a clear visual representation of each object's class and unique identif
```
=== "Inference"
```{ .py hl_lines="8 15-19 23-24" }
import numpy as np
import supervision as sv
@@ -220,7 +233,7 @@ offering a clear visual representation of each object's class and unique identif
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-tracking-ids.mp4" type="video/mp4">
</video>
## Annotate Video with Traces
### Annotate Video with Traces
Adding traces to the video involves overlaying the historical paths of the detected
objects. This feature, powered by the
@@ -229,6 +242,7 @@ allows for visualizing the trajectories of objects, helping in understanding the
movement patterns and interactions between objects in the video.
=== "Ultralytics"
```{ .py hl_lines="9 26-27" }
import numpy as np
import supervision as sv
@@ -266,6 +280,7 @@ movement patterns and interactions between objects in the video.
```
=== "Inference"
```{ .py hl_lines="9 26-27" }
import numpy as np
import supervision as sv
@@ -306,6 +321,336 @@ movement patterns and interactions between objects in the video.
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/annotate-video-with-traces.mp4" type="video/mp4">
</video>
This structured walkthrough should give a detailed pathway to annotate videos
effectively using Supervision’s various functionalities, including object tracking and
trace annotations.
## Keypoints
Models aren't limited to object detection and segmentation. Keypoint detection allows for detailed analysis of body joints and connections, especially valuable for applications like human pose estimation. This section introduces keypoint tracking. We'll walk through the steps of annotating keypoints, converting them into bounding box detections compatible with `ByteTrack`, and applying detection smoothing for enhanced stability.
To make it easier for you to follow our tutorial, let's download the video we will use as an
example. You can do this using [`supervision[assets]`](/latest/assets/) extension.
```python
from supervision.assets import download_assets, VideoAssets
download_assets(VideoAssets.SKIING)
```
<video controls muted>
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/skiing-hd.mp4" type="video/mp4">
</video>
### Keypoint Detection
First, you'll need to obtain predictions from your keypoint detection model. In this tutorial, we are using the YOLOv8 model as an example. However,
Supervision is versatile and compatible with various models. Check this [link](/latest/keypoint/core/) for guidance on how to plug in other models.
We will define a `callback` function, which will process each frame of the video by obtaining model predictions and then annotating the frame based on these predictions.
Let's immediately visualize the results with our [`EdgeAnnotator`](/latest/keypoint/annotators/#supervision.keypoint.annotators.EdgeAnnotator) and [`VertexAnnotator`](https://supervision.roboflow.com/latest/keypoint/annotators/#supervision.keypoint.annotators.VertexAnnotator).
=== "Ultralytics"
```{ .py hl_lines="5 10-11" }
import numpy as np
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8m-pose.pt")
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model(frame)[0]
key_points = sv.KeyPoints.from_ultralytics(results)
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
return vertex_annotator.annotate(
annotated_frame, key_points=key_points)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
=== "Inference"
```{ .py hl_lines="5-6 11-12" }
import numpy as np
import supervision as sv
from inference.models.utils import get_roboflow_model
model = get_roboflow_model(
model_id="yolov8m-pose-640", api_key=<ROBOFLOW API KEY>)
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model.infer(frame)[0]
key_points = sv.KeyPoints.from_inference(results)
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
return vertex_annotator.annotate(
annotated_frame, key_points=key_points)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
<video controls>
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/track-keypoints-only-keypoints.mp4" type="video/mp4">
</video>
### Convert to Detections
Keypoint tracking is currently supported via the conversion of `KeyPoints` to `Detections`. This is achieved with the [`KeyPoints.as_detections()`](/latest/keypoint/core/#supervision.keypoint.core.KeyPoints.as_detections) function.
Let's convert to detections and visualize the results with our [`BoxAnnotator`](/latest/detection/annotators/#supervision.annotators.core.BoxAnnotator).
!!! tip
You may use the `selected_keypoint_indices` argument to specify a subset of keypoints to convert. This is useful when some keypoints could be occluded. For example: a person might swing their arm, causing the elbow to be occluded by the torso sometimes.
=== "Ultralytics"
```{ .py hl_lines="8 13 19-20" }
import numpy as np
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8m-pose.pt")
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
box_annotator = sv.BoxAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model(frame)[0]
key_points = sv.KeyPoints.from_ultralytics(results)
detections = key_points.as_detections()
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
annotated_frame = vertex_annotator.annotate(
annotated_frame, key_points=key_points)
return box_annotator.annotate(
annotated_frame, detections=detections)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
=== "Inference"
```{ .py hl_lines="9 14 20-21" }
import numpy as np
import supervision as sv
from inference.models.utils import get_roboflow_model
model = get_roboflow_model(
model_id="yolov8m-pose-640", api_key=<ROBOFLOW API KEY>)
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
box_annotator = sv.BoxAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model.infer(frame)[0]
key_points = sv.KeyPoints.from_inference(results)
detections = key_points.as_detections()
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
annotated_frame = vertex_annotator.annotate(
annotated_frame, key_points=key_points)
return box_annotator.annotate(
annotated_frame, detections=detections)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
<video controls>
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/track-keypoints-converted-to-detections.mp4" type="video/mp4">
</video>
### Keypoint Tracking
Now that we have a `Detections` object, we can track it throughout the video. Utilizing Supervision’s [`sv.ByteTrack`](/latest/trackers/#supervision.tracker.byte_tracker.core.ByteTrack) functionality, each detected object is assigned a unique tracker ID, enabling the continuous following of the object's motion path across different frames. We shall visualize the result with `TraceAnnotator`.
=== "Ultralytics"
```{ .py hl_lines="10-11 17 25-26" }
import numpy as np
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8m-pose.pt")
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
box_annotator = sv.BoxAnnotator()
tracker = sv.ByteTrack()
trace_annotator = sv.TraceAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model(frame)[0]
key_points = sv.KeyPoints.from_ultralytics(results)
detections = key_points.as_detections()
detections = tracker.update_with_detections(detections)
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
annotated_frame = vertex_annotator.annotate(
annotated_frame, key_points=key_points)
annotated_frame = box_annotator.annotate(
annotated_frame, detections=detections)
return trace_annotator.annotate(
annotated_frame, detections=detections)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
=== "Inference"
```{ .py hl_lines="11-12 18 26-27" }
import numpy as np
import supervision as sv
from inference.models.utils import get_roboflow_model
model = get_roboflow_model(
model_id="yolov8m-pose-640", api_key=<ROBOFLOW API KEY>)
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
box_annotator = sv.BoxAnnotator()
tracker = sv.ByteTrack()
trace_annotator = sv.TraceAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model.infer(frame)[0]
key_points = sv.KeyPoints.from_inference(results)
detections = key_points.as_detections()
detections = tracker.update_with_detections(detections)
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
annotated_frame = vertex_annotator.annotate(
annotated_frame, key_points=key_points)
annotated_frame = box_annotator.annotate(
annotated_frame, detections=detections)
return trace_annotator.annotate(
annotated_frame, detections=detections)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
<video controls>
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/track-keypoints-with-tracking.mp4" type="video/mp4">
</video>
### Bonus: Smoothing
We could stop here as we have successfully tracked the object detected by the keypoint model. However, we can further enhance the stability of the boxes by applying [`DetectionsSmoother`](/latest/detection/tools/smoother/). This tool helps in stabilizing the boxes by smoothing the bounding box coordinates across frames. It is very simple to use:
=== "Ultralytics"
```{ .py hl_lines="11 19" }
import numpy as np
import supervision as sv
from ultralytics import YOLO
model = YOLO("yolov8m-pose.pt")
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
box_annotator = sv.BoxAnnotator()
tracker = sv.ByteTrack()
smoother = sv.DetectionsSmoother()
trace_annotator = sv.TraceAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model(frame)[0]
key_points = sv.KeyPoints.from_ultralytics(results)
detections = key_points.as_detections()
detections = tracker.update_with_detections(detections)
detections = smoother.update_with_detections(detections)
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
annotated_frame = vertex_annotator.annotate(
annotated_frame, key_points=key_points)
annotated_frame = box_annotator.annotate(
annotated_frame, detections=detections)
return trace_annotator.annotate(
annotated_frame, detections=detections)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
=== "Inference"
```{ .py hl_lines="12 20" }
import numpy as np
import supervision as sv
from inference.models.utils import get_roboflow_model
model = get_roboflow_model(
model_id="yolov8m-pose-640", api_key=<ROBOFLOW API KEY>)
edge_annotator = sv.EdgeAnnotator()
vertex_annotator = sv.VertexAnnotator()
box_annotator = sv.BoxAnnotator()
tracker = sv.ByteTrack()
smoother = sv.DetectionsSmoother()
trace_annotator = sv.TraceAnnotator()
def callback(frame: np.ndarray, _: int) -> np.ndarray:
results = model.infer(frame)[0]
key_points = sv.KeyPoints.from_inference(results)
detections = key_points.as_detections()
detections = tracker.update_with_detections(detections)
detections = smoother.update_with_detections(detections)
annotated_frame = edge_annotator.annotate(
frame.copy(), key_points=key_points)
annotated_frame = vertex_annotator.annotate(
annotated_frame, key_points=key_points)
annotated_frame = box_annotator.annotate(
annotated_frame, detections=detections)
return trace_annotator.annotate(
annotated_frame, detections=detections)
sv.process_video(
source_path="skiing.mp4",
target_path="result.mp4",
callback=callback
)
```
<video controls>
<source src="https://media.roboflow.com/supervision/video-examples/how-to/track-objects/track-keypoints-with-smoothing.mp4" type="video/mp4">
</video>
This structured walkthrough should give a detailed pathway to annotate videos effectively using Supervision’s various functionalities, including object tracking and trace annotations.
docs/keypoint/annotators.md
+1
View File
@@ -1,5 +1,6 @@
---
comments: true
status: new
---
# Annotators
docs/keypoint/core.md
+1
View File
@@ -1,5 +1,6 @@
---
comments: true
status: new
---
# Keypoint Detection
docs/metrics/common_values.md
+20
View File
@@ -0,0 +1,20 @@
---
comments: true
status: new
---
# Common Values
This page contains supplementary values, types and enums that metrics use.
<div class="md-typeset">
<h2><a href="#supervision.metrics.core.MetricTarget">MetricTarget</a></h2>
</div>
:::supervision.metrics.core.MetricTarget
<div class="md-typeset">
<h2><a href="#supervision.metrics.core.AveragingMethod">AveragingMethod</a></h2>
</div>
:::supervision.metrics.core.AveragingMethod
docs/metrics/mean_average_recall.md
+18
View File
@@ -0,0 +1,18 @@
---
comments: true
status: new
---
# Mean Average Recall
<div class="md-typeset">
<h2><a href="#supervision.metrics.mean_average_recall.MeanAverageRecall">MeanAverageRecall</a></h2>
</div>
:::supervision.metrics.mean_average_recall.MeanAverageRecall
<div class="md-typeset">
<h2><a href="#supervision.metrics.mean_average_recall.MeanAverageRecallResult">MeanAverageRecallResult</a></h2>
</div>
:::supervision.metrics.mean_average_recall.MeanAverageRecallResult
docs/metrics/precision.md
+18
View File
@@ -0,0 +1,18 @@
---
comments: true
status: new
---
# Precision
<div class="md-typeset">
<h2><a href="#supervision.metrics.precision.Precision">Precision</a></h2>
</div>
:::supervision.metrics.precision.Precision
<div class="md-typeset">
<h2><a href="#supervision.metrics.precision.PrecisionResult">PrecisionResult</a></h2>
</div>
:::supervision.metrics.precision.PrecisionResult
docs/metrics/recall.md
+18
View File
@@ -0,0 +1,18 @@
---
comments: true
status: new
---
# Recall
<div class="md-typeset">
<h2><a href="#supervision.metrics.recall.Recall">Recall</a></h2>
</div>
:::supervision.metrics.recall.Recall
<div class="md-typeset">
<h2><a href="#supervision.metrics.recall.RecallResult">RecallResult</a></h2>
</div>
:::supervision.metrics.recall.RecallResult
docs/notebooks/small-object-detection-with-sahi.ipynb
Vendored
+5 -5
View File
@@ -15,7 +15,7 @@
"\n",
"This cookbook shows how to use [Slicing Aided Hyper Inference (SAHI) ](https://arxiv.org/abs/2202.06934) for small object detection with `supervision`.\n",
"\n",
"![\"Small Object Detection\"](https://raw.githubusercontent.com/ediardo/notebooks/main/sahi/animation.gif \"Small Object Detection\")\n",
"![\"Small Object Detection\"](https://media.roboflow.com/supervision/cookbooks/sahi/animation.gif \"Small Object Detection\")\n",
"\n",
"Click the Open in Colab button to run the cookbook on Google Colab.\n",
"\n",
@@ -70,7 +70,7 @@
"\n",
"Detecting people (or their heads) is a common problem that has been addressed by many researchers in the past. In this project, we\u2019ll use an open-source public dataset and a fine-tuned model to perform inference on images.\n",
"\n",
"![Roboflow Universe](https://raw.githubusercontent.com/ediardo/notebooks/main/sahi/roboflow_universe.png \"Open source model for counting people's heads\")\n",
"![Roboflow Universe](https://media.roboflow.com/supervision/cookbooks/sahi/roboflow_universe.png \"Open source model for counting people's heads\")\n",
"\n",
"Some details about the project [\"people_counterv0 Computer Vision Project\"](https://universe.roboflow.com/sit-cx0ng/people_counterv0):\n",
"\n",
@@ -782,9 +782,9 @@
"\n",
"| Example| Observations |\n",
"|----|----|\n",
"| ![Overlapping](https://github.com/ediardo/notebooks/blob/main/sahi/overlapping_1.png?raw=true \"Overlapping\") | False Negative, Incomplete bbox |\n",
"| ![Overlapping](https://raw.githubusercontent.com/ediardo/notebooks/main/sahi/overlapping_2.png \"Overlapping\")| Double detection, Incomplete bbox|\n",
"| ![Overlapping](https://raw.githubusercontent.com/ediardo/notebooks/main/sahi/overlapping_3.png \"Overlapping\")| Incomplete bounding box|\n",
"| ![Overlapping](https://media.roboflow.com/supervision/cookbooks/sahi/overlapping_1.png \"Overlapping\") | False Negative, Incomplete bbox |\n",
"| ![Overlapping](https://media.roboflow.com/supervision/cookbooks/sahi/overlapping_2.png \"Overlapping\")| Double detection, Incomplete bbox|\n",
"| ![Overlapping](https://media.roboflow.com/supervision/cookbooks/sahi/overlapping_3.png \"Overlapping\")| Incomplete bounding box|\n",
"\n",
"## Improving Object Detection Near Boundaries with Overlapping\n",
"\n",
docs/trackers.md
+1
View File
@@ -1,5 +1,6 @@
---
comments: true
status: new
---
# ByteTrack
mkdocs.yml
+4
View File
@@ -66,7 +66,11 @@ nav:
- Utils: datasets/utils.md
- Metrics:
- mAP: metrics/mean_average_precision.md
- mAR: metrics/mean_average_recall.md
- Precision: metrics/precision.md
- Recall: metrics/recall.md
- F1 Score: metrics/f1_score.md
- Common Values: metrics/common_values.md
- Legacy Metrics: detection/metrics.md
- Utils:
- Video: utils/video.md
poetry.lock
Generated
+576 -339
View File
File diff suppressed because it is too large. Load diff
pyproject.toml
+19 -4
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "supervision"
version = "0.24.0"
version = "0.25.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 = [
@@ -9,7 +9,7 @@ maintainers = [
]
readme = "README.md"
license = "MIT"
packages = [{ include = "supervision" }]
packages = [{ include = "supervision" }, { include = "supervision/py.typed" }]
homepage = "https://github.com/roboflow/supervision"
repository = "https://github.com/roboflow/supervision"
documentation = "https://supervision.roboflow.com/latest/"
@@ -46,17 +46,32 @@ python = "^3.8"
numpy = [
{ version = ">=1.21.2,<1.23.3", python = "<=3.10" },
{ version = ">=1.23.3", python = ">3.10" },
{ version = ">=2.1.0", python = ">=3.13" },
]
scipy = [
{ version = "1.10.0", python = "<3.9" },
{ version = "^1.10.0", python = ">=3.9" },
{ version = ">=1.14.1", python = ">=3.13" },
]
# Matplotlib sub-dependency
# The 'contourpy' package is required by Matplotlib for contour plotting.
# We need to ensure compatibility with both Python 3.8 and Python 3.13.
#
# For Python 3.8 and above, we use version 1.0.7 or higher, as it is the lowest major version that supports Python 3.8.
# For Python 3.13 and above, we use version 1.3.0 or higher, as it is the first version that explicitly supports Python 3.13.
contourpy = [
{ version = ">=1.0.7", python = ">=3.8" },
{ version = ">=1.3.0", python = ">=3.13" },
]
matplotlib = ">=3.6.0"
pyyaml = ">=5.3"
defusedxml = "^0.7.1"
pillow = ">=9.4"
requests = { version = ">=2.26.0,<=2.32.3", optional = true }
tqdm = { version = ">=4.62.3,<=4.66.5", optional = true }
tqdm = { version = ">=4.62.3,<=4.67.0", optional = true }
# pandas: picked lowest major version that supports Python 3.8
pandas = { version = ">=2.0.0", optional = true }
pandas-stubs = { version = ">=2.0.0.230412", optional = true }
@@ -69,7 +84,7 @@ metrics = ["pandas", "pandas-stubs"]
[tool.poetry.group.dev.dependencies]
twine = "^5.1.1"
pytest = ">=7.2.2,<9.0.0"
wheel = ">=0.40,<0.45"
wheel = ">=0.40,<0.46"
build = ">=0.10,<1.3"
ruff = ">=0.1.0"
mypy = "^1.4.1"
supervision/annotators/core.py
+253 -116
View File
@@ -16,7 +16,7 @@ from supervision.annotators.utils import (
)
from supervision.config import CLASS_NAME_DATA_FIELD, ORIENTED_BOX_COORDINATES
from supervision.detection.core import Detections
from supervision.detection.utils import clip_boxes, mask_to_polygons
from supervision.detection.utils import clip_boxes, mask_to_polygons, spread_out_boxes
from supervision.draw.color import Color, ColorPalette
from supervision.draw.utils import draw_polygon
from supervision.geometry.core import Position
@@ -32,6 +32,8 @@ from supervision.utils.image import (
)
from supervision.utils.internal import deprecated
CV2_FONT = cv2.FONT_HERSHEY_SIMPLEX
class BoxAnnotator(BaseAnnotator):
"""
@@ -1054,6 +1056,7 @@ class LabelAnnotator(BaseAnnotator):
text_position: Position = Position.TOP_LEFT,
color_lookup: ColorLookup = ColorLookup.CLASS,
border_radius: int = 0,
smart_position: bool = False,
):
"""
Args:
@@ -1070,6 +1073,7 @@ class LabelAnnotator(BaseAnnotator):
Options are `INDEX`, `CLASS`, `TRACK`.
border_radius (int): The radius to apply round edges. If the selected
value is higher than the lower dimension, width or height, is clipped.
smart_position (bool): Spread out the labels to avoid overlapping.
"""
self.border_radius: int = border_radius
self.color: Union[Color, ColorPalette] = color
@@ -1079,6 +1083,7 @@ class LabelAnnotator(BaseAnnotator):
self.text_padding: int = text_padding
self.text_anchor: Position = text_position
self.color_lookup: ColorLookup = color_lookup
self.smart_position = smart_position
@ensure_cv2_image_for_annotation
def annotate(
@@ -1128,11 +1133,29 @@ class LabelAnnotator(BaseAnnotator):
![label-annotator-example](https://media.roboflow.com/
supervision-annotator-examples/label-annotator-example-purple.png)
"""
assert isinstance(scene, np.ndarray)
font = cv2.FONT_HERSHEY_SIMPLEX
anchors_coordinates = detections.get_anchors_coordinates(
anchor=self.text_anchor
).astype(int)
self._validate_labels(labels, detections)
labels = self._get_labels_text(detections, labels)
label_properties = self._get_label_properties(detections, labels)
if self.smart_position:
xyxy = label_properties[:, :4]
xyxy = spread_out_boxes(xyxy)
label_properties[:, :4] = xyxy
self._draw_labels(
scene=scene,
labels=labels,
label_properties=label_properties,
detections=detections,
custom_color_lookup=custom_color_lookup,
)
return scene
def _validate_labels(self, labels: Optional[List[str]], detections: Detections):
if labels is not None and len(labels) != len(detections):
raise ValueError(
f"The number of labels ({len(labels)}) does not match the "
@@ -1140,72 +1163,121 @@ class LabelAnnotator(BaseAnnotator):
f"should have exactly 1 label."
)
for detection_idx, center_coordinates in enumerate(anchors_coordinates):
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=(
self.color_lookup
if custom_color_lookup is None
else custom_color_lookup
),
)
def _get_label_properties(
self,
detections: Detections,
labels: List[str],
) -> np.ndarray:
"""
Calculate the numerical properties required to draw the labels on the image.
text_color = resolve_color(
color=self.text_color,
detections=detections,
detection_idx=detection_idx,
color_lookup=(
self.color_lookup
if custom_color_lookup is None
else custom_color_lookup
),
)
Returns:
(np.ndarray): An array of label properties, containing columns:
`min_x`, `min_y`, `max_x`, `max_y`, `padded_text_height`.
"""
label_properties = []
anchors_coordinates = detections.get_anchors_coordinates(
anchor=self.text_anchor
).astype(int)
if labels is not None:
text = labels[detection_idx]
elif CLASS_NAME_DATA_FIELD in detections.data:
text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx]
elif detections.class_id is not None:
text = str(detections.class_id[detection_idx])
else:
text = str(detection_idx)
text_w, text_h = cv2.getTextSize(
text=text,
fontFace=font,
for label, center_coords in zip(labels, anchors_coordinates):
(text_w, text_h) = cv2.getTextSize(
text=label,
fontFace=CV2_FONT,
fontScale=self.text_scale,
thickness=self.text_thickness,
)[0]
text_w_padded = text_w + 2 * self.text_padding
text_h_padded = text_h + 2 * self.text_padding
width_padded = text_w + 2 * self.text_padding
height_padded = text_h + 2 * self.text_padding
text_background_xyxy = resolve_text_background_xyxy(
center_coordinates=tuple(center_coordinates),
text_wh=(text_w_padded, text_h_padded),
center_coordinates=tuple(center_coords),
text_wh=(width_padded, height_padded),
position=self.text_anchor,
)
text_x = text_background_xyxy[0] + self.text_padding
text_y = text_background_xyxy[1] + self.text_padding + text_h
label_properties.append(
[
*text_background_xyxy,
text_h,
]
)
return np.array(label_properties).reshape(-1, 5)
@staticmethod
def _get_labels_text(
detections: Detections, custom_labels: Optional[List[str]]
) -> List[str]:
if custom_labels is not None:
return custom_labels
labels = []
for idx in range(len(detections)):
if CLASS_NAME_DATA_FIELD in detections.data:
labels.append(detections.data[CLASS_NAME_DATA_FIELD][idx])
elif detections.class_id is not None:
labels.append(str(detections.class_id[idx]))
else:
labels.append(str(idx))
return labels
def _draw_labels(
self,
scene: np.ndarray,
labels: List[str],
label_properties: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray],
) -> None:
assert len(labels) == len(label_properties) == len(detections), (
f"Number of label properties ({len(label_properties)}), "
f"labels ({len(labels)}) and detections ({len(detections)}) "
"do not match."
)
color_lookup = (
custom_color_lookup
if custom_color_lookup is not None
else self.color_lookup
)
for idx, label_property in enumerate(label_properties):
background_color = resolve_color(
color=self.color,
detections=detections,
detection_idx=idx,
color_lookup=color_lookup,
)
text_color = resolve_color(
color=self.text_color,
detections=detections,
detection_idx=idx,
color_lookup=color_lookup,
)
box_xyxy = label_property[:4]
text_height_padded = label_property[4]
self.draw_rounded_rectangle(
scene=scene,
xyxy=text_background_xyxy,
color=color.as_bgr(),
xyxy=box_xyxy,
color=background_color.as_bgr(),
border_radius=self.border_radius,
)
text_x = box_xyxy[0] + self.text_padding
text_y = box_xyxy[1] + self.text_padding + text_height_padded
cv2.putText(
img=scene,
text=text,
text=labels[idx],
org=(text_x, text_y),
fontFace=font,
fontFace=CV2_FONT,
fontScale=self.text_scale,
color=text_color.as_bgr(),
thickness=self.text_thickness,
lineType=cv2.LINE_AA,
)
return scene
@staticmethod
def draw_rounded_rectangle(
@@ -1266,6 +1338,7 @@ class RichLabelAnnotator(BaseAnnotator):
text_position: Position = Position.TOP_LEFT,
color_lookup: ColorLookup = ColorLookup.CLASS,
border_radius: int = 0,
smart_position: bool = False,
):
"""
Args:
@@ -1282,6 +1355,7 @@ class RichLabelAnnotator(BaseAnnotator):
Options are `INDEX`, `CLASS`, `TRACK`.
border_radius (int): The radius to apply round edges. If the selected
value is higher than the lower dimension, width or height, is clipped.
smart_position (bool): Spread out the labels to avoid overlapping.
"""
self.color = color
self.text_color = text_color
@@ -1289,14 +1363,8 @@ class RichLabelAnnotator(BaseAnnotator):
self.text_anchor = text_position
self.color_lookup = color_lookup
self.border_radius = border_radius
if font_path is not None:
try:
self.font = ImageFont.truetype(font_path, font_size)
except OSError:
print(f"Font path '{font_path}' not found. Using PIL's default font.")
self.font = self._load_default_font(font_size)
else:
self.font = self._load_default_font(font_size)
self.smart_position = smart_position
self.font = self._load_font(font_size, font_path)
@ensure_pil_image_for_annotation
def annotate(
@@ -1346,88 +1414,157 @@ class RichLabelAnnotator(BaseAnnotator):
"""
assert isinstance(scene, Image.Image)
self._validate_labels(labels, detections)
draw = ImageDraw.Draw(scene)
anchors_coordinates = detections.get_anchors_coordinates(
anchor=self.text_anchor
).astype(int)
labels = self._get_labels_text(detections, labels)
label_properties = self._get_label_properties(draw, detections, labels)
if self.smart_position:
xyxy = label_properties[:, :4]
xyxy = spread_out_boxes(xyxy)
label_properties[:, :4] = xyxy
self._draw_labels(
draw=draw,
labels=labels,
label_properties=label_properties,
detections=detections,
custom_color_lookup=custom_color_lookup,
)
return scene
def _validate_labels(self, labels: Optional[List[str]], detections: Detections):
if labels is not None and len(labels) != len(detections):
raise ValueError(
f"The number of labels provided ({len(labels)}) does not match the "
f"number of detections ({len(detections)}). Each detection should have "
f"a corresponding label."
)
for detection_idx, center_coordinates in enumerate(anchors_coordinates):
color = resolve_color(
color=self.color,
detections=detections,
detection_idx=detection_idx,
color_lookup=(
self.color_lookup
if custom_color_lookup is None
else custom_color_lookup
),
f"The number of labels ({len(labels)}) does not match the "
f"number of detections ({len(detections)}). Each detection "
f"should have exactly 1 label."
)
text_color = resolve_color(
color=self.text_color,
detections=detections,
detection_idx=detection_idx,
color_lookup=(
self.color_lookup
if custom_color_lookup is None
else custom_color_lookup
),
def _get_label_properties(
self, draw, detections: Detections, labels: List[str]
) -> np.ndarray:
"""
Calculate the numerical properties required to draw the labels on the image.
Returns:
(np.ndarray): An array of label properties, containing columns:
`min_x`, `min_y`, `max_x`, `max_y`, `text_left_coordinate`,
`text_top_coordinate`. The first 4 values are already padded
with `text_padding`.
"""
label_properties = []
anchor_coordinates = detections.get_anchors_coordinates(
anchor=self.text_anchor
).astype(int)
for label, center_coords in zip(labels, anchor_coordinates):
text_left, text_top, text_right, text_bottom = draw.textbbox(
(0, 0), label, font=self.font
)
text_width = text_right - text_left
text_height = text_bottom - text_top
width_padded = text_width + 2 * self.text_padding
height_padded = text_height + 2 * self.text_padding
if labels is not None:
text = labels[detection_idx]
elif CLASS_NAME_DATA_FIELD in detections.data:
text = detections.data[CLASS_NAME_DATA_FIELD][detection_idx]
elif detections.class_id is not None:
text = str(detections.class_id[detection_idx])
else:
text = str(detection_idx)
left, top, right, bottom = draw.textbbox((0, 0), text, font=self.font)
text_width = right - left
text_height = bottom - top
text_w_padded = text_width + 2 * self.text_padding
text_h_padded = text_height + 2 * self.text_padding
text_background_xyxy = resolve_text_background_xyxy(
center_coordinates=tuple(center_coordinates),
text_wh=(text_w_padded, text_h_padded),
center_coordinates=tuple(center_coords),
text_wh=(width_padded, height_padded),
position=self.text_anchor,
)
text_x = text_background_xyxy[0] + self.text_padding - left
text_y = text_background_xyxy[1] + self.text_padding - top
label_properties.append([*text_background_xyxy, text_left, text_top])
return np.array(label_properties).reshape(-1, 6)
@staticmethod
def _get_labels_text(
detections: Detections, custom_labels: Optional[List[str]]
) -> List[str]:
if custom_labels is not None:
return custom_labels
labels = []
for idx in range(len(detections)):
if CLASS_NAME_DATA_FIELD in detections.data:
labels.append(detections.data[CLASS_NAME_DATA_FIELD][idx])
elif detections.class_id is not None:
labels.append(str(detections.class_id[idx]))
else:
labels.append(str(idx))
return labels
def _draw_labels(
self,
draw,
labels: List[str],
label_properties: np.ndarray,
detections: Detections,
custom_color_lookup: Optional[np.ndarray],
) -> None:
assert len(labels) == len(label_properties) == len(detections), (
f"Number of label properties ({len(label_properties)}), "
f"labels ({len(labels)}) and detections ({len(detections)}) "
"do not match."
)
color_lookup = (
custom_color_lookup
if custom_color_lookup is not None
else self.color_lookup
)
for idx, label_property in enumerate(label_properties):
background_color = resolve_color(
color=self.color,
detections=detections,
detection_idx=idx,
color_lookup=color_lookup,
)
text_color = resolve_color(
color=self.text_color,
detections=detections,
detection_idx=idx,
color_lookup=color_lookup,
)
box_xyxy = label_property[:4]
text_left = label_property[4]
text_top = label_property[5]
label_x_position = box_xyxy[0] + self.text_padding - text_left
label_y_position = box_xyxy[1] + self.text_padding - text_top
draw.rounded_rectangle(
text_background_xyxy,
tuple(box_xyxy),
radius=self.border_radius,
fill=color.as_rgb(),
fill=background_color.as_rgb(),
outline=None,
)
draw.text(
xy=(text_x, text_y),
text=text,
xy=(label_x_position, label_y_position),
text=labels[idx],
font=self.font,
fill=text_color.as_rgb(),
)
return scene
@staticmethod
def _load_default_font(size):
"""
PIL either loads a font that accepts a size (e.g. on my machine)
or raises an error saying `load_default` does not accept arguments
(e.g. in Colab).
"""
def _load_font(font_size: int, font_path: Optional[str]):
def load_default_font(size):
try:
return ImageFont.load_default(size)
except TypeError:
return ImageFont.load_default()
if font_path is None:
return load_default_font(font_size)
try:
font = ImageFont.load_default(size)
except TypeError:
font = ImageFont.load_default()
return font
return ImageFont.truetype(font_path, font_size)
except OSError:
print(f"Font path '{font_path}' not found. Using PIL's default font.")
return load_default_font(font_size)
class IconAnnotator(BaseAnnotator):
supervision/assets/list.py
+6
View File
@@ -20,6 +20,7 @@ class VideoAssets(Enum):
| `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) |
| `SKIING` | `skiing.mp4` | [Link](https://media.roboflow.com/supervision/video-examples/skiing.mp4) |
""" # noqa: E501 // docs
VEHICLES = "vehicles.mp4"
@@ -31,6 +32,7 @@ class VideoAssets(Enum):
PEOPLE_WALKING = "people-walking.mp4"
BEACH = "beach-1.mp4"
BASKETBALL = "basketball-1.mp4"
SKIING = "skiing.mp4"
@classmethod
def list(cls):
@@ -74,4 +76,8 @@ VIDEO_ASSETS: Dict[str, Tuple[str, str]] = {
f"{BASE_VIDEO_URL}{VideoAssets.BASKETBALL.value}",
"60d94a3c7c47d16f09d342b088012ecc",
),
VideoAssets.SKIING.value: (
f"{BASE_VIDEO_URL}{VideoAssets.SKIING.value}",
"d30987cbab1bbc5934199cdd1b293119",
),
}
supervision/dataset/utils.py
+1 -1
View File
@@ -55,7 +55,7 @@ def merge_class_lists(class_lists: List[List[str]]) -> List[str]:
for class_list in class_lists:
for class_name in class_list:
unique_classes.add(class_name.lower())
unique_classes.add(class_name)
return sorted(list(unique_classes))
supervision/detection/core.py
+16
View File
@@ -32,8 +32,10 @@ from supervision.detection.utils import (
extract_ultralytics_masks,
get_data_item,
is_data_equal,
is_metadata_equal,
mask_to_xyxy,
merge_data,
merge_metadata,
process_roboflow_result,
xywh_to_xyxy,
)
@@ -125,6 +127,9 @@ 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.
metadata (Dict[str, Any]): A dictionary containing collection-level metadata
that applies to the entire set of detections. This may include information such
as the video name, camera parameters, timestamp, or other global metadata.
""" # noqa: E501 // docs
xyxy: np.ndarray
@@ -133,6 +138,7 @@ class Detections:
class_id: Optional[np.ndarray] = None
tracker_id: Optional[np.ndarray] = None
data: Dict[str, Union[np.ndarray, List]] = field(default_factory=dict)
metadata: Dict[str, Any] = field(default_factory=dict)
def __post_init__(self):
validate_detections_fields(
@@ -185,6 +191,7 @@ class Detections:
np.array_equal(self.confidence, other.confidence),
np.array_equal(self.tracker_id, other.tracker_id),
is_data_equal(self.data, other.data),
is_metadata_equal(self.metadata, other.metadata),
]
)
@@ -985,6 +992,7 @@ class Detections:
"""
empty_detections = Detections.empty()
empty_detections.data = self.data
empty_detections.metadata = self.metadata
return self == empty_detections
@classmethod
@@ -1078,6 +1086,9 @@ class Detections:
data = merge_data([d.data for d in detections_list])
metadata_list = [detections.metadata for detections in detections_list]
metadata = merge_metadata(metadata_list)
return cls(
xyxy=xyxy,
mask=mask,
@@ -1085,6 +1096,7 @@ class Detections:
class_id=class_id,
tracker_id=tracker_id,
data=data,
metadata=metadata,
)
def get_anchors_coordinates(self, anchor: Position) -> np.ndarray:
@@ -1198,6 +1210,7 @@ class Detections:
class_id=self.class_id[index] if self.class_id is not None else None,
tracker_id=self.tracker_id[index] if self.tracker_id is not None else None,
data=get_data_item(self.data, index),
metadata=self.metadata,
)
def __setitem__(self, key: str, value: Union[np.ndarray, List]):
@@ -1459,6 +1472,8 @@ def merge_inner_detection_object_pair(
else:
winning_detection = detections_2
metadata = merge_metadata([detections_1.metadata, detections_2.metadata])
return Detections(
xyxy=merged_xyxy,
mask=merged_mask,
@@ -1466,6 +1481,7 @@ def merge_inner_detection_object_pair(
class_id=winning_detection.class_id,
tracker_id=winning_detection.tracker_id,
data=winning_detection.data,
metadata=metadata,
)
supervision/detection/line_zone.py
+156 -93
View File
@@ -1,11 +1,12 @@
import math
import warnings
from collections import Counter
from collections import Counter, defaultdict, deque
from functools import lru_cache
from typing import Any, Dict, Iterable, List, Literal, Optional, Tuple
from typing import Any, Deque, Dict, Iterable, List, Literal, Optional, Tuple
import cv2
import numpy as np
import numpy.typing as npt
from supervision.config import CLASS_NAME_DATA_FIELD
from supervision.detection.core import Detections
@@ -41,6 +42,10 @@ class LineZone:
to inside.
out_count (int): The number of objects that have crossed the line from inside
to outside.
in_count_per_class (Dict[int, int]): Number of objects of each class that have
crossed the line from outside to inside.
out_count_per_class (Dict[int, int]): Number of objects of each class that have
crossed the line from inside to outside.
Example:
```python
@@ -74,6 +79,7 @@ class LineZone:
Position.BOTTOM_LEFT,
Position.BOTTOM_RIGHT,
),
minimum_crossing_threshold: int = 1,
):
"""
Args:
@@ -84,10 +90,18 @@ class LineZone:
to consider when deciding on whether the detection
has passed the line counter or not. By default, this
contains the four corners of the detection's bounding box
minimum_crossing_threshold (int): Detection needs to be seen
on the other side of the line for this many frames to be
considered as having crossed the line. This is useful when
dealing with unstable bounding boxes or when detections
may linger on the line.
"""
self.vector = Vector(start=start, end=end)
self.limits = self.calculate_region_of_interest_limits(vector=self.vector)
self.tracker_state: Dict[str, bool] = {}
self.limits = self._calculate_region_of_interest_limits(vector=self.vector)
self.crossing_history_length = max(2, minimum_crossing_threshold + 1)
self.crossing_state_history: Dict[int, Deque[bool]] = defaultdict(
lambda: deque(maxlen=self.crossing_history_length)
)
self._in_count_per_class: Counter = Counter()
self._out_count_per_class: Counter = Counter()
self.triggering_anchors = triggering_anchors
@@ -97,74 +111,20 @@ class LineZone:
@property
def in_count(self) -> int:
"""
Number of objects that have crossed the line from
outside to inside.
"""
return sum(self._in_count_per_class.values())
@property
def out_count(self) -> int:
"""
Number of objects that have crossed the line from
inside to outside.
"""
return sum(self._out_count_per_class.values())
@property
def in_count_per_class(self) -> Dict[int, int]:
"""
Number of objects of each class that have crossed
the line from outside to inside.
"""
return dict(self._in_count_per_class)
@property
def out_count_per_class(self) -> Dict[int, int]:
"""
Number of objects of each class that have crossed the line
from inside to outside.
"""
return dict(self._out_count_per_class)
@staticmethod
def calculate_region_of_interest_limits(vector: Vector) -> Tuple[Vector, Vector]:
magnitude = vector.magnitude
if magnitude == 0:
raise ValueError("The magnitude of the vector cannot be zero.")
delta_x = vector.end.x - vector.start.x
delta_y = vector.end.y - vector.start.y
unit_vector_x = delta_x / magnitude
unit_vector_y = delta_y / magnitude
perpendicular_vector_x = -unit_vector_y
perpendicular_vector_y = unit_vector_x
start_region_limit = Vector(
start=vector.start,
end=Point(
x=vector.start.x + perpendicular_vector_x,
y=vector.start.y + perpendicular_vector_y,
),
)
end_region_limit = Vector(
start=vector.end,
end=Point(
x=vector.end.x - perpendicular_vector_x,
y=vector.end.y - perpendicular_vector_y,
),
)
return start_region_limit, end_region_limit
@staticmethod
def is_point_in_limits(point: Point, limits: Tuple[Vector, Vector]) -> bool:
cross_product_1 = limits[0].cross_product(point)
cross_product_2 = limits[1].cross_product(point)
return (cross_product_1 > 0) == (cross_product_2 > 0)
def trigger(self, detections: Detections) -> Tuple[np.ndarray, np.ndarray]:
"""
Update the `in_count` and `out_count` based on the objects that cross the line.
@@ -194,6 +154,123 @@ class LineZone:
)
return crossed_in, crossed_out
self._update_class_id_to_name(detections)
in_limits, has_any_left_trigger, has_any_right_trigger = (
self._compute_anchor_sides(detections)
)
class_ids: List[Optional[int]] = (
list(detections.class_id)
if detections.class_id is not None
else [None] * len(detections)
)
for i, (class_id, tracker_id) in enumerate(
zip(class_ids, detections.tracker_id)
):
if not in_limits[i]:
continue
if has_any_left_trigger[i] and has_any_right_trigger[i]:
continue
tracker_state: bool = has_any_left_trigger[i]
crossing_history = self.crossing_state_history[tracker_id]
crossing_history.append(tracker_state)
if len(crossing_history) < self.crossing_history_length:
continue
# TODO: Account for incorrect class_id.
# Most likely this would involve indexing self.crossing_state_history
# with (tracker_id, class_id).
oldest_state = crossing_history[0]
if crossing_history.count(oldest_state) > 1:
continue
if tracker_state:
self._in_count_per_class[class_id] += 1
crossed_in[i] = True
else:
self._out_count_per_class[class_id] += 1
crossed_out[i] = True
return crossed_in, crossed_out
@staticmethod
def _calculate_region_of_interest_limits(vector: Vector) -> Tuple[Vector, Vector]:
magnitude = vector.magnitude
if magnitude == 0:
raise ValueError("The magnitude of the vector cannot be zero.")
delta_x = vector.end.x - vector.start.x
delta_y = vector.end.y - vector.start.y
unit_vector_x = delta_x / magnitude
unit_vector_y = delta_y / magnitude
perpendicular_vector_x = -unit_vector_y
perpendicular_vector_y = unit_vector_x
start_region_limit = Vector(
start=vector.start,
end=Point(
x=vector.start.x + perpendicular_vector_x,
y=vector.start.y + perpendicular_vector_y,
),
)
end_region_limit = Vector(
start=vector.end,
end=Point(
x=vector.end.x - perpendicular_vector_x,
y=vector.end.y - perpendicular_vector_y,
),
)
return start_region_limit, end_region_limit
def _compute_anchor_sides(
self, detections: Detections
) -> Tuple[npt.NDArray[np.bool_], npt.NDArray[np.bool_], npt.NDArray[np.bool_]]:
"""
Find if detections' anchors are within the limit of the line
zone and which anchors are on its left and right side.
Assumes:
* At least 1 detection is provided
* Detections have `tracker_id`
The limit is defined as the region between the two lines,
perpendicular to the line zone, and passing through its start
and end points, as shown below:
Limits:
```
| IN ↑
| |
OUT o---LINE---o OUT
| |
↓ IN |
```
Args:
detections (Detections): The detections to check.
Returns:
result (Tuple[np.ndarray, np.ndarray, np.ndarray]):
All 3 arrays are boolean arrays of shape (N, ) where N is the
number of detections. The first array, `in_limits`, indicates
if the detection's anchor is within the line zone limits.
The second array, `has_any_left_trigger`, indicates if the
detection's anchor is on the left side of the line zone.
The third array, `has_any_right_trigger`, indicates if the
detection's anchor is on the right side of the line zone.
"""
assert len(detections) > 0
assert detections.tracker_id is not None
all_anchors = np.array(
[
detections.get_anchors_coordinates(anchor)
@@ -203,52 +280,38 @@ class LineZone:
cross_products_1 = cross_product(all_anchors, self.limits[0])
cross_products_2 = cross_product(all_anchors, self.limits[1])
# Works because limit vectors are pointing in opposite directions
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)
class_ids = (
list(detections.class_id)
if detections.class_id is not None
else [None] * len(detections)
)
tracker_ids = list(detections.tracker_id)
return in_limits, has_any_left_trigger, has_any_right_trigger
if CLASS_NAME_DATA_FIELD in detections.data:
class_names = detections.data[CLASS_NAME_DATA_FIELD]
for class_id, class_name in zip(class_ids, class_names):
if class_id is None:
class_name = "No class"
self.class_id_to_name[class_id] = class_name
def _update_class_id_to_name(self, detections: Detections) -> None:
"""
Update the attribute keeping track of which class
IDs correspond to which class names.
for i, (class_ids, tracker_id) in enumerate(zip(class_ids, tracker_ids)):
if not in_limits[i]:
continue
Assumes that class_names are only provided when class_ids are.
"""
class_names = detections.data.get(CLASS_NAME_DATA_FIELD)
assert class_names is None or detections.class_id is not None
if not is_uniformly_triggered[i]:
continue
if detections.class_id is None:
return
tracker_state = has_any_left_trigger[i]
if tracker_id not in self.tracker_state:
self.tracker_state[tracker_id] = tracker_state
continue
if self.tracker_state.get(tracker_id) == tracker_state:
continue
self.tracker_state[tracker_id] = tracker_state
if tracker_state:
self._in_count_per_class[class_ids] += 1
crossed_in[i] = True
else:
self._out_count_per_class[class_ids] += 1
crossed_out[i] = True
return crossed_in, crossed_out
if class_names is None:
new_names = {class_id: str(class_id) for class_id in detections.class_id}
else:
new_names = {
class_id: class_name
for class_id, class_name in zip(detections.class_id, class_names)
}
self.class_id_to_name.update(new_names)
class LineZoneAnnotator:
supervision/detection/tools/polygon_zone.py
+30 -2
View File
@@ -17,6 +17,12 @@ class PolygonZone:
"""
A class for defining a polygon-shaped zone within a frame for detecting objects.
!!! warning
PolygonZone uses the `tracker_id`. Read
[here](/latest/trackers/) to learn how to plug
tracking into your inference pipeline.
Attributes:
polygon (np.ndarray): A polygon represented by a numpy array of shape
`(N, 2)`, containing the `x`, `y` coordinates of the points.
@@ -26,6 +32,28 @@ class PolygonZone:
(default: (sv.Position.BOTTOM_CENTER,)).
current_count (int): The current count of detected objects within the zone
mask (np.ndarray): The 2D bool mask for the polygon zone
Example:
```python
import supervision as sv
from ultralytics import YOLO
import numpy as np
import cv2
image = cv2.imread(<SOURCE_IMAGE_PATH>)
model = YOLO("yolo11s")
tracker = sv.ByteTrack()
polygon = np.array([[100, 200], [200, 100], [300, 200], [200, 300]])
polygon_zone = sv.PolygonZone(polygon=polygon)
result = model.infer(image)[0]
detections = sv.Detections.from_ultralytics(result)
detections = tracker.update_with_detections(detections)
is_detections_in_zone = polygon_zone.trigger(detections)
print(polygon_zone.current_count)
```
"""
def __init__(
@@ -88,7 +116,7 @@ class PolygonZoneAnnotator:
Attributes:
zone (PolygonZone): The polygon zone to be annotated
color (Color): The color to draw the polygon lines
color (Color): The color to draw the polygon lines, default is white
thickness (int): The thickness of the polygon lines, default is 2
text_color (Color): The color of the text on the polygon, default is black
text_scale (float): The scale of the text on the polygon, default is 0.5
@@ -104,7 +132,7 @@ class PolygonZoneAnnotator:
def __init__(
self,
zone: PolygonZone,
color: Color,
color: Color = Color.WHITE,
thickness: int = 2,
text_color: Color = Color.BLACK,
text_scale: float = 0.5,
supervision/detection/utils.py
+141 -7
View File
@@ -1,5 +1,5 @@
from itertools import chain
from typing import Dict, List, Optional, Tuple, Union
from typing import Any, Dict, List, Optional, Tuple, Union
import cv2
import numpy as np
@@ -23,10 +23,9 @@ def polygon_to_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.n
np.ndarray: The generated 2D mask, where the polygon is marked with
`1`'s and the rest is filled with `0`'s.
"""
width, height = resolution_wh
mask = np.zeros((height, width))
cv2.fillPoly(mask, [polygon], color=1)
width, height = map(int, resolution_wh)
mask = np.zeros((height, width), dtype=np.uint8)
cv2.fillPoly(mask, [polygon.astype(np.int32)], color=1)
return mask
@@ -163,9 +162,9 @@ def oriented_box_iou_batch(
boxes_true = boxes_true.reshape(-1, 4, 2)
boxes_detection = boxes_detection.reshape(-1, 4, 2)
max_height = max(boxes_true[:, :, 0].max(), boxes_detection[:, :, 0].max()) + 1
max_height = int(max(boxes_true[:, :, 0].max(), boxes_detection[:, :, 0].max()) + 1)
# adding 1 because we are 0-indexed
max_width = max(boxes_true[:, :, 1].max(), boxes_detection[:, :, 1].max()) + 1
max_width = int(max(boxes_true[:, :, 1].max(), boxes_detection[:, :, 1].max()) + 1)
mask_true = np.zeros((boxes_true.shape[0], max_height, max_width))
for i, box_true in enumerate(boxes_true):
@@ -808,12 +807,36 @@ def is_data_equal(data_a: Dict[str, np.ndarray], data_b: Dict[str, np.ndarray])
)
def is_metadata_equal(metadata_a: Dict[str, Any], metadata_b: Dict[str, Any]) -> bool:
"""
Compares the metadata payloads of two Detections instances.
Args:
metadata_a, metadata_b: The metadata payloads of the instances.
Returns:
True if the metadata payloads are equal, False otherwise.
"""
return set(metadata_a.keys()) == set(metadata_b.keys()) and all(
np.array_equal(metadata_a[key], metadata_b[key])
if (
isinstance(metadata_a[key], np.ndarray)
and isinstance(metadata_b[key], np.ndarray)
)
else metadata_a[key] == metadata_b[key]
for key in metadata_a
)
def merge_data(
data_list: List[Dict[str, Union[npt.NDArray[np.generic], List]]],
) -> Dict[str, Union[npt.NDArray[np.generic], List]]:
"""
Merges the data payloads of a list of Detections instances.
Warning: Assumes that empty detections were filtered-out before passing data to
this function.
Args:
data_list: The data payloads of the Detections instances. Each data payload
is a dictionary with the same keys, and the values are either lists or
@@ -866,6 +889,61 @@ def merge_data(
return merged_data
def merge_metadata(metadata_list: List[Dict[str, Any]]) -> Dict[str, Any]:
"""
Merge metadata from a list of metadata dictionaries.
This function combines the metadata dictionaries. If a key appears in more than one
dictionary, the values must be identical for the merge to succeed.
Warning: Assumes that empty detections were filtered-out before passing metadata to
this function.
Args:
metadata_list (List[Dict[str, Any]]): A list of metadata dictionaries to merge.
Returns:
Dict[str, Any]: A single merged metadata dictionary.
Raises:
ValueError: If there are conflicting values for the same key or if
dictionaries have different keys.
"""
if not metadata_list:
return {}
all_keys_sets = [set(metadata.keys()) for metadata in metadata_list]
if not all(keys_set == all_keys_sets[0] for keys_set in all_keys_sets):
raise ValueError("All metadata dictionaries must have the same keys to merge.")
merged_metadata: Dict[str, Any] = {}
for metadata in metadata_list:
for key, value in metadata.items():
if key not in merged_metadata:
merged_metadata[key] = value
continue
other_value = merged_metadata[key]
if isinstance(value, np.ndarray) and isinstance(other_value, np.ndarray):
if not np.array_equal(merged_metadata[key], value):
raise ValueError(
f"Conflicting metadata for key: '{key}': "
"{type(value)}, {type(other_value)}."
)
elif isinstance(value, np.ndarray) or isinstance(other_value, np.ndarray):
# Since [] == np.array([]).
raise ValueError(
f"Conflicting metadata for key: '{key}': "
"{type(value)}, {type(other_value)}."
)
else:
print("hm")
if merged_metadata[key] != value:
raise ValueError(f"Conflicting metadata for key: '{key}'.")
return merged_metadata
def get_data_item(
data: Dict[str, Union[np.ndarray, List]],
index: Union[int, slice, List[int], np.ndarray],
@@ -1039,3 +1117,59 @@ def cross_product(anchors: np.ndarray, vector: Vector) -> np.ndarray:
)
vector_start = np.array([vector.start.x, vector.start.y])
return np.cross(vector_at_zero, anchors - vector_start)
def spread_out_boxes(
xyxy: np.ndarray,
max_iterations: int = 100,
) -> np.ndarray:
"""
Spread out boxes that overlap with each other.
Args:
xyxy: Numpy array of shape (N, 4) where N is the number of boxes.
max_iterations: Maximum number of iterations to run the algorithm for.
"""
if len(xyxy) == 0:
return xyxy
xyxy_padded = pad_boxes(xyxy, px=1)
for _ in range(max_iterations):
# NxN
iou = box_iou_batch(xyxy_padded, xyxy_padded)
np.fill_diagonal(iou, 0)
if np.all(iou == 0):
break
overlap_mask = iou > 0
# Nx2
centers = (xyxy_padded[:, :2] + xyxy_padded[:, 2:]) / 2
# NxNx2
delta_centers = centers[:, np.newaxis, :] - centers[np.newaxis, :, :]
delta_centers *= overlap_mask[:, :, np.newaxis]
# Nx2
delta_sum = np.sum(delta_centers, axis=1)
delta_magnitude = np.linalg.norm(delta_sum, axis=1, keepdims=True)
direction_vectors = np.divide(
delta_sum,
delta_magnitude,
out=np.zeros_like(delta_sum),
where=delta_magnitude != 0,
)
force_vectors = np.sum(iou, axis=1)
force_vectors = force_vectors[:, np.newaxis] * direction_vectors
force_vectors *= 10
force_vectors[(force_vectors > 0) & (force_vectors < 2)] = 2
force_vectors[(force_vectors < 0) & (force_vectors > -2)] = -2
force_vectors = force_vectors.astype(int)
xyxy_padded[:, [0, 1]] += force_vectors
xyxy_padded[:, [2, 3]] += force_vectors
return pad_boxes(xyxy_padded, px=-1)
supervision/draw/utils.py
+18 -8
View File
@@ -9,7 +9,11 @@ from supervision.geometry.core import Point, Rect
def draw_line(
scene: np.ndarray, start: Point, end: Point, color: Color, thickness: int = 2
scene: np.ndarray,
start: Point,
end: Point,
color: Color = Color.ROBOFLOW,
thickness: int = 2,
) -> np.ndarray:
"""
Draws a line on a given scene.
@@ -18,7 +22,7 @@ def draw_line(
scene (np.ndarray): The scene on which the line will be drawn
start (Point): The starting point of the line
end (Point): The end point of the line
color (Color): The color of the line
color (Color): The color of the line, defaults to Color.ROBOFLOW
thickness (int): The thickness of the line
Returns:
@@ -35,7 +39,7 @@ def draw_line(
def draw_rectangle(
scene: np.ndarray, rect: Rect, color: Color, thickness: int = 2
scene: np.ndarray, rect: Rect, color: Color = Color.ROBOFLOW, thickness: int = 2
) -> np.ndarray:
"""
Draws a rectangle on an image.
@@ -60,7 +64,7 @@ def draw_rectangle(
def draw_filled_rectangle(
scene: np.ndarray, rect: Rect, color: Color, opacity: float = 1
scene: np.ndarray, rect: Rect, color: Color = Color.ROBOFLOW, opacity: float = 1
) -> np.ndarray:
"""
Draws a filled rectangle on an image.
@@ -151,14 +155,17 @@ def draw_rounded_rectangle(
def draw_polygon(
scene: np.ndarray, polygon: np.ndarray, color: Color, thickness: int = 2
scene: np.ndarray,
polygon: np.ndarray,
color: Color = Color.ROBOFLOW,
thickness: int = 2,
) -> np.ndarray:
"""Draw a polygon on a scene.
Parameters:
scene (np.ndarray): The scene to draw the polygon on.
polygon (np.ndarray): The polygon to be drawn, given as a list of vertices.
color (Color): The color of the polygon.
color (Color): The color of the polygon. Defaults to Color.ROBOFLOW.
thickness (int): The thickness of the polygon lines, by default 2.
Returns:
@@ -171,14 +178,17 @@ def draw_polygon(
def draw_filled_polygon(
scene: np.ndarray, polygon: np.ndarray, color: Color, opacity: float = 1
scene: np.ndarray,
polygon: np.ndarray,
color: Color = Color.ROBOFLOW,
opacity: float = 1,
) -> np.ndarray:
"""Draw a filled polygon on a scene.
Parameters:
scene (np.ndarray): The scene to draw the polygon on.
polygon (np.ndarray): The polygon to be drawn, given as a list of vertices.
color (Color): The color of the polygon.
color (Color): The color of the polygon. Defaults to Color.ROBOFLOW.
opacity (float): The opacity of polygon when drawn on the scene.
Returns:
supervision/geometry/utils.py
+6
View File
@@ -16,6 +16,9 @@ def get_polygon_center(polygon: np.ndarray) -> Point:
Point: The center of the polygon, represented as a
Point object with x and y attributes.
Raises:
ValueError: If the polygon has no vertices.
Examples:
```python
import numpy as np
@@ -30,6 +33,9 @@ def get_polygon_center(polygon: np.ndarray) -> Point:
# This is one of the 3 candidate algorithms considered for centroid calculation.
# For a more detailed discussion, see PR #1084 and commit eb33176
if len(polygon) == 0:
raise ValueError("Polygon must have at least one vertex.")
shift_polygon = np.roll(polygon, -1, axis=0)
signed_areas = np.cross(polygon, shift_polygon) / 2
if signed_areas.sum() == 0:
supervision/keypoint/annotators.py
+9 -2
View File
@@ -5,10 +5,11 @@ from typing import List, Optional, Tuple, Union
import cv2
import numpy as np
from supervision import Rect, pad_boxes
from supervision.annotators.base import ImageType
from supervision.detection.utils import pad_boxes, spread_out_boxes
from supervision.draw.color import Color
from supervision.draw.utils import draw_rounded_rectangle
from supervision.geometry.core import Rect
from supervision.keypoint.core import KeyPoints
from supervision.keypoint.skeletons import SKELETONS_BY_VERTEX_COUNT
from supervision.utils.conversion import ensure_cv2_image_for_annotation
@@ -201,6 +202,7 @@ class VertexLabelAnnotator:
text_thickness: int = 1,
text_padding: int = 10,
border_radius: int = 0,
smart_position: bool = False,
):
"""
Args:
@@ -215,6 +217,7 @@ class VertexLabelAnnotator:
text_padding (int): The padding around the text.
border_radius (int): The radius of the rounded corners of the
boxes. Set to a high value to produce circles.
smart_position (bool): Spread out the labels to avoid overlap.
"""
self.border_radius: int = border_radius
self.color: Union[Color, List[Color]] = color
@@ -222,6 +225,7 @@ class VertexLabelAnnotator:
self.text_scale: float = text_scale
self.text_thickness: int = text_thickness
self.text_padding: int = text_padding
self.smart_position = smart_position
def annotate(
self,
@@ -356,9 +360,12 @@ class VertexLabelAnnotator:
for anchor, label in zip(anchors, labels)
]
)
xyxy_padded = pad_boxes(xyxy=xyxy, px=self.text_padding)
if self.smart_position:
xyxy_padded = spread_out_boxes(xyxy_padded)
xyxy = pad_boxes(xyxy=xyxy_padded, px=-self.text_padding)
for text, color, text_color, box, box_padded in zip(
labels, colors, text_colors, xyxy, xyxy_padded
):
supervision/keypoint/core.py
+74 -1
View File
@@ -2,12 +2,13 @@ from __future__ import annotations
from contextlib import suppress
from dataclasses import dataclass, field
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union
from typing import Any, Dict, Iterable, Iterator, List, Optional, Tuple, Union
import numpy as np
import numpy.typing as npt
from supervision.config import CLASS_NAME_DATA_FIELD
from supervision.detection.core import Detections
from supervision.detection.utils import get_data_item, is_data_equal
from supervision.validators import validate_keypoints_fields
@@ -612,3 +613,75 @@ class KeyPoints:
```
"""
return cls(xy=np.empty((0, 0, 2), dtype=np.float32))
def is_empty(self) -> bool:
"""
Returns `True` if the `KeyPoints` object is considered empty.
"""
empty_keypoints = KeyPoints.empty()
empty_keypoints.data = self.data
return self == empty_keypoints
def as_detections(
self, selected_keypoint_indices: Optional[Iterable[int]] = None
) -> Detections:
"""
Convert a KeyPoints object to a Detections object. This
approximates the bounding box of the detected object by
taking the bounding box that fits all keypoints.
Arguments:
selected_keypoint_indices (Optional[Iterable[int]]): The
indices of the keypoints to include in the bounding box
calculation. This helps focus on a subset of keypoints,
e.g. when some are occluded. Captures all keypoints by default.
Returns:
detections (Detections): The converted detections object.
Example:
```python
keypoints = sv.KeyPoints.from_inference(...)
detections = keypoints.as_detections()
```
"""
if self.is_empty():
return Detections.empty()
detections_list = []
for i, xy in enumerate(self.xy):
if selected_keypoint_indices:
xy = xy[selected_keypoint_indices]
# [0, 0] used by some frameworks to indicate missing keypoints
xy = xy[~np.all(xy == 0, axis=1)]
if len(xy) == 0:
xyxy = np.array([[0, 0, 0, 0]], dtype=np.float32)
else:
x_min = xy[:, 0].min()
x_max = xy[:, 0].max()
y_min = xy[:, 1].min()
y_max = xy[:, 1].max()
xyxy = np.array([[x_min, y_min, x_max, y_max]], dtype=np.float32)
if self.confidence is None:
confidence = None
else:
confidence = self.confidence[i]
if selected_keypoint_indices:
confidence = confidence[selected_keypoint_indices]
confidence = np.array([confidence.mean()], dtype=np.float32)
detections_list.append(
Detections(
xyxy=xyxy,
confidence=confidence,
)
)
detections = Detections.merge(detections_list)
detections.class_id = self.class_id
detections.data = self.data
detections = detections[detections.area > 0]
return detections
supervision/metrics/__init__.py
+6 -1
View File
@@ -1,5 +1,4 @@
from supervision.metrics.core import (
CLASS_ID_NONE,
AveragingMethod,
Metric,
MetricTarget,
@@ -9,6 +8,12 @@ from supervision.metrics.mean_average_precision import (
MeanAveragePrecision,
MeanAveragePrecisionResult,
)
from supervision.metrics.mean_average_recall import (
MeanAverageRecall,
MeanAverageRecallResult,
)
from supervision.metrics.precision import Precision, PrecisionResult
from supervision.metrics.recall import Recall, RecallResult
from supervision.metrics.utils.object_size import (
ObjectSizeCategory,
get_detection_size_category,
supervision/metrics/core.py
+14 -15
View File
@@ -4,9 +4,6 @@ from abc import ABC, abstractmethod
from enum import Enum
from typing import Any
CLASS_ID_NONE = -1
"""Used by metrics module as class ID, when none is present"""
class Metric(ABC):
"""
@@ -40,9 +37,10 @@ class MetricTarget(Enum):
"""
Specifies what type of detection is used to compute the metric.
* BOXES: xyxy bounding boxes
* MASKS: Binary masks
* ORIENTED_BOUNDING_BOXES: Oriented bounding boxes (OBB)
Attributes:
BOXES: xyxy bounding boxes
MASKS: Binary masks
ORIENTED_BOUNDING_BOXES: Oriented bounding boxes (OBB)
"""
BOXES = "boxes"
@@ -57,15 +55,16 @@ class AveragingMethod(Enum):
Suppose, before returning the final result, a metric is computed for each class.
How do you combine those to get the final number?
* MACRO: Calculate the metric for each class and average the results. The simplest
averaging method, but it does not take class imbalance into account.
* MICRO: Calculate the metric globally by counting the total true positives, false
positives, and false negatives. Micro averaging is useful when you want to give
more importance to classes with more samples. It's also more appropriate if you
have an imbalance in the number of instances per class.
* WEIGHTED: Calculate the metric for each class and average the results, weighted by
the number of true instances of each class. Use weighted averaging if you want
to take class imbalance into account.
Attributes:
MACRO: Calculate the metric for each class and average the results. The simplest
averaging method, but it does not take class imbalance into account.
MICRO: Calculate the metric globally by counting the total true positives, false
positives, and false negatives. Micro averaging is useful when you want to
give more importance to classes with more samples. It's also more
appropriate if you have an imbalance in the number of instances per class.
WEIGHTED: Calculate the metric for each class and average the results, weighted
by the number of true instances of each class. Use weighted averaging if
you want to take class imbalance into account.
"""
MACRO = "macro"
supervision/metrics/f1_score.py
+118 -15
View File
@@ -9,7 +9,11 @@ from matplotlib import pyplot as plt
from supervision.config import ORIENTED_BOX_COORDINATES
from supervision.detection.core import Detections
from supervision.detection.utils import box_iou_batch, mask_iou_batch
from supervision.detection.utils import (
box_iou_batch,
mask_iou_batch,
oriented_box_iou_batch,
)
from supervision.draw.color import LEGACY_COLOR_PALETTE
from supervision.metrics.core import AveragingMethod, Metric, MetricTarget
from supervision.metrics.utils.object_size import (
@@ -23,23 +27,75 @@ if TYPE_CHECKING:
class F1Score(Metric):
"""
F1 Score is a metric used to evaluate object detection models. It is the harmonic
mean of precision and recall, calculated at different IoU thresholds.
In simple terms, F1 Score is a measure of a model's balance between precision and
recall (accuracy and completeness), calculated as:
`F1 = 2 * (precision * recall) / (precision + recall)`
Example:
```python
import supervision as sv
from supervision.metrics import F1Score
predictions = sv.Detections(...)
targets = sv.Detections(...)
f1_metric = F1Score()
f1_result = f1_metric.update(predictions, targets).compute()
print(f1_result.f1_50)
# 0.7618
print(f1_result)
# F1ScoreResult:
# Metric target: MetricTarget.BOXES
# Averaging method: AveragingMethod.WEIGHTED
# F1 @ 50: 0.7618
# F1 @ 75: 0.7487
# F1 @ thresh: [0.76175 0.76068 0.76068]
# IoU thresh: [0.5 0.55 0.6 ...]
# F1 per class:
# 0: [0.70968 0.70968 0.70968 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
f1_result.plot()
```
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/f1_plot_example.png\
){ align=center width="800" }
"""
def __init__(
self,
metric_target: MetricTarget = MetricTarget.BOXES,
averaging_method: AveragingMethod = AveragingMethod.WEIGHTED,
):
self._metric_target = metric_target
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
raise NotImplementedError(
"F1 score is not implemented for oriented bounding boxes."
)
"""
Initialize the F1Score metric.
Args:
metric_target (MetricTarget): The type of detection data to use.
averaging_method (AveragingMethod): The averaging method used to compute the
F1 scores. Determines how the F1 scores are aggregated across classes.
"""
self._metric_target = metric_target
self.averaging_method = averaging_method
self._predictions_list: List[Detections] = []
self._targets_list: List[Detections] = []
def reset(self) -> None:
"""
Reset the metric to its initial state, clearing all stored data.
"""
self._predictions_list = []
self._targets_list = []
@@ -48,6 +104,16 @@ class F1Score(Metric):
predictions: Union[Detections, List[Detections]],
targets: Union[Detections, List[Detections]],
) -> F1Score:
"""
Add new predictions and targets to the metric, but do not compute the result.
Args:
predictions (Union[Detections, List[Detections]]): The predicted detections.
targets (Union[Detections, List[Detections]]): The target detections.
Returns:
(F1Score): The updated metric instance.
"""
if not isinstance(predictions, list):
predictions = [predictions]
if not isinstance(targets, list):
@@ -65,6 +131,13 @@ class F1Score(Metric):
return self
def compute(self) -> F1ScoreResult:
"""
Calculate the F1 score metric based on the stored predictions and ground-truth
data, at different IoU thresholds.
Returns:
(F1ScoreResult): The F1 score metric result.
"""
result = self._compute(self._predictions_list, self._targets_list)
small_predictions, small_targets = self._filter_predictions_and_targets_by_size(
@@ -112,8 +185,12 @@ class F1Score(Metric):
iou = box_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.MASKS:
iou = mask_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
iou = oriented_box_iou_batch(
target_contents, prediction_contents
)
else:
raise NotImplementedError(
raise ValueError(
"Unsupported metric target for IoU calculation"
)
@@ -312,12 +389,22 @@ class F1Score(Metric):
return (
detections.mask
if detections.mask is not None
else np.empty((0, 0, 0), dtype=bool)
else self._make_empty_content()
)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
if obb := detections.data.get(ORIENTED_BOX_COORDINATES):
return np.ndarray(obb, dtype=np.float32)
return np.empty((0, 8), dtype=np.float32)
obb = detections.data.get(ORIENTED_BOX_COORDINATES)
if obb is not None and len(obb) > 0:
return np.array(obb, dtype=np.float32)
return self._make_empty_content()
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _make_empty_content(self) -> np.ndarray:
if self._metric_target == MetricTarget.BOXES:
return np.empty((0, 4), dtype=np.float32)
if self._metric_target == MetricTarget.MASKS:
return np.empty((0, 0, 0), dtype=bool)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
return np.empty((0, 4, 2), dtype=np.float32)
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _filter_detections_by_size(
@@ -373,7 +460,6 @@ class F1ScoreResult:
The results of the F1 score metric calculation.
Defaults to `0` if no detections or targets were provided.
Provides a custom `__str__` method for pretty printing.
Attributes:
metric_target (MetricTarget): the type of data used for the metric -
@@ -390,11 +476,11 @@ class F1ScoreResult:
matched_classes (np.ndarray): the class IDs of all matched classes.
Corresponds to the rows of `f1_per_class`.
small_objects (Optional[F1ScoreResult]): the F1 metric results
for small objects.
for small objects (area < 32²).
medium_objects (Optional[F1ScoreResult]): the F1 metric results
for medium objects.
for medium objects (32² ≤ area < 96²).
large_objects (Optional[F1ScoreResult]): the F1 metric results
for large objects.
for large objects (area ≥ 96²).
"""
metric_target: MetricTarget
@@ -424,6 +510,19 @@ class F1ScoreResult:
Example:
```python
print(f1_result)
# F1ScoreResult:
# Metric target: MetricTarget.BOXES
# Averaging method: AveragingMethod.WEIGHTED
# F1 @ 50: 0.7618
# F1 @ 75: 0.7487
# F1 @ thresh: [0.76175 0.76068 0.76068]
# IoU thresh: [0.5 0.55 0.6 ...]
# F1 per class:
# 0: [0.70968 0.70968 0.70968 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
```
"""
out_str = (
@@ -487,6 +586,10 @@ class F1ScoreResult:
def plot(self):
"""
Plot the F1 results.
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/f1_plot_example.png\
){ align=center width="800" }
"""
labels = ["F1@50", "F1@75"]
supervision/metrics/mean_average_precision.py
+98 -32
View File
@@ -9,7 +9,11 @@ from matplotlib import pyplot as plt
from supervision.config import ORIENTED_BOX_COORDINATES
from supervision.detection.core import Detections
from supervision.detection.utils import box_iou_batch, mask_iou_batch
from supervision.detection.utils import (
box_iou_batch,
mask_iou_batch,
oriented_box_iou_batch,
)
from supervision.draw.color import LEGACY_COLOR_PALETTE
from supervision.metrics.core import Metric, MetricTarget
from supervision.metrics.utils.object_size import (
@@ -23,6 +27,48 @@ if TYPE_CHECKING:
class MeanAveragePrecision(Metric):
"""
Mean Average Precision (mAP) is a metric used to evaluate object detection models.
It is the average of the precision-recall curves at different IoU thresholds.
Example:
```python
import supervision as sv
from supervision.metrics import MeanAveragePrecision
predictions = sv.Detections(...)
targets = sv.Detections(...)
map_metric = MeanAveragePrecision()
map_result = map_metric.update(predictions, targets).compute()
print(map_result.map50_95)
# 0.4674
print(map_result)
# MeanAveragePrecisionResult:
# Metric target: MetricTarget.BOXES
# Class agnostic: False
# mAP @ 50:95: 0.4674
# mAP @ 50: 0.5048
# mAP @ 75: 0.4796
# mAP scores: [0.50485 0.50377 0.50377 ...]
# IoU thresh: [0.5 0.55 0.6 ...]
# AP per class:
# 0: [0.67699 0.67699 0.67699 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
map_result.plot()
```
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/mAP_plot_example.png\
){ align=center width="800" }
"""
def __init__(
self,
metric_target: MetricTarget = MetricTarget.BOXES,
@@ -36,17 +82,15 @@ class MeanAveragePrecision(Metric):
class_agnostic (bool): Whether to treat all data as a single class.
"""
self._metric_target = metric_target
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
raise NotImplementedError(
"Mean Average Precision is not implemented for oriented bounding boxes."
)
self._class_agnostic = class_agnostic
self._predictions_list: List[Detections] = []
self._targets_list: List[Detections] = []
def reset(self) -> None:
"""
Reset the metric to its initial state, clearing all stored data.
"""
self._predictions_list = []
self._targets_list = []
@@ -76,6 +120,15 @@ class MeanAveragePrecision(Metric):
f" targets ({len(targets)}) during the update must be the same."
)
if self._class_agnostic:
predictions = deepcopy(predictions)
targets = deepcopy(targets)
for prediction in predictions:
prediction.class_id[:] = -1
for target in targets:
target.class_id[:] = -1
self._predictions_list.extend(predictions)
self._targets_list.extend(targets)
@@ -86,26 +139,10 @@ class MeanAveragePrecision(Metric):
) -> MeanAveragePrecisionResult:
"""
Calculate Mean Average Precision based on predicted and ground-truth
detections at different thresholds.
detections at different thresholds.
Returns:
(MeanAveragePrecisionResult): New instance of MeanAveragePrecision.
Example:
```python
import supervision as sv
from supervision.metrics import MeanAveragePrecision
predictions = sv.Detections(...)
targets = sv.Detections(...)
map_metric = MeanAveragePrecision()
map_result = map_metric.update(predictions, targets).compute()
print(map_result)
print(map_result.map50_95)
map_result.plot()
```
(MeanAveragePrecisionResult): The Mean Average Precision result.
"""
result = self._compute(self._predictions_list, self._targets_list)
@@ -172,14 +209,19 @@ class MeanAveragePrecision(Metric):
iou = box_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.MASKS:
iou = mask_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
iou = oriented_box_iou_batch(
target_contents, prediction_contents
)
else:
raise NotImplementedError(
raise ValueError(
"Unsupported metric target for IoU calculation"
)
matches = self._match_detection_batch(
predictions.class_id, targets.class_id, iou, iou_thresholds
)
stats.append(
(
matches,
@@ -203,6 +245,7 @@ class MeanAveragePrecision(Metric):
return MeanAveragePrecisionResult(
metric_target=self._metric_target,
is_class_agnostic=self._class_agnostic,
mAP_scores=mAP_scores,
iou_thresholds=iou_thresholds,
matched_classes=unique_classes,
@@ -230,7 +273,7 @@ class MeanAveragePrecision(Metric):
for r, p in zip(recall[::-1], precision[::-1]):
precision_levels[recall_levels <= r] = p
average_precision = (1 / 100 * precision_levels).sum()
average_precision = (1 / 101 * precision_levels).sum()
return average_precision
@staticmethod
@@ -332,8 +375,9 @@ class MeanAveragePrecision(Metric):
else self._make_empty_content()
)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
if obb := detections.data.get(ORIENTED_BOX_COORDINATES):
return np.ndarray(obb, dtype=np.float32)
obb = detections.data.get(ORIENTED_BOX_COORDINATES)
if obb is not None and len(obb) > 0:
return np.array(obb, dtype=np.float32)
return self._make_empty_content()
raise ValueError(f"Invalid metric target: {self._metric_target}")
@@ -343,7 +387,7 @@ class MeanAveragePrecision(Metric):
if self._metric_target == MetricTarget.MASKS:
return np.empty((0, 0, 0), dtype=bool)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
return np.empty((0, 8), dtype=np.float32)
return np.empty((0, 4, 2), dtype=np.float32)
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _filter_detections_by_size(
@@ -383,6 +427,8 @@ class MeanAveragePrecisionResult:
Attributes:
metric_target (MetricTarget): the type of data used for the metric -
boxes, masks or oriented bounding boxes.
class_agnostic (bool): When computing class-agnostic results, class ID
is set to `-1`.
mAP_map50_95 (float): the mAP score at IoU thresholds from `0.5` to `0.95`.
mAP_map50 (float): the mAP score at IoU threshold of `0.5`.
mAP_map75 (float): the mAP score at IoU threshold of `0.75`.
@@ -394,14 +440,15 @@ class MeanAveragePrecisionResult:
matched_classes (np.ndarray): the class IDs of all matched classes.
Corresponds to the rows of `ap_per_class`.
small_objects (Optional[MeanAveragePrecisionResult]): the mAP results
for small objects.
for small objects (area < 32²).
medium_objects (Optional[MeanAveragePrecisionResult]): the mAP results
for medium objects.
for medium objects (32² ≤ area < 96²).
large_objects (Optional[MeanAveragePrecisionResult]): the mAP results
for large objects.
for large objects (area ≥ 96²).
"""
metric_target: MetricTarget
is_class_agnostic: bool
@property
def map50_95(self) -> float:
@@ -430,12 +477,27 @@ class MeanAveragePrecisionResult:
Example:
```python
print(map_result)
# MeanAveragePrecisionResult:
# Metric target: MetricTarget.BOXES
# Class agnostic: False
# mAP @ 50:95: 0.4674
# mAP @ 50: 0.5048
# mAP @ 75: 0.4796
# mAP scores: [0.50485 0.50377 0.50377 ...]
# IoU thresh: [0.5 0.55 0.6 ...]
# AP per class:
# 0: [0.67699 0.67699 0.67699 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
```
"""
out_str = (
f"{self.__class__.__name__}:\n"
f"Metric target: {self.metric_target}\n"
f"Class agnostic: {self.is_class_agnostic}\n"
f"mAP @ 50:95: {self.map50_95:.4f}\n"
f"mAP @ 50: {self.map50:.4f}\n"
f"mAP @ 75: {self.map75:.4f}\n"
@@ -500,6 +562,10 @@ class MeanAveragePrecisionResult:
def plot(self):
"""
Plot the mAP results.
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/mAP_plot_example.png\
){ align=center width="800" }
"""
labels = ["mAP@50:95", "mAP@50", "mAP@75"]
supervision/metrics/mean_average_recall.py
+697
View File
@@ -0,0 +1,697 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
import numpy as np
from matplotlib import pyplot as plt
from supervision.config import ORIENTED_BOX_COORDINATES
from supervision.detection.core import Detections
from supervision.detection.utils import (
box_iou_batch,
mask_iou_batch,
oriented_box_iou_batch,
)
from supervision.draw.color import LEGACY_COLOR_PALETTE
from supervision.metrics.core import Metric, MetricTarget
from supervision.metrics.utils.object_size import (
ObjectSizeCategory,
get_detection_size_category,
)
from supervision.metrics.utils.utils import ensure_pandas_installed
if TYPE_CHECKING:
import pandas as pd
class MeanAverageRecall(Metric):
"""
Mean Average Recall (mAR) measures how well the model detects
and retrieves relevant objects by averaging recall over multiple
IoU thresholds, classes and detection limits.
Intuitively, while Recall measures the ability to find all relevant
objects, mAR narrows down how many detections are considered for each
class. For example, mAR @ 100 considers the top 100 highest confidence
detections for each class. mAR @ 1 considers only the highest
confidence detection for each class.
Example:
```python
import supervision as sv
from supervision.metrics import MeanAverageRecall
predictions = sv.Detections(...)
targets = sv.Detections(...)
map_metric = MeanAverageRecall()
map_result = map_metric.update(predictions, targets).compute()
print(mar_results.mar_at_100)
# 0.5241
print(mar_results)
# MeanAverageRecallResult:
# Metric target: MetricTarget.BOXES
# mAR @ 1: 0.1362
# mAR @ 10: 0.4239
# mAR @ 100: 0.5241
# max detections: [1 10 100]
# IoU thresh: [0.5 0.55 0.6 ...]
# mAR per class:
# 0: [0.78571 0.78571 0.78571 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
mar_results.plot()
```
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/mAR_plot_example.png\
){ align=center width="800" }
"""
def __init__(
self,
metric_target: MetricTarget = MetricTarget.BOXES,
):
"""
Initialize the Mean Average Recall metric.
Args:
metric_target (MetricTarget): The type of detection data to use.
"""
self._metric_target = metric_target
self._predictions_list: List[Detections] = []
self._targets_list: List[Detections] = []
self.max_detections = np.array([1, 10, 100])
def reset(self) -> None:
"""
Reset the metric to its initial state, clearing all stored data.
"""
self._predictions_list = []
self._targets_list = []
def update(
self,
predictions: Union[Detections, List[Detections]],
targets: Union[Detections, List[Detections]],
) -> MeanAverageRecall:
"""
Add new predictions and targets to the metric, but do not compute the result.
Args:
predictions (Union[Detections, List[Detections]]): The predicted detections.
targets (Union[Detections, List[Detections]]): The target detections.
Returns:
(Recall): The updated metric instance.
"""
if not isinstance(predictions, list):
predictions = [predictions]
if not isinstance(targets, list):
targets = [targets]
if len(predictions) != len(targets):
raise ValueError(
f"The number of predictions ({len(predictions)}) and"
f" targets ({len(targets)}) during the update must be the same."
)
self._predictions_list.extend(predictions)
self._targets_list.extend(targets)
return self
def compute(self) -> MeanAverageRecallResult:
"""
Calculate the Mean Average Recall metric based on the stored predictions
and ground-truth, at different IoU thresholds and maximum detection counts.
Returns:
(MeanAverageRecallResult): The Mean Average Recall metric result.
"""
result = self._compute(self._predictions_list, self._targets_list)
small_predictions, small_targets = self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.SMALL
)
result.small_objects = self._compute(small_predictions, small_targets)
medium_predictions, medium_targets = (
self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.MEDIUM
)
)
result.medium_objects = self._compute(medium_predictions, medium_targets)
large_predictions, large_targets = self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.LARGE
)
result.large_objects = self._compute(large_predictions, large_targets)
return result
def _compute(
self, predictions_list: List[Detections], targets_list: List[Detections]
) -> MeanAverageRecallResult:
iou_thresholds = np.linspace(0.5, 0.95, 10)
stats = []
for predictions, targets in zip(predictions_list, targets_list):
prediction_contents = self._detections_content(predictions)
target_contents = self._detections_content(targets)
if len(targets) > 0:
if len(predictions) == 0:
stats.append(
(
np.zeros((0, iou_thresholds.size), dtype=bool),
np.zeros((0,), dtype=np.float32),
np.zeros((0,), dtype=int),
targets.class_id,
)
)
else:
if self._metric_target == MetricTarget.BOXES:
iou = box_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.MASKS:
iou = mask_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
iou = oriented_box_iou_batch(
target_contents, prediction_contents
)
else:
raise ValueError(
"Unsupported metric target for IoU calculation"
)
matches = self._match_detection_batch(
predictions.class_id, targets.class_id, iou, iou_thresholds
)
stats.append(
(
matches,
predictions.confidence,
predictions.class_id,
targets.class_id,
)
)
if not stats:
return MeanAverageRecallResult(
metric_target=self._metric_target,
recall_scores=np.zeros(iou_thresholds.shape[0]),
recall_per_class=np.zeros((0, iou_thresholds.shape[0])),
max_detections=self.max_detections,
iou_thresholds=iou_thresholds,
matched_classes=np.array([], dtype=int),
small_objects=None,
medium_objects=None,
large_objects=None,
)
concatenated_stats = [np.concatenate(items, 0) for items in zip(*stats)]
recall_scores_per_k, recall_per_class, unique_classes = (
self._compute_average_recall_for_classes(*concatenated_stats)
)
return MeanAverageRecallResult(
metric_target=self._metric_target,
recall_scores=recall_scores_per_k,
recall_per_class=recall_per_class,
max_detections=self.max_detections,
iou_thresholds=iou_thresholds,
matched_classes=unique_classes,
small_objects=None,
medium_objects=None,
large_objects=None,
)
def _compute_average_recall_for_classes(
self,
matches: np.ndarray,
prediction_confidence: np.ndarray,
prediction_class_ids: np.ndarray,
true_class_ids: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
sorted_indices = np.argsort(-prediction_confidence)
matches = matches[sorted_indices]
prediction_class_ids = prediction_class_ids[sorted_indices]
unique_classes, class_counts = np.unique(true_class_ids, return_counts=True)
recalls_at_k = []
for max_detections in self.max_detections:
# Shape: PxTh,P,C,C -> CxThx3
confusion_matrix = self._compute_confusion_matrix(
matches,
prediction_class_ids,
unique_classes,
class_counts,
max_detections=max_detections,
)
# Shape: CxThx3 -> CxTh
recall_per_class = self._compute_recall(confusion_matrix)
recalls_at_k.append(recall_per_class)
# Shape: KxCxTh -> KxC
recalls_at_k = np.array(recalls_at_k)
average_recall_per_class = np.mean(recalls_at_k, axis=2)
# Shape: KxC -> K
recall_scores = np.mean(average_recall_per_class, axis=1)
return recall_scores, recall_per_class, unique_classes
@staticmethod
def _match_detection_batch(
predictions_classes: np.ndarray,
target_classes: np.ndarray,
iou: np.ndarray,
iou_thresholds: np.ndarray,
) -> np.ndarray:
num_predictions, num_iou_levels = (
predictions_classes.shape[0],
iou_thresholds.shape[0],
)
correct = np.zeros((num_predictions, num_iou_levels), dtype=bool)
correct_class = target_classes[:, None] == predictions_classes
for i, iou_level in enumerate(iou_thresholds):
matched_indices = np.where((iou >= iou_level) & correct_class)
if matched_indices[0].shape[0]:
combined_indices = np.stack(matched_indices, axis=1)
iou_values = iou[matched_indices][:, None]
matches = np.hstack([combined_indices, iou_values])
if matched_indices[0].shape[0] > 1:
matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
correct[matches[:, 1].astype(int), i] = True
return correct
@staticmethod
def _compute_confusion_matrix(
sorted_matches: np.ndarray,
sorted_prediction_class_ids: np.ndarray,
unique_classes: np.ndarray,
class_counts: np.ndarray,
max_detections: Optional[int] = None,
) -> np.ndarray:
"""
Compute the confusion matrix for each class and IoU threshold.
Assumes the matches and prediction_class_ids are sorted by confidence
in descending order.
Args:
sorted_matches: np.ndarray, bool, shape (P, Th), that is True
if the prediction is a true positive at the given IoU threshold.
sorted_prediction_class_ids: np.ndarray, int, shape (P,), containing
the class id for each prediction.
unique_classes: np.ndarray, int, shape (C,), containing the unique
class ids.
class_counts: np.ndarray, int, shape (C,), containing the number
of true instances for each class.
max_detections: Optional[int], the maximum number of detections to
consider for each class. Extra detections are considered false
positives. By default, all detections are considered.
Returns:
np.ndarray, shape (C, Th, 3), containing the true positives, false
positives, and false negatives for each class and IoU threshold.
"""
num_thresholds = sorted_matches.shape[1]
num_classes = unique_classes.shape[0]
confusion_matrix = np.zeros((num_classes, num_thresholds, 3))
for class_idx, class_id in enumerate(unique_classes):
is_class = sorted_prediction_class_ids == class_id
num_true = class_counts[class_idx]
num_predictions = is_class.sum()
if num_predictions == 0:
true_positives = np.zeros(num_thresholds)
false_positives = np.zeros(num_thresholds)
false_negatives = np.full(num_thresholds, num_true)
elif num_true == 0:
true_positives = np.zeros(num_thresholds)
false_positives = np.full(num_thresholds, num_predictions)
false_negatives = np.zeros(num_thresholds)
else:
limited_matches = sorted_matches[is_class][slice(max_detections)]
true_positives = limited_matches.sum(0)
false_positives = (1 - limited_matches).sum(0)
false_negatives = num_true - true_positives
false_negatives = num_true - true_positives
confusion_matrix[class_idx] = np.stack(
[true_positives, false_positives, false_negatives], axis=1
)
return confusion_matrix
@staticmethod
def _compute_recall(confusion_matrix: np.ndarray) -> np.ndarray:
"""
Broadcastable function, computing the recall from the confusion matrix.
Arguments:
confusion_matrix: np.ndarray, shape (N, ..., 3), where the last dimension
contains the true positives, false positives, and false negatives.
Returns:
np.ndarray, shape (N, ...), containing the recall for each element.
"""
if not confusion_matrix.shape[-1] == 3:
raise ValueError(
f"Confusion matrix must have shape (..., 3), got "
f"{confusion_matrix.shape}"
)
true_positives = confusion_matrix[..., 0]
false_negatives = confusion_matrix[..., 2]
denominator = true_positives + false_negatives
recall = np.where(denominator == 0, 0, true_positives / denominator)
return recall
def _detections_content(self, detections: Detections) -> np.ndarray:
"""Return boxes, masks or oriented bounding boxes from detections."""
if self._metric_target == MetricTarget.BOXES:
return detections.xyxy
if self._metric_target == MetricTarget.MASKS:
return (
detections.mask
if detections.mask is not None
else self._make_empty_content()
)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
obb = detections.data.get(ORIENTED_BOX_COORDINATES)
if obb is not None and len(obb) > 0:
return np.array(obb, dtype=np.float32)
return self._make_empty_content()
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _make_empty_content(self) -> np.ndarray:
if self._metric_target == MetricTarget.BOXES:
return np.empty((0, 4), dtype=np.float32)
if self._metric_target == MetricTarget.MASKS:
return np.empty((0, 0, 0), dtype=bool)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
return np.empty((0, 4, 2), dtype=np.float32)
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _filter_detections_by_size(
self, detections: Detections, size_category: ObjectSizeCategory
) -> Detections:
"""Return a copy of detections with contents filtered by object size."""
new_detections = deepcopy(detections)
if detections.is_empty() or size_category == ObjectSizeCategory.ANY:
return new_detections
sizes = get_detection_size_category(new_detections, self._metric_target)
size_mask = sizes == size_category.value
new_detections.xyxy = new_detections.xyxy[size_mask]
if new_detections.mask is not None:
new_detections.mask = new_detections.mask[size_mask]
if new_detections.class_id is not None:
new_detections.class_id = new_detections.class_id[size_mask]
if new_detections.confidence is not None:
new_detections.confidence = new_detections.confidence[size_mask]
if new_detections.tracker_id is not None:
new_detections.tracker_id = new_detections.tracker_id[size_mask]
if new_detections.data is not None:
for key, value in new_detections.data.items():
new_detections.data[key] = np.array(value)[size_mask]
return new_detections
def _filter_predictions_and_targets_by_size(
self,
predictions_list: List[Detections],
targets_list: List[Detections],
size_category: ObjectSizeCategory,
) -> Tuple[List[Detections], List[Detections]]:
new_predictions_list = []
new_targets_list = []
for predictions, targets in zip(predictions_list, targets_list):
new_predictions_list.append(
self._filter_detections_by_size(predictions, size_category)
)
new_targets_list.append(
self._filter_detections_by_size(targets, size_category)
)
return new_predictions_list, new_targets_list
@dataclass
class MeanAverageRecallResult:
# """
# The results of the recall metric calculation.
# Defaults to `0` if no detections or targets were provided.
# Attributes:
# metric_target (MetricTarget): the type of data used for the metric -
# boxes, masks or oriented bounding boxes.
# averaging_method (AveragingMethod): the averaging method used to compute the
# recall. Determines how the recall is aggregated across classes.
# recall_at_50 (float): the recall at IoU threshold of `0.5`.
# recall_at_75 (float): the recall at IoU threshold of `0.75`.
# recall_scores (np.ndarray): the recall scores at each IoU threshold.
# Shape: `(num_iou_thresholds,)`
# recall_per_class (np.ndarray): the recall scores per class and IoU threshold.
# Shape: `(num_target_classes, num_iou_thresholds)`
# iou_thresholds (np.ndarray): the IoU thresholds used in the calculations.
# matched_classes (np.ndarray): the class IDs of all matched classes.
# Corresponds to the rows of `recall_per_class`.
# small_objects (Optional[RecallResult]): the Recall metric results
# for small objects.
# medium_objects (Optional[RecallResult]): the Recall metric results
# for medium objects.
# large_objects (Optional[RecallResult]): the Recall metric results
# for large objects.
# """
"""
The results of the Mean Average Recall metric calculation.
Defaults to `0` if no detections or targets were provided.
Attributes:
metric_target (MetricTarget): the type of data used for the metric -
boxes, masks or oriented bounding boxes.
mAR_at_1 (float): the Mean Average Recall, when considering only the top
highest confidence detection for each class.
mAR_at_10 (float): the Mean Average Recall, when considering top 10
highest confidence detections for each class.
mAR_at_100 (float): the Mean Average Recall, when considering top 100
highest confidence detections for each class.
recall_per_class (np.ndarray): the recall scores per class and IoU threshold.
Shape: `(num_target_classes, num_iou_thresholds)`
max_detections (np.ndarray): the array with maximum number of detections
considered.
iou_thresholds (np.ndarray): the IoU thresholds used in the calculations.
matched_classes (np.ndarray): the class IDs of all matched classes.
Corresponds to the rows of `recall_per_class`.
small_objects (Optional[MeanAverageRecallResult]): the Mean Average Recall
metric results for small objects (area < 32²).
medium_objects (Optional[MeanAverageRecallResult]): the Mean Average Recall
metric results for medium objects (32² ≤ area < 96²).
large_objects (Optional[MeanAverageRecallResult]): the Mean Average Recall
metric results for large objects (area ≥ 96²).
"""
metric_target: MetricTarget
@property
def mAR_at_1(self) -> float:
return self.recall_scores[0]
@property
def mAR_at_10(self) -> float:
return self.recall_scores[1]
@property
def mAR_at_100(self) -> float:
return self.recall_scores[2]
recall_scores: np.ndarray
recall_per_class: np.ndarray
max_detections: np.ndarray
iou_thresholds: np.ndarray
matched_classes: np.ndarray
small_objects: Optional[MeanAverageRecallResult]
medium_objects: Optional[MeanAverageRecallResult]
large_objects: Optional[MeanAverageRecallResult]
def __str__(self) -> str:
"""
Format as a pretty string.
Example:
```python
# MeanAverageRecallResult:
# Metric target: MetricTarget.BOXES
# mAR @ 1: 0.1362
# mAR @ 10: 0.4239
# mAR @ 100: 0.5241
# max detections: [1 10 100]
# IoU thresh: [0.5 0.55 0.6 ...]
# mAR per class:
# 0: [0.78571 0.78571 0.78571 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
```
"""
out_str = (
f"{self.__class__.__name__}:\n"
f"Metric target: {self.metric_target}\n"
f"mAR @ 1: {self.mAR_at_1:.4f}\n"
f"mAR @ 10: {self.mAR_at_10:.4f}\n"
f"mAR @ 100: {self.mAR_at_100:.4f}\n"
f"max detections: {self.max_detections}\n"
f"IoU thresh: {self.iou_thresholds}\n"
f"mAR per class:\n"
)
if self.recall_per_class.size == 0:
out_str += " No results\n"
for class_id, recall_of_class in zip(
self.matched_classes, self.recall_per_class
):
out_str += f" {class_id}: {recall_of_class}\n"
indent = " "
if self.small_objects is not None:
indented = indent + str(self.small_objects).replace("\n", f"\n{indent}")
out_str += f"\nSmall objects:\n{indented}"
if self.medium_objects is not None:
indented = indent + str(self.medium_objects).replace("\n", f"\n{indent}")
out_str += f"\nMedium objects:\n{indented}"
if self.large_objects is not None:
indented = indent + str(self.large_objects).replace("\n", f"\n{indent}")
out_str += f"\nLarge objects:\n{indented}"
return out_str
def to_pandas(self) -> "pd.DataFrame":
"""
Convert the result to a pandas DataFrame.
Returns:
(pd.DataFrame): The result as a DataFrame.
"""
ensure_pandas_installed()
import pandas as pd
pandas_data = {
"mAR @ 1": self.mAR_at_1,
"mAR @ 10": self.mAR_at_10,
"mAR @ 100": self.mAR_at_100,
}
if self.small_objects is not None:
small_objects_df = self.small_objects.to_pandas()
for key, value in small_objects_df.items():
pandas_data[f"small_objects_{key}"] = value
if self.medium_objects is not None:
medium_objects_df = self.medium_objects.to_pandas()
for key, value in medium_objects_df.items():
pandas_data[f"medium_objects_{key}"] = value
if self.large_objects is not None:
large_objects_df = self.large_objects.to_pandas()
for key, value in large_objects_df.items():
pandas_data[f"large_objects_{key}"] = value
return pd.DataFrame(pandas_data, index=[0])
def plot(self):
"""
Plot the Mean Average Recall results.
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/mAR_plot_example.png\
){ align=center width="800" }
"""
labels = ["mAR @ 1", "mAR @ 10", "mAR @ 100"]
values = [self.mAR_at_1, self.mAR_at_10, self.mAR_at_100]
colors = [LEGACY_COLOR_PALETTE[0]] * 3
if self.small_objects is not None:
small_objects = self.small_objects
labels += ["Small: mAR @ 1", "Small: mAR @ 10", "Small: mAR @ 100"]
values += [
small_objects.mAR_at_1,
small_objects.mAR_at_10,
small_objects.mAR_at_100,
]
colors += [LEGACY_COLOR_PALETTE[3]] * 3
if self.medium_objects is not None:
medium_objects = self.medium_objects
labels += ["Medium: mAR @ 1", "Medium: mAR @ 10", "Medium: mAR @ 100"]
values += [
medium_objects.mAR_at_1,
medium_objects.mAR_at_10,
medium_objects.mAR_at_100,
]
colors += [LEGACY_COLOR_PALETTE[2]] * 3
if self.large_objects is not None:
large_objects = self.large_objects
labels += ["Large: mAR @ 1", "Large: mAR @ 10", "Large: mAR @ 100"]
values += [
large_objects.mAR_at_1,
large_objects.mAR_at_10,
large_objects.mAR_at_100,
]
colors += [LEGACY_COLOR_PALETTE[4]] * 3
plt.rcParams["font.family"] = "monospace"
_, ax = plt.subplots(figsize=(10, 6))
ax.set_ylim(0, 1)
ax.set_ylabel("Value", fontweight="bold")
title = (
f"Mean Average Recall, by Object Size"
f"\n(target: {self.metric_target.value})"
)
ax.set_title(title, fontweight="bold")
x_positions = range(len(labels))
bars = ax.bar(x_positions, values, color=colors, align="center")
ax.set_xticks(x_positions)
ax.set_xticklabels(labels, rotation=45, ha="right")
for bar in bars:
y_value = bar.get_height()
ax.text(
bar.get_x() + bar.get_width() / 2,
y_value + 0.02,
f"{y_value:.2f}",
ha="center",
va="bottom",
)
plt.rcParams["font.family"] = "sans-serif"
plt.tight_layout()
plt.show()
supervision/metrics/precision.py
+653
View File
@@ -0,0 +1,653 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
import numpy as np
from matplotlib import pyplot as plt
from supervision.config import ORIENTED_BOX_COORDINATES
from supervision.detection.core import Detections
from supervision.detection.utils import (
box_iou_batch,
mask_iou_batch,
oriented_box_iou_batch,
)
from supervision.draw.color import LEGACY_COLOR_PALETTE
from supervision.metrics.core import AveragingMethod, Metric, MetricTarget
from supervision.metrics.utils.object_size import (
ObjectSizeCategory,
get_detection_size_category,
)
from supervision.metrics.utils.utils import ensure_pandas_installed
if TYPE_CHECKING:
import pandas as pd
class Precision(Metric):
"""
Precision is a metric used to evaluate object detection models. It is the ratio of
true positive detections to the total number of predicted detections. We calculate
it at different IoU thresholds.
In simple terms, Precision is a measure of a model's accuracy, calculated as:
`Precision = TP / (TP + FP)`
Here, `TP` is the number of true positives (correct detections), and `FP` is the
number of false positive detections (detected, but incorrectly).
Example:
```python
import supervision as sv
from supervision.metrics import Precision
predictions = sv.Detections(...)
targets = sv.Detections(...)
precision_metric = Precision()
precision_result = precision_metric.update(predictions, targets).compute()
print(precision_result.precision_at_50)
# 0.8099
print(precision_result)
# PrecisionResult:
# Metric target: MetricTarget.BOXES
# Averaging method: AveragingMethod.WEIGHTED
# P @ 50: 0.8099
# P @ 75: 0.7969
# P @ thresh: [0.80992 0.80905 0.80905 ...]
# IoU thresh: [0.5 0.55 0.6 ...]
# Precision per class:
# 0: [0.64706 0.64706 0.64706 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
print(precision_result.small_objects.precision_at_50)
```
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/precision_plot_example.png\
){ align=center width="800" }
"""
def __init__(
self,
metric_target: MetricTarget = MetricTarget.BOXES,
averaging_method: AveragingMethod = AveragingMethod.WEIGHTED,
):
"""
Initialize the Precision metric.
Args:
metric_target (MetricTarget): The type of detection data to use.
averaging_method (AveragingMethod): The averaging method used to compute the
precision. Determines how the precision is aggregated across classes.
"""
self._metric_target = metric_target
self.averaging_method = averaging_method
self._predictions_list: List[Detections] = []
self._targets_list: List[Detections] = []
def reset(self) -> None:
"""
Reset the metric to its initial state, clearing all stored data.
"""
self._predictions_list = []
self._targets_list = []
def update(
self,
predictions: Union[Detections, List[Detections]],
targets: Union[Detections, List[Detections]],
) -> Precision:
"""
Add new predictions and targets to the metric, but do not compute the result.
Args:
predictions (Union[Detections, List[Detections]]): The predicted detections.
targets (Union[Detections, List[Detections]]): The target detections.
Returns:
(Precision): The updated metric instance.
"""
if not isinstance(predictions, list):
predictions = [predictions]
if not isinstance(targets, list):
targets = [targets]
if len(predictions) != len(targets):
raise ValueError(
f"The number of predictions ({len(predictions)}) and"
f" targets ({len(targets)}) during the update must be the same."
)
self._predictions_list.extend(predictions)
self._targets_list.extend(targets)
return self
def compute(self) -> PrecisionResult:
"""
Calculate the precision metric based on the stored predictions and ground-truth
data, at different IoU thresholds.
Returns:
(PrecisionResult): The precision metric result.
"""
result = self._compute(self._predictions_list, self._targets_list)
small_predictions, small_targets = self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.SMALL
)
result.small_objects = self._compute(small_predictions, small_targets)
medium_predictions, medium_targets = (
self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.MEDIUM
)
)
result.medium_objects = self._compute(medium_predictions, medium_targets)
large_predictions, large_targets = self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.LARGE
)
result.large_objects = self._compute(large_predictions, large_targets)
return result
def _compute(
self, predictions_list: List[Detections], targets_list: List[Detections]
) -> PrecisionResult:
iou_thresholds = np.linspace(0.5, 0.95, 10)
stats = []
for predictions, targets in zip(predictions_list, targets_list):
prediction_contents = self._detections_content(predictions)
target_contents = self._detections_content(targets)
if len(targets) > 0:
if len(predictions) == 0:
stats.append(
(
np.zeros((0, iou_thresholds.size), dtype=bool),
np.zeros((0,), dtype=np.float32),
np.zeros((0,), dtype=int),
targets.class_id,
)
)
else:
if self._metric_target == MetricTarget.BOXES:
iou = box_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.MASKS:
iou = mask_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
iou = oriented_box_iou_batch(
target_contents, prediction_contents
)
else:
raise ValueError(
"Unsupported metric target for IoU calculation"
)
matches = self._match_detection_batch(
predictions.class_id, targets.class_id, iou, iou_thresholds
)
stats.append(
(
matches,
predictions.confidence,
predictions.class_id,
targets.class_id,
)
)
if not stats:
return PrecisionResult(
metric_target=self._metric_target,
averaging_method=self.averaging_method,
precision_scores=np.zeros(iou_thresholds.shape[0]),
precision_per_class=np.zeros((0, iou_thresholds.shape[0])),
iou_thresholds=iou_thresholds,
matched_classes=np.array([], dtype=int),
small_objects=None,
medium_objects=None,
large_objects=None,
)
concatenated_stats = [np.concatenate(items, 0) for items in zip(*stats)]
precision_scores, precision_per_class, unique_classes = (
self._compute_precision_for_classes(*concatenated_stats)
)
return PrecisionResult(
metric_target=self._metric_target,
averaging_method=self.averaging_method,
precision_scores=precision_scores,
precision_per_class=precision_per_class,
iou_thresholds=iou_thresholds,
matched_classes=unique_classes,
small_objects=None,
medium_objects=None,
large_objects=None,
)
def _compute_precision_for_classes(
self,
matches: np.ndarray,
prediction_confidence: np.ndarray,
prediction_class_ids: np.ndarray,
true_class_ids: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
sorted_indices = np.argsort(-prediction_confidence)
matches = matches[sorted_indices]
prediction_class_ids = prediction_class_ids[sorted_indices]
unique_classes, class_counts = np.unique(true_class_ids, return_counts=True)
# Shape: PxTh,P,C,C -> CxThx3
confusion_matrix = self._compute_confusion_matrix(
matches, prediction_class_ids, unique_classes, class_counts
)
# Shape: CxThx3 -> CxTh
precision_per_class = self._compute_precision(confusion_matrix)
# Shape: CxTh -> Th
if self.averaging_method == AveragingMethod.MACRO:
precision_scores = np.mean(precision_per_class, axis=0)
elif self.averaging_method == AveragingMethod.MICRO:
confusion_matrix_merged = confusion_matrix.sum(0)
precision_scores = self._compute_precision(confusion_matrix_merged)
elif self.averaging_method == AveragingMethod.WEIGHTED:
class_counts = class_counts.astype(np.float32)
precision_scores = np.average(
precision_per_class, axis=0, weights=class_counts
)
return precision_scores, precision_per_class, unique_classes
@staticmethod
def _match_detection_batch(
predictions_classes: np.ndarray,
target_classes: np.ndarray,
iou: np.ndarray,
iou_thresholds: np.ndarray,
) -> np.ndarray:
num_predictions, num_iou_levels = (
predictions_classes.shape[0],
iou_thresholds.shape[0],
)
correct = np.zeros((num_predictions, num_iou_levels), dtype=bool)
correct_class = target_classes[:, None] == predictions_classes
for i, iou_level in enumerate(iou_thresholds):
matched_indices = np.where((iou >= iou_level) & correct_class)
if matched_indices[0].shape[0]:
combined_indices = np.stack(matched_indices, axis=1)
iou_values = iou[matched_indices][:, None]
matches = np.hstack([combined_indices, iou_values])
if matched_indices[0].shape[0] > 1:
matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
correct[matches[:, 1].astype(int), i] = True
return correct
@staticmethod
def _compute_confusion_matrix(
sorted_matches: np.ndarray,
sorted_prediction_class_ids: np.ndarray,
unique_classes: np.ndarray,
class_counts: np.ndarray,
) -> np.ndarray:
"""
Compute the confusion matrix for each class and IoU threshold.
Assumes the matches and prediction_class_ids are sorted by confidence
in descending order.
Arguments:
sorted_matches: np.ndarray, bool, shape (P, Th), that is True
if the prediction is a true positive at the given IoU threshold.
sorted_prediction_class_ids: np.ndarray, int, shape (P,), containing
the class id for each prediction.
unique_classes: np.ndarray, int, shape (C,), containing the unique
class ids.
class_counts: np.ndarray, int, shape (C,), containing the number
of true instances for each class.
Returns:
np.ndarray, shape (C, Th, 3), containing the true positives, false
positives, and false negatives for each class and IoU threshold.
"""
num_thresholds = sorted_matches.shape[1]
num_classes = unique_classes.shape[0]
confusion_matrix = np.zeros((num_classes, num_thresholds, 3))
for class_idx, class_id in enumerate(unique_classes):
is_class = sorted_prediction_class_ids == class_id
num_true = class_counts[class_idx]
num_predictions = is_class.sum()
if num_predictions == 0:
true_positives = np.zeros(num_thresholds)
false_positives = np.zeros(num_thresholds)
false_negatives = np.full(num_thresholds, num_true)
elif num_true == 0:
true_positives = np.zeros(num_thresholds)
false_positives = np.full(num_thresholds, num_predictions)
false_negatives = np.zeros(num_thresholds)
else:
true_positives = sorted_matches[is_class].sum(0)
false_positives = (1 - sorted_matches[is_class]).sum(0)
false_negatives = num_true - true_positives
confusion_matrix[class_idx] = np.stack(
[true_positives, false_positives, false_negatives], axis=1
)
return confusion_matrix
@staticmethod
def _compute_precision(confusion_matrix: np.ndarray) -> np.ndarray:
"""
Broadcastable function, computing the precision from the confusion matrix.
Arguments:
confusion_matrix: np.ndarray, shape (N, ..., 3), where the last dimension
contains the true positives, false positives, and false negatives.
Returns:
np.ndarray, shape (N, ...), containing the precision for each element.
"""
if not confusion_matrix.shape[-1] == 3:
raise ValueError(
f"Confusion matrix must have shape (..., 3), got "
f"{confusion_matrix.shape}"
)
true_positives = confusion_matrix[..., 0]
false_positives = confusion_matrix[..., 1]
denominator = true_positives + false_positives
precision = np.where(denominator == 0, 0, true_positives / denominator)
return precision
def _detections_content(self, detections: Detections) -> np.ndarray:
"""Return boxes, masks or oriented bounding boxes from detections."""
if self._metric_target == MetricTarget.BOXES:
return detections.xyxy
if self._metric_target == MetricTarget.MASKS:
return (
detections.mask
if detections.mask is not None
else self._make_empty_content()
)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
obb = detections.data.get(ORIENTED_BOX_COORDINATES)
if obb is not None and len(obb) > 0:
return np.array(obb, dtype=np.float32)
return self._make_empty_content()
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _make_empty_content(self) -> np.ndarray:
if self._metric_target == MetricTarget.BOXES:
return np.empty((0, 4), dtype=np.float32)
if self._metric_target == MetricTarget.MASKS:
return np.empty((0, 0, 0), dtype=bool)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
return np.empty((0, 4, 2), dtype=np.float32)
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _filter_detections_by_size(
self, detections: Detections, size_category: ObjectSizeCategory
) -> Detections:
"""Return a copy of detections with contents filtered by object size."""
new_detections = deepcopy(detections)
if detections.is_empty() or size_category == ObjectSizeCategory.ANY:
return new_detections
sizes = get_detection_size_category(new_detections, self._metric_target)
size_mask = sizes == size_category.value
new_detections.xyxy = new_detections.xyxy[size_mask]
if new_detections.mask is not None:
new_detections.mask = new_detections.mask[size_mask]
if new_detections.class_id is not None:
new_detections.class_id = new_detections.class_id[size_mask]
if new_detections.confidence is not None:
new_detections.confidence = new_detections.confidence[size_mask]
if new_detections.tracker_id is not None:
new_detections.tracker_id = new_detections.tracker_id[size_mask]
if new_detections.data is not None:
for key, value in new_detections.data.items():
new_detections.data[key] = np.array(value)[size_mask]
return new_detections
def _filter_predictions_and_targets_by_size(
self,
predictions_list: List[Detections],
targets_list: List[Detections],
size_category: ObjectSizeCategory,
) -> Tuple[List[Detections], List[Detections]]:
"""
Filter predictions and targets by object size category.
"""
new_predictions_list = []
new_targets_list = []
for predictions, targets in zip(predictions_list, targets_list):
new_predictions_list.append(
self._filter_detections_by_size(predictions, size_category)
)
new_targets_list.append(
self._filter_detections_by_size(targets, size_category)
)
return new_predictions_list, new_targets_list
@dataclass
class PrecisionResult:
"""
The results of the precision metric calculation.
Defaults to `0` if no detections or targets were provided.
Attributes:
metric_target (MetricTarget): the type of data used for the metric -
boxes, masks or oriented bounding boxes.
averaging_method (AveragingMethod): the averaging method used to compute the
precision. Determines how the precision is aggregated across classes.
precision_at_50 (float): the precision at IoU threshold of `0.5`.
precision_at_75 (float): the precision at IoU threshold of `0.75`.
precision_scores (np.ndarray): the precision scores at each IoU threshold.
Shape: `(num_iou_thresholds,)`
precision_per_class (np.ndarray): the precision scores per class and
IoU threshold. Shape: `(num_target_classes, num_iou_thresholds)`
iou_thresholds (np.ndarray): the IoU thresholds used in the calculations.
matched_classes (np.ndarray): the class IDs of all matched classes.
Corresponds to the rows of `precision_per_class`.
small_objects (Optional[PrecisionResult]): the Precision metric results
for small objects (area < 32²).
medium_objects (Optional[PrecisionResult]): the Precision metric results
for medium objects (32² ≤ area < 96²).
large_objects (Optional[PrecisionResult]): the Precision metric results
for large objects (area ≥ 96²).
"""
metric_target: MetricTarget
averaging_method: AveragingMethod
@property
def precision_at_50(self) -> float:
return self.precision_scores[0]
@property
def precision_at_75(self) -> float:
return self.precision_scores[5]
precision_scores: np.ndarray
precision_per_class: np.ndarray
iou_thresholds: np.ndarray
matched_classes: np.ndarray
small_objects: Optional[PrecisionResult]
medium_objects: Optional[PrecisionResult]
large_objects: Optional[PrecisionResult]
def __str__(self) -> str:
"""
Format as a pretty string.
Example:
```python
print(precision_result)
# PrecisionResult:
# Metric target: MetricTarget.BOXES
# Averaging method: AveragingMethod.WEIGHTED
# P @ 50: 0.8099
# P @ 75: 0.7969
# P @ thresh: [0.80992 0.80905 0.80905 ...]
# IoU thresh: [0.5 0.55 0.6 ...]
# Precision per class:
# 0: [0.64706 0.64706 0.64706 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
```
"""
out_str = (
f"{self.__class__.__name__}:\n"
f"Metric target: {self.metric_target}\n"
f"Averaging method: {self.averaging_method}\n"
f"P @ 50: {self.precision_at_50:.4f}\n"
f"P @ 75: {self.precision_at_75:.4f}\n"
f"P @ thresh: {self.precision_scores}\n"
f"IoU thresh: {self.iou_thresholds}\n"
f"Precision per class:\n"
)
if self.precision_per_class.size == 0:
out_str += " No results\n"
for class_id, precision_of_class in zip(
self.matched_classes, self.precision_per_class
):
out_str += f" {class_id}: {precision_of_class}\n"
indent = " "
if self.small_objects is not None:
indented = indent + str(self.small_objects).replace("\n", f"\n{indent}")
out_str += f"\nSmall objects:\n{indented}"
if self.medium_objects is not None:
indented = indent + str(self.medium_objects).replace("\n", f"\n{indent}")
out_str += f"\nMedium objects:\n{indented}"
if self.large_objects is not None:
indented = indent + str(self.large_objects).replace("\n", f"\n{indent}")
out_str += f"\nLarge objects:\n{indented}"
return out_str
def to_pandas(self) -> "pd.DataFrame":
"""
Convert the result to a pandas DataFrame.
Returns:
(pd.DataFrame): The result as a DataFrame.
"""
ensure_pandas_installed()
import pandas as pd
pandas_data = {
"P@50": self.precision_at_50,
"P@75": self.precision_at_75,
}
if self.small_objects is not None:
small_objects_df = self.small_objects.to_pandas()
for key, value in small_objects_df.items():
pandas_data[f"small_objects_{key}"] = value
if self.medium_objects is not None:
medium_objects_df = self.medium_objects.to_pandas()
for key, value in medium_objects_df.items():
pandas_data[f"medium_objects_{key}"] = value
if self.large_objects is not None:
large_objects_df = self.large_objects.to_pandas()
for key, value in large_objects_df.items():
pandas_data[f"large_objects_{key}"] = value
return pd.DataFrame(pandas_data, index=[0])
def plot(self):
"""
Plot the precision results.
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/precision_plot_example.png\
){ align=center width="800" }
"""
labels = ["Precision@50", "Precision@75"]
values = [self.precision_at_50, self.precision_at_75]
colors = [LEGACY_COLOR_PALETTE[0]] * 2
if self.small_objects is not None:
small_objects = self.small_objects
labels += ["Small: P@50", "Small: P@75"]
values += [small_objects.precision_at_50, small_objects.precision_at_75]
colors += [LEGACY_COLOR_PALETTE[3]] * 2
if self.medium_objects is not None:
medium_objects = self.medium_objects
labels += ["Medium: P@50", "Medium: P@75"]
values += [medium_objects.precision_at_50, medium_objects.precision_at_75]
colors += [LEGACY_COLOR_PALETTE[2]] * 2
if self.large_objects is not None:
large_objects = self.large_objects
labels += ["Large: P@50", "Large: P@75"]
values += [large_objects.precision_at_50, large_objects.precision_at_75]
colors += [LEGACY_COLOR_PALETTE[4]] * 2
plt.rcParams["font.family"] = "monospace"
_, ax = plt.subplots(figsize=(10, 6))
ax.set_ylim(0, 1)
ax.set_ylabel("Value", fontweight="bold")
title = (
f"Precision, by Object Size"
f"\n(target: {self.metric_target.value},"
f" averaging: {self.averaging_method.value})"
)
ax.set_title(title, fontweight="bold")
x_positions = range(len(labels))
bars = ax.bar(x_positions, values, color=colors, align="center")
ax.set_xticks(x_positions)
ax.set_xticklabels(labels, rotation=45, ha="right")
for bar in bars:
y_value = bar.get_height()
ax.text(
bar.get_x() + bar.get_width() / 2,
y_value + 0.02,
f"{y_value:.2f}",
ha="center",
va="bottom",
)
plt.rcParams["font.family"] = "sans-serif"
plt.tight_layout()
plt.show()
supervision/metrics/recall.py
+652
View File
@@ -0,0 +1,652 @@
from __future__ import annotations
from copy import deepcopy
from dataclasses import dataclass
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
import numpy as np
from matplotlib import pyplot as plt
from supervision.config import ORIENTED_BOX_COORDINATES
from supervision.detection.core import Detections
from supervision.detection.utils import (
box_iou_batch,
mask_iou_batch,
oriented_box_iou_batch,
)
from supervision.draw.color import LEGACY_COLOR_PALETTE
from supervision.metrics.core import AveragingMethod, Metric, MetricTarget
from supervision.metrics.utils.object_size import (
ObjectSizeCategory,
get_detection_size_category,
)
from supervision.metrics.utils.utils import ensure_pandas_installed
if TYPE_CHECKING:
import pandas as pd
class Recall(Metric):
"""
Recall is a metric used to evaluate object detection models. It is the ratio of
true positive detections to the total number of ground truth instances. We calculate
it at different IoU thresholds.
In simple terms, Recall is a measure of a model's completeness, calculated as:
`Recall = TP / (TP + FN)`
Here, `TP` is the number of true positives (correct detections), and `FN` is the
number of false negatives (missed detections).
Example:
```python
import supervision as sv
from supervision.metrics import Recall
predictions = sv.Detections(...)
targets = sv.Detections(...)
recall_metric = Recall()
recall_result = recall_metric.update(predictions, targets).compute()
print(recall_result.recall_at_50)
# 0.7615
print(recall_result)
# RecallResult:
# Metric target: MetricTarget.BOXES
# Averaging method: AveragingMethod.WEIGHTED
# R @ 50: 0.7615
# R @ 75: 0.7462
# R @ thresh: [0.76151 0.76011 0.76011 0.75732 ...]
# IoU thresh: [0.5 0.55 0.6 ...]
# Recall per class:
# 0: [0.78571 0.78571 0.78571 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
recall_result.plot()
```
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/recall_plot_example.png\
){ align=center width="800" }
"""
def __init__(
self,
metric_target: MetricTarget = MetricTarget.BOXES,
averaging_method: AveragingMethod = AveragingMethod.WEIGHTED,
):
"""
Initialize the Recall metric.
Args:
metric_target (MetricTarget): The type of detection data to use.
averaging_method (AveragingMethod): The averaging method used to compute the
recall. Determines how the recall is aggregated across classes.
"""
self._metric_target = metric_target
self.averaging_method = averaging_method
self._predictions_list: List[Detections] = []
self._targets_list: List[Detections] = []
def reset(self) -> None:
"""
Reset the metric to its initial state, clearing all stored data.
"""
self._predictions_list = []
self._targets_list = []
def update(
self,
predictions: Union[Detections, List[Detections]],
targets: Union[Detections, List[Detections]],
) -> Recall:
"""
Add new predictions and targets to the metric, but do not compute the result.
Args:
predictions (Union[Detections, List[Detections]]): The predicted detections.
targets (Union[Detections, List[Detections]]): The target detections.
Returns:
(Recall): The updated metric instance.
"""
if not isinstance(predictions, list):
predictions = [predictions]
if not isinstance(targets, list):
targets = [targets]
if len(predictions) != len(targets):
raise ValueError(
f"The number of predictions ({len(predictions)}) and"
f" targets ({len(targets)}) during the update must be the same."
)
self._predictions_list.extend(predictions)
self._targets_list.extend(targets)
return self
def compute(self) -> RecallResult:
"""
Calculate the recall metric based on the stored predictions and ground-truth
data, at different IoU thresholds.
Returns:
(RecallResult): The recall metric result.
"""
result = self._compute(self._predictions_list, self._targets_list)
small_predictions, small_targets = self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.SMALL
)
result.small_objects = self._compute(small_predictions, small_targets)
medium_predictions, medium_targets = (
self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.MEDIUM
)
)
result.medium_objects = self._compute(medium_predictions, medium_targets)
large_predictions, large_targets = self._filter_predictions_and_targets_by_size(
self._predictions_list, self._targets_list, ObjectSizeCategory.LARGE
)
result.large_objects = self._compute(large_predictions, large_targets)
return result
def _compute(
self, predictions_list: List[Detections], targets_list: List[Detections]
) -> RecallResult:
iou_thresholds = np.linspace(0.5, 0.95, 10)
stats = []
for predictions, targets in zip(predictions_list, targets_list):
prediction_contents = self._detections_content(predictions)
target_contents = self._detections_content(targets)
if len(targets) > 0:
if len(predictions) == 0:
stats.append(
(
np.zeros((0, iou_thresholds.size), dtype=bool),
np.zeros((0,), dtype=np.float32),
np.zeros((0,), dtype=int),
targets.class_id,
)
)
else:
if self._metric_target == MetricTarget.BOXES:
iou = box_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.MASKS:
iou = mask_iou_batch(target_contents, prediction_contents)
elif self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
iou = oriented_box_iou_batch(
target_contents, prediction_contents
)
else:
raise ValueError(
"Unsupported metric target for IoU calculation"
)
matches = self._match_detection_batch(
predictions.class_id, targets.class_id, iou, iou_thresholds
)
stats.append(
(
matches,
predictions.confidence,
predictions.class_id,
targets.class_id,
)
)
if not stats:
return RecallResult(
metric_target=self._metric_target,
averaging_method=self.averaging_method,
recall_scores=np.zeros(iou_thresholds.shape[0]),
recall_per_class=np.zeros((0, iou_thresholds.shape[0])),
iou_thresholds=iou_thresholds,
matched_classes=np.array([], dtype=int),
small_objects=None,
medium_objects=None,
large_objects=None,
)
concatenated_stats = [np.concatenate(items, 0) for items in zip(*stats)]
recall_scores, recall_per_class, unique_classes = (
self._compute_recall_for_classes(*concatenated_stats)
)
return RecallResult(
metric_target=self._metric_target,
averaging_method=self.averaging_method,
recall_scores=recall_scores,
recall_per_class=recall_per_class,
iou_thresholds=iou_thresholds,
matched_classes=unique_classes,
small_objects=None,
medium_objects=None,
large_objects=None,
)
def _compute_recall_for_classes(
self,
matches: np.ndarray,
prediction_confidence: np.ndarray,
prediction_class_ids: np.ndarray,
true_class_ids: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
sorted_indices = np.argsort(-prediction_confidence)
matches = matches[sorted_indices]
prediction_class_ids = prediction_class_ids[sorted_indices]
unique_classes, class_counts = np.unique(true_class_ids, return_counts=True)
# Shape: PxTh,P,C,C -> CxThx3
confusion_matrix = self._compute_confusion_matrix(
matches, prediction_class_ids, unique_classes, class_counts
)
# Shape: CxThx3 -> CxTh
recall_per_class = self._compute_recall(confusion_matrix)
# Shape: CxTh -> Th
if self.averaging_method == AveragingMethod.MACRO:
recall_scores = np.mean(recall_per_class, axis=0)
elif self.averaging_method == AveragingMethod.MICRO:
confusion_matrix_merged = confusion_matrix.sum(0)
recall_scores = self._compute_recall(confusion_matrix_merged)
elif self.averaging_method == AveragingMethod.WEIGHTED:
class_counts = class_counts.astype(np.float32)
recall_scores = np.average(recall_per_class, axis=0, weights=class_counts)
return recall_scores, recall_per_class, unique_classes
@staticmethod
def _match_detection_batch(
predictions_classes: np.ndarray,
target_classes: np.ndarray,
iou: np.ndarray,
iou_thresholds: np.ndarray,
) -> np.ndarray:
num_predictions, num_iou_levels = (
predictions_classes.shape[0],
iou_thresholds.shape[0],
)
correct = np.zeros((num_predictions, num_iou_levels), dtype=bool)
correct_class = target_classes[:, None] == predictions_classes
for i, iou_level in enumerate(iou_thresholds):
matched_indices = np.where((iou >= iou_level) & correct_class)
if matched_indices[0].shape[0]:
combined_indices = np.stack(matched_indices, axis=1)
iou_values = iou[matched_indices][:, None]
matches = np.hstack([combined_indices, iou_values])
if matched_indices[0].shape[0] > 1:
matches = matches[matches[:, 2].argsort()[::-1]]
matches = matches[np.unique(matches[:, 1], return_index=True)[1]]
matches = matches[np.unique(matches[:, 0], return_index=True)[1]]
correct[matches[:, 1].astype(int), i] = True
return correct
@staticmethod
def _compute_confusion_matrix(
sorted_matches: np.ndarray,
sorted_prediction_class_ids: np.ndarray,
unique_classes: np.ndarray,
class_counts: np.ndarray,
) -> np.ndarray:
"""
Compute the confusion matrix for each class and IoU threshold.
Assumes the matches and prediction_class_ids are sorted by confidence
in descending order.
Arguments:
sorted_matches: np.ndarray, bool, shape (P, Th), that is True
if the prediction is a true positive at the given IoU threshold.
sorted_prediction_class_ids: np.ndarray, int, shape (P,), containing
the class id for each prediction.
unique_classes: np.ndarray, int, shape (C,), containing the unique
class ids.
class_counts: np.ndarray, int, shape (C,), containing the number
of true instances for each class.
Returns:
np.ndarray, shape (C, Th, 3), containing the true positives, false
positives, and false negatives for each class and IoU threshold.
"""
num_thresholds = sorted_matches.shape[1]
num_classes = unique_classes.shape[0]
confusion_matrix = np.zeros((num_classes, num_thresholds, 3))
for class_idx, class_id in enumerate(unique_classes):
is_class = sorted_prediction_class_ids == class_id
num_true = class_counts[class_idx]
num_predictions = is_class.sum()
if num_predictions == 0:
true_positives = np.zeros(num_thresholds)
false_positives = np.zeros(num_thresholds)
false_negatives = np.full(num_thresholds, num_true)
elif num_true == 0:
true_positives = np.zeros(num_thresholds)
false_positives = np.full(num_thresholds, num_predictions)
false_negatives = np.zeros(num_thresholds)
else:
true_positives = sorted_matches[is_class].sum(0)
false_positives = (1 - sorted_matches[is_class]).sum(0)
false_negatives = num_true - true_positives
confusion_matrix[class_idx] = np.stack(
[true_positives, false_positives, false_negatives], axis=1
)
return confusion_matrix
@staticmethod
def _compute_recall(confusion_matrix: np.ndarray) -> np.ndarray:
"""
Broadcastable function, computing the recall from the confusion matrix.
Arguments:
confusion_matrix: np.ndarray, shape (N, ..., 3), where the last dimension
contains the true positives, false positives, and false negatives.
Returns:
np.ndarray, shape (N, ...), containing the recall for each element.
"""
if not confusion_matrix.shape[-1] == 3:
raise ValueError(
f"Confusion matrix must have shape (..., 3), got "
f"{confusion_matrix.shape}"
)
true_positives = confusion_matrix[..., 0]
false_negatives = confusion_matrix[..., 2]
denominator = true_positives + false_negatives
recall = np.where(denominator == 0, 0, true_positives / denominator)
return recall
def _detections_content(self, detections: Detections) -> np.ndarray:
"""Return boxes, masks or oriented bounding boxes from detections."""
if self._metric_target == MetricTarget.BOXES:
return detections.xyxy
if self._metric_target == MetricTarget.MASKS:
return (
detections.mask
if detections.mask is not None
else self._make_empty_content()
)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
obb = detections.data.get(ORIENTED_BOX_COORDINATES)
if obb is not None and len(obb) > 0:
return np.array(obb, dtype=np.float32)
return self._make_empty_content()
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _make_empty_content(self) -> np.ndarray:
if self._metric_target == MetricTarget.BOXES:
return np.empty((0, 4), dtype=np.float32)
if self._metric_target == MetricTarget.MASKS:
return np.empty((0, 0, 0), dtype=bool)
if self._metric_target == MetricTarget.ORIENTED_BOUNDING_BOXES:
return np.empty((0, 4, 2), dtype=np.float32)
raise ValueError(f"Invalid metric target: {self._metric_target}")
def _filter_detections_by_size(
self, detections: Detections, size_category: ObjectSizeCategory
) -> Detections:
"""Return a copy of detections with contents filtered by object size."""
new_detections = deepcopy(detections)
if detections.is_empty() or size_category == ObjectSizeCategory.ANY:
return new_detections
sizes = get_detection_size_category(new_detections, self._metric_target)
size_mask = sizes == size_category.value
new_detections.xyxy = new_detections.xyxy[size_mask]
if new_detections.mask is not None:
new_detections.mask = new_detections.mask[size_mask]
if new_detections.class_id is not None:
new_detections.class_id = new_detections.class_id[size_mask]
if new_detections.confidence is not None:
new_detections.confidence = new_detections.confidence[size_mask]
if new_detections.tracker_id is not None:
new_detections.tracker_id = new_detections.tracker_id[size_mask]
if new_detections.data is not None:
for key, value in new_detections.data.items():
new_detections.data[key] = np.array(value)[size_mask]
return new_detections
def _filter_predictions_and_targets_by_size(
self,
predictions_list: List[Detections],
targets_list: List[Detections],
size_category: ObjectSizeCategory,
) -> Tuple[List[Detections], List[Detections]]:
"""
Filter predictions and targets by object size category.
"""
new_predictions_list = []
new_targets_list = []
for predictions, targets in zip(predictions_list, targets_list):
new_predictions_list.append(
self._filter_detections_by_size(predictions, size_category)
)
new_targets_list.append(
self._filter_detections_by_size(targets, size_category)
)
return new_predictions_list, new_targets_list
@dataclass
class RecallResult:
"""
The results of the recall metric calculation.
Defaults to `0` if no detections or targets were provided.
Attributes:
metric_target (MetricTarget): the type of data used for the metric -
boxes, masks or oriented bounding boxes.
averaging_method (AveragingMethod): the averaging method used to compute the
recall. Determines how the recall is aggregated across classes.
recall_at_50 (float): the recall at IoU threshold of `0.5`.
recall_at_75 (float): the recall at IoU threshold of `0.75`.
recall_scores (np.ndarray): the recall scores at each IoU threshold.
Shape: `(num_iou_thresholds,)`
recall_per_class (np.ndarray): the recall scores per class and IoU threshold.
Shape: `(num_target_classes, num_iou_thresholds)`
iou_thresholds (np.ndarray): the IoU thresholds used in the calculations.
matched_classes (np.ndarray): the class IDs of all matched classes.
Corresponds to the rows of `recall_per_class`.
small_objects (Optional[RecallResult]): the Recall metric results
for small objects (area < 32²).
medium_objects (Optional[RecallResult]): the Recall metric results
for medium objects (32² ≤ area < 96²).
large_objects (Optional[RecallResult]): the Recall metric results
for large objects (area ≥ 96²).
"""
metric_target: MetricTarget
averaging_method: AveragingMethod
@property
def recall_at_50(self) -> float:
return self.recall_scores[0]
@property
def recall_at_75(self) -> float:
return self.recall_scores[5]
recall_scores: np.ndarray
recall_per_class: np.ndarray
iou_thresholds: np.ndarray
matched_classes: np.ndarray
small_objects: Optional[RecallResult]
medium_objects: Optional[RecallResult]
large_objects: Optional[RecallResult]
def __str__(self) -> str:
"""
Format as a pretty string.
Example:
```python
print(recall_result)
# RecallResult:
# Metric target: MetricTarget.BOXES
# Averaging method: AveragingMethod.WEIGHTED
# R @ 50: 0.7615
# R @ 75: 0.7462
# R @ thresh: [0.76151 0.76011 0.76011 0.75732 ...]
# IoU thresh: [0.5 0.55 0.6 ...]
# Recall per class:
# 0: [0.78571 0.78571 0.78571 ...]
# ...
# Small objects: ...
# Medium objects: ...
# Large objects: ...
```
"""
out_str = (
f"{self.__class__.__name__}:\n"
f"Metric target: {self.metric_target}\n"
f"Averaging method: {self.averaging_method}\n"
f"R @ 50: {self.recall_at_50:.4f}\n"
f"R @ 75: {self.recall_at_75:.4f}\n"
f"R @ thresh: {self.recall_scores}\n"
f"IoU thresh: {self.iou_thresholds}\n"
f"Recall per class:\n"
)
if self.recall_per_class.size == 0:
out_str += " No results\n"
for class_id, recall_of_class in zip(
self.matched_classes, self.recall_per_class
):
out_str += f" {class_id}: {recall_of_class}\n"
indent = " "
if self.small_objects is not None:
indented = indent + str(self.small_objects).replace("\n", f"\n{indent}")
out_str += f"\nSmall objects:\n{indented}"
if self.medium_objects is not None:
indented = indent + str(self.medium_objects).replace("\n", f"\n{indent}")
out_str += f"\nMedium objects:\n{indented}"
if self.large_objects is not None:
indented = indent + str(self.large_objects).replace("\n", f"\n{indent}")
out_str += f"\nLarge objects:\n{indented}"
return out_str
def to_pandas(self) -> "pd.DataFrame":
"""
Convert the result to a pandas DataFrame.
Returns:
(pd.DataFrame): The result as a DataFrame.
"""
ensure_pandas_installed()
import pandas as pd
pandas_data = {
"R@50": self.recall_at_50,
"R@75": self.recall_at_75,
}
if self.small_objects is not None:
small_objects_df = self.small_objects.to_pandas()
for key, value in small_objects_df.items():
pandas_data[f"small_objects_{key}"] = value
if self.medium_objects is not None:
medium_objects_df = self.medium_objects.to_pandas()
for key, value in medium_objects_df.items():
pandas_data[f"medium_objects_{key}"] = value
if self.large_objects is not None:
large_objects_df = self.large_objects.to_pandas()
for key, value in large_objects_df.items():
pandas_data[f"large_objects_{key}"] = value
return pd.DataFrame(pandas_data, index=[0])
def plot(self):
"""
Plot the recall results.
![example_plot](\
https://media.roboflow.com/supervision-docs/metrics/recall_plot_example.png\
){ align=center width="800" }
"""
labels = ["Recall@50", "Recall@75"]
values = [self.recall_at_50, self.recall_at_75]
colors = [LEGACY_COLOR_PALETTE[0]] * 2
if self.small_objects is not None:
small_objects = self.small_objects
labels += ["Small: R@50", "Small: R@75"]
values += [small_objects.recall_at_50, small_objects.recall_at_75]
colors += [LEGACY_COLOR_PALETTE[3]] * 2
if self.medium_objects is not None:
medium_objects = self.medium_objects
labels += ["Medium: R@50", "Medium: R@75"]
values += [medium_objects.recall_at_50, medium_objects.recall_at_75]
colors += [LEGACY_COLOR_PALETTE[2]] * 2
if self.large_objects is not None:
large_objects = self.large_objects
labels += ["Large: R@50", "Large: R@75"]
values += [large_objects.recall_at_50, large_objects.recall_at_75]
colors += [LEGACY_COLOR_PALETTE[4]] * 2
plt.rcParams["font.family"] = "monospace"
_, ax = plt.subplots(figsize=(10, 6))
ax.set_ylim(0, 1)
ax.set_ylabel("Value", fontweight="bold")
title = (
f"Recall, by Object Size"
f"\n(target: {self.metric_target.value},"
f" averaging: {self.averaging_method.value})"
)
ax.set_title(title, fontweight="bold")
x_positions = range(len(labels))
bars = ax.bar(x_positions, values, color=colors, align="center")
ax.set_xticks(x_positions)
ax.set_xticklabels(labels, rotation=45, ha="right")
for bar in bars:
y_value = bar.get_height()
ax.text(
bar.get_x() + bar.get_width() / 2,
y_value + 0.02,
f"{y_value:.2f}",
ha="center",
va="bottom",
)
plt.rcParams["font.family"] = "sans-serif"
plt.tight_layout()
plt.show()
supervision/py.typed
View File
Whitespace-only changes.
supervision/tracker/byte_tracker/basetrack.py
-63
View File
@@ -1,63 +0,0 @@
from collections import OrderedDict
from enum import Enum
import numpy as np
class TrackState(Enum):
New = 0
Tracked = 1
Lost = 2
Removed = 3
class BaseTrack:
_count = 0
def __init__(self):
self.track_id = 0
self.is_activated = False
self.state = TrackState.New
self.history = OrderedDict()
self.features = []
self.curr_feature = None
self.score = 0
self.start_frame = 0
self.frame_id = 0
self.time_since_update = 0
# multi-camera
self.location = (np.inf, np.inf)
@property
def end_frame(self) -> int:
return self.frame_id
@staticmethod
def next_id() -> int:
BaseTrack._count += 1
return BaseTrack._count
@staticmethod
def reset_counter():
BaseTrack._count = 0
BaseTrack.track_id = 0
BaseTrack.start_frame = 0
BaseTrack.frame_id = 0
BaseTrack.time_since_update = 0
def activate(self, *args):
raise NotImplementedError
def predict(self):
raise NotImplementedError
def update(self, *args, **kwargs):
raise NotImplementedError
def mark_lost(self):
self.state = TrackState.Lost
def mark_removed(self):
self.state = TrackState.Removed
supervision/tracker/byte_tracker/core.py
+49 -204
View File
@@ -5,186 +5,9 @@ import numpy as np
from supervision.detection.core import Detections
from supervision.detection.utils import box_iou_batch
from supervision.tracker.byte_tracker import matching
from supervision.tracker.byte_tracker.basetrack import BaseTrack, TrackState
from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter
class STrack(BaseTrack):
shared_kalman = KalmanFilter()
_external_count = 0
def __init__(self, tlwh, score, class_ids, minimum_consecutive_frames):
# wait activate
self._tlwh = np.asarray(tlwh, dtype=np.float32)
self.kalman_filter = None
self.mean, self.covariance = None, None
self.is_activated = False
self.score = score
self.class_ids = class_ids
self.tracklet_len = 0
self.external_track_id = -1
self.minimum_consecutive_frames = minimum_consecutive_frames
def predict(self):
mean_state = self.mean.copy()
if self.state != TrackState.Tracked:
mean_state[7] = 0
self.mean, self.covariance = self.kalman_filter.predict(
mean_state, self.covariance
)
@staticmethod
def multi_predict(stracks):
if len(stracks) > 0:
multi_mean = []
multi_covariance = []
for i, st in enumerate(stracks):
multi_mean.append(st.mean.copy())
multi_covariance.append(st.covariance)
if st.state != TrackState.Tracked:
multi_mean[i][7] = 0
multi_mean, multi_covariance = STrack.shared_kalman.multi_predict(
np.asarray(multi_mean), np.asarray(multi_covariance)
)
for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
stracks[i].mean = mean
stracks[i].covariance = cov
def activate(self, kalman_filter, frame_id):
"""Start a new tracklet"""
self.kalman_filter = kalman_filter
self.internal_track_id = self.next_id()
self.mean, self.covariance = self.kalman_filter.initiate(
self.tlwh_to_xyah(self._tlwh)
)
self.tracklet_len = 0
self.state = TrackState.Tracked
if frame_id == 1:
self.is_activated = True
if self.minimum_consecutive_frames == 1:
self.external_track_id = self.next_external_id()
self.frame_id = frame_id
self.start_frame = frame_id
def re_activate(self, new_track, frame_id, new_id=False):
self.mean, self.covariance = self.kalman_filter.update(
self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
)
self.tracklet_len = 0
self.state = TrackState.Tracked
self.frame_id = frame_id
if new_id:
self.internal_track_id = self.next_id()
self.score = new_track.score
def update(self, new_track, frame_id):
"""
Update a matched track
:type new_track: STrack
:type frame_id: int
:type update_feature: bool
:return:
"""
self.frame_id = frame_id
self.tracklet_len += 1
new_tlwh = new_track.tlwh
self.mean, self.covariance = self.kalman_filter.update(
self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh)
)
self.state = TrackState.Tracked
if self.tracklet_len == self.minimum_consecutive_frames:
self.is_activated = True
if self.external_track_id == -1:
self.external_track_id = self.next_external_id()
self.score = new_track.score
@property
def tlwh(self):
"""Get current position in bounding box format `(top left x, top left y,
width, height)`.
"""
if self.mean is None:
return self._tlwh.copy()
ret = self.mean[:4].copy()
ret[2] *= ret[3]
ret[:2] -= ret[2:] / 2
return ret
@property
def tlbr(self):
"""Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
`(top left, bottom right)`.
"""
ret = self.tlwh.copy()
ret[2:] += ret[:2]
return ret
@staticmethod
def tlwh_to_xyah(tlwh):
"""Convert bounding box to format `(center x, center y, aspect ratio,
height)`, where the aspect ratio is `width / height`.
"""
ret = np.asarray(tlwh).copy()
ret[:2] += ret[2:] / 2
ret[2] /= ret[3]
return ret
def to_xyah(self):
return self.tlwh_to_xyah(self.tlwh)
@staticmethod
def next_external_id():
STrack._external_count += 1
return STrack._external_count
@staticmethod
def reset_external_counter():
STrack._external_count = 0
@staticmethod
def tlbr_to_tlwh(tlbr):
ret = np.asarray(tlbr).copy()
ret[2:] -= ret[:2]
return ret
@staticmethod
def tlwh_to_tlbr(tlwh):
ret = np.asarray(tlwh).copy()
ret[2:] += ret[:2]
return ret
def __repr__(self):
return "OT_{}_({}-{})".format(
self.internal_track_id, self.start_frame, self.end_frame
)
def detections2boxes(detections: Detections) -> np.ndarray:
"""
Convert Supervision Detections to numpy tensors for further computation.
Args:
detections (Detections): Detections/Targets in the format of sv.Detections.
Returns:
(np.ndarray): Detections as numpy tensors as in
`(x_min, y_min, x_max, y_max, confidence, class_id)` order.
"""
return np.hstack(
(
detections.xyxy,
detections.confidence[:, np.newaxis],
detections.class_id[:, np.newaxis],
)
)
from supervision.tracker.byte_tracker.single_object_track import STrack, TrackState
from supervision.tracker.byte_tracker.utils import IdCounter
class ByteTrack:
@@ -230,11 +53,17 @@ class ByteTrack:
self.max_time_lost = int(frame_rate / 30.0 * lost_track_buffer)
self.minimum_consecutive_frames = minimum_consecutive_frames
self.kalman_filter = KalmanFilter()
self.shared_kalman = KalmanFilter()
self.tracked_tracks: List[STrack] = []
self.lost_tracks: List[STrack] = []
self.removed_tracks: List[STrack] = []
# Warning, possible bug: If you also set internal_id to start at 1,
# all traces will be connected across objects.
self.internal_id_counter = IdCounter()
self.external_id_counter = IdCounter(start_id=1)
def update_with_detections(self, detections: Detections) -> Detections:
"""
Updates the tracker with the provided detections and returns the updated
@@ -274,8 +103,12 @@ class ByteTrack:
)
```
"""
tensors = detections2boxes(detections=detections)
tensors = np.hstack(
(
detections.xyxy,
detections.confidence[:, np.newaxis],
)
)
tracks = self.update_with_tensors(tensors=tensors)
if len(tracks) > 0:
@@ -301,7 +134,7 @@ class ByteTrack:
return detections
def reset(self):
def reset(self) -> None:
"""
Resets the internal state of the ByteTrack tracker.
@@ -311,11 +144,11 @@ class ByteTrack:
ensuring the tracker starts with a clean state for each new video.
"""
self.frame_id = 0
self.tracked_tracks: List[STrack] = []
self.lost_tracks: List[STrack] = []
self.removed_tracks: List[STrack] = []
BaseTrack.reset_counter()
STrack.reset_external_counter()
self.internal_id_counter.reset()
self.external_id_counter.reset()
self.tracked_tracks = []
self.lost_tracks = []
self.removed_tracks = []
def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
"""
@@ -333,7 +166,6 @@ class ByteTrack:
lost_stracks = []
removed_stracks = []
class_ids = tensors[:, 5]
scores = tensors[:, 4]
bboxes = tensors[:, :4]
@@ -347,14 +179,18 @@ class ByteTrack:
scores_keep = scores[remain_inds]
scores_second = scores[inds_second]
class_ids_keep = class_ids[remain_inds]
class_ids_second = class_ids[inds_second]
if len(dets) > 0:
"""Detections"""
detections = [
STrack(STrack.tlbr_to_tlwh(tlbr), s, c, self.minimum_consecutive_frames)
for (tlbr, s, c) in zip(dets, scores_keep, class_ids_keep)
STrack(
STrack.tlbr_to_tlwh(tlbr),
score_keep,
self.minimum_consecutive_frames,
self.shared_kalman,
self.internal_id_counter,
self.external_id_counter,
)
for (tlbr, score_keep) in zip(dets, scores_keep)
]
else:
detections = []
@@ -372,7 +208,7 @@ class ByteTrack:
""" Step 2: First association, with high score detection boxes"""
strack_pool = joint_tracks(tracked_stracks, self.lost_tracks)
# Predict the current location with KF
STrack.multi_predict(strack_pool)
STrack.multi_predict(strack_pool, self.shared_kalman)
dists = matching.iou_distance(strack_pool, detections)
dists = matching.fuse_score(dists, detections)
@@ -387,7 +223,7 @@ class ByteTrack:
track.update(detections[idet], self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
track.re_activate(det, self.frame_id)
refind_stracks.append(track)
""" Step 3: Second association, with low score detection boxes"""
@@ -395,8 +231,15 @@ class ByteTrack:
if len(dets_second) > 0:
"""Detections"""
detections_second = [
STrack(STrack.tlbr_to_tlwh(tlbr), s, c, self.minimum_consecutive_frames)
for (tlbr, s, c) in zip(dets_second, scores_second, class_ids_second)
STrack(
STrack.tlbr_to_tlwh(tlbr),
score_second,
self.minimum_consecutive_frames,
self.shared_kalman,
self.internal_id_counter,
self.external_id_counter,
)
for (tlbr, score_second) in zip(dets_second, scores_second)
]
else:
detections_second = []
@@ -416,13 +259,13 @@ class ByteTrack:
track.update(det, self.frame_id)
activated_starcks.append(track)
else:
track.re_activate(det, self.frame_id, new_id=False)
track.re_activate(det, self.frame_id)
refind_stracks.append(track)
for it in u_track:
track = r_tracked_stracks[it]
if not track.state == TrackState.Lost:
track.mark_lost()
track.state = TrackState.Lost
lost_stracks.append(track)
"""Deal with unconfirmed tracks, usually tracks with only one beginning frame"""
@@ -438,7 +281,7 @@ class ByteTrack:
activated_starcks.append(unconfirmed[itracked])
for it in u_unconfirmed:
track = unconfirmed[it]
track.mark_removed()
track.state = TrackState.Removed
removed_stracks.append(track)
""" Step 4: Init new stracks"""
@@ -450,8 +293,8 @@ class ByteTrack:
activated_starcks.append(track)
""" Step 5: Update state"""
for track in self.lost_tracks:
if self.frame_id - track.end_frame > self.max_time_lost:
track.mark_removed()
if self.frame_id - track.frame_id > self.max_time_lost:
track.state = TrackState.Removed
removed_stracks.append(track)
self.tracked_tracks = [
@@ -497,7 +340,7 @@ def joint_tracks(
return result
def sub_tracks(track_list_a: List, track_list_b: List) -> List[int]:
def sub_tracks(track_list_a: List[STrack], track_list_b: List[STrack]) -> List[int]:
"""
Returns a list of tracks from track_list_a after removing any tracks
that share the same internal_track_id with tracks in track_list_b.
@@ -518,7 +361,9 @@ def sub_tracks(track_list_a: List, track_list_b: List) -> List[int]:
return list(tracks.values())
def remove_duplicate_tracks(tracks_a: List, tracks_b: List) -> Tuple[List, List]:
def remove_duplicate_tracks(
tracks_a: List[STrack], tracks_b: List[STrack]
) -> Tuple[List[STrack], List[STrack]]:
pairwise_distance = matching.iou_distance(tracks_a, tracks_b)
matching_pairs = np.where(pairwise_distance < 0.15)
supervision/tracker/byte_tracker/matching.py
+10 -5
View File
@@ -1,10 +1,15 @@
from typing import List, Tuple
from __future__ import annotations
from typing import TYPE_CHECKING, List, Tuple
import numpy as np
from scipy.optimize import linear_sum_assignment
from supervision.detection.utils import box_iou_batch
if TYPE_CHECKING:
from supervision.tracker.byte_tracker.core import STrack
def indices_to_matches(
cost_matrix: np.ndarray, indices: np.ndarray, thresh: float
@@ -20,7 +25,7 @@ def indices_to_matches(
def linear_assignment(
cost_matrix: np.ndarray, thresh: float
) -> [np.ndarray, Tuple[int], Tuple[int, int]]:
) -> Tuple[np.ndarray, Tuple[int], Tuple[int, int]]:
if cost_matrix.size == 0:
return (
np.empty((0, 2), dtype=int),
@@ -35,7 +40,7 @@ def linear_assignment(
return indices_to_matches(cost_matrix, indices, thresh)
def iou_distance(atracks: List, btracks: List) -> np.ndarray:
def iou_distance(atracks: List[STrack], btracks: List[STrack]) -> np.ndarray:
if (len(atracks) > 0 and isinstance(atracks[0], np.ndarray)) or (
len(btracks) > 0 and isinstance(btracks[0], np.ndarray)
):
@@ -53,11 +58,11 @@ def iou_distance(atracks: List, btracks: List) -> np.ndarray:
return cost_matrix
def fuse_score(cost_matrix: np.ndarray, detections: List) -> np.ndarray:
def fuse_score(cost_matrix: np.ndarray, stracks: List[STrack]) -> np.ndarray:
if cost_matrix.size == 0:
return cost_matrix
iou_sim = 1 - cost_matrix
det_scores = np.array([det.score for det in detections])
det_scores = np.array([strack.score for strack in stracks])
det_scores = np.expand_dims(det_scores, axis=0).repeat(cost_matrix.shape[0], axis=0)
fuse_sim = iou_sim * det_scores
fuse_cost = 1 - fuse_sim
supervision/tracker/byte_tracker/single_object_track.py
+178
View File
@@ -0,0 +1,178 @@
from __future__ import annotations
from enum import Enum
from typing import List
import numpy as np
import numpy.typing as npt
from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter
from supervision.tracker.byte_tracker.utils import IdCounter
class TrackState(Enum):
New = 0
Tracked = 1
Lost = 2
Removed = 3
class STrack:
def __init__(
self,
tlwh: npt.NDArray[np.float32],
score: npt.NDArray[np.float32],
minimum_consecutive_frames: int,
shared_kalman: KalmanFilter,
internal_id_counter: IdCounter,
external_id_counter: IdCounter,
):
self.state = TrackState.New
self.is_activated = False
self.start_frame = 0
self.frame_id = 0
self._tlwh = np.asarray(tlwh, dtype=np.float32)
self.kalman_filter = None
self.shared_kalman = shared_kalman
self.mean, self.covariance = None, None
self.is_activated = False
self.score = score
self.tracklet_len = 0
self.minimum_consecutive_frames = minimum_consecutive_frames
self.internal_id_counter = internal_id_counter
self.external_id_counter = external_id_counter
self.internal_track_id = self.internal_id_counter.NO_ID
self.external_track_id = self.external_id_counter.NO_ID
def predict(self) -> None:
mean_state = self.mean.copy()
if self.state != TrackState.Tracked:
mean_state[7] = 0
self.mean, self.covariance = self.kalman_filter.predict(
mean_state, self.covariance
)
@staticmethod
def multi_predict(stracks: List[STrack], shared_kalman: KalmanFilter) -> None:
if len(stracks) > 0:
multi_mean = []
multi_covariance = []
for i, st in enumerate(stracks):
multi_mean.append(st.mean.copy())
multi_covariance.append(st.covariance)
if st.state != TrackState.Tracked:
multi_mean[i][7] = 0
multi_mean, multi_covariance = shared_kalman.multi_predict(
np.asarray(multi_mean), np.asarray(multi_covariance)
)
for i, (mean, cov) in enumerate(zip(multi_mean, multi_covariance)):
stracks[i].mean = mean
stracks[i].covariance = cov
def activate(self, kalman_filter: KalmanFilter, frame_id: int) -> None:
"""Start a new tracklet"""
self.kalman_filter = kalman_filter
self.internal_track_id = self.internal_id_counter.new_id()
self.mean, self.covariance = self.kalman_filter.initiate(
self.tlwh_to_xyah(self._tlwh)
)
self.tracklet_len = 0
self.state = TrackState.Tracked
if frame_id == 1:
self.is_activated = True
if self.minimum_consecutive_frames == 1:
self.external_track_id = self.external_id_counter.new_id()
self.frame_id = frame_id
self.start_frame = frame_id
def re_activate(self, new_track: STrack, frame_id: int) -> None:
self.mean, self.covariance = self.kalman_filter.update(
self.mean, self.covariance, self.tlwh_to_xyah(new_track.tlwh)
)
self.tracklet_len = 0
self.state = TrackState.Tracked
self.frame_id = frame_id
self.score = new_track.score
def update(self, new_track: STrack, frame_id: int) -> None:
"""
Update a matched track
:type new_track: STrack
:type frame_id: int
:type update_feature: bool
:return:
"""
self.frame_id = frame_id
self.tracklet_len += 1
new_tlwh = new_track.tlwh
self.mean, self.covariance = self.kalman_filter.update(
self.mean, self.covariance, self.tlwh_to_xyah(new_tlwh)
)
self.state = TrackState.Tracked
if self.tracklet_len == self.minimum_consecutive_frames:
self.is_activated = True
if self.external_track_id == self.external_id_counter.NO_ID:
self.external_track_id = self.external_id_counter.new_id()
self.score = new_track.score
@property
def tlwh(self) -> npt.NDArray[np.float32]:
"""Get current position in bounding box format `(top left x, top left y,
width, height)`.
"""
if self.mean is None:
return self._tlwh.copy()
ret = self.mean[:4].copy()
ret[2] *= ret[3]
ret[:2] -= ret[2:] / 2
return ret
@property
def tlbr(self) -> npt.NDArray[np.float32]:
"""Convert bounding box to format `(min x, min y, max x, max y)`, i.e.,
`(top left, bottom right)`.
"""
ret = self.tlwh.copy()
ret[2:] += ret[:2]
return ret
@staticmethod
def tlwh_to_xyah(tlwh) -> npt.NDArray[np.float32]:
"""Convert bounding box to format `(center x, center y, aspect ratio,
height)`, where the aspect ratio is `width / height`.
"""
ret = np.asarray(tlwh).copy()
ret[:2] += ret[2:] / 2
ret[2] /= ret[3]
return ret
def to_xyah(self) -> npt.NDArray[np.float32]:
return self.tlwh_to_xyah(self.tlwh)
@staticmethod
def tlbr_to_tlwh(tlbr) -> npt.NDArray[np.float32]:
ret = np.asarray(tlbr).copy()
ret[2:] -= ret[:2]
return ret
@staticmethod
def tlwh_to_tlbr(tlwh) -> npt.NDArray[np.float32]:
ret = np.asarray(tlwh).copy()
ret[2:] += ret[:2]
return ret
def __repr__(self) -> str:
return "OT_{}_({}-{})".format(
self.internal_track_id, self.start_frame, self.frame_id
)
supervision/tracker/byte_tracker/utils.py
+18
View File
@@ -0,0 +1,18 @@
class IdCounter:
def __init__(self, start_id: int = 0):
self.start_id = start_id
if self.start_id <= self.NO_ID:
raise ValueError(f"start_id must be greater than {self.NO_ID}")
self.reset()
def reset(self) -> None:
self._id = self.start_id
def new_id(self) -> int:
returned_id = self._id
self._id += 1
return returned_id
@property
def NO_ID(self) -> int:
return -1
test/detection/test_core.py
+205 -2
View File
@@ -106,6 +106,26 @@ TEST_DET_DIFFERENT_DATA = Detections(
"never_seen_key": [9],
},
)
TEST_DET_WITH_METADATA = Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
)
TEST_DET_WITH_METADATA_2 = Detections(
xyxy=np.array([[30, 30, 40, 40]]),
class_id=np.array([2]),
metadata={"source": "camera1"},
)
TEST_DET_NO_METADATA = Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
)
TEST_DET_DIFFERENT_METADATA = Detections(
xyxy=np.array([[50, 50, 60, 60]]),
class_id=np.array([3]),
metadata={"source": "camera2"},
)
@pytest.mark.parametrize(
@@ -258,6 +278,11 @@ def test_getitem(
TEST_DET_1,
DoesNotRaise(),
), # Single detection and empty-array fields
(
[TEST_DET_ZERO_LENGTH, TEST_DET_ZERO_LENGTH],
TEST_DET_ZERO_LENGTH,
DoesNotRaise(),
), # Zero-length fields across all Detections
(
[
TEST_DET_1,
@@ -287,12 +312,190 @@ def test_getitem(
Detections.empty(),
],
mock_detections(
xyxy=[[10, 10, 20, 20]],
xyxy=np.array([[10, 10, 20, 20]]),
class_id=[1],
mask=[np.zeros((4, 4), dtype=bool)],
),
DoesNotRaise(),
), # Segmentation + Empty
# Metadata
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
),
Detections.empty(),
],
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
),
DoesNotRaise(),
), # Metadata merge with empty detections
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
),
Detections(xyxy=np.array([[30, 30, 40, 40]]), class_id=np.array([2])),
],
None,
pytest.raises(ValueError),
), # Empty and non-empty metadata
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
)
],
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
),
DoesNotRaise(),
), # Single detection with metadata
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
),
Detections(
xyxy=np.array([[30, 30, 40, 40]]),
class_id=np.array([2]),
metadata={"source": "camera1"},
),
],
Detections(
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]]),
class_id=np.array([1, 2]),
metadata={"source": "camera1"},
),
DoesNotRaise(),
), # Multiple metadata entries with identical values
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
class_id=np.array([1]),
metadata={"source": "camera1"},
),
Detections(
xyxy=np.array([[50, 50, 60, 60]]),
class_id=np.array([3]),
metadata={"source": "camera2"},
),
],
None,
pytest.raises(ValueError),
), # Different metadata values
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
metadata={"source": "camera1", "resolution": "1080p"},
),
Detections(
xyxy=np.array([[30, 30, 40, 40]]),
metadata={"source": "camera1", "resolution": "1080p"},
),
],
Detections(
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]]),
metadata={"source": "camera1", "resolution": "1080p"},
),
DoesNotRaise(),
), # Large metadata with multiple identical entries
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]), metadata={"source": "camera1"}
),
Detections(
xyxy=np.array([[30, 30, 40, 40]]), metadata={"source": ["camera1"]}
),
],
None,
pytest.raises(ValueError),
), # Inconsistent types in metadata values
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]), metadata={"source": "camera1"}
),
Detections(
xyxy=np.array([[30, 30, 40, 40]]), metadata={"location": "indoor"}
),
],
None,
pytest.raises(ValueError),
), # Metadata key mismatch
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
metadata={
"source": "camera1",
"settings": {"resolution": "1080p", "fps": 30},
},
),
Detections(
xyxy=np.array([[30, 30, 40, 40]]),
metadata={
"source": "camera1",
"settings": {"resolution": "1080p", "fps": 30},
},
),
],
Detections(
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]]),
metadata={
"source": "camera1",
"settings": {"resolution": "1080p", "fps": 30},
},
),
DoesNotRaise(),
), # multi-field metadata
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
metadata={"calibration_matrix": np.array([[1, 0], [0, 1]])},
),
Detections(
xyxy=np.array([[30, 30, 40, 40]]),
metadata={"calibration_matrix": np.array([[1, 0], [0, 1]])},
),
],
Detections(
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]]),
metadata={"calibration_matrix": np.array([[1, 0], [0, 1]])},
),
DoesNotRaise(),
), # Identical 2D numpy arrays in metadata
(
[
Detections(
xyxy=np.array([[10, 10, 20, 20]]),
metadata={"calibration_matrix": np.array([[1, 0], [0, 1]])},
),
Detections(
xyxy=np.array([[30, 30, 40, 40]]),
metadata={"calibration_matrix": np.array([[2, 0], [0, 2]])},
),
],
None,
pytest.raises(ValueError),
), # Mismatching 2D numpy arrays in metadata
],
)
def test_merge(
@@ -302,7 +505,7 @@ def test_merge(
) -> None:
with exception:
result = Detections.merge(detections_list=detections_list)
assert result == expected_result
assert result == expected_result, f"Expected: {expected_result}, Got: {result}"
@pytest.mark.parametrize(
test/detection/test_line_counter.py
+300 -1
View File
@@ -69,7 +69,7 @@ def test_calculate_region_of_interest_limits(
exception: Exception,
) -> None:
with exception:
result = LineZone.calculate_region_of_interest_limits(vector=vector)
result = LineZone._calculate_region_of_interest_limits(vector=vector)
assert result == expected_result
@@ -303,6 +303,19 @@ def test_line_zone_one_detection_default_anchors(
[False, True, False, True],
[False, False, True, False],
),
( # Scrape line, left side, center anchor (along line point)
Vector(Point(0, 0), Point(10, 0)),
[
[-2, 4, 2, 6],
[-2, 4 - 10, 2, 6 - 10],
[-2, 4, 2, 6],
[-2, 4 - 10, 2, 6 - 10],
[-2, 4 - 10, 2, 6 - 10],
],
[Position.CENTER],
[False, True, False, True, False],
[False, False, True, False, False],
),
( # Scrape line, right side, corner anchors
Vector(Point(0, 0), Point(10, 0)),
[
@@ -477,3 +490,289 @@ def test_line_zone_multiple_detections(
assert crossed_in_list == expected_crossed_in
assert crossed_out_list == expected_crossed_out
@pytest.mark.parametrize(
"vector, xyxy_sequence, triggering_anchors, minimum_crossing_threshold, "
"expected_crossed_in, expected_crossed_out",
[
( # Detection lingers around line, all crosses counted
Vector(Point(0, 0), Point(10, 0)),
[
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4 - 10, 3, 6],
],
[
Position.TOP_LEFT,
],
1,
[False, True, False, True, False],
[False, False, True, False, False],
),
( # Detection lingers around line, only final cross counted
Vector(Point(0, 0), Point(10, 0)),
[
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4 - 10, 3, 6],
],
[
Position.TOP_LEFT,
],
2,
[False, False, False, False, True],
[False, False, False, False, False],
),
( # Detection lingers around line for a long time
Vector(Point(0, 0), Point(10, 0)),
[
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4 - 10, 3, 6],
],
[
Position.TOP_LEFT,
],
2,
[False] * 12 + [True],
[False] * 13,
),
( # Detection lingers around line, longer cycle
Vector(Point(0, 0), Point(10, 0)),
[
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4, 3, 6],
[2, 4, 3, 6],
[2, 4, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4 - 10, 3, 6],
[2, 4 - 10, 3, 6],
],
[
Position.TOP_LEFT,
],
4,
[False] * 8 + [True],
[False] * 9,
),
],
)
def test_line_zone_one_detection_long_horizon(
vector: Vector,
xyxy_sequence: List[List[float]],
triggering_anchors: List[Position],
minimum_crossing_threshold: int,
expected_crossed_in: List[bool],
expected_crossed_out: List[bool],
) -> None:
line_zone = LineZone(
start=vector.start,
end=vector.end,
triggering_anchors=triggering_anchors,
minimum_crossing_threshold=minimum_crossing_threshold,
)
crossed_in_list = []
crossed_out_list = []
for i, bbox in enumerate(xyxy_sequence):
detections = mock_detections(
xyxy=[bbox],
tracker_id=[0],
)
crossed_in, crossed_out = line_zone.trigger(detections)
crossed_in_list.append(crossed_in[0])
crossed_out_list.append(crossed_out[0])
assert (
crossed_in_list == expected_crossed_in
), f"expected {expected_crossed_in}, got {crossed_in_list}"
assert (
crossed_out_list == expected_crossed_out
), f"expected {expected_crossed_out}, got {crossed_out_list}"
@pytest.mark.parametrize(
"vector, xyxy_sequence, anchors, minimum_crossing_threshold, "
"expected_crossed_in, expected_crossed_out, expected_count_in, "
"expected_count_out, exception",
[
( # One stays, one crosses, one disappears before crossing
Vector(Point(0, 0), Point(10, 0)),
[
[[4, 4, 6, 6], [4, 4, 6, 6], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
],
[
Position.TOP_LEFT,
],
1,
[
[False, False, False],
[False, True, False],
[False, False],
[False, True],
[False, False],
],
[
[False, False, False],
[False, False, False],
[False, True],
[False, False],
[False, False],
],
[0, 1, 1, 2, 2],
[0, 0, 1, 1, 1],
DoesNotRaise(),
),
( # One stays, one crosses, one disappears immediately after crossing
Vector(Point(0, 0), Point(10, 0)),
[
[[4, 4, 6, 6], [4, 4, 6, 6], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
],
[
Position.TOP_LEFT,
],
1,
[
[False, False, False],
[False, True, False],
[False, False, True],
[False, True],
[False, False],
],
[
[False, False, False],
[False, False, False],
[False, True, False],
[False, False],
[False, False],
],
[0, 1, 2, 3, 3],
[0, 0, 1, 1, 1],
DoesNotRaise(),
),
( # One stays, one crosses, one disappears before crossing
Vector(Point(0, 0), Point(10, 0)),
[
[[4, 4, 6, 6], [4, 4, 6, 6], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
],
[
Position.TOP_LEFT,
],
2,
[
[False, False, False],
[False, False, False],
[False, False],
[False, False],
[False, True],
],
[
[False, False, False],
[False, False, False],
[False, False],
[False, False],
[False, False],
],
[0, 0, 0, 0, 1],
[0, 0, 0, 0, 0],
DoesNotRaise(),
),
( # One stays, one crosses, one disappears immediately after crossing
Vector(Point(0, 0), Point(10, 0)),
[
[[4, 4, 6, 6], [4, 4, 6, 6], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10], [4, 4, 6, 6]],
[[4, 4, 6, 6], [4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
[[4, 4, 6, 6], [4, 4 - 10, 6, 6 - 10]],
],
[
Position.TOP_LEFT,
],
2,
[
[False, False, False],
[False, False, False],
[False, False, False],
[False, False],
[False, True],
],
[
[False, False, False],
[False, False, False],
[False, False, False],
[False, False],
[False, False],
],
[0, 0, 0, 0, 1],
[0, 0, 0, 0, 0],
DoesNotRaise(),
),
],
)
def test_line_zone_long_horizon_disappearing_detections(
vector: Vector,
xyxy_sequence: List[List[Optional[List[float]]]],
anchors: List[Position],
minimum_crossing_threshold: int,
expected_crossed_in: List[List[bool]],
expected_crossed_out: List[List[bool]],
expected_count_in: List[int],
expected_count_out: List[int],
exception: Exception,
) -> None:
with exception:
line_zone = LineZone(
start=vector.start,
end=vector.end,
triggering_anchors=anchors,
minimum_crossing_threshold=minimum_crossing_threshold,
)
crossed_in_list = []
crossed_out_list = []
count_in_list = []
count_out_list = []
for bboxes in xyxy_sequence:
detections = mock_detections(
xyxy=bboxes,
tracker_id=[i for i in range(0, len(bboxes))],
)
crossed_in, crossed_out = line_zone.trigger(detections)
crossed_in_list.append(list(crossed_in))
crossed_out_list.append(list(crossed_out))
count_in_list.append(line_zone.in_count)
count_out_list.append(line_zone.out_count)
assert crossed_in_list == expected_crossed_in
assert crossed_out_list == expected_crossed_out
assert count_in_list == expected_count_in
assert count_out_list == expected_count_out
test/detection/test_utils.py
+161
View File
@@ -14,6 +14,7 @@ from supervision.detection.utils import (
filter_polygons_by_area,
get_data_item,
merge_data,
merge_metadata,
move_boxes,
process_roboflow_result,
scale_boxes,
@@ -1138,3 +1139,163 @@ def test_xywh_to_xyxy(xywh: np.ndarray, expected_result: np.ndarray) -> None:
def test_xcycwh_to_xyxy(xcycwh: np.ndarray, expected_result: np.ndarray) -> None:
result = xcycwh_to_xyxy(xcycwh)
np.testing.assert_array_equal(result, expected_result)
@pytest.mark.parametrize(
"metadata_list, expected_result, exception",
[
# Identical metadata with a single key
([{"key1": "value1"}, {"key1": "value1"}], {"key1": "value1"}, DoesNotRaise()),
# Identical metadata with multiple keys
(
[
{"key1": "value1", "key2": "value2"},
{"key1": "value1", "key2": "value2"},
],
{"key1": "value1", "key2": "value2"},
DoesNotRaise(),
),
# Conflicting values for the same key
([{"key1": "value1"}, {"key1": "value2"}], None, pytest.raises(ValueError)),
# Different sets of keys across dictionaries
([{"key1": "value1"}, {"key2": "value2"}], None, pytest.raises(ValueError)),
# Empty metadata list
([], {}, DoesNotRaise()),
# Empty metadata dictionaries
([{}, {}], {}, DoesNotRaise()),
# Different declaration order for keys
(
[
{"key1": "value1", "key2": "value2"},
{"key2": "value2", "key1": "value1"},
],
{"key1": "value1", "key2": "value2"},
DoesNotRaise(),
),
# Nested metadata dictionaries
(
[{"key1": {"sub_key": "sub_value"}}, {"key1": {"sub_key": "sub_value"}}],
{"key1": {"sub_key": "sub_value"}},
DoesNotRaise(),
),
# Large metadata dictionaries with many keys
(
[
{f"key{i}": f"value{i}" for i in range(100)},
{f"key{i}": f"value{i}" for i in range(100)},
],
{f"key{i}": f"value{i}" for i in range(100)},
DoesNotRaise(),
),
# Mixed types in list metadata values
(
[{"key1": ["value1", 2, True]}, {"key1": ["value1", 2, True]}],
{"key1": ["value1", 2, True]},
DoesNotRaise(),
),
# Identical lists across metadata dictionaries
(
[{"key1": [1, 2, 3]}, {"key1": [1, 2, 3]}],
{"key1": [1, 2, 3]},
DoesNotRaise(),
),
# Identical numpy arrays across metadata dictionaries
(
[{"key1": np.array([1, 2, 3])}, {"key1": np.array([1, 2, 3])}],
{"key1": np.array([1, 2, 3])},
DoesNotRaise(),
),
# Identical numpy arrays across metadata dictionaries, different datatype
(
[
{"key1": np.array([1, 2, 3], dtype=np.int32)},
{"key1": np.array([1, 2, 3], dtype=np.int64)},
],
{"key1": np.array([1, 2, 3])},
DoesNotRaise(),
),
# Conflicting lists for the same key
([{"key1": [1, 2, 3]}, {"key1": [4, 5, 6]}], None, pytest.raises(ValueError)),
# Conflicting numpy arrays for the same key
(
[{"key1": np.array([1, 2, 3])}, {"key1": np.array([4, 5, 6])}],
None,
pytest.raises(ValueError),
),
# Mixed data types: list and numpy array for the same key
(
[{"key1": [1, 2, 3]}, {"key1": np.array([1, 2, 3])}],
None,
pytest.raises(ValueError),
),
# Empty lists and numpy arrays for the same key
([{"key1": []}, {"key1": np.array([])}], None, pytest.raises(ValueError)),
# Identical multi-dimensional lists across metadata dictionaries
(
[{"key1": [[1, 2], [3, 4]]}, {"key1": [[1, 2], [3, 4]]}],
{"key1": [[1, 2], [3, 4]]},
DoesNotRaise(),
),
# Identical multi-dimensional numpy arrays across metadata dictionaries
(
[
{"key1": np.arange(4).reshape(2, 2)},
{"key1": np.arange(4).reshape(2, 2)},
],
{"key1": np.arange(4).reshape(2, 2)},
DoesNotRaise(),
),
# Conflicting multi-dimensional lists for the same key
(
[{"key1": [[1, 2], [3, 4]]}, {"key1": [[5, 6], [7, 8]]}],
None,
pytest.raises(ValueError),
),
# Conflicting multi-dimensional numpy arrays for the same key
(
[
{"key1": np.arange(4).reshape(2, 2)},
{"key1": np.arange(4, 8).reshape(2, 2)},
],
None,
pytest.raises(ValueError),
),
# Mixed types with multi-dimensional list and array for the same key
(
[{"key1": [[1, 2], [3, 4]]}, {"key1": np.arange(4).reshape(2, 2)}],
None,
pytest.raises(ValueError),
),
# Identical higher-dimensional (3D) numpy arrays across
# metadata dictionaries
(
[
{"key1": np.arange(8).reshape(2, 2, 2)},
{"key1": np.arange(8).reshape(2, 2, 2)},
],
{"key1": np.arange(8).reshape(2, 2, 2)},
DoesNotRaise(),
),
# Differently-shaped higher-dimensional (3D) numpy arrays
# across metadata dictionaries
(
[
{"key1": np.arange(8).reshape(2, 2, 2)},
{"key1": np.arange(8).reshape(4, 1, 2)},
],
None,
pytest.raises(ValueError),
),
],
)
def test_merge_metadata(metadata_list, expected_result, exception):
with exception:
result = merge_metadata(metadata_list)
if expected_result is None:
assert result is None, f"Expected an error, but got a result {result}"
for key, value in result.items():
assert key in expected_result
if isinstance(value, np.ndarray):
np.testing.assert_array_equal(value, expected_result[key])
else:
assert value == expected_result[key]
test/tracker/__init__.py
View File
Whitespace-only changes.
test/tracker/test_byte_tracker.py
+40
View File
@@ -0,0 +1,40 @@
from typing import List
import numpy as np
import pytest
import supervision as sv
@pytest.mark.parametrize(
"detections, expected_results",
[
(
[
sv.Detections(
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]]),
class_id=np.array([1, 1]),
confidence=np.array([1, 1]),
),
sv.Detections(
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]]),
class_id=np.array([1, 1]),
confidence=np.array([1, 1]),
),
],
sv.Detections(
xyxy=np.array([[10, 10, 20, 20], [30, 30, 40, 40]]),
class_id=np.array([1, 1]),
confidence=np.array([1, 1]),
tracker_id=np.array([1, 2]),
),
),
],
)
def test_byte_tracker(
detections: List[sv.Detections],
expected_results: sv.Detections,
) -> None:
byte_tracker = sv.ByteTrack()
tracked_detections = [byte_tracker.update_with_detections(d) for d in detections]
assert tracked_detections[-1] == expected_results
test/utils/test_internal.py
+21 -2
View File
@@ -121,7 +121,15 @@ class MockDataclass:
(
Detections.empty(),
False,
{"xyxy", "class_id", "confidence", "mask", "tracker_id", "data"},
{
"xyxy",
"class_id",
"confidence",
"mask",
"tracker_id",
"data",
"metadata",
},
DoesNotRaise(),
),
(
@@ -134,6 +142,7 @@ class MockDataclass:
"mask",
"tracker_id",
"data",
"metadata",
"area",
"box_area",
},
@@ -149,6 +158,7 @@ class MockDataclass:
"mask",
"tracker_id",
"data",
"metadata",
},
DoesNotRaise(),
),
@@ -169,13 +179,22 @@ class MockDataclass:
"mask",
"tracker_id",
"data",
"metadata",
},
DoesNotRaise(),
),
(
Detections.empty(),
False,
{"xyxy", "class_id", "confidence", "mask", "tracker_id", "data"},
{
"xyxy",
"class_id",
"confidence",
"mask",
"tracker_id",
"data",
"metadata",
},
DoesNotRaise(),
),
],
tox.ini
+1 -1
View File
@@ -1,5 +1,5 @@
[tox]
envlist = py38,py39,py310,py311,py312
envlist = py38,py39,py310,py311,py312,py313
[testenv]
changedir = test