Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
1b449294c4 | ||
|
|
36f95d060d | ||
|
|
07c344e050 | ||
|
|
f3bdb6efa2 | ||
|
|
f4f9e7ddc7 | ||
|
|
53f2b4fd45 | ||
|
|
c7582f7c07 | ||
|
|
3e8906d830 | ||
|
|
6521f9f4ed | ||
|
|
560a35b75f | ||
|
|
405a4c73f9 | ||
|
|
87fd10c366 | ||
|
|
8141cb357f | ||
|
|
87fb336eb6 | ||
|
|
28d279bec3 | ||
|
|
1b0bcea5a6 | ||
|
|
507e000bca | ||
|
|
f7bf766a60 | ||
|
|
1e0f18f792 | ||
|
|
be5c1b1aa1 | ||
|
|
c35a7714d2 | ||
|
|
a8cc5a68bb | ||
|
|
de4056192e | ||
|
|
d337a03d28 | ||
|
|
00b67588d8 | ||
|
|
e10e0f21df | ||
|
|
4fb6760d2c | ||
|
|
de7ed5e115 | ||
|
|
d929b86c9f | ||
|
|
0b1627f835 | ||
|
|
20b1b21774 | ||
|
|
e5eaf26137 | ||
|
|
be3e67949c | ||
|
|
2587fdfefd | ||
|
|
08a88a8c0c | ||
|
|
181c0b833b | ||
|
|
11601e3f2d | ||
|
|
1b8ff617c3 | ||
|
|
5bc83bf6cc | ||
|
|
dd80753b2b | ||
|
|
b2bb69bebb | ||
|
|
9c081b0e11 | ||
|
|
7f613c6d27 | ||
|
|
9736eba9e6 | ||
|
|
53dfeac19b | ||
|
|
5c25e58e39 | ||
|
|
52949f2fd1 | ||
|
|
56c511b5d8 | ||
|
|
62014186a2 | ||
|
|
6cd7cf0209 | ||
|
|
a39c92016f | ||
|
|
4e88e26d91 | ||
|
|
484e53324e | ||
|
|
e5c4d8e3b9 | ||
|
|
52ef3fadcb | ||
|
|
472ddd8d0e | ||
|
|
90b23430d2 | ||
|
|
e91f0cb496 | ||
|
|
fa6238dbc4 | ||
|
|
1a7306c4ae | ||
|
|
3d4326c79b | ||
|
|
ca212b70c8 | ||
|
|
0008b33eb1 | ||
|
|
77190f287c | ||
|
|
121d4f31bb | ||
|
|
7c97e87e79 | ||
|
|
9640f1d756 | ||
|
|
2419d42224 | ||
|
|
c6517de143 | ||
|
|
2706304a73 | ||
|
|
8538c26a22 | ||
|
|
a0c60fe043 | ||
|
|
01d8b0316c | ||
|
|
6356a8ac3d | ||
|
|
bbd9a955f8 | ||
|
|
7b2bcf00ad | ||
|
|
775a3f2795 | ||
|
|
5dae02e93e | ||
|
|
666427df08 | ||
|
|
78411b14b5 | ||
|
|
69ae5c6cd9 | ||
|
|
94d08bc9b5 | ||
|
|
693465d74d | ||
|
|
92413d3cc2 | ||
|
|
8b764218eb | ||
|
|
117ba2fcc0 | ||
|
|
1d28b8ae26 | ||
|
|
cc2e08ac72 | ||
|
|
ba03fefbd8 | ||
|
|
4cd9f05fab | ||
|
|
8066125f03 | ||
|
|
aa772dde8d | ||
|
|
1b5e6db196 | ||
|
|
fc4bb27aa8 | ||
|
|
ace6ce363f | ||
|
|
f0cef99de7 | ||
|
|
c4d296fbc6 | ||
|
|
9c4555a800 | ||
|
|
4b423fd402 | ||
|
|
27efbc044d | ||
|
|
2a2c259964 | ||
|
|
c705d50e74 | ||
|
|
d5b8a3aa1a | ||
|
|
d9f6c5c73c | ||
|
|
65a93f8287 | ||
|
|
1bcf8346dc | ||
|
|
aba6ff57a9 | ||
|
|
b5e1342038 | ||
|
|
64d223b76d | ||
|
|
21968e53c6 | ||
|
|
9b77b2a1c6 | ||
|
|
ddabe04409 |
@@ -22,8 +22,5 @@ jobs:
|
||||
python -m pip install --upgrade pip
|
||||
pip install -e .
|
||||
pip install pytest
|
||||
pip install isort
|
||||
pip install flake8
|
||||
pip install "black==22.3.0"
|
||||
- name: 🧪 Test
|
||||
run: "python -m pytest ./test"
|
||||
|
||||
@@ -69,12 +69,12 @@ repos:
|
||||
|
||||
|
||||
- repo: https://github.com/psf/black
|
||||
rev: 23.7.0
|
||||
rev: 23.9.1
|
||||
hooks:
|
||||
- id: black
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.0.286
|
||||
rev: v0.0.291
|
||||
hooks:
|
||||
- id: ruff
|
||||
args: [--fix, --exit-non-zero-on-fix]
|
||||
|
||||
@@ -273,19 +273,13 @@ array([
|
||||
|
||||
</details>
|
||||
|
||||
## 🛠️ built with supervision
|
||||
|
||||
Did you build something cool using supervision? [Let us know!](https://github.com/roboflow/supervision/discussions/categories/built-with-supervision)
|
||||
|
||||
https://user-images.githubusercontent.com/26109316/207858600-ee862b22-0353-440b-ad85-caa0c4777904.mp4
|
||||
|
||||
## 🎬 tutorials
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/ae1ca38e-40b7-4b35-8582-e8ea5de3806e" alt="Accelerate Image Annotation with SAM and Grounding DINO" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/oEQYStnF2l8" title="Accelerate Image Annotation with SAM and Grounding DINO"><strong>Accelerate Image Annotation with SAM and Grounding DINO</strong></a>
|
||||
<div><strong>Created: 20 Apr 2023</strong> | <strong>Updated: 20 Apr 2023</strong></div>
|
||||
<br/> Discover how to speed up your image annotation process using Grounding DINO and Segment Anything Model (SAM). Learn how to convert object detection datasets into instance segmentation datasets, and see the potential of using these models to automatically annotate your datasets for real-time detectors like YOLOv8... </p>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><img src="https://github.com/roboflow/supervision/assets/26109316/54afdf1c-218c-4451-8f12-627fb85f1682" alt="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><strong>Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking</strong></a>
|
||||
<div><strong>Created: 6 Sep 2023</strong> | <strong>Updated: 6 Sep 2023</strong></div>
|
||||
<br/> In this video, we explore real-time traffic analysis using YOLOv8 and ByteTrack to detect and track vehicles on aerial images. Harnessing the power of Python and Supervision, we delve deep into assigning cars to specific entry zones and understanding their direction of movement. By visualizing their paths, we gain insights into traffic flow across bustling roundabouts... </p>
|
||||
|
||||
<br/>
|
||||
|
||||
@@ -295,6 +289,14 @@ https://user-images.githubusercontent.com/26109316/207858600-ee862b22-0353-440b-
|
||||
<div><strong>Created: 11 Apr 2023</strong> | <strong>Updated: 11 Apr 2023</strong></div>
|
||||
<br/> Discover the incredible potential of Meta AI's Segment Anything Model (SAM)! We dive into SAM, an efficient and promptable model for image segmentation, which has revolutionized computer vision tasks. With over 1 billion masks on 11M licensed and privacy-respecting images, SAM's zero-shot performance is often competitive with or even superior to prior fully supervised results... </p>
|
||||
|
||||
## 💜 built with supervision
|
||||
|
||||
Did you build something cool using supervision? [Let us know!](https://github.com/roboflow/supervision/discussions/categories/built-with-supervision)
|
||||
|
||||
https://user-images.githubusercontent.com/26109316/207858600-ee862b22-0353-440b-ad85-caa0c4777904.mp4
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/c9436828-9fbf-4c25-ae8c-60e9c81b3900
|
||||
|
||||
## 📚 documentation
|
||||
|
||||
Visit our [documentation](https://roboflow.github.io/supervision) page to learn how supervision can help you build computer vision applications faster and more reliably.
|
||||
|
||||
@@ -1,7 +1,19 @@
|
||||
## BoxAnnotator
|
||||
## BoundingBoxAnnotator
|
||||
|
||||
:::supervision.detection.annotate.BoxAnnotator
|
||||
:::supervision.annotators.core.BoundingBoxAnnotator
|
||||
|
||||
## MaskAnnotator
|
||||
|
||||
:::supervision.detection.annotate.MaskAnnotator
|
||||
:::supervision.annotators.core.MaskAnnotator
|
||||
|
||||
## EllipseAnnotator
|
||||
|
||||
:::supervision.annotators.core.EllipseAnnotator
|
||||
|
||||
## BoxCornerAnnotator
|
||||
|
||||
:::supervision.annotators.core.BoxCornerAnnotator
|
||||
|
||||
## LabelAnnotator
|
||||
|
||||
:::supervision.annotators.core.LabelAnnotator
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
## Position
|
||||
|
||||
:::supervision.geometry.core.Position
|
||||
@@ -16,7 +16,8 @@ def process_video(
|
||||
model = YOLO(source_weights_path)
|
||||
|
||||
tracker = sv.ByteTrack()
|
||||
box_annotator = sv.BoxAnnotator()
|
||||
box_annotator = sv.BoundingBoxAnnotator()
|
||||
label_annotator = sv.LabelAnnotator(classes=model.names)
|
||||
frame_generator = sv.get_video_frames_generator(source_path=source_video_path)
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=source_video_path)
|
||||
|
||||
@@ -28,16 +29,15 @@ def process_video(
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections = tracker.update_with_detections(detections)
|
||||
|
||||
labels = [
|
||||
f"#{tracker_id} {model.model.names[class_id]}"
|
||||
for _, _, _, class_id, tracker_id in detections
|
||||
]
|
||||
|
||||
annotated_frame = box_annotator.annotate(
|
||||
scene=frame.copy(), detections=detections, labels=labels
|
||||
scene=frame.copy(), detections=detections
|
||||
)
|
||||
|
||||
sink.write_frame(frame=annotated_frame)
|
||||
annotated_labeled_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
|
||||
sink.write_frame(frame=annotated_labeled_frame)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
data/
|
||||
@@ -0,0 +1,44 @@
|
||||
## 👋 hello
|
||||
|
||||
This script performs traffic flow analysis using YOLOv8, an object-detection method and ByteTrack, a simple yet effective online multi-object tracking method. It uses the supervision package for multiple tasks such as tracking, annotations, etc.
|
||||
|
||||
https://github.com/roboflow/supervision/assets/26109316/c9436828-9fbf-4c25-ae8c-60e9c81b3900
|
||||
|
||||
## 💻 install
|
||||
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/traffic_analysis
|
||||
```
|
||||
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
- download `traffic_analysis.pt` and `traffic_analysis.mov` files
|
||||
|
||||
```bash
|
||||
./setup.sh
|
||||
```
|
||||
|
||||
## ⚙️ run
|
||||
|
||||
```bash
|
||||
python script.py \
|
||||
--source_weights_path data/traffic_analysis.pt \
|
||||
--source_video_path data/traffic_analysis.mov \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5 \
|
||||
--target_video_path data/traffic_analysis_result.mov
|
||||
```
|
||||
@@ -0,0 +1,4 @@
|
||||
supervision>=0.15.0rc1
|
||||
tqdm
|
||||
ultralytics
|
||||
gdown
|
||||
@@ -0,0 +1,219 @@
|
||||
import argparse
|
||||
from typing import Dict, List, Set, Tuple
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from tqdm import tqdm
|
||||
from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
|
||||
COLORS = sv.ColorPalette.default()
|
||||
|
||||
ZONE_IN_POLYGONS = [
|
||||
np.array([[592, 282], [900, 282], [900, 82], [592, 82]]),
|
||||
np.array([[950, 860], [1250, 860], [1250, 1060], [950, 1060]]),
|
||||
np.array([[592, 582], [592, 860], [392, 860], [392, 582]]),
|
||||
np.array([[1250, 282], [1250, 530], [1450, 530], [1450, 282]]),
|
||||
]
|
||||
|
||||
ZONE_OUT_POLYGONS = [
|
||||
np.array([[950, 282], [1250, 282], [1250, 82], [950, 82]]),
|
||||
np.array([[592, 860], [900, 860], [900, 1060], [592, 1060]]),
|
||||
np.array([[592, 282], [592, 550], [392, 550], [392, 282]]),
|
||||
np.array([[1250, 860], [1250, 560], [1450, 560], [1450, 860]]),
|
||||
]
|
||||
|
||||
|
||||
class DetectionsManager:
|
||||
def __init__(self) -> None:
|
||||
self.tracker_id_to_zone_id: Dict[int, int] = {}
|
||||
self.counts: Dict[int, Dict[int, Set[int]]] = {}
|
||||
|
||||
def update(
|
||||
self,
|
||||
detections_all: sv.Detections,
|
||||
detections_in_zones: List[sv.Detections],
|
||||
detections_out_zones: List[sv.Detections],
|
||||
) -> sv.Detections:
|
||||
for zone_in_id, detections_in_zone in enumerate(detections_in_zones):
|
||||
for tracker_id in detections_in_zone.tracker_id:
|
||||
self.tracker_id_to_zone_id.setdefault(tracker_id, zone_in_id)
|
||||
|
||||
for zone_out_id, detections_out_zone in enumerate(detections_out_zones):
|
||||
for tracker_id in detections_out_zone.tracker_id:
|
||||
if tracker_id in self.tracker_id_to_zone_id:
|
||||
zone_in_id = self.tracker_id_to_zone_id[tracker_id]
|
||||
self.counts.setdefault(zone_out_id, {})
|
||||
self.counts[zone_out_id].setdefault(zone_in_id, set())
|
||||
self.counts[zone_out_id][zone_in_id].add(tracker_id)
|
||||
|
||||
detections_all.class_id = np.vectorize(
|
||||
lambda x: self.tracker_id_to_zone_id.get(x, -1)
|
||||
)(detections_all.tracker_id)
|
||||
return detections_all[detections_all.class_id != -1]
|
||||
|
||||
|
||||
def initiate_polygon_zones(
|
||||
polygons: List[np.ndarray],
|
||||
frame_resolution_wh: Tuple[int, int],
|
||||
triggering_position: sv.Position = sv.Position.CENTER,
|
||||
) -> List[sv.PolygonZone]:
|
||||
return [
|
||||
sv.PolygonZone(
|
||||
polygon=polygon,
|
||||
frame_resolution_wh=frame_resolution_wh,
|
||||
triggering_position=triggering_position,
|
||||
)
|
||||
for polygon in polygons
|
||||
]
|
||||
|
||||
|
||||
class VideoProcessor:
|
||||
def __init__(
|
||||
self,
|
||||
source_weights_path: str,
|
||||
source_video_path: str,
|
||||
target_video_path: str = None,
|
||||
confidence_threshold: float = 0.3,
|
||||
iou_threshold: float = 0.7,
|
||||
) -> None:
|
||||
self.conf_threshold = confidence_threshold
|
||||
self.iou_threshold = iou_threshold
|
||||
self.source_video_path = source_video_path
|
||||
self.target_video_path = target_video_path
|
||||
|
||||
self.model = YOLO(source_weights_path)
|
||||
self.tracker = sv.ByteTrack()
|
||||
|
||||
self.video_info = sv.VideoInfo.from_video_path(source_video_path)
|
||||
self.zones_in = initiate_polygon_zones(
|
||||
ZONE_IN_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
|
||||
)
|
||||
self.zones_out = initiate_polygon_zones(
|
||||
ZONE_OUT_POLYGONS, self.video_info.resolution_wh, sv.Position.CENTER
|
||||
)
|
||||
|
||||
self.box_annotator = sv.BoxAnnotator(color=COLORS)
|
||||
self.trace_annotator = sv.TraceAnnotator(
|
||||
color=COLORS, position=sv.Position.CENTER, trace_length=100, thickness=2
|
||||
)
|
||||
self.detections_manager = DetectionsManager()
|
||||
|
||||
def process_video(self):
|
||||
frame_generator = sv.get_video_frames_generator(
|
||||
source_path=self.source_video_path
|
||||
)
|
||||
|
||||
if self.target_video_path:
|
||||
with sv.VideoSink(self.target_video_path, self.video_info) as sink:
|
||||
for frame in tqdm(frame_generator, total=self.video_info.total_frames):
|
||||
annotated_frame = self.process_frame(frame)
|
||||
sink.write_frame(annotated_frame)
|
||||
else:
|
||||
for frame in tqdm(frame_generator, total=self.video_info.total_frames):
|
||||
annotated_frame = self.process_frame(frame)
|
||||
cv2.imshow("Processed Video", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
|
||||
def annotate_frame(
|
||||
self, frame: np.ndarray, detections: sv.Detections
|
||||
) -> np.ndarray:
|
||||
annotated_frame = frame.copy()
|
||||
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
|
||||
annotated_frame = sv.draw_polygon(
|
||||
annotated_frame, zone_in.polygon, COLORS.colors[i]
|
||||
)
|
||||
annotated_frame = sv.draw_polygon(
|
||||
annotated_frame, zone_out.polygon, COLORS.colors[i]
|
||||
)
|
||||
|
||||
labels = [f"#{tracker_id}" for tracker_id in detections.tracker_id]
|
||||
annotated_frame = self.trace_annotator.annotate(annotated_frame, detections)
|
||||
annotated_frame = self.box_annotator.annotate(
|
||||
annotated_frame, detections, labels
|
||||
)
|
||||
|
||||
for zone_out_id, zone_out in enumerate(self.zones_out):
|
||||
zone_center = sv.get_polygon_center(polygon=zone_out.polygon)
|
||||
if zone_out_id in self.detections_manager.counts:
|
||||
counts = self.detections_manager.counts[zone_out_id]
|
||||
for i, zone_in_id in enumerate(counts):
|
||||
count = len(self.detections_manager.counts[zone_out_id][zone_in_id])
|
||||
text_anchor = sv.Point(x=zone_center.x, y=zone_center.y + 40 * i)
|
||||
annotated_frame = sv.draw_text(
|
||||
scene=annotated_frame,
|
||||
text=str(count),
|
||||
text_anchor=text_anchor,
|
||||
background_color=COLORS.colors[zone_in_id],
|
||||
)
|
||||
|
||||
return annotated_frame
|
||||
|
||||
def process_frame(self, frame: np.ndarray) -> np.ndarray:
|
||||
results = self.model(
|
||||
frame, verbose=False, conf=self.conf_threshold, iou=self.iou_threshold
|
||||
)[0]
|
||||
detections = sv.Detections.from_ultralytics(results)
|
||||
detections.class_id = np.zeros(len(detections))
|
||||
detections = self.tracker.update_with_detections(detections)
|
||||
|
||||
detections_in_zones = []
|
||||
detections_out_zones = []
|
||||
|
||||
for i, (zone_in, zone_out) in enumerate(zip(self.zones_in, self.zones_out)):
|
||||
detections_in_zone = detections[zone_in.trigger(detections=detections)]
|
||||
detections_in_zones.append(detections_in_zone)
|
||||
detections_out_zone = detections[zone_out.trigger(detections=detections)]
|
||||
detections_out_zones.append(detections_out_zone)
|
||||
|
||||
detections = self.detections_manager.update(
|
||||
detections, detections_in_zones, detections_out_zones
|
||||
)
|
||||
return self.annotate_frame(frame, detections)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Traffic Flow Analysis with YOLO and ByteTrack"
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--source_weights_path",
|
||||
required=True,
|
||||
help="Path to the source weights file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
default=None,
|
||||
help="Path to the target video file (output)",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--confidence_threshold",
|
||||
default=0.3,
|
||||
help="Confidence threshold for the model",
|
||||
type=float,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
processor = VideoProcessor(
|
||||
source_weights_path=args.source_weights_path,
|
||||
source_video_path=args.source_video_path,
|
||||
target_video_path=args.target_video_path,
|
||||
confidence_threshold=args.confidence_threshold,
|
||||
iou_threshold=args.iou_threshold,
|
||||
)
|
||||
processor.process_video()
|
||||
Executable
+17
@@ -0,0 +1,17 @@
|
||||
#!/bin/bash
|
||||
|
||||
# Get the directory where the script is located
|
||||
DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
|
||||
|
||||
# Check if 'data' directory does not exist and then create it
|
||||
if [[ ! -e $DIR/data ]]; then
|
||||
mkdir "$DIR/data"
|
||||
else
|
||||
echo "'data' directory already exists."
|
||||
fi
|
||||
|
||||
# Download the traffic_analysis.mov file from Google Drive
|
||||
gdown -O "$DIR/data/traffic_analysis.mov" "https://drive.google.com/uc?id=1qadBd7lgpediafCpL_yedGjQPk-FLK-W"
|
||||
|
||||
# Download the traffic_analysis.pt file from Google Drive
|
||||
gdown -O "$DIR/data/traffic_analysis.pt" "https://drive.google.com/uc?id=1y-IfToCjRXa3ZdC1JpnKRopC7mcQW-5z"
|
||||
Generated
+38
-24
@@ -493,6 +493,7 @@ files = [
|
||||
{file = "contourpy-1.1.0-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:18a64814ae7bce73925131381603fff0116e2df25230dfc80d6d690aa6e20b37"},
|
||||
{file = "contourpy-1.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:90c81f22b4f572f8a2110b0b741bb64e5a6427e0a198b2cdc1fbaf85f352a3aa"},
|
||||
{file = "contourpy-1.1.0-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:53cc3a40635abedbec7f1bde60f8c189c49e84ac180c665f2cd7c162cc454baa"},
|
||||
{file = "contourpy-1.1.0-cp310-cp310-win32.whl", hash = "sha256:9b2dd2ca3ac561aceef4c7c13ba654aaa404cf885b187427760d7f7d4c57cff8"},
|
||||
{file = "contourpy-1.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:1f795597073b09d631782e7245016a4323cf1cf0b4e06eef7ea6627e06a37ff2"},
|
||||
{file = "contourpy-1.1.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:0b7b04ed0961647691cfe5d82115dd072af7ce8846d31a5fac6c142dcce8b882"},
|
||||
{file = "contourpy-1.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:27bc79200c742f9746d7dd51a734ee326a292d77e7d94c8af6e08d1e6c15d545"},
|
||||
@@ -501,6 +502,7 @@ files = [
|
||||
{file = "contourpy-1.1.0-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e5cec36c5090e75a9ac9dbd0ff4a8cf7cecd60f1b6dc23a374c7d980a1cd710e"},
|
||||
{file = "contourpy-1.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1f0cbd657e9bde94cd0e33aa7df94fb73c1ab7799378d3b3f902eb8eb2e04a3a"},
|
||||
{file = "contourpy-1.1.0-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:181cbace49874f4358e2929aaf7ba84006acb76694102e88dd15af861996c16e"},
|
||||
{file = "contourpy-1.1.0-cp311-cp311-win32.whl", hash = "sha256:edb989d31065b1acef3828a3688f88b2abb799a7db891c9e282df5ec7e46221b"},
|
||||
{file = "contourpy-1.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:fb3b7d9e6243bfa1efb93ccfe64ec610d85cfe5aec2c25f97fbbd2e58b531256"},
|
||||
{file = "contourpy-1.1.0-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:bcb41692aa09aeb19c7c213411854402f29f6613845ad2453d30bf421fe68fed"},
|
||||
{file = "contourpy-1.1.0-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:5d123a5bc63cd34c27ff9c7ac1cd978909e9c71da12e05be0231c608048bb2ae"},
|
||||
@@ -509,6 +511,7 @@ files = [
|
||||
{file = "contourpy-1.1.0-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:317267d915490d1e84577924bd61ba71bf8681a30e0d6c545f577363157e5e94"},
|
||||
{file = "contourpy-1.1.0-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d551f3a442655f3dcc1285723f9acd646ca5858834efeab4598d706206b09c9f"},
|
||||
{file = "contourpy-1.1.0-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:e7a117ce7df5a938fe035cad481b0189049e8d92433b4b33aa7fc609344aafa1"},
|
||||
{file = "contourpy-1.1.0-cp38-cp38-win32.whl", hash = "sha256:108dfb5b3e731046a96c60bdc46a1a0ebee0760418951abecbe0fc07b5b93b27"},
|
||||
{file = "contourpy-1.1.0-cp38-cp38-win_amd64.whl", hash = "sha256:d4f26b25b4f86087e7d75e63212756c38546e70f2a92d2be44f80114826e1cd4"},
|
||||
{file = "contourpy-1.1.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:bc00bb4225d57bff7ebb634646c0ee2a1298402ec10a5fe7af79df9a51c1bfd9"},
|
||||
{file = "contourpy-1.1.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:189ceb1525eb0655ab8487a9a9c41f42a73ba52d6789754788d1883fb06b2d8a"},
|
||||
@@ -517,6 +520,7 @@ files = [
|
||||
{file = "contourpy-1.1.0-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:143dde50520a9f90e4a2703f367cf8ec96a73042b72e68fcd184e1279962eb6f"},
|
||||
{file = "contourpy-1.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e94bef2580e25b5fdb183bf98a2faa2adc5b638736b2c0a4da98691da641316a"},
|
||||
{file = "contourpy-1.1.0-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:ed614aea8462735e7d70141374bd7650afd1c3f3cb0c2dbbcbe44e14331bf002"},
|
||||
{file = "contourpy-1.1.0-cp39-cp39-win32.whl", hash = "sha256:71551f9520f008b2950bef5f16b0e3587506ef4f23c734b71ffb7b89f8721999"},
|
||||
{file = "contourpy-1.1.0-cp39-cp39-win_amd64.whl", hash = "sha256:438ba416d02f82b692e371858143970ed2eb6337d9cdbbede0d8ad9f3d7dd17d"},
|
||||
{file = "contourpy-1.1.0-pp38-pypy38_pp73-macosx_10_9_x86_64.whl", hash = "sha256:a698c6a7a432789e587168573a864a7ea374c6be8d4f31f9d87c001d5a843493"},
|
||||
{file = "contourpy-1.1.0-pp38-pypy38_pp73-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:397b0ac8a12880412da3551a8cb5a187d3298a72802b45a3bd1805e204ad8439"},
|
||||
@@ -539,34 +543,34 @@ test-no-images = ["pytest", "pytest-cov", "wurlitzer"]
|
||||
|
||||
[[package]]
|
||||
name = "cryptography"
|
||||
version = "41.0.2"
|
||||
version = "41.0.4"
|
||||
description = "cryptography is a package which provides cryptographic recipes and primitives to Python developers."
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "cryptography-41.0.2-cp37-abi3-macosx_10_12_universal2.whl", hash = "sha256:01f1d9e537f9a15b037d5d9ee442b8c22e3ae11ce65ea1f3316a41c78756b711"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:079347de771f9282fbfe0e0236c716686950c19dee1b76240ab09ce1624d76d7"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:439c3cc4c0d42fa999b83ded80a9a1fb54d53c58d6e59234cfe97f241e6c781d"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f14ad275364c8b4e525d018f6716537ae7b6d369c094805cae45300847e0894f"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:84609ade00a6ec59a89729e87a503c6e36af98ddcd566d5f3be52e29ba993182"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:49c3222bb8f8e800aead2e376cbef687bc9e3cb9b58b29a261210456a7783d83"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:d73f419a56d74fef257955f51b18d046f3506270a5fd2ac5febbfa259d6c0fa5"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:2a034bf7d9ca894720f2ec1d8b7b5832d7e363571828037f9e0c4f18c1b58a58"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-win32.whl", hash = "sha256:d124682c7a23c9764e54ca9ab5b308b14b18eba02722b8659fb238546de83a76"},
|
||||
{file = "cryptography-41.0.2-cp37-abi3-win_amd64.whl", hash = "sha256:9c3fe6534d59d071ee82081ca3d71eed3210f76ebd0361798c74abc2bcf347d4"},
|
||||
{file = "cryptography-41.0.2-pp310-pypy310_pp73-macosx_10_12_x86_64.whl", hash = "sha256:a719399b99377b218dac6cf547b6ec54e6ef20207b6165126a280b0ce97e0d2a"},
|
||||
{file = "cryptography-41.0.2-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:182be4171f9332b6741ee818ec27daff9fb00349f706629f5cbf417bd50e66fd"},
|
||||
{file = "cryptography-41.0.2-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:7a9a3bced53b7f09da251685224d6a260c3cb291768f54954e28f03ef14e3766"},
|
||||
{file = "cryptography-41.0.2-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:f0dc40e6f7aa37af01aba07277d3d64d5a03dc66d682097541ec4da03cc140ee"},
|
||||
{file = "cryptography-41.0.2-pp38-pypy38_pp73-macosx_10_12_x86_64.whl", hash = "sha256:674b669d5daa64206c38e507808aae49904c988fa0a71c935e7006a3e1e83831"},
|
||||
{file = "cryptography-41.0.2-pp38-pypy38_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:7af244b012711a26196450d34f483357e42aeddb04128885d95a69bd8b14b69b"},
|
||||
{file = "cryptography-41.0.2-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:9b6d717393dbae53d4e52684ef4f022444fc1cce3c48c38cb74fca29e1f08eaa"},
|
||||
{file = "cryptography-41.0.2-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:192255f539d7a89f2102d07d7375b1e0a81f7478925b3bc2e0549ebf739dae0e"},
|
||||
{file = "cryptography-41.0.2-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:f772610fe364372de33d76edcd313636a25684edb94cee53fd790195f5989d14"},
|
||||
{file = "cryptography-41.0.2-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:b332cba64d99a70c1e0836902720887fb4529ea49ea7f5462cf6640e095e11d2"},
|
||||
{file = "cryptography-41.0.2-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:9a6673c1828db6270b76b22cc696f40cde9043eb90373da5c2f8f2158957f42f"},
|
||||
{file = "cryptography-41.0.2-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:342f3767e25876751e14f8459ad85e77e660537ca0a066e10e75df9c9e9099f0"},
|
||||
{file = "cryptography-41.0.2.tar.gz", hash = "sha256:7d230bf856164de164ecb615ccc14c7fc6de6906ddd5b491f3af90d3514c925c"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-macosx_10_12_universal2.whl", hash = "sha256:80907d3faa55dc5434a16579952ac6da800935cd98d14dbd62f6f042c7f5e839"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-macosx_10_12_x86_64.whl", hash = "sha256:35c00f637cd0b9d5b6c6bd11b6c3359194a8eba9c46d4e875a3660e3b400005f"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cecfefa17042941f94ab54f769c8ce0fe14beff2694e9ac684176a2535bf9714"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e40211b4923ba5a6dc9769eab704bdb3fbb58d56c5b336d30996c24fcf12aadb"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:23a25c09dfd0d9f28da2352503b23e086f8e78096b9fd585d1d14eca01613e13"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2ed09183922d66c4ec5fdaa59b4d14e105c084dd0febd27452de8f6f74704143"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-musllinux_1_1_aarch64.whl", hash = "sha256:5a0f09cefded00e648a127048119f77bc2b2ec61e736660b5789e638f43cc397"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-musllinux_1_1_x86_64.whl", hash = "sha256:9eeb77214afae972a00dee47382d2591abe77bdae166bda672fb1e24702a3860"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-win32.whl", hash = "sha256:3b224890962a2d7b57cf5eeb16ccaafba6083f7b811829f00476309bce2fe0fd"},
|
||||
{file = "cryptography-41.0.4-cp37-abi3-win_amd64.whl", hash = "sha256:c880eba5175f4307129784eca96f4e70b88e57aa3f680aeba3bab0e980b0f37d"},
|
||||
{file = "cryptography-41.0.4-pp310-pypy310_pp73-macosx_10_12_x86_64.whl", hash = "sha256:004b6ccc95943f6a9ad3142cfabcc769d7ee38a3f60fb0dddbfb431f818c3a67"},
|
||||
{file = "cryptography-41.0.4-pp310-pypy310_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:86defa8d248c3fa029da68ce61fe735432b047e32179883bdb1e79ed9bb8195e"},
|
||||
{file = "cryptography-41.0.4-pp310-pypy310_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:37480760ae08065437e6573d14be973112c9e6dcaf5f11d00147ee74f37a3829"},
|
||||
{file = "cryptography-41.0.4-pp310-pypy310_pp73-win_amd64.whl", hash = "sha256:b5f4dfe950ff0479f1f00eda09c18798d4f49b98f4e2006d644b3301682ebdca"},
|
||||
{file = "cryptography-41.0.4-pp38-pypy38_pp73-macosx_10_12_x86_64.whl", hash = "sha256:7e53db173370dea832190870e975a1e09c86a879b613948f09eb49324218c14d"},
|
||||
{file = "cryptography-41.0.4-pp38-pypy38_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:5b72205a360f3b6176485a333256b9bcd48700fc755fef51c8e7e67c4b63e3ac"},
|
||||
{file = "cryptography-41.0.4-pp38-pypy38_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:93530900d14c37a46ce3d6c9e6fd35dbe5f5601bf6b3a5c325c7bffc030344d9"},
|
||||
{file = "cryptography-41.0.4-pp38-pypy38_pp73-win_amd64.whl", hash = "sha256:efc8ad4e6fc4f1752ebfb58aefece8b4e3c4cae940b0994d43649bdfce8d0d4f"},
|
||||
{file = "cryptography-41.0.4-pp39-pypy39_pp73-macosx_10_12_x86_64.whl", hash = "sha256:c3391bd8e6de35f6f1140e50aaeb3e2b3d6a9012536ca23ab0d9c35ec18c8a91"},
|
||||
{file = "cryptography-41.0.4-pp39-pypy39_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:0d9409894f495d465fe6fda92cb70e8323e9648af912d5b9141d616df40a87b8"},
|
||||
{file = "cryptography-41.0.4-pp39-pypy39_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:8ac4f9ead4bbd0bc8ab2d318f97d85147167a488be0e08814a37eb2f439d5cf6"},
|
||||
{file = "cryptography-41.0.4-pp39-pypy39_pp73-win_amd64.whl", hash = "sha256:047c4603aeb4bbd8db2756e38f5b8bd7e94318c047cfe4efeb5d715e08b49311"},
|
||||
{file = "cryptography-41.0.4.tar.gz", hash = "sha256:7febc3094125fc126a7f6fb1f420d0da639f3f32cb15c8ff0dc3997c4549f51a"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -1451,6 +1455,16 @@ files = [
|
||||
{file = "MarkupSafe-2.1.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:5bbe06f8eeafd38e5d0a4894ffec89378b6c6a625ff57e3028921f8ff59318ac"},
|
||||
{file = "MarkupSafe-2.1.3-cp311-cp311-win32.whl", hash = "sha256:dd15ff04ffd7e05ffcb7fe79f1b98041b8ea30ae9234aed2a9168b5797c3effb"},
|
||||
{file = "MarkupSafe-2.1.3-cp311-cp311-win_amd64.whl", hash = "sha256:134da1eca9ec0ae528110ccc9e48041e0828d79f24121a1a146161103c76e686"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:f698de3fd0c4e6972b92290a45bd9b1536bffe8c6759c62471efaa8acb4c37bc"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:aa57bd9cf8ae831a362185ee444e15a93ecb2e344c8e52e4d721ea3ab6ef1823"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ffcc3f7c66b5f5b7931a5aa68fc9cecc51e685ef90282f4a82f0f5e9b704ad11"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47d4f1c5f80fc62fdd7777d0d40a2e9dda0a05883ab11374334f6c4de38adffd"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1f67c7038d560d92149c060157d623c542173016c4babc0c1913cca0564b9939"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:9aad3c1755095ce347e26488214ef77e0485a3c34a50c5a5e2471dff60b9dd9c"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:14ff806850827afd6b07a5f32bd917fb7f45b046ba40c57abdb636674a8b559c"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8f9293864fe09b8149f0cc42ce56e3f0e54de883a9de90cd427f191c346eb2e1"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-win32.whl", hash = "sha256:715d3562f79d540f251b99ebd6d8baa547118974341db04f5ad06d5ea3eb8007"},
|
||||
{file = "MarkupSafe-2.1.3-cp312-cp312-win_amd64.whl", hash = "sha256:1b8dd8c3fd14349433c79fa8abeb573a55fc0fdd769133baac1f5e07abf54aeb"},
|
||||
{file = "MarkupSafe-2.1.3-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:8e254ae696c88d98da6555f5ace2279cf7cd5b3f52be2b5cf97feafe883b58d2"},
|
||||
{file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:cb0932dc158471523c9637e807d9bfb93e06a95cbf010f1a38b98623b929ef2b"},
|
||||
{file = "MarkupSafe-2.1.3-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9402b03f1a1b4dc4c19845e5c749e3ab82d5078d16a2a4c2cd2df62d57bb0707"},
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "supervision"
|
||||
version = "0.14.0"
|
||||
version = "0.15.0rc2"
|
||||
description = "A set of easy-to-use utils that will come in handy in any Computer Vision project"
|
||||
authors = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
maintainers = ["Piotr Skalski <piotr.skalski92@gmail.com>"]
|
||||
|
||||
@@ -6,14 +6,20 @@ try:
|
||||
except importlib_metadata.PackageNotFoundError:
|
||||
__version__ = "development"
|
||||
|
||||
|
||||
from supervision.annotators.core import (
|
||||
BoundingBoxAnnotator,
|
||||
BoxCornerAnnotator,
|
||||
EllipseAnnotator,
|
||||
LabelAnnotator,
|
||||
MaskAnnotator,
|
||||
)
|
||||
from supervision.classification.core import Classifications
|
||||
from supervision.dataset.core import (
|
||||
BaseDataset,
|
||||
ClassificationDataset,
|
||||
DetectionDataset,
|
||||
)
|
||||
from supervision.detection.annotate import BoxAnnotator, MaskAnnotator
|
||||
from supervision.detection.annotate import BoxAnnotator, TraceAnnotator
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.detection.line_counter import LineZone, LineZoneAnnotator
|
||||
from supervision.detection.tools.inference_slicer import InferenceSlicer
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from enum import Enum
|
||||
from typing import Union
|
||||
|
||||
import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
|
||||
|
||||
class ColorMap(Enum):
|
||||
"""
|
||||
Enum for annotator color mapping.
|
||||
"""
|
||||
|
||||
INDEX = "index"
|
||||
CLASS = "class"
|
||||
TRACK = "track"
|
||||
|
||||
|
||||
def resolve_color_idx(
|
||||
detections: Detections, detection_idx: int, color_map: ColorMap = ColorMap.CLASS
|
||||
) -> int:
|
||||
if detection_idx >= len(detections):
|
||||
raise ValueError(
|
||||
f"Detection index {detection_idx}"
|
||||
f"is out of bounds for detections of length {len(detections)}"
|
||||
)
|
||||
|
||||
if color_map == ColorMap.INDEX:
|
||||
return detection_idx
|
||||
elif color_map == ColorMap.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:
|
||||
if detections.tracker_id is None:
|
||||
raise ValueError(
|
||||
"Could not resolve color by track because"
|
||||
"Detections do not have tracker_id"
|
||||
)
|
||||
return detections.tracker_id[detection_idx]
|
||||
|
||||
|
||||
def resolve_color(color: Union[Color, ColorPalette], idx: int) -> Color:
|
||||
if isinstance(color, ColorPalette):
|
||||
return color.by_idx(idx)
|
||||
else:
|
||||
return color
|
||||
|
||||
|
||||
class BaseAnnotator(ABC):
|
||||
@abstractmethod
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
pass
|
||||
@@ -0,0 +1,478 @@
|
||||
from typing import List, Tuple, Union
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from supervision.annotators.base import (
|
||||
BaseAnnotator,
|
||||
ColorMap,
|
||||
resolve_color,
|
||||
resolve_color_idx,
|
||||
)
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
|
||||
class BoundingBoxAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing bounding boxes on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
color_map: str = "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`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> 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.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> bounding_box_annotator = sv.BoundingBoxAnnotator()
|
||||
>>> annotated_frame = bounding_box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(x1, y1),
|
||||
pt2=(x2, y2),
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class MaskAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing masks on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
opacity: float = 0.5,
|
||||
color_map: str = "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`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.opacity = opacity
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> 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.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> mask_annotator = sv.MaskAnnotator()
|
||||
>>> annotated_frame = mask_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
if detections.mask is None:
|
||||
return scene
|
||||
|
||||
for detection_idx in np.flip(np.argsort(detections.area)):
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
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()
|
||||
|
||||
scene = np.where(
|
||||
np.expand_dims(mask, axis=-1),
|
||||
np.uint8(self.opacity * colored_mask + (1 - self.opacity) * scene),
|
||||
scene,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class EllipseAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing ellipses on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
start_angle: int = -45,
|
||||
end_angle: int = 235,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
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`.
|
||||
"""
|
||||
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)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> 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.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> ellipse_annotator = sv.EllipseAnnotator()
|
||||
>>> annotated_frame = ellipse_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
|
||||
center = (int((x1 + x2) / 2), y2)
|
||||
width = x2 - x1
|
||||
cv2.ellipse(
|
||||
scene,
|
||||
center=center,
|
||||
axes=(int(width), int(0.35 * width)),
|
||||
angle=0.0,
|
||||
startAngle=self.start_angle,
|
||||
endAngle=self.end_angle,
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
lineType=cv2.LINE_4,
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class BoxCornerAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing box corners on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 4,
|
||||
corner_length: int = 25,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
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`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.corner_length: int = corner_length
|
||||
self.color_map: ColorMap = ColorMap(color_map)
|
||||
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> 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.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> corner_annotator = sv.BoxCornerAnnotator()
|
||||
>>> annotated_frame = corner_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
for detection_idx in range(len(detections)):
|
||||
x1, y1, x2, y2 = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
corners = [(x1, y1), (x2, y1), (x1, y2), (x2, y2)]
|
||||
|
||||
for x, y in corners:
|
||||
x_end = x + self.corner_length if x == x1 else x - self.corner_length
|
||||
cv2.line(
|
||||
scene, (x, y), (x_end, y), color.as_bgr(), thickness=self.thickness
|
||||
)
|
||||
|
||||
y_end = y + self.corner_length if y == y1 else y - self.corner_length
|
||||
cv2.line(
|
||||
scene, (x, y), (x, y_end), color.as_bgr(), thickness=self.thickness
|
||||
)
|
||||
return scene
|
||||
|
||||
|
||||
class LabelAnnotator:
|
||||
"""
|
||||
A class for annotating labels on an image using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
text_color: Color = Color.black(),
|
||||
text_scale: float = 0.5,
|
||||
text_thickness: int = 1,
|
||||
text_padding: int = 10,
|
||||
text_position: Position = Position.TOP_LEFT,
|
||||
color_map: str = "class",
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating the text background.
|
||||
text_color (Color): The color to use for the text.
|
||||
text_scale (float): Font scale for the text.
|
||||
text_thickness (int): Thickness of the text characters.
|
||||
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`.
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.text_color: Color = text_color
|
||||
self.text_scale: float = text_scale
|
||||
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)
|
||||
|
||||
@staticmethod
|
||||
def resolve_text_background_xyxy(
|
||||
detection_xyxy: Tuple[int, int, int, int],
|
||||
text_wh: Tuple[int, int],
|
||||
text_padding: int,
|
||||
position: Position,
|
||||
) -> Tuple[int, int, int, int]:
|
||||
padded_text_wh = (text_wh[0] + 2 * text_padding, text_wh[1] + 2 * text_padding)
|
||||
x1, y1, x2, y2 = detection_xyxy
|
||||
center_x = (x1 + x2) // 2
|
||||
center_y = (y1 + y2) // 2
|
||||
|
||||
if position == Position.TOP_LEFT:
|
||||
return x1, y1 - padded_text_wh[1], x1 + padded_text_wh[0], y1
|
||||
elif position == Position.TOP_RIGHT:
|
||||
return x2 - padded_text_wh[0], y1 - padded_text_wh[1], x2, y1
|
||||
elif position == Position.TOP_CENTER:
|
||||
return (
|
||||
center_x - padded_text_wh[0] // 2,
|
||||
y1 - padded_text_wh[1],
|
||||
center_x + padded_text_wh[0] // 2,
|
||||
y1,
|
||||
)
|
||||
elif position == Position.CENTER:
|
||||
return (
|
||||
center_x - padded_text_wh[0] // 2,
|
||||
center_y - padded_text_wh[1] // 2,
|
||||
center_x + padded_text_wh[0] // 2,
|
||||
center_y + padded_text_wh[1] // 2,
|
||||
)
|
||||
elif position == Position.BOTTOM_LEFT:
|
||||
return x1, y2, x1 + padded_text_wh[0], y2 + padded_text_wh[1]
|
||||
elif position == Position.BOTTOM_RIGHT:
|
||||
return x2 - padded_text_wh[0], y2, x2, y2 + padded_text_wh[1]
|
||||
elif position == Position.BOTTOM_CENTER:
|
||||
return (
|
||||
center_x - padded_text_wh[0] // 2,
|
||||
y2,
|
||||
center_x + padded_text_wh[0] // 2,
|
||||
y2 + padded_text_wh[1],
|
||||
)
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
labels: List[str] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with labels based on the provided detections.
|
||||
|
||||
Args:
|
||||
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.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> label_annotator = sv.LabelAnnotator(text_position=sv.Position.CENTER)
|
||||
>>> annotated_frame = label_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||

|
||||
"""
|
||||
font = cv2.FONT_HERSHEY_SIMPLEX
|
||||
for detection_idx in range(len(detections)):
|
||||
detection_xyxy = detections.xyxy[detection_idx].astype(int)
|
||||
idx = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=self.color_map,
|
||||
)
|
||||
color = resolve_color(color=self.color, idx=idx)
|
||||
text = (
|
||||
f"{detections.class_id[detection_idx]}"
|
||||
if (labels is None or len(detections) != len(labels))
|
||||
else labels[detection_idx]
|
||||
)
|
||||
text_wh = cv2.getTextSize(
|
||||
text=text,
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
thickness=self.text_thickness,
|
||||
)[0]
|
||||
|
||||
text_background_xyxy = self.resolve_text_background_xyxy(
|
||||
detection_xyxy=detection_xyxy,
|
||||
text_wh=text_wh,
|
||||
text_padding=self.text_padding,
|
||||
position=self.text_position,
|
||||
)
|
||||
|
||||
text_x = text_background_xyxy[0] + self.text_padding
|
||||
text_y = text_background_xyxy[1] + self.text_padding + text_wh[1]
|
||||
|
||||
cv2.rectangle(
|
||||
img=scene,
|
||||
pt1=(text_background_xyxy[0], text_background_xyxy[1]),
|
||||
pt2=(text_background_xyxy[2], text_background_xyxy[3]),
|
||||
color=color.as_bgr(),
|
||||
thickness=cv2.FILLED,
|
||||
)
|
||||
cv2.putText(
|
||||
img=scene,
|
||||
text=text,
|
||||
org=(text_x, text_y),
|
||||
fontFace=font,
|
||||
fontScale=self.text_scale,
|
||||
color=self.text_color.as_rgb(),
|
||||
thickness=self.text_thickness,
|
||||
lineType=cv2.LINE_AA,
|
||||
)
|
||||
return scene
|
||||
@@ -5,6 +5,7 @@ import numpy as np
|
||||
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.draw.color import Color, ColorPalette
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
|
||||
class BoxAnnotator:
|
||||
@@ -146,44 +147,100 @@ class BoxAnnotator:
|
||||
return scene
|
||||
|
||||
|
||||
class MaskAnnotator:
|
||||
class Trace:
|
||||
def __init__(
|
||||
self,
|
||||
max_size: Optional[int] = None,
|
||||
start_frame_id: int = 0,
|
||||
anchor: Position = Position.CENTER,
|
||||
) -> None:
|
||||
self.current_frame_id = start_frame_id
|
||||
self.max_size = max_size
|
||||
self.anchor = anchor
|
||||
|
||||
self.frame_id = np.array([], dtype=int)
|
||||
self.xy = np.empty((0, 2), dtype=np.float32)
|
||||
self.tracker_id = np.array([], dtype=int)
|
||||
|
||||
def put(self, detections: Detections) -> None:
|
||||
frame_id = np.full(len(detections), self.current_frame_id, dtype=int)
|
||||
self.frame_id = np.concatenate([self.frame_id, frame_id])
|
||||
self.xy = np.concatenate(
|
||||
[self.xy, detections.get_anchor_coordinates(self.anchor)]
|
||||
)
|
||||
self.tracker_id = np.concatenate([self.tracker_id, detections.tracker_id])
|
||||
|
||||
unique_frame_id = np.unique(self.frame_id)
|
||||
|
||||
if 0 < self.max_size < len(unique_frame_id):
|
||||
max_allowed_frame_id = self.current_frame_id - self.max_size + 1
|
||||
filtering_mask = self.frame_id >= max_allowed_frame_id
|
||||
self.frame_id = self.frame_id[filtering_mask]
|
||||
self.xy = self.xy[filtering_mask]
|
||||
self.tracker_id = self.tracker_id[filtering_mask]
|
||||
|
||||
self.current_frame_id += 1
|
||||
|
||||
def get(self, tracker_id: int) -> np.ndarray:
|
||||
return self.xy[self.tracker_id == tracker_id]
|
||||
|
||||
|
||||
class TraceAnnotator:
|
||||
"""
|
||||
A class for overlaying masks on an image using detections provided.
|
||||
A class for drawing trace paths on an image based on detection coordinates.
|
||||
|
||||
Attributes:
|
||||
color (Union[Color, ColorPalette]): The color to fill the mask,
|
||||
can be a single color or a color palette
|
||||
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. 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`.
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
position: Optional[Position] = Position.CENTER,
|
||||
trace_length: int = 30,
|
||||
thickness: int = 2,
|
||||
):
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.position = position
|
||||
self.trace = Trace(max_size=trace_length)
|
||||
self.thickness = thickness
|
||||
|
||||
def annotate(
|
||||
self, scene: np.ndarray, detections: Detections, opacity: float = 0.5
|
||||
) -> np.ndarray:
|
||||
def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray:
|
||||
"""
|
||||
Overlays the masks on the given image based on the provided detections,
|
||||
with a specified opacity.
|
||||
Draws trace paths on the frame based on the detection coordinates provided.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image on which the masks will be overlaid
|
||||
detections (Detections): The detections for which the
|
||||
masks will be overlaid
|
||||
opacity (float): The opacity of the masks, between 0 and 1, default is 0.5
|
||||
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.
|
||||
|
||||
Returns:
|
||||
np.ndarray: The image with the masks overlaid
|
||||
"""
|
||||
if detections.mask is None:
|
||||
return scene
|
||||
np.ndarray: The image with the trace paths drawn on it.
|
||||
|
||||
for i in np.flip(np.argsort(detections.area)):
|
||||
class_id = (
|
||||
detections.class_id[i] if detections.class_id is not None else None
|
||||
)
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> trace_annotator = sv.TraceAnnotator()
|
||||
>>> annotated_frame = trace_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
"""
|
||||
self.trace.put(detections)
|
||||
|
||||
for i, (xyxy, mask, confidence, class_id, tracker_id) in enumerate(detections):
|
||||
class_id = detections.class_id[i] if class_id is not None else None
|
||||
idx = class_id if class_id is not None else i
|
||||
color = (
|
||||
self.color.by_idx(idx)
|
||||
@@ -191,14 +248,13 @@ class MaskAnnotator:
|
||||
else self.color
|
||||
)
|
||||
|
||||
mask = detections.mask[i]
|
||||
colored_mask = np.zeros_like(scene, dtype=np.uint8)
|
||||
colored_mask[:] = color.as_bgr()
|
||||
|
||||
scene = np.where(
|
||||
np.expand_dims(mask, axis=-1),
|
||||
np.uint8(opacity * colored_mask + (1 - opacity) * scene),
|
||||
scene,
|
||||
)
|
||||
|
||||
xy = self.trace.get(tracker_id=tracker_id)
|
||||
if len(xy) > 1:
|
||||
scene = cv2.polylines(
|
||||
scene,
|
||||
[xy.astype(np.int32)],
|
||||
False,
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
return scene
|
||||
|
||||
@@ -78,7 +78,7 @@ class Detections:
|
||||
"""
|
||||
|
||||
xyxy: np.ndarray
|
||||
mask: np.Optional[np.ndarray] = None
|
||||
mask: Optional[np.ndarray] = None
|
||||
confidence: Optional[np.ndarray] = None
|
||||
class_id: Optional[np.ndarray] = None
|
||||
tracker_id: Optional[np.ndarray] = None
|
||||
@@ -178,6 +178,7 @@ class Detections:
|
||||
```
|
||||
"""
|
||||
yolov5_detections_predictions = yolov5_results.pred[0].cpu().cpu().numpy()
|
||||
|
||||
return cls(
|
||||
xyxy=yolov5_detections_predictions[:, :4],
|
||||
confidence=yolov5_detections_predictions[:, 4],
|
||||
@@ -216,6 +217,7 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_yolov8(result)
|
||||
```
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=yolov8_results.boxes.xyxy.cpu().numpy(),
|
||||
confidence=yolov8_results.boxes.conf.cpu().numpy(),
|
||||
@@ -255,6 +257,7 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_ultralytics(result)
|
||||
```
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=ultralytics_results.boxes.xyxy.cpu().numpy(),
|
||||
confidence=ultralytics_results.boxes.conf.cpu().numpy(),
|
||||
@@ -293,6 +296,9 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_yolo_nas(result)
|
||||
```
|
||||
"""
|
||||
if np.asarray(yolo_nas_results.bboxes_xyxy).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
return cls(
|
||||
xyxy=yolo_nas_results.prediction.bboxes_xyxy,
|
||||
confidence=yolo_nas_results.prediction.confidence,
|
||||
@@ -320,14 +326,16 @@ class Detections:
|
||||
|
||||
>>> yolo_pipeline = Pipeline.create(
|
||||
... task="yolo",
|
||||
... model_stub = "zoo:cv/detection/yolov5-l/pytorch/" \
|
||||
... model_path = "zoo:cv/detection/yolov5-l/pytorch/" \
|
||||
... "ultralytics/coco/pruned80_quant-none"
|
||||
>>> pipeline_outputs = yolo_pipeline(SOURCE_IMAGE_PATH,
|
||||
... iou_thres=0.6, conf_thres=0.001)
|
||||
>>> result = list(pipeline_outputs.boxes[0])
|
||||
>>> detections = sv.Detections.from_yolo_nas(result)
|
||||
>>> detections = sv.Detections.from_deepsparse(result)
|
||||
```
|
||||
"""
|
||||
if np.asarray(deepsparse_results.boxes[0]).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
return cls(
|
||||
xyxy=np.array(deepsparse_results.boxes[0]),
|
||||
confidence=np.array(deepsparse_results.scores[0]),
|
||||
@@ -360,6 +368,7 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_mmdet(mmdet_result)
|
||||
```
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=mmdet_results.pred_instances.bboxes.cpu().numpy(),
|
||||
confidence=mmdet_results.pred_instances.scores.cpu().numpy(),
|
||||
@@ -375,6 +384,7 @@ class Detections:
|
||||
Returns:
|
||||
Detections: A new Detections object.
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=transformers_results["boxes"].cpu().numpy(),
|
||||
confidence=transformers_results["scores"].cpu().numpy(),
|
||||
@@ -411,6 +421,7 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_detectron2(result)
|
||||
```
|
||||
"""
|
||||
|
||||
return cls(
|
||||
xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
|
||||
confidence=detectron2_results["instances"].scores.cpu().numpy(),
|
||||
@@ -461,7 +472,11 @@ class Detections:
|
||||
xyxy, confidence, class_id, masks = process_roboflow_result(
|
||||
roboflow_result=roboflow_result, class_list=class_list
|
||||
)
|
||||
return Detections(
|
||||
|
||||
if np.asarray(xyxy).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
return cls(
|
||||
xyxy=xyxy,
|
||||
confidence=confidence,
|
||||
class_id=class_id,
|
||||
@@ -504,7 +519,11 @@ class Detections:
|
||||
xywh = np.array([mask["bbox"] for mask in sorted_generated_masks])
|
||||
mask = np.array([mask["segmentation"] for mask in sorted_generated_masks])
|
||||
|
||||
return Detections(xyxy=xywh_to_xyxy(boxes_xywh=xywh), mask=mask)
|
||||
if np.asarray(xywh).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
xyxy = xywh_to_xyxy(boxes_xywh=xywh)
|
||||
return cls(xyxy=xyxy, mask=mask)
|
||||
|
||||
@classmethod
|
||||
def from_paddledet(cls, paddledet_result) -> Detections:
|
||||
@@ -538,6 +557,10 @@ class Detections:
|
||||
>>> detections = sv.Detections.from_paddledet(paddledet_result)
|
||||
```
|
||||
"""
|
||||
|
||||
if np.asarray(paddledet_result["bbox"][:, 2:6]).shape[0] == 0:
|
||||
return cls.empty()
|
||||
|
||||
return cls(
|
||||
xyxy=paddledet_result["bbox"][:, 2:6],
|
||||
confidence=paddledet_result["bbox"][:, 1],
|
||||
@@ -621,15 +644,23 @@ class Detections:
|
||||
|
||||
def get_anchor_coordinates(self, anchor: Position) -> np.ndarray:
|
||||
"""
|
||||
Returns the bounding box coordinates for a specific anchor.
|
||||
Calculates and returns the coordinates of a specific anchor point
|
||||
within the bounding boxes defined by the `xyxy` attribute. The anchor
|
||||
point can be any of the predefined positions in the `Position` enum,
|
||||
such as `CENTER`, `CENTER_LEFT`, `BOTTOM_RIGHT`, etc.
|
||||
|
||||
Args:
|
||||
anchor (Position): Position of bounding box anchor
|
||||
for which to return the coordinates.
|
||||
anchor (Position): An enum specifying the position of the anchor point
|
||||
within the bounding box. Supported positions are defined in the
|
||||
`Position` enum.
|
||||
|
||||
Returns:
|
||||
np.ndarray: An array of shape `(n, 2)` containing the bounding
|
||||
box anchor coordinates in format `[x, y]`.
|
||||
np.ndarray: An array of shape `(n, 2)`, where `n` is the number of bounding
|
||||
boxes. Each row contains the `[x, y]` coordinates of the specified
|
||||
anchor point for the corresponding bounding box.
|
||||
|
||||
Raises:
|
||||
ValueError: If the provided `anchor` is not supported.
|
||||
"""
|
||||
if anchor == Position.CENTER:
|
||||
return np.array(
|
||||
@@ -638,10 +669,36 @@ class Detections:
|
||||
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
|
||||
]
|
||||
).transpose()
|
||||
elif anchor == Position.CENTER_LEFT:
|
||||
return np.array(
|
||||
[
|
||||
self.xyxy[:, 0],
|
||||
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
|
||||
]
|
||||
).transpose()
|
||||
elif anchor == Position.CENTER_RIGHT:
|
||||
return np.array(
|
||||
[
|
||||
self.xyxy[:, 2],
|
||||
(self.xyxy[:, 1] + self.xyxy[:, 3]) / 2,
|
||||
]
|
||||
).transpose()
|
||||
elif anchor == Position.BOTTOM_CENTER:
|
||||
return np.array(
|
||||
[(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 3]]
|
||||
).transpose()
|
||||
elif anchor == Position.BOTTOM_LEFT:
|
||||
return np.array([self.xyxy[:, 0], self.xyxy[:, 3]]).transpose()
|
||||
elif anchor == Position.BOTTOM_RIGHT:
|
||||
return np.array([self.xyxy[:, 2], self.xyxy[:, 3]]).transpose()
|
||||
elif anchor == Position.TOP_CENTER:
|
||||
return np.array(
|
||||
[(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 1]]
|
||||
).transpose()
|
||||
elif anchor == Position.TOP_LEFT:
|
||||
return np.array([self.xyxy[:, 0], self.xyxy[:, 1]]).transpose()
|
||||
elif anchor == Position.TOP_RIGHT:
|
||||
return np.array([self.xyxy[:, 2], self.xyxy[:, 1]]).transpose()
|
||||
|
||||
raise ValueError(f"{anchor} is not supported.")
|
||||
|
||||
|
||||
@@ -23,9 +23,10 @@ def move_detections(detections: Detections, offset: np.array) -> Detections:
|
||||
class InferenceSlicer:
|
||||
"""
|
||||
InferenceSlicer performs slicing-based inference for small target detection. This
|
||||
method, often referred to as Slicing Adaptive Inference (SAHI), involves dividing a
|
||||
larger image into smaller slices, performing inference on each slice, and then
|
||||
merging the detections.
|
||||
method, often referred to as
|
||||
[Slicing Adaptive Inference (SAHI)](https://ieeexplore.ieee.org/document/9897990),
|
||||
involves dividing a larger image into smaller slices, performing inference on each
|
||||
slice, and then merging the detections.
|
||||
|
||||
Attributes:
|
||||
slice_wh (Tuple[int, int]): Dimensions of each slice in the format
|
||||
|
||||
@@ -63,7 +63,7 @@ class PolygonZone:
|
||||
clipped_detections.get_anchor_coordinates(anchor=self.triggering_position)
|
||||
).astype(int)
|
||||
is_in_zone = self.mask[clipped_anchors[:, 1], clipped_anchors[:, 0]]
|
||||
self.current_count = np.sum(is_in_zone)
|
||||
self.current_count = int(np.sum(is_in_zone))
|
||||
return is_in_zone.astype(bool)
|
||||
|
||||
|
||||
|
||||
@@ -6,8 +6,19 @@ from typing import Tuple
|
||||
|
||||
|
||||
class Position(Enum):
|
||||
"""
|
||||
Enum representing the position of an anchor point.
|
||||
"""
|
||||
|
||||
CENTER = "CENTER"
|
||||
CENTER_LEFT = "CENTER_LEFT"
|
||||
CENTER_RIGHT = "CENTER_RIGHT"
|
||||
TOP_CENTER = "TOP_CENTER"
|
||||
TOP_LEFT = "TOP_LEFT"
|
||||
TOP_RIGHT = "TOP_RIGHT"
|
||||
BOTTOM_LEFT = "BOTTOM_LEFT"
|
||||
BOTTOM_CENTER = "BOTTOM_CENTER"
|
||||
BOTTOM_RIGHT = "BOTTOM_RIGHT"
|
||||
|
||||
@classmethod
|
||||
def list(cls):
|
||||
|
||||
@@ -65,6 +65,7 @@ class VideoSink:
|
||||
target_path (str): The path to the output file where the video will be saved.
|
||||
video_info (VideoInfo): Information about the video resolution, fps,
|
||||
and total frame count.
|
||||
codec (str): FOURCC code for video format
|
||||
|
||||
Example:
|
||||
```python
|
||||
@@ -73,20 +74,26 @@ class VideoSink:
|
||||
>>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4')
|
||||
|
||||
>>> with sv.VideoSink(target_path='target_video.mp4',
|
||||
... video_info=video_info) as sink:
|
||||
... video_info=video_info,
|
||||
codec='H264') as sink:
|
||||
... for frame in get_video_frames_generator(source_path='source_video.mp4',
|
||||
... stride=2):
|
||||
... sink.write_frame(frame=frame)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, target_path: str, video_info: VideoInfo):
|
||||
def __init__(self, target_path: str, video_info: VideoInfo, codec: str = "mp4v"):
|
||||
self.target_path = target_path
|
||||
self.video_info = video_info
|
||||
self.__fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
self.__codec = codec
|
||||
self.__writer = None
|
||||
|
||||
def __enter__(self):
|
||||
try:
|
||||
self.__fourcc = cv2.VideoWriter_fourcc(*self.__codec)
|
||||
except TypeError as e:
|
||||
print(str(e) + ". Defaulting to mp4v...")
|
||||
self.__fourcc = cv2.VideoWriter_fourcc(*"mp4v")
|
||||
self.__writer = cv2.VideoWriter(
|
||||
self.target_path,
|
||||
self.__fourcc,
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
from contextlib import ExitStack as DoesNotRaise
|
||||
from test.utils import mock_detections
|
||||
from typing import Optional
|
||||
|
||||
import pytest
|
||||
|
||||
from supervision.annotators.core import ColorMap, resolve_color_idx
|
||||
from supervision.detection.core import Detections
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, detection_idx, color_map, expected_result, exception",
|
||||
[
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.INDEX,
|
||||
0,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; index mapping
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.CLASS,
|
||||
5,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; class mapping
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
0,
|
||||
ColorMap.TRACK,
|
||||
2,
|
||||
DoesNotRaise(),
|
||||
), # multiple detections; track mapping
|
||||
(
|
||||
Detections.empty(),
|
||||
0,
|
||||
ColorMap.INDEX,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # no detections; index mapping; out of bounds
|
||||
(
|
||||
mock_detections(
|
||||
xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]],
|
||||
class_id=[5, 3],
|
||||
tracker_id=[2, 6],
|
||||
),
|
||||
2,
|
||||
ColorMap.INDEX,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; index mapping; out of bounds
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
ColorMap.CLASS,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; class mapping; no class_id
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
0,
|
||||
ColorMap.TRACK,
|
||||
None,
|
||||
pytest.raises(ValueError),
|
||||
), # multiple detections; class mapping; no track_id
|
||||
],
|
||||
)
|
||||
def test_resolve_color_idx(
|
||||
detections: Detections,
|
||||
detection_idx: int,
|
||||
color_map: ColorMap,
|
||||
expected_result: Optional[int],
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
with exception:
|
||||
result = resolve_color_idx(
|
||||
detections=detections,
|
||||
detection_idx=detection_idx,
|
||||
color_map=color_map,
|
||||
)
|
||||
assert result == expected_result
|
||||
@@ -5,7 +5,8 @@ from typing import List, Optional, Union
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from supervision import Detections
|
||||
from supervision.detection.core import Detections
|
||||
from supervision.geometry.core import Position
|
||||
|
||||
PREDICTIONS = np.array(
|
||||
[
|
||||
@@ -190,3 +191,85 @@ def test_merge(
|
||||
with exception:
|
||||
result = Detections.merge(detections_list=detections_list)
|
||||
assert result == expected_result
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"detections, anchor, expected_result, exception",
|
||||
[
|
||||
(
|
||||
Detections.empty(),
|
||||
Position.CENTER,
|
||||
np.empty((0, 2), dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # empty detections
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20]]),
|
||||
Position.CENTER,
|
||||
np.array([[15, 15]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # single detection; center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.CENTER,
|
||||
np.array([[15, 15], [25, 25]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.CENTER_LEFT,
|
||||
np.array([[10, 15], [20, 25]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; center left anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.CENTER_RIGHT,
|
||||
np.array([[20, 15], [30, 25]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; center right anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.TOP_CENTER,
|
||||
np.array([[15, 10], [25, 20]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; top center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.TOP_LEFT,
|
||||
np.array([[10, 10], [20, 20]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; top left anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.TOP_RIGHT,
|
||||
np.array([[20, 10], [30, 20]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; top right anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.BOTTOM_CENTER,
|
||||
np.array([[15, 20], [25, 30]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; bottom center anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.BOTTOM_LEFT,
|
||||
np.array([[10, 20], [20, 30]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; bottom left anchor
|
||||
(
|
||||
mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]),
|
||||
Position.BOTTOM_RIGHT,
|
||||
np.array([[20, 20], [30, 30]], dtype=np.float32),
|
||||
DoesNotRaise(),
|
||||
), # two detections; bottom right anchor
|
||||
],
|
||||
)
|
||||
def test_get_anchor_coordinates(
|
||||
detections: Detections,
|
||||
anchor: Position,
|
||||
expected_result: np.ndarray,
|
||||
exception: Exception,
|
||||
) -> None:
|
||||
result = detections.get_anchor_coordinates(anchor)
|
||||
with exception:
|
||||
assert np.array_equal(result, expected_result)
|
||||
|
||||
Reference in New Issue
Block a user