Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2ba6a22d69 | ||
|
|
3de2efea12 | ||
|
|
f7145aa6aa | ||
|
|
b12d73a883 | ||
|
|
3b531249bc | ||
|
|
624a2e4c10 | ||
|
|
4c063acd2b | ||
|
|
0b805b7a04 | ||
|
|
a8081bec05 | ||
|
|
68df3604ea | ||
|
|
3224032542 | ||
|
|
1005de429c | ||
|
|
1f1edfc8f8 | ||
|
|
338d62b2cb | ||
|
|
17f7949b96 | ||
|
|
5ceb61a185 | ||
|
|
27d6832507 | ||
|
|
8acc0a9ed8 | ||
|
|
97d9472f42 | ||
|
|
f3fa12d169 | ||
|
|
9b5bba1422 | ||
|
|
a337ee591a | ||
|
|
a1da604798 | ||
|
|
1d4c3f3a68 | ||
|
|
a72ad42628 | ||
|
|
be5748c086 | ||
|
|
9e05a545c7 | ||
|
|
f0142344ae | ||
|
|
d596d10b3f | ||
|
|
86134993c8 | ||
|
|
8bab49157d | ||
|
|
4ae9131f0a | ||
|
|
69894c82af | ||
|
|
e009402934 | ||
|
|
351ef464d9 | ||
|
|
949e326559 | ||
|
|
0a7e116d7f | ||
|
|
5ce65a2fd5 | ||
|
|
73ae555e10 | ||
|
|
98c0b8c2a5 | ||
|
|
7b7a734f2c | ||
|
|
3e6c1799ef | ||
|
|
d83547ec80 | ||
|
|
7e8d44ac35 | ||
|
|
2e9b2a1bcb | ||
|
|
5df9730b3f | ||
|
|
4d955c8531 | ||
|
|
bbc9374f06 | ||
|
|
921256548e | ||
|
|
b68c019fe1 | ||
|
|
055d06b56f | ||
|
|
5274147856 | ||
|
|
2b8c31a481 | ||
|
|
0bd9862f8a | ||
|
|
587ed25f49 | ||
|
|
df6b628a30 | ||
|
|
f0e7c0c847 |
@@ -1,28 +0,0 @@
|
||||
attrs==23.1.0
|
||||
certifi==2023.7.22
|
||||
charset-normalizer==2.0.12
|
||||
cycler==0.12.1
|
||||
exceptiongroup==1.1.3
|
||||
fonttools==4.43.1
|
||||
idna==3.4
|
||||
iniconfig==2.0.0
|
||||
kiwisolver==1.4.5
|
||||
matplotlib==3.5.0
|
||||
numpy==1.21.2
|
||||
opencv-python==4.5.5.64
|
||||
packaging==23.2
|
||||
Pillow==10.1.0
|
||||
pluggy==1.3.0
|
||||
pyparsing==3.1.1
|
||||
pytest==7.2.0
|
||||
python-dateutil==2.8.2
|
||||
PyYAML==5.3
|
||||
requests==2.26.0
|
||||
scipy==1.9.0
|
||||
setuptools-scm==8.0.4
|
||||
six==1.16.0
|
||||
tomli==2.0.1
|
||||
tqdm==4.62.3
|
||||
typing_extensions==4.8.0
|
||||
urllib3==1.26.18
|
||||
defusedxml == 0.7.1
|
||||
@@ -20,13 +20,39 @@ jobs:
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
check-latest: true
|
||||
cache: 'pip'
|
||||
cache-dependency-path: '**/requirements-dev-min.txt'
|
||||
|
||||
- name: 📦 Install dependencies
|
||||
run: |
|
||||
python -m pip install --upgrade pip
|
||||
pip install -r .github/requirements-dev-min.txt
|
||||
pip install \
|
||||
attrs==23.1.0 \
|
||||
certifi==2023.7.22 \
|
||||
charset-normalizer==2.0.12 \
|
||||
cycler==0.12.1 \
|
||||
exceptiongroup==1.1.3 \
|
||||
fonttools==4.43.1 \
|
||||
idna==3.4 \
|
||||
iniconfig==2.0.0 \
|
||||
kiwisolver==1.4.5 \
|
||||
matplotlib==3.5.0 \
|
||||
numpy==1.21.2 \
|
||||
opencv-python==4.5.5.64 \
|
||||
packaging==23.2 \
|
||||
Pillow==10.1.0 \
|
||||
pluggy==1.3.0 \
|
||||
pyparsing==3.1.1 \
|
||||
pytest==7.2.0 \
|
||||
python-dateutil==2.8.2 \
|
||||
PyYAML==5.3 \
|
||||
requests==2.26.0 \
|
||||
scipy==1.9.0 \
|
||||
setuptools-scm==8.0.4 \
|
||||
six==1.16.0 \
|
||||
tomli==2.0.1 \
|
||||
tqdm==4.62.3 \
|
||||
typing_extensions==4.8.0 \
|
||||
urllib3==1.26.18 \
|
||||
defusedxml==0.7.1
|
||||
|
||||
- name: 🧪 Test
|
||||
run: "python -m pytest ./test"
|
||||
|
||||
@@ -46,7 +46,7 @@ repos:
|
||||
|
||||
|
||||
- repo: https://github.com/astral-sh/ruff-pre-commit
|
||||
rev: v0.1.9
|
||||
rev: v0.1.13
|
||||
hooks:
|
||||
- id: ruff
|
||||
args: [--fix, --exit-non-zero-on-fix]
|
||||
|
||||
@@ -216,18 +216,18 @@ Supervision provides a set of [utils](https://supervision.roboflow.com/datasets/
|
||||
## 🎬 tutorials
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><img src="https://github.com/roboflow/supervision/assets/26109316/54afdf1c-218c-4451-8f12-627fb85f1682" alt="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/4Q3ut7vqD5o" title="Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking"><strong>Traffic Analysis with YOLOv8 and ByteTrack - Vehicle Detection and Tracking</strong></a>
|
||||
<div><strong>Created: 6 Sep 2023</strong> | <strong>Updated: 6 Sep 2023</strong></div>
|
||||
<br/> In this video, we explore real-time traffic analysis using YOLOv8 and ByteTrack to detect and track vehicles on aerial images. Harnessing the power of Python and Supervision, we delve deep into assigning cars to specific entry zones and understanding their direction of movement. By visualizing their paths, we gain insights into traffic flow across bustling roundabouts... </p>
|
||||
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/61a444c8-b135-48ce-b979-2a5ab47c5a91" alt="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/uWP6UjDeZvY" title="Speed Estimation & Vehicle Tracking | Computer Vision | Open Source"><strong>Speed Estimation & Vehicle Tracking | Computer Vision | Open Source</strong></a>
|
||||
<div><strong>Created: 11 Jan 2024</strong> | <strong>Updated: 11 Jan 2024</strong></div>
|
||||
<br/> Learn how to track and estimate the speed of vehicles using YOLO, ByteTrack, and Roboflow Inference. This comprehensive tutorial covers object detection, multi-object tracking, filtering detections, perspective transformation, speed estimation, visualization improvements, and more.</p>
|
||||
|
||||
<br/>
|
||||
|
||||
<p align="left">
|
||||
<a href="https://youtu.be/D-D6ZmadzPE" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><img src="https://github.com/SkalskiP/SkalskiP/assets/26109316/6913ff11-53c6-4341-8d90-eaff3023c3fd" alt="SAM - Segment Anything Model by Meta AI: Complete Guide" width="300px" align="left" /></a>
|
||||
<a href="https://youtu.be/D-D6ZmadzPE" title="SAM - Segment Anything Model by Meta AI: Complete Guide"><strong>SAM - Segment Anything Model by Meta AI: Complete Guide</strong></a>
|
||||
<div><strong>Created: 11 Apr 2023</strong> | <strong>Updated: 11 Apr 2023</strong></div>
|
||||
<br/> Discover the incredible potential of Meta AI's Segment Anything Model (SAM)! We dive into SAM, an efficient and promptable model for image segmentation, which has revolutionized computer vision tasks. With over 1 billion masks on 11M licensed and privacy-respecting images, SAM's zero-shot performance is often competitive with or even superior to prior fully supervised results... </p>
|
||||
<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>
|
||||
|
||||
## 💜 built with supervision
|
||||
|
||||
|
||||
@@ -19,6 +19,27 @@
|
||||
|
||||
</div>
|
||||
|
||||
=== "RoundBox"
|
||||
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> round_box_annotator = sv.RoundBoxAnnotator()
|
||||
>>> annotated_frame = round_box_annotator.annotate(
|
||||
... scene=image.copy(),
|
||||
... detections=detections
|
||||
... )
|
||||
```
|
||||
|
||||
<div class="result" markdown>
|
||||
|
||||
{ align=center width="800" }
|
||||
|
||||
</div>
|
||||
|
||||
=== "BoxCorner"
|
||||
|
||||
```python
|
||||
@@ -391,6 +412,10 @@
|
||||
|
||||
:::supervision.annotators.core.TraceAnnotator
|
||||
|
||||
## RoundBoxAnnotator
|
||||
|
||||
:::supervision.annotators.core.RoundBoxAnnotator
|
||||
|
||||
## ColorLookup
|
||||
|
||||
:::supervision.annotators.utils.ColorLookup
|
||||
|
||||
@@ -7,6 +7,7 @@ interfaces with diverse applications.
|
||||
- [tracking](./tracking) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [count people in zone](./count_people_in_zone) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [traffic analysis](./traffic_analysis) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [speed estimation](./speed_estimation) by [@SkalskiP](https://github.com/SkalskiP)
|
||||
- [heatmap and track](./heatmap_and_track/) by [@HinePo](https://github.com/HinePo)
|
||||
|
||||
## contributing
|
||||
|
||||
@@ -0,0 +1,2 @@
|
||||
data/
|
||||
venv*/
|
||||
@@ -0,0 +1,111 @@
|
||||
# speed estimation
|
||||
|
||||
## 👋 hello
|
||||
|
||||
This example performs speed estimation analysis using various object-detection models
|
||||
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/0542fd3c-bb5f-475e-b96c-793560abeb18
|
||||
|
||||
## 💻 install
|
||||
|
||||
> [!NOTE]
|
||||
> YOLO-NAS is compatible with Python versions up to and including Python 3.10.
|
||||
|
||||
- clone repository and navigate to example directory
|
||||
|
||||
```bash
|
||||
git clone https://github.com/roboflow/supervision.git
|
||||
cd supervision/examples/speed_estimation
|
||||
```
|
||||
|
||||
- setup python environment and activate it [optional]
|
||||
|
||||
```bash
|
||||
python3.10 -m venv venv
|
||||
source venv/bin/activate
|
||||
```
|
||||
|
||||
- install required dependencies
|
||||
|
||||
```bash
|
||||
pip install -r requirements.txt
|
||||
```
|
||||
|
||||
- download `vehicles.mp4` file
|
||||
|
||||
```bash
|
||||
python3.10 video_downloader.py
|
||||
```
|
||||
|
||||
## 🛠️ script arguments
|
||||
|
||||
- `--roboflow_api_key` (optional): The API key for Roboflow services. If not provided
|
||||
directly, the script tries to fetch it from the `ROBOFLOW_API_KEY` environment
|
||||
variable. Follow [this guide](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key)
|
||||
to acquire your `API KEY`.
|
||||
- `--model_id` (optional): Designates the Roboflow model ID to be used. The default
|
||||
value is `"yolov8x-1280"`.
|
||||
|
||||
- `--source_weights_path`: Required. Specifies the path to the YOLO model's weights
|
||||
file, which is essential for the object detection process. This file contains the
|
||||
data that the model uses to identify objects in the video.
|
||||
- `--source_video_path`: Required. The path to the source video file that will be
|
||||
analyzed. This is the input video on which traffic flow analysis will be performed.
|
||||
- `--target_video_path`: The path to save the output video with
|
||||
annotations. If not specified, the processed video will be displayed in real-time
|
||||
without being saved.
|
||||
- `--confidence_threshold` (optional): Sets the confidence threshold for the YOLO
|
||||
model to filter detections. Default is `0.3`. This determines how confident the
|
||||
model should be to recognize an object in the video.
|
||||
- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold
|
||||
for the model. Default is 0.7. This value is used to manage object detection
|
||||
accuracy, particularly in distinguishing between different objects.
|
||||
|
||||
## ⚙️ run
|
||||
|
||||
- yolo-nas
|
||||
|
||||
```bash
|
||||
python yolo_nas_example.py \
|
||||
--source_video_path data/vehicles.mp4 \
|
||||
--target_video_path data/vehicles-result.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
- inference
|
||||
|
||||
```bash
|
||||
python inference_example.py \
|
||||
--roboflow_api_key <ROBOFLOW API KEY> \
|
||||
--source_video_path data/vehicles.mp4 \
|
||||
--target_video_path data/vehicles-result.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
- ultralytics
|
||||
|
||||
```bash
|
||||
python ultralytics_example.py \
|
||||
--source_video_path data/vehicles.mp4 \
|
||||
--target_video_path data/vehicles-result.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
|
||||
## © license
|
||||
|
||||
This demo integrates two main components, each with its own licensing:
|
||||
|
||||
- ultralytics: The object detection model used in this demo, YOLOv8, is distributed
|
||||
under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE).
|
||||
You can find more details about this license here.
|
||||
|
||||
- supervision: The analytics code that powers the zone-based analysis in this demo is
|
||||
based on the Supervision library, which is licensed under the
|
||||
[MIT license](https://github.com/roboflow/supervision/blob/develop/LICENSE.md). This
|
||||
makes the Supervision part of the code fully open source and freely usable in your
|
||||
projects.
|
||||
@@ -0,0 +1,167 @@
|
||||
import argparse
|
||||
import os
|
||||
from collections import defaultdict, deque
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from inference.models.utils import get_roboflow_model
|
||||
|
||||
import supervision as sv
|
||||
|
||||
SOURCE = np.array([[1252, 787], [2298, 803], [5039, 2159], [-550, 2159]])
|
||||
|
||||
TARGET_WIDTH = 25
|
||||
TARGET_HEIGHT = 250
|
||||
|
||||
TARGET = np.array(
|
||||
[
|
||||
[0, 0],
|
||||
[TARGET_WIDTH - 1, 0],
|
||||
[TARGET_WIDTH - 1, TARGET_HEIGHT - 1],
|
||||
[0, TARGET_HEIGHT - 1],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class ViewTransformer:
|
||||
def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
|
||||
source = source.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
self.m = cv2.getPerspectiveTransform(source, target)
|
||||
|
||||
def transform_points(self, points: np.ndarray) -> np.ndarray:
|
||||
reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
|
||||
transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
|
||||
return transformed_points.reshape(-1, 2)
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Vehicle Speed Estimation using Inference and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model_id",
|
||||
default="yolov8x-640",
|
||||
help="Roboflow model ID",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--roboflow_api_key",
|
||||
default=None,
|
||||
help="Roboflow API KEY",
|
||||
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",
|
||||
required=True,
|
||||
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
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_arguments()
|
||||
|
||||
api_key = args.roboflow_api_key
|
||||
api_key = os.environ.get("ROBOFLOW_API_KEY", api_key)
|
||||
if api_key is None:
|
||||
raise ValueError(
|
||||
"Roboflow API key is missing. Please provide it as an argument or set the "
|
||||
"ROBOFLOW_API_KEY environment variable."
|
||||
)
|
||||
args.roboflow_api_key = api_key
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
|
||||
model = get_roboflow_model(model_id=args.model_id, api_key=args.roboflow_api_key)
|
||||
|
||||
byte_track = sv.ByteTrack(
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
text_position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
trace_annotator = sv.TraceAnnotator(
|
||||
thickness=thickness,
|
||||
trace_length=video_info.fps * 2,
|
||||
position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in frame_generator:
|
||||
results = model.infer(frame)[0]
|
||||
detections = sv.Detections.from_inference(results)
|
||||
detections = detections[detections.confidence > args.confidence_threshold]
|
||||
detections = detections[polygon_zone.trigger(detections)]
|
||||
detections = detections.with_nms(threshold=args.iou_threshold)
|
||||
detections = byte_track.update_with_detections(detections=detections)
|
||||
|
||||
points = detections.get_anchors_coordinates(
|
||||
anchor=sv.Position.BOTTOM_CENTER
|
||||
)
|
||||
points = view_transformer.transform_points(points=points).astype(int)
|
||||
|
||||
for tracker_id, [_, y] in zip(detections.tracker_id, points):
|
||||
coordinates[tracker_id].append(y)
|
||||
|
||||
labels = []
|
||||
for tracker_id in detections.tracker_id:
|
||||
if len(coordinates[tracker_id]) < video_info.fps / 2:
|
||||
labels.append(f"#{tracker_id}")
|
||||
else:
|
||||
coordinate_start = coordinates[tracker_id][-1]
|
||||
coordinate_end = coordinates[tracker_id][0]
|
||||
distance = abs(coordinate_start - coordinate_end)
|
||||
time = len(coordinates[tracker_id]) / video_info.fps
|
||||
speed = distance / time * 3.6
|
||||
labels.append(f"#{tracker_id} {int(speed)} km/h")
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels
|
||||
)
|
||||
|
||||
sink.write_frame(annotated_frame)
|
||||
cv2.imshow("frame", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,6 @@
|
||||
supervision==0.18.0rc1
|
||||
tqdm
|
||||
requests
|
||||
ultralytics
|
||||
super-gradients
|
||||
inference
|
||||
@@ -0,0 +1,145 @@
|
||||
import argparse
|
||||
from collections import defaultdict, deque
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from ultralytics import YOLO
|
||||
|
||||
import supervision as sv
|
||||
|
||||
SOURCE = np.array([[1252, 787], [2298, 803], [5039, 2159], [-550, 2159]])
|
||||
|
||||
TARGET_WIDTH = 25
|
||||
TARGET_HEIGHT = 250
|
||||
|
||||
TARGET = np.array(
|
||||
[
|
||||
[0, 0],
|
||||
[TARGET_WIDTH - 1, 0],
|
||||
[TARGET_WIDTH - 1, TARGET_HEIGHT - 1],
|
||||
[0, TARGET_HEIGHT - 1],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class ViewTransformer:
|
||||
def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
|
||||
source = source.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
self.m = cv2.getPerspectiveTransform(source, target)
|
||||
|
||||
def transform_points(self, points: np.ndarray) -> np.ndarray:
|
||||
reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
|
||||
transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
|
||||
return transformed_points.reshape(-1, 2)
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Vehicle Speed Estimation using Ultralytics and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
required=True,
|
||||
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
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_arguments()
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
|
||||
model = YOLO("yolov8x.pt")
|
||||
|
||||
byte_track = sv.ByteTrack(
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
text_position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
trace_annotator = sv.TraceAnnotator(
|
||||
thickness=thickness,
|
||||
trace_length=video_info.fps * 2,
|
||||
position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in frame_generator:
|
||||
result = model(frame)[0]
|
||||
detections = sv.Detections.from_ultralytics(result)
|
||||
detections = detections[detections.confidence > args.confidence_threshold]
|
||||
detections = detections[polygon_zone.trigger(detections)]
|
||||
detections = detections.with_nms(threshold=args.iou_threshold)
|
||||
detections = byte_track.update_with_detections(detections=detections)
|
||||
|
||||
points = detections.get_anchors_coordinates(
|
||||
anchor=sv.Position.BOTTOM_CENTER
|
||||
)
|
||||
points = view_transformer.transform_points(points=points).astype(int)
|
||||
|
||||
for tracker_id, [_, y] in zip(detections.tracker_id, points):
|
||||
coordinates[tracker_id].append(y)
|
||||
|
||||
labels = []
|
||||
for tracker_id in detections.tracker_id:
|
||||
if len(coordinates[tracker_id]) < video_info.fps / 2:
|
||||
labels.append(f"#{tracker_id}")
|
||||
else:
|
||||
coordinate_start = coordinates[tracker_id][-1]
|
||||
coordinate_end = coordinates[tracker_id][0]
|
||||
distance = abs(coordinate_start - coordinate_end)
|
||||
time = len(coordinates[tracker_id]) / video_info.fps
|
||||
speed = distance / time * 3.6
|
||||
labels.append(f"#{tracker_id} {int(speed)} km/h")
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels
|
||||
)
|
||||
|
||||
sink.write_frame(annotated_frame)
|
||||
cv2.imshow("frame", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
@@ -0,0 +1,8 @@
|
||||
import os
|
||||
|
||||
from supervision.assets import VideoAssets, download_assets
|
||||
|
||||
if not os.path.exists("data"):
|
||||
os.makedirs("data")
|
||||
os.chdir("data")
|
||||
download_assets(VideoAssets.VEHICLES)
|
||||
@@ -0,0 +1,145 @@
|
||||
import argparse
|
||||
from collections import defaultdict, deque
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from super_gradients.common.object_names import Models
|
||||
from super_gradients.training import models
|
||||
|
||||
import supervision as sv
|
||||
|
||||
SOURCE = np.array([[1252, 787], [2298, 803], [5039, 2159], [-550, 2159]])
|
||||
|
||||
TARGET_WIDTH = 25
|
||||
TARGET_HEIGHT = 250
|
||||
|
||||
TARGET = np.array(
|
||||
[
|
||||
[0, 0],
|
||||
[TARGET_WIDTH - 1, 0],
|
||||
[TARGET_WIDTH - 1, TARGET_HEIGHT - 1],
|
||||
[0, TARGET_HEIGHT - 1],
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
class ViewTransformer:
|
||||
def __init__(self, source: np.ndarray, target: np.ndarray) -> None:
|
||||
source = source.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
self.m = cv2.getPerspectiveTransform(source, target)
|
||||
|
||||
def transform_points(self, points: np.ndarray) -> np.ndarray:
|
||||
reshaped_points = points.reshape(-1, 1, 2).astype(np.float32)
|
||||
transformed_points = cv2.perspectiveTransform(reshaped_points, self.m)
|
||||
return transformed_points.reshape(-1, 2)
|
||||
|
||||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Vehicle Speed Estimation using YOLO-NAS and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
required=True,
|
||||
help="Path to the source video file",
|
||||
type=str,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target_video_path",
|
||||
required=True,
|
||||
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
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_arguments()
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
|
||||
model = models.get(Models.YOLO_NAS_L, pretrained_weights="coco")
|
||||
|
||||
byte_track = sv.ByteTrack(
|
||||
frame_rate=video_info.fps, track_thresh=args.confidence_threshold
|
||||
)
|
||||
|
||||
thickness = sv.calculate_dynamic_line_thickness(
|
||||
resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
text_scale = sv.calculate_dynamic_text_scale(resolution_wh=video_info.resolution_wh)
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
text_thickness=thickness,
|
||||
text_position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
trace_annotator = sv.TraceAnnotator(
|
||||
thickness=thickness,
|
||||
trace_length=video_info.fps * 2,
|
||||
position=sv.Position.BOTTOM_CENTER,
|
||||
)
|
||||
|
||||
frame_generator = sv.get_video_frames_generator(source_path=args.source_video_path)
|
||||
|
||||
polygon_zone = sv.PolygonZone(
|
||||
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
|
||||
)
|
||||
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
|
||||
|
||||
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
|
||||
|
||||
with sv.VideoSink(args.target_video_path, video_info) as sink:
|
||||
for frame in frame_generator:
|
||||
result = model.predict(frame)[0]
|
||||
detections = sv.Detections.from_yolo_nas(result)
|
||||
detections = detections[polygon_zone.trigger(detections)]
|
||||
detections = detections.with_nms(threshold=args.iou_threshold)
|
||||
detections = byte_track.update_with_detections(detections=detections)
|
||||
|
||||
points = detections.get_anchors_coordinates(
|
||||
anchor=sv.Position.BOTTOM_CENTER
|
||||
)
|
||||
points = view_transformer.transform_points(points=points).astype(int)
|
||||
|
||||
for tracker_id, [_, y] in zip(detections.tracker_id, points):
|
||||
coordinates[tracker_id].append(y)
|
||||
|
||||
labels = []
|
||||
for tracker_id in detections.tracker_id:
|
||||
if len(coordinates[tracker_id]) < video_info.fps / 2:
|
||||
labels.append(f"#{tracker_id}")
|
||||
else:
|
||||
coordinate_start = coordinates[tracker_id][-1]
|
||||
coordinate_end = coordinates[tracker_id][0]
|
||||
distance = abs(coordinate_start - coordinate_end)
|
||||
time = len(coordinates[tracker_id]) / video_info.fps
|
||||
speed = distance / time * 3.6
|
||||
labels.append(f"#{tracker_id} {int(speed)} km/h")
|
||||
|
||||
annotated_frame = frame.copy()
|
||||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections
|
||||
)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
scene=annotated_frame, detections=detections, labels=labels
|
||||
)
|
||||
|
||||
sink.write_frame(annotated_frame)
|
||||
cv2.imshow("frame", annotated_frame)
|
||||
if cv2.waitKey(1) & 0xFF == ord("q"):
|
||||
break
|
||||
cv2.destroyAllWindows()
|
||||
@@ -1 +1,2 @@
|
||||
data/
|
||||
venv/
|
||||
|
||||
Generated
+35
-35
@@ -1151,13 +1151,13 @@ trio = ["async_generator", "trio"]
|
||||
|
||||
[[package]]
|
||||
name = "jinja2"
|
||||
version = "3.1.2"
|
||||
version = "3.1.3"
|
||||
description = "A very fast and expressive template engine."
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "Jinja2-3.1.2-py3-none-any.whl", hash = "sha256:6088930bfe239f0e6710546ab9c19c9ef35e29792895fed6e6e31a023a182a61"},
|
||||
{file = "Jinja2-3.1.2.tar.gz", hash = "sha256:31351a702a408a9e7595a8fc6150fc3f43bb6bf7e319770cbc0db9df9437e852"},
|
||||
{file = "Jinja2-3.1.3-py3-none-any.whl", hash = "sha256:7d6d50dd97d52cbc355597bd845fabfbac3f551e1f99619e39a35ce8c370b5fa"},
|
||||
{file = "Jinja2-3.1.3.tar.gz", hash = "sha256:ac8bd6544d4bb2c9792bf3a159e80bba8fda7f07e81bc3aed565432d5925ba90"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -1856,13 +1856,13 @@ mkdocs = ">=1.1"
|
||||
|
||||
[[package]]
|
||||
name = "mkdocs-material"
|
||||
version = "9.5.3"
|
||||
version = "9.5.4"
|
||||
description = "Documentation that simply works"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "mkdocs_material-9.5.3-py3-none-any.whl", hash = "sha256:76c93a8525cceb0b395b9cedab3428bf518cf6439adef2b940f1c1574b775d89"},
|
||||
{file = "mkdocs_material-9.5.3.tar.gz", hash = "sha256:5899219f422f0a6de784232d9d40374416302ffae3c160cacc72969fcc1ee372"},
|
||||
{file = "mkdocs_material-9.5.4-py3-none-any.whl", hash = "sha256:efd7cc8ae03296d728da9bd38f4db8b07ab61f9738a0cbd0dfaf2a15a50e7343"},
|
||||
{file = "mkdocs_material-9.5.4.tar.gz", hash = "sha256:3d196ee67fad16b2df1a458d650a8ac1890294eaae368d26cee71bc24ad41c40"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
@@ -3310,28 +3310,28 @@ files = [
|
||||
|
||||
[[package]]
|
||||
name = "ruff"
|
||||
version = "0.1.11"
|
||||
version = "0.1.13"
|
||||
description = "An extremely fast Python linter and code formatter, written in Rust."
|
||||
optional = false
|
||||
python-versions = ">=3.7"
|
||||
files = [
|
||||
{file = "ruff-0.1.11-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:a7f772696b4cdc0a3b2e527fc3c7ccc41cdcb98f5c80fdd4f2b8c50eb1458196"},
|
||||
{file = "ruff-0.1.11-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:934832f6ed9b34a7d5feea58972635c2039c7a3b434fe5ba2ce015064cb6e955"},
|
||||
{file = "ruff-0.1.11-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ea0d3e950e394c4b332bcdd112aa566010a9f9c95814844a7468325290aabfd9"},
|
||||
{file = "ruff-0.1.11-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9bd4025b9c5b429a48280785a2b71d479798a69f5c2919e7d274c5f4b32c3607"},
|
||||
{file = "ruff-0.1.11-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e1ad00662305dcb1e987f5ec214d31f7d6a062cae3e74c1cbccef15afd96611d"},
|
||||
{file = "ruff-0.1.11-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:4b077ce83f47dd6bea1991af08b140e8b8339f0ba8cb9b7a484c30ebab18a23f"},
|
||||
{file = "ruff-0.1.11-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:c4a88efecec23c37b11076fe676e15c6cdb1271a38f2b415e381e87fe4517f18"},
|
||||
{file = "ruff-0.1.11-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:5b25093dad3b055667730a9b491129c42d45e11cdb7043b702e97125bcec48a1"},
|
||||
{file = "ruff-0.1.11-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:231d8fb11b2cc7c0366a326a66dafc6ad449d7fcdbc268497ee47e1334f66f77"},
|
||||
{file = "ruff-0.1.11-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:09c415716884950080921dd6237767e52e227e397e2008e2bed410117679975b"},
|
||||
{file = "ruff-0.1.11-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:0f58948c6d212a6b8d41cd59e349751018797ce1727f961c2fa755ad6208ba45"},
|
||||
{file = "ruff-0.1.11-py3-none-musllinux_1_2_i686.whl", hash = "sha256:190a566c8f766c37074d99640cd9ca3da11d8deae2deae7c9505e68a4a30f740"},
|
||||
{file = "ruff-0.1.11-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:6464289bd67b2344d2a5d9158d5eb81025258f169e69a46b741b396ffb0cda95"},
|
||||
{file = "ruff-0.1.11-py3-none-win32.whl", hash = "sha256:9b8f397902f92bc2e70fb6bebfa2139008dc72ae5177e66c383fa5426cb0bf2c"},
|
||||
{file = "ruff-0.1.11-py3-none-win_amd64.whl", hash = "sha256:eb85ee287b11f901037a6683b2374bb0ec82928c5cbc984f575d0437979c521a"},
|
||||
{file = "ruff-0.1.11-py3-none-win_arm64.whl", hash = "sha256:97ce4d752f964ba559c7023a86e5f8e97f026d511e48013987623915431c7ea9"},
|
||||
{file = "ruff-0.1.11.tar.gz", hash = "sha256:f9d4d88cb6eeb4dfe20f9f0519bd2eaba8119bde87c3d5065c541dbae2b5a2cb"},
|
||||
{file = "ruff-0.1.13-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:e3fd36e0d48aeac672aa850045e784673449ce619afc12823ea7868fcc41d8ba"},
|
||||
{file = "ruff-0.1.13-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:9fb6b3b86450d4ec6a6732f9f60c4406061b6851c4b29f944f8c9d91c3611c7a"},
|
||||
{file = "ruff-0.1.13-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:b13ba5d7156daaf3fd08b6b993360a96060500aca7e307d95ecbc5bb47a69296"},
|
||||
{file = "ruff-0.1.13-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:9ebb40442f7b531e136d334ef0851412410061e65d61ca8ce90d894a094feb22"},
|
||||
{file = "ruff-0.1.13-py3-none-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:226b517f42d59a543d6383cfe03cccf0091e3e0ed1b856c6824be03d2a75d3b6"},
|
||||
{file = "ruff-0.1.13-py3-none-manylinux_2_17_ppc64.manylinux2014_ppc64.whl", hash = "sha256:5f0312ba1061e9b8c724e9a702d3c8621e3c6e6c2c9bd862550ab2951ac75c16"},
|
||||
{file = "ruff-0.1.13-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:2f59bcf5217c661254bd6bc42d65a6fd1a8b80c48763cb5c2293295babd945dd"},
|
||||
{file = "ruff-0.1.13-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:e6894b00495e00c27b6ba61af1fc666f17de6140345e5ef27dd6e08fb987259d"},
|
||||
{file = "ruff-0.1.13-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9a1600942485c6e66119da294c6294856b5c86fd6df591ce293e4a4cc8e72989"},
|
||||
{file = "ruff-0.1.13-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:ee3febce7863e231a467f90e681d3d89210b900d49ce88723ce052c8761be8c7"},
|
||||
{file = "ruff-0.1.13-py3-none-musllinux_1_2_armv7l.whl", hash = "sha256:dcaab50e278ff497ee4d1fe69b29ca0a9a47cd954bb17963628fa417933c6eb1"},
|
||||
{file = "ruff-0.1.13-py3-none-musllinux_1_2_i686.whl", hash = "sha256:f57de973de4edef3ad3044d6a50c02ad9fc2dff0d88587f25f1a48e3f72edf5e"},
|
||||
{file = "ruff-0.1.13-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7a36fa90eb12208272a858475ec43ac811ac37e91ef868759770b71bdabe27b6"},
|
||||
{file = "ruff-0.1.13-py3-none-win32.whl", hash = "sha256:a623349a505ff768dad6bd57087e2461be8db58305ebd5577bd0e98631f9ae69"},
|
||||
{file = "ruff-0.1.13-py3-none-win_amd64.whl", hash = "sha256:f988746e3c3982bea7f824c8fa318ce7f538c4dfefec99cd09c8770bd33e6539"},
|
||||
{file = "ruff-0.1.13-py3-none-win_arm64.whl", hash = "sha256:6bbbc3042075871ec17f28864808540a26f0f79a4478c357d3e3d2284e832998"},
|
||||
{file = "ruff-0.1.13.tar.gz", hash = "sha256:e261f1baed6291f434ffb1d5c6bd8051d1c2a26958072d38dfbec39b3dda7352"},
|
||||
]
|
||||
|
||||
[[package]]
|
||||
@@ -3543,30 +3543,30 @@ files = [
|
||||
|
||||
[[package]]
|
||||
name = "tox"
|
||||
version = "4.11.4"
|
||||
version = "4.12.0"
|
||||
description = "tox is a generic virtualenv management and test command line tool"
|
||||
optional = false
|
||||
python-versions = ">=3.8"
|
||||
files = [
|
||||
{file = "tox-4.11.4-py3-none-any.whl", hash = "sha256:2adb83d68f27116812b69aa36676a8d6a52249cb0d173649de0e7d0c2e3e7229"},
|
||||
{file = "tox-4.11.4.tar.gz", hash = "sha256:73a7240778fabf305aeb05ab8ea26e575e042ab5a18d71d0ed13e343a51d6ce1"},
|
||||
{file = "tox-4.12.0-py3-none-any.whl", hash = "sha256:c94bf5852ba41f3d9f1e3470ccf3390e0b7bdc938095be3cd96dce25ab5062a0"},
|
||||
{file = "tox-4.12.0.tar.gz", hash = "sha256:76adc53a3baff7bde80d6ad7f63235735cfc5bc42e8cb6fccfbf62cb5ffd4d92"},
|
||||
]
|
||||
|
||||
[package.dependencies]
|
||||
cachetools = ">=5.3.1"
|
||||
cachetools = ">=5.3.2"
|
||||
chardet = ">=5.2"
|
||||
colorama = ">=0.4.6"
|
||||
filelock = ">=3.12.3"
|
||||
packaging = ">=23.1"
|
||||
platformdirs = ">=3.10"
|
||||
filelock = ">=3.13.1"
|
||||
packaging = ">=23.2"
|
||||
platformdirs = ">=4.1"
|
||||
pluggy = ">=1.3"
|
||||
pyproject-api = ">=1.6.1"
|
||||
tomli = {version = ">=2.0.1", markers = "python_version < \"3.11\""}
|
||||
virtualenv = ">=20.24.3"
|
||||
virtualenv = ">=20.25"
|
||||
|
||||
[package.extras]
|
||||
docs = ["furo (>=2023.8.19)", "sphinx (>=7.2.4)", "sphinx-argparse-cli (>=1.11.1)", "sphinx-autodoc-typehints (>=1.24)", "sphinx-copybutton (>=0.5.2)", "sphinx-inline-tabs (>=2023.4.21)", "sphinxcontrib-towncrier (>=0.2.1a0)", "towncrier (>=23.6)"]
|
||||
testing = ["build[virtualenv] (>=0.10)", "covdefaults (>=2.3)", "detect-test-pollution (>=1.1.1)", "devpi-process (>=1)", "diff-cover (>=7.7)", "distlib (>=0.3.7)", "flaky (>=3.7)", "hatch-vcs (>=0.3)", "hatchling (>=1.18)", "psutil (>=5.9.5)", "pytest (>=7.4)", "pytest-cov (>=4.1)", "pytest-mock (>=3.11.1)", "pytest-xdist (>=3.3.1)", "re-assert (>=1.1)", "time-machine (>=2.12)", "wheel (>=0.41.2)"]
|
||||
docs = ["furo (>=2023.9.10)", "sphinx (>=7.2.6)", "sphinx-argparse-cli (>=1.11.1)", "sphinx-autodoc-typehints (>=1.25.2)", "sphinx-copybutton (>=0.5.2)", "sphinx-inline-tabs (>=2023.4.21)", "sphinxcontrib-towncrier (>=0.2.1a0)", "towncrier (>=23.11)"]
|
||||
testing = ["build[virtualenv] (>=1.0.3)", "covdefaults (>=2.3)", "detect-test-pollution (>=1.2)", "devpi-process (>=1)", "diff-cover (>=8.0.2)", "distlib (>=0.3.8)", "flaky (>=3.7)", "hatch-vcs (>=0.4)", "hatchling (>=1.21)", "psutil (>=5.9.7)", "pytest (>=7.4.4)", "pytest-cov (>=4.1)", "pytest-mock (>=3.12)", "pytest-xdist (>=3.5)", "re-assert (>=1.1)", "time-machine (>=2.13)", "wheel (>=0.42)"]
|
||||
|
||||
[[package]]
|
||||
name = "tqdm"
|
||||
@@ -3825,4 +3825,4 @@ desktop = ["opencv-python"]
|
||||
[metadata]
|
||||
lock-version = "2.0"
|
||||
python-versions = "^3.8"
|
||||
content-hash = "0dd573d54dcb8bc2627062330d52e90293f39a90e3f49b7afda7ef31abec3abf"
|
||||
content-hash = "b9f8ce8318216bf537b94e5a41bc9182fc35fc49346beaedfd8f39bc7598fa09"
|
||||
|
||||
+2
-2
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "supervision"
|
||||
version = "0.18.0rc1"
|
||||
version = "0.18.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>"]
|
||||
@@ -54,7 +54,7 @@ pytest = "^7.2.2"
|
||||
wheel = ">=0.40,<0.43"
|
||||
notebook = ">=6.5.3,<8.0.0"
|
||||
build = ">=0.10,<1.1"
|
||||
ruff = ">=0.0.280,<0.1.12"
|
||||
ruff = ">=0.1.0"
|
||||
isort = "^5.12.0"
|
||||
mypy = "^1.4.1"
|
||||
pre-commit = "^3.3.3"
|
||||
|
||||
@@ -20,6 +20,7 @@ from supervision.annotators.core import (
|
||||
MaskAnnotator,
|
||||
PixelateAnnotator,
|
||||
PolygonAnnotator,
|
||||
RoundBoxAnnotator,
|
||||
TraceAnnotator,
|
||||
TriangleAnnotator,
|
||||
)
|
||||
|
||||
@@ -1368,3 +1368,131 @@ class TriangleAnnotator(BaseAnnotator):
|
||||
cv2.fillPoly(scene, [vertices], color.as_bgr())
|
||||
|
||||
return scene
|
||||
|
||||
|
||||
class RoundBoxAnnotator(BaseAnnotator):
|
||||
"""
|
||||
A class for drawing bounding boxes with round edges on an image
|
||||
using provided detections.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
color: Union[Color, ColorPalette] = ColorPalette.default(),
|
||||
thickness: int = 2,
|
||||
color_lookup: ColorLookup = ColorLookup.CLASS,
|
||||
roundness: float = 0.6,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
color (Union[Color, ColorPalette]): The color or color palette to use for
|
||||
annotating detections.
|
||||
thickness (int): Thickness of the bounding box lines.
|
||||
color_lookup (str): Strategy for mapping colors to annotations.
|
||||
Options are `INDEX`, `CLASS`, `TRACK`.
|
||||
roundness (float): Percent of roundness for edges of bounding box.
|
||||
Value must be float 0 < roundness <= 1.0
|
||||
By default roundness percent is calculated based on smaller side
|
||||
length (width or height).
|
||||
"""
|
||||
self.color: Union[Color, ColorPalette] = color
|
||||
self.thickness: int = thickness
|
||||
self.color_lookup: ColorLookup = color_lookup
|
||||
if not 0 < roundness <= 1.0:
|
||||
raise ValueError("roundness attribute must be float between (0, 1.0]")
|
||||
self.roundness: float = roundness
|
||||
|
||||
def annotate(
|
||||
self,
|
||||
scene: np.ndarray,
|
||||
detections: Detections,
|
||||
custom_color_lookup: Optional[np.ndarray] = None,
|
||||
) -> np.ndarray:
|
||||
"""
|
||||
Annotates the given scene with bounding boxes with rounded edges
|
||||
based on the provided detections.
|
||||
|
||||
Args:
|
||||
scene (np.ndarray): The image where rounded 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:
|
||||
The annotated image.
|
||||
|
||||
Example:
|
||||
```python
|
||||
>>> import supervision as sv
|
||||
|
||||
>>> image = ...
|
||||
>>> detections = sv.Detections(...)
|
||||
|
||||
>>> round_box_annotator = sv.RoundBoxAnnotator()
|
||||
>>> annotated_frame = round_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)
|
||||
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,
|
||||
)
|
||||
|
||||
radius = (
|
||||
int((x2 - x1) // 2 * self.roundness)
|
||||
if abs(x1 - x2) < abs(y1 - y2)
|
||||
else int((y2 - y1) // 2 * self.roundness)
|
||||
)
|
||||
|
||||
circle_coordinates = [
|
||||
((x1 + radius), (y1 + radius)),
|
||||
((x2 - radius), (y1 + radius)),
|
||||
((x2 - radius), (y2 - radius)),
|
||||
((x1 + radius), (y2 - radius)),
|
||||
]
|
||||
|
||||
line_coordinates = [
|
||||
((x1 + radius, y1), (x2 - radius, y1)),
|
||||
((x2, y1 + radius), (x2, y2 - radius)),
|
||||
((x1 + radius, y2), (x2 - radius, y2)),
|
||||
((x1, y1 + radius), (x1, y2 - radius)),
|
||||
]
|
||||
|
||||
start_angles = (180, 270, 0, 90)
|
||||
end_angles = (270, 360, 90, 180)
|
||||
|
||||
for center_coordinates, line, start_angle, end_angle in zip(
|
||||
circle_coordinates, line_coordinates, start_angles, end_angles
|
||||
):
|
||||
cv2.ellipse(
|
||||
img=scene,
|
||||
center=center_coordinates,
|
||||
axes=(radius, radius),
|
||||
angle=0,
|
||||
startAngle=start_angle,
|
||||
endAngle=end_angle,
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
cv2.line(
|
||||
img=scene,
|
||||
pt1=line[0],
|
||||
pt2=line[1],
|
||||
color=color.as_bgr(),
|
||||
thickness=self.thickness,
|
||||
)
|
||||
|
||||
return scene
|
||||
|
||||
@@ -54,7 +54,9 @@ class LineZone:
|
||||
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):
|
||||
for i, (xyxy, tracker_id) in enumerate(
|
||||
zip(detections.xyxy, detections.tracker_id)
|
||||
):
|
||||
if tracker_id is None:
|
||||
continue
|
||||
|
||||
|
||||
Reference in New Issue
Block a user