Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
f34993c202 | ||
|
|
f44a9376ec | ||
|
|
48752b432b | ||
|
|
c1797ffd81 | ||
|
|
a640e527d0 | ||
|
|
d894c7814b | ||
|
|
815c43beb8 | ||
|
|
2fed8ae0c3 | ||
|
|
586ff5b3fb | ||
|
|
8b73c12bcc | ||
|
|
67d1323daa | ||
|
|
0e6d69cf85 | ||
|
|
a5c211ff61 | ||
|
|
5f5b4da845 | ||
|
|
be16b84340 | ||
|
|
08af18ad41 | ||
|
|
caf504c0c9 | ||
|
|
893a19a991 | ||
|
|
b35a6161dd | ||
|
|
aabca4a4b5 | ||
|
|
57ba059f1c | ||
|
|
743f3281a5 | ||
|
|
737cb37249 | ||
|
|
78a4e73f25 | ||
|
|
dc9c1151c5 | ||
|
|
5f82aac9e6 | ||
|
|
d66a539779 | ||
|
|
a037a4a404 | ||
|
|
871968a916 | ||
|
|
5ca30f5898 | ||
|
|
54bb79c220 | ||
|
|
809d8946e6 | ||
|
|
e2305e6f57 | ||
|
|
e173510117 | ||
|
|
e209c7436c | ||
|
|
25a45aa423 | ||
|
|
b9d1d78f31 | ||
|
|
142fe02959 | ||
|
|
6815814aec | ||
|
|
1b6613d4a3 | ||
|
|
fabfe0783a | ||
|
|
2c4e7589b0 | ||
|
|
0abe36e079 | ||
|
|
61809f9329 | ||
|
|
69f3d10e42 | ||
|
|
62c45548a6 | ||
|
|
00e209744f | ||
|
|
8dcfb42751 | ||
|
|
a5a853e141 | ||
|
|
034c2bd9a8 | ||
|
|
a7449ba56f | ||
|
|
8f302c0646 | ||
|
|
1ddb78b662 | ||
|
|
db4afe9885 | ||
|
|
72a780b74a | ||
|
|
b2e9a668b2 | ||
|
|
57ef24bcb3 | ||
|
|
eda6909c91 | ||
|
|
9a9bf62fe0 | ||
|
|
1ab50cdb3c | ||
|
|
f78c282b7e | ||
|
|
3d843d710d | ||
|
|
75db51aa90 | ||
|
|
fe8a8260dc | ||
|
|
bf725ef0f7 | ||
|
|
338c589032 | ||
|
|
4baf570b28 | ||
|
|
1324cda600 | ||
|
|
9ded70c8da | ||
|
|
8cda753334 | ||
|
|
a14159c150 | ||
|
|
ae00a9d65f | ||
|
|
c51159b598 | ||
|
|
476a31fa79 | ||
|
|
e1e455f99c | ||
|
|
5b5e0eb88d | ||
|
|
d0139f5293 | ||
|
|
bed7ed45b4 | ||
|
|
35142fc755 | ||
|
|
f041f0ddcc | ||
|
|
253a136f43 | ||
|
|
9cefff8627 | ||
|
|
f2af290166 | ||
|
|
988da89bf4 | ||
|
|
b4a8859616 | ||
|
|
9fff77a6e3 | ||
|
|
e7f95e0761 | ||
|
|
2254e59bdf | ||
|
|
6a0482ded5 | ||
|
|
082f248159 | ||
|
|
227615b766 | ||
|
|
baff75332f | ||
|
|
434a695324 |
@@ -0,0 +1,17 @@
|
||||
contourpy==1.1.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
cycler==0.11.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
fonttools==4.41.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
importlib-resources==6.0.0 ; python_version >= "3.8" and python_version < "3.10"
|
||||
kiwisolver==1.4.4 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
matplotlib==3.7.2 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
numpy==1.24.4 ; python_version >= "3.8" and python_version < "3.12"
|
||||
opencv-python-headless==4.8.0.74 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
packaging==23.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
pillow==10.0.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
pyparsing==3.0.9 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
python-dateutil==2.8.2 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
pyyaml==6.0.1 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
scipy==1.10.1 ; python_version >= "3.8" and python_version < "3.12"
|
||||
six==1.16.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
zipp==3.16.2 ; python_version >= "3.8" and python_version < "3.10"
|
||||
pytest==7.4.0 ; python_version >= "3.8" and python_full_version < "3.12.0"
|
||||
@@ -0,0 +1,33 @@
|
||||
|
||||
name: Clear cache
|
||||
|
||||
on:
|
||||
schedule:
|
||||
- cron: '0 0 1 * *'
|
||||
workflow_dispatch:
|
||||
|
||||
permissions:
|
||||
actions: write
|
||||
|
||||
jobs:
|
||||
clear-cache:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- name: Clear cache
|
||||
uses: actions/github-script@v6
|
||||
with:
|
||||
script: |
|
||||
console.log("About to clear")
|
||||
const caches = await github.rest.actions.getActionsCacheList({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
})
|
||||
for (const cache of caches.data.actions_caches) {
|
||||
console.log(cache)
|
||||
github.rest.actions.deleteActionsCacheById({
|
||||
owner: context.repo.owner,
|
||||
repo: context.repo.repo,
|
||||
cache_id: cache.id,
|
||||
})
|
||||
}
|
||||
console.log("Clear completed")
|
||||
@@ -15,12 +15,18 @@ jobs:
|
||||
uses: actions/checkout@v3
|
||||
- name: 🐍 Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v4
|
||||
# id based on python version
|
||||
id: python-setup
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
- name: 🦾 Install dependencies
|
||||
check-latest: true
|
||||
cache: 'pip'
|
||||
cache-dependency-path: '**/requirements-dev.txt'
|
||||
|
||||
- name: 📦 Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -e .
|
||||
pip install pytest
|
||||
pip install -r .github/requirements-dev.txt
|
||||
|
||||
- name: 🧪 Test
|
||||
run: "python -m pytest ./test"
|
||||
|
||||
@@ -7,7 +7,7 @@ ci:
|
||||
|
||||
repos:
|
||||
- repo: https://github.com/pre-commit/pre-commit-hooks
|
||||
rev: v4.4.0
|
||||
rev: v4.5.0
|
||||
hooks:
|
||||
- id: end-of-file-fixer
|
||||
- id: trailing-whitespace
|
||||
@@ -34,11 +34,6 @@ repos:
|
||||
# name: Upgrade code
|
||||
# args: [--py38-plus]
|
||||
|
||||
- repo: https://github.com/PyCQA/isort
|
||||
rev: 5.12.0
|
||||
hooks:
|
||||
- id: isort
|
||||
name: Sort imports
|
||||
|
||||
- repo: https://github.com/PyCQA/flake8
|
||||
rev: 6.1.0
|
||||
|
||||
@@ -19,6 +19,7 @@
|
||||
[](https://github.com/roboflow/supervision/blob/main/LICENSE.md)
|
||||
[](https://badge.fury.io/py/supervision)
|
||||
[](https://colab.research.google.com/github/roboflow/supervision/blob/main/demo.ipynb)
|
||||
[](https://huggingface.co/spaces/Roboflow/Annotators)
|
||||
[](https://discord.gg/GbfgXGJ8Bk)
|
||||
|
||||
</div>
|
||||
@@ -363,5 +364,4 @@ We love your input! Please see our [contributing guide](https://github.com/robof
|
||||
</a>
|
||||
</a>
|
||||
</div>
|
||||
|
||||
</div>
|
||||
|
||||
+78
-6
@@ -82,6 +82,27 @@
|
||||
|
||||
</div>
|
||||
|
||||
=== "Dot"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> dot_annotator = sv.DotAnnotator()
|
||||
>>> annotated_frame = dot_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "Ellipse"
|
||||
|
||||
```python
|
||||
@@ -191,15 +212,25 @@
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
>>> model = YOLO('yolov8x.pt')
|
||||
|
||||
>>> trace_annotator = sv.TraceAnnotator()
|
||||
>>> annotated_frame = trace_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='...')
|
||||
>>> frames_generator = get_video_frames_generator(source_path='...')
|
||||
>>> tracker = sv.ByteTrack()
|
||||
|
||||
>>> with sv.VideoSink(target_path='...', video_info=video_info) as sink:
|
||||
... for frame in frames_generator:
|
||||
... result = model(frame)[0]
|
||||
... detections = sv.Detections.from_ultralytics(result)
|
||||
... detections = tracker.update_with_detections(detections)
|
||||
... annotated_frame = trace_annotator.annotate(
|
||||
... scene=frame.copy(),
|
||||
... detections=detections)
|
||||
... sink.write_frame(frame=annotated_frame)
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
@@ -208,6 +239,35 @@
|
||||
|
||||
</div>
|
||||
|
||||
=== "HeatMap"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> model = YOLO('yolov8x.pt')
|
||||
|
||||
>>> heat_map_annotator = sv.HeatMapAnnotator()
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='...')
|
||||
>>> frames_generator = get_video_frames_generator(source_path='...')
|
||||
|
||||
>>> with sv.VideoSink(target_path='...', video_info=video_info) as sink:
|
||||
... for frame in frames_generator:
|
||||
... result = model(frame)[0]
|
||||
... detections = sv.Detections.from_ultralytics(result)
|
||||
... annotated_frame = heat_map_annotator.annotate(
|
||||
... scene=frame.copy(),
|
||||
... detections=detections)
|
||||
... sink.write_frame(frame=annotated_frame)
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
## BoundingBoxAnnotator
|
||||
|
||||
:::supervision.annotators.core.BoundingBoxAnnotator
|
||||
@@ -224,6 +284,10 @@
|
||||
|
||||
:::supervision.annotators.core.CircleAnnotator
|
||||
|
||||
## DotAnnotator
|
||||
|
||||
:::supervision.annotators.core.DotAnnotator
|
||||
|
||||
## EllipseAnnotator
|
||||
|
||||
:::supervision.annotators.core.EllipseAnnotator
|
||||
@@ -232,6 +296,10 @@
|
||||
|
||||
:::supervision.annotators.core.HaloAnnotator
|
||||
|
||||
## HeatMapAnnotator
|
||||
|
||||
:::supervision.annotators.core.HeatMapAnnotator
|
||||
|
||||
## MaskAnnotator
|
||||
|
||||
:::supervision.annotators.core.MaskAnnotator
|
||||
@@ -247,3 +315,7 @@
|
||||
## TraceAnnotator
|
||||
|
||||
:::supervision.annotators.core.TraceAnnotator
|
||||
|
||||
## ColorLookup
|
||||
|
||||
:::supervision.annotators.utils.ColorLookup
|
||||
|
||||
@@ -1,3 +1,47 @@
|
||||
### 0.16.0 <small>October 19, 2023</small>
|
||||
|
||||
- Added [#422](https://github.com/roboflow/supervision/pull/422): [`sv.BoxMaskAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoxMaskAnnotator) allowing to annotate images and videos with mox masks.
|
||||
|
||||
- Added [#433](https://github.com/roboflow/supervision/pull/433): [`sv.HaloAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.HaloAnnotator) allowing to annotate images and videos with halo effect.
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> halo_annotator = sv.HaloAnnotator()
|
||||
>>> annotated_frame = halo_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
- Added [#466](https://github.com/roboflow/supervision/pull/466): [`sv.HeatMapAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.HeatMapAnnotator) allowing to annotate videos with heat maps.
|
||||
|
||||
- Added [#492](https://github.com/roboflow/supervision/pull/492): [`sv.DotAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.DotAnnotator) allowing to annotate images and videos with dots.
|
||||
|
||||
- Added [#449](https://github.com/roboflow/supervision/pull/449): [`sv.draw_image`](https://supervision.roboflow.com/draw/utils/#supervision.draw.utils.draw_image) allowing to draw an image onto a given scene with specified opacity and dimensions.
|
||||
|
||||
- Added [#280](https://github.com/roboflow/supervision/pull/280): [`sv.FPSMonitor`](https://supervision.roboflow.com/utils/video/#supervision.utils.video.FPSMonitor) for monitoring frames per second (FPS) to benchmark latency.
|
||||
|
||||
- Added [#454](https://github.com/roboflow/supervision/pull/454): 🤗 Hugging Face Annotators [space](https://huggingface.co/spaces/Roboflow/Annotators).
|
||||
|
||||
- Changed [#482](https://github.com/roboflow/supervision/pull/482): [`sv.LineZone.tigger`](https://supervision.roboflow.com/detection/tools/line_zone/#supervision.detection.line_counter.LineZone.trigger) now return `Tuple[np.ndarray, np.ndarray]`. The first array indicates which detections have crossed the line from outside to inside. The second array indicates which detections have crossed the line from inside to outside.
|
||||
|
||||
- Changed [#465](https://github.com/roboflow/supervision/pull/465): Annotator argument name from `color_map: str` to `color_lookup: ColorLookup` enum to increase type safety.
|
||||
|
||||
- Changed [#426](https://github.com/roboflow/supervision/pull/426): [`sv.MaskAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.MaskAnnotator) allowing 2x faster annotation.
|
||||
|
||||
- Fixed [#477](https://github.com/roboflow/supervision/pull/477): Poetry env definition allowing proper local installation.
|
||||
|
||||
- Fixed [#430](https://github.com/roboflow/supervision/pull/430): [`sv.ByteTrack`](https://supervision.roboflow.com/trackers/#supervision.tracker.byte_tracker.core.ByteTrack) to return `np.array([], dtype=int)` when `svDetections` is empty.
|
||||
|
||||
!!! warning
|
||||
|
||||
`sv.Detections.from_yolov8` and `sv.Classifications.from_yolov8` as those are now replaced by [`sv.Detections.from_ultralytics`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_ultralytics) and [`sv.Classifications.from_ultralytics`](https://supervision.roboflow.com/classification/core/#supervision.classification.core.Classifications.from_ultralytics).
|
||||
|
||||
|
||||
### 0.15.0 <small>October 5, 2023</small>
|
||||
|
||||
- Added [#170](https://github.com/roboflow/supervision/pull/170): [`sv.BoundingBoxAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoundingBoxAnnotator) allowing to annotate images and videos with bounding boxes.
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
## LineZone
|
||||
|
||||
:::supervision.detection.line_counter.LineZone
|
||||
@@ -0,0 +1,7 @@
|
||||
## Color
|
||||
|
||||
:::supervision.draw.color.Color
|
||||
|
||||
## ColorPalette
|
||||
|
||||
:::supervision.draw.color.ColorPalette
|
||||
@@ -17,3 +17,7 @@
|
||||
## draw_text
|
||||
|
||||
:::supervision.draw.utils.draw_text
|
||||
|
||||
## draw_image
|
||||
|
||||
:::supervision.draw.utils.draw_image
|
||||
|
||||
@@ -1 +1,78 @@
|
||||
🚧 Page under construction.
|
||||
With Supervision, you can easily [annotate](https://supervision.roboflow.com/annotators/) predictions obtained from a variety of object detection and segmentation models. This document outlines how to run inference using the [Ultralytics](https://github.com/ultralytics/ultralytics) YOLOv8 model, load these predictions into Supervision, and annotate the image.
|
||||
|
||||
## Run Inference
|
||||
|
||||
First, you'll need to obtain predictions from your object detection or segmentation model.
|
||||
```python
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread("image.jpg")
|
||||
results = model(image)[0]
|
||||
```
|
||||
|
||||
## Load Predictions into Supervision
|
||||
|
||||
Now that we have predictions from a model, we can load them into Supervision. We can do so using the [`sv.Detections.from_ultralytics`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_ultralytics) method, which accepts model results from both detection and segmentation models.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
import supervision as sv
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread("image.jpg")
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
```
|
||||
|
||||
You can conveniently load predictions from other computer vision frameworks and libraries using:
|
||||
|
||||
- [`from_deepsparse`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_deepsparse) ([Deepsparse](https://github.com/neuralmagic/deepsparse))
|
||||
- [`from_detectron2`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_detectron2) ([Detectron2](https://github.com/facebookresearch/detectron2))
|
||||
- [`from_mmdetection`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_mmdetection) ([MMDetection](https://github.com/open-mmlab/mmdetection))
|
||||
- [`from_roboflow`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_roboflow) ([Roboflow Inference](https://github.com/roboflow/inference))
|
||||
- [`from_sam`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_sam) ([Segment Anything Model](https://github.com/facebookresearch/segment-anything))
|
||||
- [`from_transformers`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_transformers) ([HuggingFace Transformers](https://github.com/huggingface/transformers))
|
||||
- [`from_yolo_nas`](https://supervision.roboflow.com/detection/core/#supervision.detection.core.Detections.from_yolo_nas) ([YOLO-NAS](https://github.com/Deci-AI/super-gradients/blob/master/YOLONAS.md))
|
||||
|
||||
|
||||
## Annotate Image
|
||||
|
||||
Finally, we can annotate the image with the predictions. Since we are working with an object detection model, we will use the [`sv.BoundingBoxAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoundingBoxAnnotator) and [`sv.LabelAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.LabelAnnotator) classes. If you are running the segmentation model [`sv.MaskAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.MaskAnnotator) is a drop-in replacement for [`sv.BoundingBoxAnnotator`](https://supervision.roboflow.com/annotators/#supervision.annotators.core.BoundingBoxAnnotator) that will allow you to draw masks instead of boxes.
|
||||
|
||||
```python
|
||||
import cv2
|
||||
from ultralytics import YOLO
|
||||
import supervision as sv
|
||||
|
||||
model = YOLO("yolov8n.pt")
|
||||
image = cv2.imread("image.jpg")
|
||||
results = model(image)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator()
|
||||
|
||||
labels = [
|
||||
results.names[class_id]
|
||||
for class_id
|
||||
in detections.class_id
|
||||
]
|
||||
|
||||
annotated_image = bounding_box_annotator.annotate(
|
||||
scene=image, detections=detections)
|
||||
annotated_image = label_annotator.annotate(
|
||||
scene=annotated_image, detections=detections, labels=labels)
|
||||
```
|
||||
|
||||

|
||||
|
||||
## Display Annotated Image
|
||||
|
||||
To display the annotated image in Jupyter Notebook or Google Colab, use the [`sv.plot_image`](https://supervision.roboflow.com/utils/notebook/#supervision.utils.notebook.plot_image) function.
|
||||
|
||||
```python
|
||||
sv.plot_image(annotated_image)
|
||||
```
|
||||
|
||||
+2
-1
@@ -46,6 +46,7 @@ You can install `supervision` with pip in a
|
||||
# setup python environment and activate it
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
pip install --upgrade pip
|
||||
|
||||
# headless install
|
||||
pip install -e "."
|
||||
@@ -62,7 +63,7 @@ You can install `supervision` with pip in a
|
||||
cd supervision
|
||||
|
||||
# setup python environment and activate it
|
||||
poetry env use python 3.10
|
||||
poetry env use python3.10
|
||||
poetry shell
|
||||
|
||||
# headless install
|
||||
|
||||
@@ -0,0 +1,31 @@
|
||||
## Supervision Examples
|
||||
|
||||
This repository is packed with real-world use-cases, provided through Python scripts or
|
||||
interactive notebooks. Browse through to understand how the Supervision library
|
||||
interfaces with diverse applications.
|
||||
|
||||
| **Title** | **Contributor** |
|
||||
|:--------------------------------------:|:----------------------------------------:|
|
||||
| [Tracking](./tracking) | [@SkalskiP](https://github.com/SkalskiP) |
|
||||
| [Traffic Analysis](./traffic_analysis) | [@SkalskiP](https://github.com/SkalskiP) |
|
||||
|
||||
|
||||
## Guidelines for Contributing
|
||||
|
||||
We welcome contributions from the community in the form of examples, applications, and
|
||||
guides. To contribute, please follow these steps:
|
||||
|
||||
1. Create a pull request (PR) with the `[Example]` prefix in the title, adding your
|
||||
project folder to the `examples/` directory in the repository.
|
||||
2. Confirm your project aligns with the following standards:
|
||||
- Incorporates the `supervision` package.
|
||||
- Provides a `README.md` file, detailing the instructions to execute the project.
|
||||
- Showcases visual results, demonstrating the app's functionality.
|
||||
- Avoids adding large assets or dependencies unless absolutely necessary.
|
||||
- The contributor is expected to provide support for issues related to their
|
||||
examples.
|
||||
- In case the presented model has licensing complications, kindly specify them to
|
||||
circumvent potential misunderstandings.
|
||||
|
||||
For inquiries or concerns about these prerequisites, feel free to raise a PR. We are
|
||||
committed to assist and guide you.
|
||||
+3
-1
@@ -1,5 +1,5 @@
|
||||
site_name: Supervision
|
||||
site_url: https://roboflow.github.io/supervision
|
||||
site_url: https://supervision.roboflow.com/
|
||||
site_author: Roboflow
|
||||
site_description: A set of easy-to-use utils that will come in handy in any Computer Vision project
|
||||
repo_name: roboflow/supervision
|
||||
@@ -39,6 +39,7 @@ nav:
|
||||
- Core: detection/core.md
|
||||
- Utils: detection/utils.md
|
||||
- Tools:
|
||||
- Line Zone: detection/tools/line_zone.md
|
||||
- Polygon Zone: detection/tools/polygon_zone.md
|
||||
- Inference Slicer: detection/tools/inference_slicer.md
|
||||
- Annotators: annotators.md
|
||||
@@ -47,6 +48,7 @@ nav:
|
||||
- Metrics:
|
||||
- Object Detection: metrics/detection.md
|
||||
- Draw:
|
||||
- Color: draw/color.md
|
||||
- Utils: draw/utils.md
|
||||
- Utils:
|
||||
- Video: utils/video.md
|
||||
|
||||
Generated
+3
-3
@@ -3358,13 +3358,13 @@ dev = ["flake8", "flake8-annotations", "flake8-bandit", "flake8-bugbear", "flake
|
||||
|
||||
[[package]]
|
||||
name = "urllib3"
|
||||
version = "2.0.6"
|
||||
version = "2.0.7"
|
||||
description = "HTTP library with thread-safe connection pooling, file post, and more."
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "urllib3-2.0.6-py3-none-any.whl", hash = "sha256:7a7c7003b000adf9e7ca2a377c9688bbc54ed41b985789ed576570342a375cd2"},
|
||||
{file = "urllib3-2.0.6.tar.gz", hash = "sha256:b19e1a85d206b56d7df1d5e683df4a7725252a964e3993648dd0fb5a1c157564"},
|
||||
{file = "urllib3-2.0.7-py3-none-any.whl", hash = "sha256:fdb6d215c776278489906c2f8916e6e7d4f5a9b602ccbcfdf7f016fc8da0596e"},
|
||||
{file = "urllib3-2.0.7.tar.gz", hash = "sha256:c97dfde1f7bd43a71c8d2a58e369e9b2bf692d1334ea9f9cae55add7d0dd0f84"},
|
||||
]
|
||||
|
||||
[package.extras]
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "supervision"
|
||||
version = "0.16.0rc2"
|
||||
version = "0.16.0"
|
||||
description = "A set of easy-to-use utils that will come in handy in any Computer Vision project"
|
||||
authors = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
maintainers = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
@@ -176,7 +176,7 @@ max-args = 20
|
||||
include-package-data = false
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
exclude = ["docs*", "test*"]
|
||||
exclude = ["docs*", "test*","examples*"]
|
||||
|
||||
[build-system]
|
||||
requires = ["poetry-core"]
|
||||
|
||||
+11
-2
@@ -12,13 +12,15 @@ from supervision.annotators.core import (
|
||||
BoxCornerAnnotator,
|
||||
BoxMaskAnnotator,
|
||||
CircleAnnotator,
|
||||
DotAnnotator,
|
||||
EllipseAnnotator,
|
||||
HaloAnnotator,
|
||||
HeatMapAnnotator,
|
||||
LabelAnnotator,
|
||||
MaskAnnotator,
|
||||
TraceAnnotator,
|
||||
)
|
||||
from supervision.annotators.utils import ColorMap
|
||||
from supervision.annotators.utils import ColorLookup
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.core import (
|
||||
BaseDataset,
|
||||
@@ -40,7 +42,14 @@ from supervision.detection.utils import (
|
||||
polygon_to_xyxy,
|
||||
)
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text
|
||||
from supervision.draw.utils import (
|
||||
draw_filled_rectangle,
|
||||
draw_image,
|
||||
draw_line,
|
||||
draw_polygon,
|
||||
draw_rectangle,
|
||||
draw_text,
|
||||
)
|
||||
from supervision.geometry.core import Point, Position, Rect
|
||||
from supervision.geometry.utils import get_polygon_center
|
||||
from supervision.metrics.detection import ConfusionMatrix, MeanAveragePrecision
|
||||
|
||||
+334
-93
@@ -5,12 +5,7 @@ import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.annotators.base import BaseAnnotator
|
||||
from supervision.annotators.utils import (
|
||||
ColorMap,
|
||||
Trace,
|
||||
resolve_color,
|
||||
resolve_color_idx,
|
||||
)
|
||||
from supervision.annotators.utils import ColorLookup, Trace, resolve_color
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.geometry.core import Position
|
||||
@@ -25,27 +20,34 @@ class BoundingBoxAnnotator(BaseAnnotator):
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the bounding box lines.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with bounding boxes based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where bounding boxes will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -69,12 +71,14 @@ class BoundingBoxAnnotator(BaseAnnotator):
|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
@@ -94,27 +98,34 @@ class MaskAnnotator(BaseAnnotator):
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
opacity: float = 0.5,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.opacity = opacity
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with masks based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where masks will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -140,12 +151,14 @@ class MaskAnnotator(BaseAnnotator):
|
||||
return scene
|
||||
|
||||
for detection_idx in np.flip(np.argsort(detections.area)):
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
mask = detections.mask[detection_idx]
|
||||
colored_mask = np.zeros_like(scene, dtype=np.uint8)
|
||||
colored_mask[:] = color.as_bgr()
|
||||
@@ -165,27 +178,34 @@ class BoxMaskAnnotator(BaseAnnotator):
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
opacity: float = 0.5,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
self.opacity = opacity
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with box masks based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where bounding boxes will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -210,12 +230,14 @@ class BoxMaskAnnotator(BaseAnnotator):
|
||||
mask_image = scene.copy()
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
@@ -239,30 +261,37 @@ class HaloAnnotator(BaseAnnotator):
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
opacity: float = 0.8,
|
||||
kernel_size: int = 40,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
opacity (float): Opacity of the overlay mask. Must be between `0` and `1`.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
kernel_size (int): The size of the average pooling kernel used for creating
|
||||
the halo.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.opacity = opacity
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
self.kernel_size: int = kernel_size
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with halos based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where masks will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -292,12 +321,14 @@ class HaloAnnotator(BaseAnnotator):
|
||||
)
|
||||
|
||||
for detection_idx in np.flip(np.argsort(detections.area)):
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
mask = detections.mask[detection_idx]
|
||||
fmask = np.logical_or(fmask, mask)
|
||||
color_bgr = color.as_bgr()
|
||||
@@ -323,7 +354,7 @@ class EllipseAnnotator(BaseAnnotator):
|
||||
thickness: int = 2,
|
||||
start_angle: int = -45,
|
||||
end_angle: int = 235,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -332,22 +363,29 @@ class EllipseAnnotator(BaseAnnotator):
|
||||
thickness (int): Thickness of the ellipse lines.
|
||||
start_angle (int): Starting angle of the ellipse.
|
||||
end_angle (int): Ending angle of the ellipse.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.start_angle: int = start_angle
|
||||
self.end_angle: int = end_angle
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with ellipses based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where ellipses will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -371,13 +409,14 @@ class EllipseAnnotator(BaseAnnotator):
|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
|
||||
center = (int((x1 + x2) / 2), y2)
|
||||
width = x2 - x1
|
||||
cv2.ellipse(
|
||||
@@ -404,7 +443,7 @@ class BoxCornerAnnotator(BaseAnnotator):
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 4,
|
||||
corner_length: int = 15,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -412,21 +451,28 @@ class BoxCornerAnnotator(BaseAnnotator):
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the corner lines.
|
||||
corner_length (int): Length of each corner line.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.corner_length: int = corner_length
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with box corners based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where box corners will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -450,12 +496,14 @@ class BoxCornerAnnotator(BaseAnnotator):
|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
|
||||
|
||||
for x, y in corners:
|
||||
@@ -479,26 +527,27 @@ class CircleAnnotator(BaseAnnotator):
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 4,
|
||||
color_map: str = "class",
|
||||
thickness: int = 2,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the circle line.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with circles based on the provided detections.
|
||||
@@ -506,6 +555,8 @@ class CircleAnnotator(BaseAnnotator):
|
||||
Args:
|
||||
scene (np.ndarray): The image where box corners will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -532,19 +583,14 @@ class CircleAnnotator(BaseAnnotator):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
center = ((x1 + x2) // 2, (y1 + y2) // 2)
|
||||
distance = sqrt((x1 - center[0]) ** 2 + (y1 - center[1]) ** 2)
|
||||
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
if isinstance(self.color, ColorPalette)
|
||||
else self.color
|
||||
)
|
||||
|
||||
cv2.circle(
|
||||
img=scene,
|
||||
center=center,
|
||||
@@ -556,6 +602,83 @@ class CircleAnnotator(BaseAnnotator):
|
||||
return scene
|
||||
|
||||
|
||||
class DotAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing dots on an image at specific coordinates based on provided
|
||||
detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
radius: int = 4,
|
||||
position: Position = Position.CENTER,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
radius (int): Radius of the drawn dots.
|
||||
position (Position): The anchor position for placing the dot.
|
||||
color_lookup (ColorLookup): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.radius: int = radius
|
||||
self.position: Position = position
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with dots based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where dots will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> dot_annotator = sv.DotAnnotator()
|
||||
>>> annotated_frame = dot_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
xy = detections.get_anchor_coordinates(anchor=self.position)
|
||||
for detection_idx in range(len(detections)):
|
||||
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,
|
||||
)
|
||||
center = (int(xy[detection_idx, 0]), int(xy[detection_idx, 1]))
|
||||
cv2.circle(scene, center, self.radius, color.as_bgr(), -1)
|
||||
return scene
|
||||
|
||||
|
||||
class LabelAnnotator:
|
||||
"""
|
||||
A class for annotating labels on an image using provided detections.
|
||||
@@ -569,7 +692,7 @@ class LabelAnnotator:
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
text_position: Position = Position.TOP_LEFT,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
@@ -581,8 +704,8 @@ class LabelAnnotator:
|
||||
text_padding (int): Padding around the text within its background box.
|
||||
text_position (Position): Position of the text relative to the detection.
|
||||
Possible values are defined in the `Position` enum.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.text_color: Color = text_color
|
||||
@@ -590,7 +713,7 @@ class LabelAnnotator:
|
||||
self.text_thickness: int = text_thickness
|
||||
self.text_padding: int = text_padding
|
||||
self.text_position: Position = text_position
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
@staticmethod
|
||||
def resolve_text_background_xyxy(
|
||||
@@ -639,6 +762,7 @@ class LabelAnnotator:
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
labels: List[str] = None,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with labels based on the provided detections.
|
||||
@@ -647,6 +771,8 @@ class LabelAnnotator:
|
||||
scene (np.ndarray): The image where labels will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
labels (List[str]): Optional. Custom labels for each detection.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
@@ -671,12 +797,14 @@ class LabelAnnotator:
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for detection_idx in range(len(detections)):
|
||||
detection_xyxy = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
text = (
|
||||
f"{detections.class_id[detection_idx]}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
@@ -787,30 +915,35 @@ class TraceAnnotator:
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
position: Optional[Position] = Position.CENTER,
|
||||
position: Position = Position.CENTER,
|
||||
trace_length: int = 30,
|
||||
thickness: int = 2,
|
||||
color_map: str = "class",
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color to draw the trace, can be
|
||||
a single color or a color palette.
|
||||
position (Optional[Position]): The position of the trace.
|
||||
position (Position): The position of the trace.
|
||||
Defaults to `CENTER`.
|
||||
trace_length (int): The maximum length of the trace in terms of historical
|
||||
points. Defaults to `30`.
|
||||
thickness (int): The thickness of the trace lines. Defaults to `2`.
|
||||
color_map (str): Strategy for mapping colors to annotations.
|
||||
Options are `index`, `class`, or `track`.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACE`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.position = position
|
||||
self.trace = Trace(max_size=trace_length)
|
||||
self.thickness = thickness
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draws trace paths on the frame based on the detection coordinates provided.
|
||||
|
||||
@@ -818,6 +951,8 @@ class TraceAnnotator:
|
||||
scene (np.ndarray): The image on which the traces will be drawn.
|
||||
detections (Detections): The detections which include coordinates for
|
||||
which the traces will be drawn.
|
||||
custom_color_lookup (Optional[np.ndarray]): Custom color lookup array.
|
||||
Allows to override the default color mapping strategy.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the trace paths drawn on it.
|
||||
@@ -825,15 +960,25 @@ class TraceAnnotator:
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
>>> model = YOLO('yolov8x.pt')
|
||||
|
||||
>>> trace_annotator = sv.TraceAnnotator()
|
||||
>>> annotated_frame = trace_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='...')
|
||||
>>> frames_generator = sv.get_video_frames_generator(source_path='...')
|
||||
>>> tracker = sv.ByteTrack()
|
||||
|
||||
>>> with sv.VideoSink(target_path='...', video_info=video_info) as sink:
|
||||
... for frame in frames_generator:
|
||||
... result = model(frame)[0]
|
||||
... detections = sv.Detections.from_ultralytics(result)
|
||||
... detections = tracker.update_with_detections(detections)
|
||||
... annotated_frame = trace_annotator.annotate(
|
||||
... scene=frame.copy(),
|
||||
... detections=detections)
|
||||
... sink.write_frame(frame=annotated_frame)
|
||||
```
|
||||
|
||||
):
|
||||
tracker_id = int(detections.tracker_id[detection_idx])
|
||||
idx = resolve_color_idx(
|
||||
color = resolve_color(
|
||||
color=self.color,
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
color_lookup=self.color_lookup
|
||||
if custom_color_lookup is None
|
||||
else custom_color_lookup,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
xy = self.trace.get(tracker_id=tracker_id)
|
||||
if len(xy) > 1:
|
||||
scene = cv2.polylines(
|
||||
@@ -859,3 +1006,97 @@ class TraceAnnotator:
|
||||
thickness=self.thickness,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class HeatMapAnnotator:
|
||||
"""
|
||||
A class for drawing heatmaps on an image based on provided detections.
|
||||
Heat accumulates over time and is drawn as a semi-transparent overlay
|
||||
of blurred circles.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
position: Position = Position.BOTTOM_CENTER,
|
||||
opacity: float = 0.2,
|
||||
radius: int = 40,
|
||||
kernel_size: int = 25,
|
||||
top_hue: int = 0,
|
||||
low_hue: int = 125,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
position (Position): The position of the heatmap. Defaults to
|
||||
`BOTTOM_CENTER`.
|
||||
opacity (float): Opacity of the overlay mask, between 0 and 1.
|
||||
radius (int): Radius of the heat circle.
|
||||
kernel_size (int): Kernel size for blurring the heatmap.
|
||||
top_hue (int): Hue at the top of the heatmap. Defaults to 0 (red).
|
||||
low_hue (int): Hue at the bottom of the heatmap. Defaults to 125 (blue).
|
||||
"""
|
||||
self.position = position
|
||||
self.opacity = opacity
|
||||
self.radius = radius
|
||||
self.kernel_size = kernel_size
|
||||
self.heat_mask = None
|
||||
self.top_hue = top_hue
|
||||
self.low_hue = low_hue
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Annotates the scene with a heatmap based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where the heatmap will be drawn.
|
||||
detections (Detections): Object detections to annotate.
|
||||
|
||||
Returns:
|
||||
np.ndarray: Annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
>>> from ultralytics import YOLO
|
||||
|
||||
>>> model = YOLO('yolov8x.pt')
|
||||
|
||||
>>> heat_map_annotator = sv.HeatMapAnnotator()
|
||||
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='...')
|
||||
>>> frames_generator = get_video_frames_generator(source_path='...')
|
||||
|
||||
>>> with sv.VideoSink(target_path='...', video_info=video_info) as sink:
|
||||
... for frame in frames_generator:
|
||||
... result = model(frame)[0]
|
||||
... detections = sv.Detections.from_ultralytics(result)
|
||||
... annotated_frame = heat_map_annotator.annotate(
|
||||
... scene=frame.copy(),
|
||||
... detections=detections)
|
||||
... sink.write_frame(frame=annotated_frame)
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
|
||||
if self.heat_mask is None:
|
||||
self.heat_mask = np.zeros(scene.shape[:2])
|
||||
mask = np.zeros(scene.shape[:2])
|
||||
for xy in detections.get_anchor_coordinates(self.position):
|
||||
cv2.circle(mask, (int(xy[0]), int(xy[1])), self.radius, 1, -1)
|
||||
self.heat_mask = mask + self.heat_mask
|
||||
temp = self.heat_mask.copy()
|
||||
temp = self.low_hue - temp / temp.max() * (self.low_hue - self.top_hue)
|
||||
temp = temp.astype(np.uint8)
|
||||
if self.kernel_size is not None:
|
||||
temp = cv2.blur(temp, (self.kernel_size, self.kernel_size))
|
||||
hsv = np.zeros(scene.shape)
|
||||
hsv[..., 0] = temp
|
||||
hsv[..., 1] = 255
|
||||
hsv[..., 2] = 255
|
||||
temp = cv2.cvtColor(hsv.astype(np.uint8), cv2.COLOR_HSV2BGR)
|
||||
mask = cv2.cvtColor(self.heat_mask.astype(np.uint8), cv2.COLOR_GRAY2BGR) > 0
|
||||
scene[mask] = cv2.addWeighted(temp, self.opacity, scene, 1 - self.opacity, 0)[
|
||||
mask
|
||||
]
|
||||
return scene
|
||||
|
||||
@@ -8,18 +8,29 @@ from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
|
||||
class ColorMap(Enum):
|
||||
class ColorLookup(Enum):
|
||||
"""
|
||||
Enum for annotator color mapping.
|
||||
Enumeration class to define strategies for mapping colors to annotations.
|
||||
|
||||
This enum supports three different lookup strategies:
|
||||
- `INDEX`: Colors are determined by the index of the detection within the scene.
|
||||
- `CLASS`: Colors are determined by the class label of the detected object.
|
||||
- `TRACK`: Colors are determined by the tracking identifier of the object.
|
||||
"""
|
||||
|
||||
INDEX = "index"
|
||||
CLASS = "class"
|
||||
TRACK = "track"
|
||||
|
||||
@classmethod
|
||||
def list(cls):
|
||||
return list(map(lambda c: c.value, cls))
|
||||
|
||||
|
||||
def resolve_color_idx(
|
||||
detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS
|
||||
detections: Detections,
|
||||
detection_idx: int,
|
||||
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
|
||||
) -> int:
|
||||
if detection_idx >= len(detections):
|
||||
raise ValueError(
|
||||
@@ -27,16 +38,23 @@ def resolve_color_idx(
|
||||
f"is out of bounds for detections of length {len(detections)}"
|
||||
)
|
||||
|
||||
if color_map == ColorMap.INDEX:
|
||||
if isinstance(color_lookup, np.ndarray):
|
||||
if len(color_lookup) != len(detections):
|
||||
raise ValueError(
|
||||
f"Length of color lookup {len(color_lookup)}"
|
||||
f"does not match length of detections {len(detections)}"
|
||||
)
|
||||
return color_lookup[detection_idx]
|
||||
elif color_lookup == ColorLookup.INDEX:
|
||||
return detection_idx
|
||||
elif color_map == ColorMap.CLASS:
|
||||
elif color_lookup == ColorLookup.CLASS:
|
||||
if detections.class_id is None:
|
||||
raise ValueError(
|
||||
"Could not resolve color by class because"
|
||||
"Detections do not have class_id"
|
||||
)
|
||||
return detections.class_id[detection_idx]
|
||||
elif color_map == ColorMap.TRACK:
|
||||
elif color_lookup == ColorLookup.TRACK:
|
||||
if detections.tracker_id is None:
|
||||
raise ValueError(
|
||||
"Could not resolve color by track because"
|
||||
@@ -45,12 +63,26 @@ def resolve_color_idx(
|
||||
return detections.tracker_id[detection_idx]
|
||||
|
||||
|
||||
def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color:
|
||||
def get_color_by_index(color: Union[Color, ColorPalette], idx: int) -> Color:
|
||||
if isinstance(color, ColorPalette):
|
||||
return color.by_idx(idx)
|
||||
return color
|
||||
|
||||
|
||||
def resolve_color(
|
||||
color: Union[Color, ColorPalette],
|
||||
detections: Detections,
|
||||
detection_idx: int,
|
||||
color_lookup: Union[ColorLookup, np.ndarray] = ColorLookup.CLASS,
|
||||
) -> Color:
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_lookup=color_lookup,
|
||||
)
|
||||
return get_color_by_index(color=color, idx=idx)
|
||||
|
||||
|
||||
class Trace:
|
||||
def __init__(
|
||||
self,
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from typing import Dict, Optional
|
||||
from typing import Dict, Optional, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@@ -10,37 +10,54 @@ from supervision.geometry.core import Point, Rect, Vector
|
||||
|
||||
class LineZone:
|
||||
"""
|
||||
Count the number of objects that cross a line.
|
||||
This class is responsible for counting the number of objects that cross a
|
||||
predefined line.
|
||||
|
||||
!!! warning
|
||||
|
||||
LineZone utilizes the `tracker_id`. Read
|
||||
[here](https://supervision.roboflow.com/trackers/) to learn how to plug
|
||||
tracking into your inference pipeline.
|
||||
|
||||
Attributes:
|
||||
in_count (int): The number of objects that have crossed the line from outside
|
||||
to inside.
|
||||
out_count (int): The number of objects that have crossed the line from inside
|
||||
to outside.
|
||||
"""
|
||||
|
||||
def __init__(self, start: Point, end: Point):
|
||||
"""
|
||||
Initialize a LineCounter object.
|
||||
|
||||
Attributes:
|
||||
Args:
|
||||
start (Point): The starting point of the line.
|
||||
end (Point): The ending point of the line.
|
||||
|
||||
"""
|
||||
self.vector = Vector(start=start, end=end)
|
||||
self.tracker_state: Dict[str, bool] = {}
|
||||
self.in_count: int = 0
|
||||
self.out_count: int = 0
|
||||
|
||||
def trigger(self, detections: Detections):
|
||||
def trigger(self, detections: Detections) -> Tuple[np.ndarray, np.ndarray]:
|
||||
"""
|
||||
Update the in_count and out_count for the detections that cross the line.
|
||||
Update the `in_count` and `out_count` based on the objects that cross the line.
|
||||
|
||||
Attributes:
|
||||
detections (Detections): The detections for which to update the counts.
|
||||
Args:
|
||||
detections (Detections): A list of detections for which to update the
|
||||
counts.
|
||||
|
||||
Returns:
|
||||
A tuple of two boolean NumPy arrays. The first array indicates which
|
||||
detections have crossed the line from outside to inside. The second
|
||||
array indicates which detections have crossed the line from inside to
|
||||
outside.
|
||||
"""
|
||||
for xyxy, _, confidence, class_id, tracker_id in detections:
|
||||
# handle detections with no tracker_id
|
||||
crossed_in = np.full(len(detections), False)
|
||||
crossed_out = np.full(len(detections), False)
|
||||
|
||||
for i, (xyxy, _, confidence, class_id, tracker_id) in enumerate(detections):
|
||||
if tracker_id is None:
|
||||
continue
|
||||
|
||||
# we check if all four anchors of bbox are on the same side of vector
|
||||
x1, y1, x2, y2 = xyxy
|
||||
anchors = [
|
||||
Point(x=x1, y=y1),
|
||||
@@ -50,25 +67,27 @@ class LineZone:
|
||||
]
|
||||
triggers = [self.vector.is_in(point=anchor) for anchor in anchors]
|
||||
|
||||
# detection is partially in and partially out
|
||||
if len(set(triggers)) == 2:
|
||||
continue
|
||||
|
||||
tracker_state = triggers[0]
|
||||
# handle new detection
|
||||
|
||||
if tracker_id not in self.tracker_state:
|
||||
self.tracker_state[tracker_id] = tracker_state
|
||||
continue
|
||||
|
||||
# handle detection on the same side of the line
|
||||
if self.tracker_state.get(tracker_id) == tracker_state:
|
||||
continue
|
||||
|
||||
self.tracker_state[tracker_id] = tracker_state
|
||||
if tracker_state:
|
||||
self.in_count += 1
|
||||
crossed_in[i] = True
|
||||
else:
|
||||
self.out_count += 1
|
||||
crossed_out[i] = True
|
||||
|
||||
return crossed_in, crossed_out
|
||||
|
||||
|
||||
class LineZoneAnnotator:
|
||||
|
||||
+71
-21
@@ -34,21 +34,35 @@ def _validate_color_hex(color_hex: str):
|
||||
|
||||
@dataclass
|
||||
class Color:
|
||||
"""
|
||||
Represents a color in RGB format.
|
||||
|
||||
Attributes:
|
||||
r (int): Red channel.
|
||||
g (int): Green channel.
|
||||
b (int): Blue channel.
|
||||
"""
|
||||
|
||||
r: int
|
||||
g: int
|
||||
b: int
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, color_hex: str):
|
||||
def from_hex(cls, color_hex: str) -> Color:
|
||||
"""
|
||||
Creates a Color instance from a color hex string
|
||||
Create a Color instance from a hex string.
|
||||
|
||||
:param color_hex: str : The color hex string in the format
|
||||
of "fff", "ffffff", "#fff", or "#ffffff"
|
||||
:return: Color : A Color instance representing the color
|
||||
Args:
|
||||
color_hex (str): Hex string of the color.
|
||||
|
||||
Returns:
|
||||
Color: Instance representing the color.
|
||||
|
||||
Example:
|
||||
color = Color.from_hex('#ff00ff')
|
||||
```
|
||||
>>> Color.from_hex('#ff00ff')
|
||||
Color(r=255, g=0, b=255)
|
||||
```
|
||||
"""
|
||||
_validate_color_hex(color_hex)
|
||||
color_hex = color_hex.lstrip("#")
|
||||
@@ -59,17 +73,31 @@ class Color:
|
||||
|
||||
def as_rgb(self) -> Tuple[int, int, int]:
|
||||
"""
|
||||
Returns the color as a tuple of integers in the RGB format
|
||||
Returns the color as an RGB tuple.
|
||||
|
||||
:return: Tuple[int, int, int] : The color in the RGB format
|
||||
Returns:
|
||||
Tuple[int, int, int]: RGB tuple.
|
||||
|
||||
Example:
|
||||
```
|
||||
>>> color.as_rgb()
|
||||
(255, 0, 255)
|
||||
```
|
||||
"""
|
||||
return self.r, self.g, self.b
|
||||
|
||||
def as_bgr(self) -> Tuple[int, int, int]:
|
||||
"""
|
||||
Returns the color as a tuple of integers in the BGR format
|
||||
Returns the color as a BGR tuple.
|
||||
|
||||
:return: Tuple[int, int, int] : The color in the BGR format
|
||||
Returns:
|
||||
Tuple[int, int, int]: BGR tuple.
|
||||
|
||||
Example:
|
||||
```
|
||||
>>> color.as_bgr()
|
||||
(255, 0, 255)
|
||||
```
|
||||
"""
|
||||
return self.b, self.g, self.r
|
||||
|
||||
@@ -100,33 +128,55 @@ class ColorPalette:
|
||||
|
||||
@classmethod
|
||||
def default(cls) -> ColorPalette:
|
||||
"""
|
||||
Returns a default color palette.
|
||||
|
||||
Returns:
|
||||
ColorPalette: A ColorPalette instance with default colors.
|
||||
|
||||
Example:
|
||||
```
|
||||
>>> ColorPalette.default()
|
||||
ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...])
|
||||
```
|
||||
"""
|
||||
return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE)
|
||||
|
||||
@classmethod
|
||||
def from_hex(cls, color_hex_list: List[str]):
|
||||
def from_hex(cls, color_hex_list: List[str]) -> ColorPalette:
|
||||
"""
|
||||
Creates a ColorPalette instance from a list of color hex strings
|
||||
Create a ColorPalette instance from a list of hex strings.
|
||||
|
||||
:param color_hex_list: List[str] : A list of color hex strings in the
|
||||
format of "fff", "ffffff", "#fff", or "#ffffff"
|
||||
:return: ColorPalette : A ColorPalette instance representing the color palette
|
||||
Args:
|
||||
color_hex_list (List[str]): List of color hex strings.
|
||||
|
||||
Returns:
|
||||
ColorPalette: A ColorPalette instance.
|
||||
|
||||
Example:
|
||||
color_palette = ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff'])
|
||||
```
|
||||
>>> ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff'])
|
||||
ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...])
|
||||
```
|
||||
"""
|
||||
colors = [Color.from_hex(color_hex) for color_hex in color_hex_list]
|
||||
return cls(colors)
|
||||
|
||||
def by_idx(self, idx: int) -> Color:
|
||||
"""
|
||||
Returns the color at a given index in the color palette.
|
||||
Return the color at a given index in the palette.
|
||||
|
||||
:param idx: int : The index of the color in the color palette
|
||||
:return: Color : The color at the given index
|
||||
Args:
|
||||
idx (int): Index of the color in the palette.
|
||||
|
||||
Returns:
|
||||
Color: Color at the given index.
|
||||
|
||||
Example:
|
||||
color_palette = ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff'])
|
||||
color = color_palette.by_idx(1)
|
||||
```
|
||||
>>> color_palette.by_idx(1)
|
||||
Color(r=0, g=255, b=0)
|
||||
```
|
||||
"""
|
||||
if idx < 0:
|
||||
raise ValueError("idx argument should not be negative")
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
from typing import Optional
|
||||
import os
|
||||
from typing import Optional, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
@@ -168,3 +169,65 @@ def draw_text(
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
def draw_image(
|
||||
scene: np.ndarray, image: Union[str, np.ndarray], opacity: float, rect: Rect
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Draws an image onto a given scene with specified opacity and dimensions.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): Background image where the new image will be drawn.
|
||||
image (Union[str, np.ndarray]): Image to draw.
|
||||
opacity (float): Opacity of the image to be drawn.
|
||||
rect (Rect): Rectangle specifying where to draw the image.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The updated scene.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: If the image path does not exist.
|
||||
ValueError: For invalid opacity or rectangle dimensions.
|
||||
"""
|
||||
|
||||
# Validate and load image
|
||||
if isinstance(image, str):
|
||||
if not os.path.exists(image):
|
||||
raise FileNotFoundError(f"Image path ('{image}') does not exist.")
|
||||
image = cv2.imread(image, cv2.IMREAD_UNCHANGED)
|
||||
|
||||
# Validate opacity
|
||||
if not 0.0 <= opacity <= 1.0:
|
||||
raise ValueError("Opacity must be between 0.0 and 1.0.")
|
||||
|
||||
# Validate rectangle dimensions
|
||||
if (
|
||||
rect.x < 0
|
||||
or rect.y < 0
|
||||
or rect.x + rect.width > scene.shape[1]
|
||||
or rect.y + rect.height > scene.shape[0]
|
||||
):
|
||||
raise ValueError("Invalid rectangle dimensions.")
|
||||
|
||||
# Resize and isolate alpha channel
|
||||
image = cv2.resize(image, (rect.width, rect.height))
|
||||
alpha_channel = (
|
||||
image[:, :, 3]
|
||||
if image.shape[2] == 4
|
||||
else np.ones((rect.height, rect.width), dtype=image.dtype) * 255
|
||||
)
|
||||
alpha_scaled = cv2.convertScaleAbs(alpha_channel * opacity)
|
||||
|
||||
# Perform blending
|
||||
scene_roi = scene[rect.y : rect.y + rect.height, rect.x : rect.x + rect.width]
|
||||
alpha_float = alpha_scaled.astype(np.float32) / 255.0
|
||||
blended_roi = cv2.convertScaleAbs(
|
||||
(1 - alpha_float[..., np.newaxis]) * scene_roi
|
||||
+ alpha_float[..., np.newaxis] * image[:, :, :3]
|
||||
)
|
||||
|
||||
# Update the scene
|
||||
scene[rect.y : rect.y + rect.height, rect.x : rect.x + rect.width] = blended_roi
|
||||
|
||||
return scene
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
from typing import List, Optional, Tuple
|
||||
from typing import List, Tuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
@@ -11,7 +11,7 @@ from supervision.tracker.byte_tracker.kalman_filter import KalmanFilter
|
||||
class STrack(BaseTrack):
|
||||
shared_kalman = KalmanFilter()
|
||||
|
||||
def __init__(self, tlwh, score, class_ids, mask: Optional[np.array] = None):
|
||||
def __init__(self, tlwh, score, class_ids):
|
||||
# wait activate
|
||||
self._tlwh = np.asarray(tlwh, dtype=np.float32)
|
||||
self.kalman_filter = None
|
||||
@@ -21,7 +21,6 @@ class STrack(BaseTrack):
|
||||
self.score = score
|
||||
self.class_ids = class_ids
|
||||
self.tracklet_len = 0
|
||||
self.mask = mask
|
||||
|
||||
def predict(self):
|
||||
mean_state = self.mean.copy()
|
||||
@@ -75,7 +74,6 @@ class STrack(BaseTrack):
|
||||
if new_id:
|
||||
self.track_id = self.next_id()
|
||||
self.score = new_track.score
|
||||
self.mask = new_track.mask
|
||||
|
||||
def update(self, new_track, frame_id):
|
||||
"""
|
||||
@@ -97,8 +95,6 @@ class STrack(BaseTrack):
|
||||
|
||||
self.score = new_track.score
|
||||
|
||||
self.mask = new_track.mask
|
||||
|
||||
@property
|
||||
def tlwh(self):
|
||||
"""Get current position in bounding box format `(top left x, top left y,
|
||||
@@ -235,8 +231,9 @@ class ByteTrack:
|
||||
... )
|
||||
```
|
||||
"""
|
||||
|
||||
tracks = self.update_with_tensors(
|
||||
tensors=detections2boxes(detections=detections), masks=detections.mask
|
||||
tensors=detections2boxes(detections=detections)
|
||||
)
|
||||
detections = Detections.empty()
|
||||
if len(tracks) > 0:
|
||||
@@ -252,22 +249,17 @@ class ByteTrack:
|
||||
detections.confidence = np.array(
|
||||
[t.score for t in tracks], dtype=np.float32
|
||||
)
|
||||
detections.mask = np.array([t.mask for t in tracks], dtype=bool)
|
||||
|
||||
else:
|
||||
detections.tracker_id = np.array([], dtype=int)
|
||||
|
||||
return detections
|
||||
|
||||
def update_with_tensors(
|
||||
self, tensors: np.ndarray, masks: Optional[np.array] = None
|
||||
) -> List[STrack]:
|
||||
def update_with_tensors(self, tensors: np.ndarray) -> List[STrack]:
|
||||
"""
|
||||
Updates the tracker with the provided tensors and returns the updated tracks.
|
||||
|
||||
Parameters:
|
||||
tensors: The new tensors to update with.
|
||||
masks: The new masks associated to new tensors
|
||||
|
||||
Returns:
|
||||
List[STrack]: Updated tracks.
|
||||
@@ -289,12 +281,6 @@ class ByteTrack:
|
||||
inds_second = np.logical_and(inds_low, inds_high)
|
||||
dets_second = bboxes[inds_second]
|
||||
dets = bboxes[remain_inds]
|
||||
if masks is not None:
|
||||
masks_keep = masks[remain_inds]
|
||||
masks_second = masks[inds_second]
|
||||
else:
|
||||
masks_keep = np.array([None] * len(remain_inds))
|
||||
masks_second = np.array([None] * len(inds_second))
|
||||
scores_keep = scores[remain_inds]
|
||||
scores_second = scores[inds_second]
|
||||
|
||||
@@ -304,10 +290,8 @@ class ByteTrack:
|
||||
if len(dets) > 0:
|
||||
"""Detections"""
|
||||
detections = [
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m)
|
||||
for (tlbr, s, c, m) in zip(
|
||||
dets, scores_keep, class_ids_keep, masks_keep
|
||||
)
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
|
||||
for (tlbr, s, c) in zip(dets, scores_keep, class_ids_keep)
|
||||
]
|
||||
else:
|
||||
detections = []
|
||||
@@ -347,10 +331,8 @@ class ByteTrack:
|
||||
if len(dets_second) > 0:
|
||||
"""Detections"""
|
||||
detections_second = [
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c, m)
|
||||
for (tlbr, s, c, m) in zip(
|
||||
dets_second, scores_second, class_ids_second, masks_second
|
||||
)
|
||||
STrack(STrack.tlbr_to_tlwh(tlbr), s, c)
|
||||
for (tlbr, s, c) in zip(dets_second, scores_second, class_ids_second)
|
||||
]
|
||||
else:
|
||||
detections_second = []
|
||||
|
||||
@@ -2,14 +2,15 @@ from contextlib import ExitStack as DoesNotRaise
|
||||
from test.utils import mock_detections
|
||||
from typing import Optional
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision.annotators.utils import ColorMap, resolve_color_idx
|
||||
from supervision.annotators.utils import ColorLookup, resolve_color_idx
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, detection_idx, color_map, expected_result, exception",
|
||||
"detections, detection_idx, color_lookup, expected_result, exception",
|
||||
[
|
||||
(
|
||||
mock_detections(
|
||||
@@ -18,10 +19,10 @@ from supervision.detection.core import Detections
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.INDEX,
|
||||
ColorLookup.INDEX,
|
||||
0,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; index mapping
|
||||
), # multiple detections; index lookup
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
@@ -29,10 +30,10 @@ from supervision.detection.core import Detections
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.CLASS,
|
||||
ColorLookup.CLASS,
|
||||
5,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; class mapping
|
||||
), # multiple detections; class lookup
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
@@ -40,17 +41,17 @@ from supervision.detection.core import Detections
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.TRACK,
|
||||
ColorLookup.TRACK,
|
||||
2,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; track mapping
|
||||
), # multiple detections; track lookup
|
||||
(
|
||||
Detections.empty(),
|
||||
0,
|
||||
ColorMap.INDEX,
|
||||
ColorLookup.INDEX,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # no detections; index mapping; out of bounds
|
||||
), # no detections; index lookup; out of bounds
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
@@ -58,30 +59,44 @@ from supervision.detection.core import Detections
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
2,
|
||||
ColorMap.INDEX,
|
||||
ColorLookup.INDEX,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; index mapping; out of bounds
|
||||
), # multiple detections; index lookup; out of bounds
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
ColorMap.CLASS,
|
||||
ColorLookup.CLASS,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; class mapping; no class_id
|
||||
), # multiple detections; class lookup; no class_id
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
ColorMap.TRACK,
|
||||
ColorLookup.TRACK,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; class mapping; no track_id
|
||||
), # multiple detections; class lookup; no track_id
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
np.array([1, 0]),
|
||||
1,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; custom lookup; correct length
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
np.array([1]),
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; custom lookup; wrong length
|
||||
],
|
||||
)
|
||||
def test_resolve_color_idx(
|
||||
detections: Detections,
|
||||
detection_idx: int,
|
||||
color_map: ColorMap,
|
||||
color_lookup: ColorLookup,
|
||||
expected_result: Optional[int],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
@@ -89,6 +104,6 @@ def test_resolve_color_idx(
|
||||
result = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=color_map,
|
||||
color_lookup=color_lookup,
|
||||
)
|
||||
assert result == expected_result
|
||||
|
||||
Reference in New Issue
Block a user