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

Merge pull request #714 from roboflow/example/speed_estimation 

example/speed_estimation 🚗 💨💨💨
This commit is contained in:
Piotr Skalski 2024-01-10 12:10:52 +01:00 committed by GitHub
parent a1da604798 a337ee591a
commit 9b5bba1422
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
9 changed files with 586 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
examples/README.md
View File

@ -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

2
examples/speed_estimation/.gitignore vendored Normal file
View File

@ -0,0 +1,2 @@
data/
venv*/

111
examples/speed_estimation/README.md Normal file
View File

@ -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.

167
examples/speed_estimation/inference_example.py Normal file
View File

@ -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()

6
examples/speed_estimation/requirements.txt Normal file
View File

@ -0,0 +1,6 @@
supervision==0.18.0rc1
tqdm
requests
ultralytics
super-gradients
inference

145
examples/speed_estimation/ultralytics_example.py Normal file
View File

@ -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()

8
examples/speed_estimation/video_downloader.py Normal file
View File

@ -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)

145
examples/speed_estimation/yolo_nas_example.py Normal file
View File

@ -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
examples/traffic_analysis/.gitignore vendored
View File

@ -1 +1,2 @@
data/
venv/