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

fix(pre_commit): 🎨 auto format pre-commit hooks

This commit is contained in:
pre-commit-ci[bot] 2024-01-08 12:03:30 +00:00
parent 69894c82af
commit 4ae9131f0a
7 changed files with 104 additions and 121 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
examples/speed_estimation/.gitignore vendored
View File

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

10
examples/speed_estimation/README.md
View File

@ -2,7 +2,7 @@
## 👋 hello
This example performs speed estimation analysis using various object-detection models
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.
@ -27,13 +27,13 @@ supervision package for multiple tasks such as tracking, annotations, etc.
```bash
pip install -r requirements.txt
```
- download `vehicles.mp4` file
```bash
./setup.sh
```
## 🛠️ script arguments
- `--roboflow_api_key` (optional): The API key for Roboflow services. If not provided
@ -69,7 +69,7 @@ supervision package for multiple tasks such as tracking, annotations, etc.
--confidence_threshold 0.3 \
--iou_threshold 0.5
```
- inference
```bash
@ -103,4 +103,4 @@ This demo integrates two main components, each with its own licensing:
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.
projects.

67
examples/speed_estimation/inference_example.py
View File

@ -8,26 +8,22 @@ from inference.models.utils import get_roboflow_model
import supervision as sv
SOURCE = np.array([
[1252, 787],
[2298, 803],
[5039, 2159],
[-550, 2159]
])
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]
])
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)
@ -74,10 +70,7 @@ def parse_arguments() -> argparse.Namespace:
type=float,
)
parser.add_argument(
"--iou_threshold",
default=0.7,
help="IOU threshold for the model",
type=float
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
)
return parser.parse_args()
@ -99,29 +92,30 @@ if __name__ == "__main__":
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)
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)
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)
text_position=sv.Position.BOTTOM_CENTER,
)
trace_annotator = sv.TraceAnnotator(
thickness=thickness,
trace_length=video_info.fps * 2,
position=sv.Position.BOTTOM_CENTER)
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)
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
)
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
@ -135,7 +129,9 @@ if __name__ == "__main__":
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 = 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):
@ -155,15 +151,14 @@ if __name__ == "__main__":
annotated_frame = frame.copy()
annotated_frame = trace_annotator.annotate(
scene=annotated_frame,
detections=detections)
scene=annotated_frame, detections=detections
)
annotated_frame = bounding_box_annotator.annotate(
scene=annotated_frame,
detections=detections)
scene=annotated_frame, detections=detections
)
annotated_frame = label_annotator.annotate(
scene=annotated_frame,
detections=detections,
labels=labels)
scene=annotated_frame, detections=detections, labels=labels
)
sink.write_frame(annotated_frame)
cv2.imshow("frame", annotated_frame)

2
examples/speed_estimation/requirements.txt
View File

@ -3,4 +3,4 @@ supervision
ultralytics
super-gradients
inference
inference

69
examples/speed_estimation/ultralytics_example.py
View File

@ -7,26 +7,22 @@ from ultralytics import YOLO
import supervision as sv
SOURCE = np.array([
[1252, 787],
[2298, 803],
[5039, 2159],
[-550, 2159]
])
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]
])
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)
@ -61,10 +57,7 @@ def parse_arguments() -> argparse.Namespace:
type=float,
)
parser.add_argument(
"--iou_threshold",
default=0.7,
help="IOU threshold for the model",
type=float
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
)
return parser.parse_args()
@ -74,32 +67,33 @@ if __name__ == "__main__":
args = parse_arguments()
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
model = YOLO('yolov8x.pt')
model = YOLO("yolov8x.pt")
byte_track = sv.ByteTrack(
frame_rate=video_info.fps,
track_thresh=args.confidence_threshold)
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)
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)
text_position=sv.Position.BOTTOM_CENTER,
)
trace_annotator = sv.TraceAnnotator(
thickness=thickness,
trace_length=video_info.fps * 2,
position=sv.Position.BOTTOM_CENTER)
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)
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
)
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
@ -113,7 +107,9 @@ if __name__ == "__main__":
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 = 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):
@ -133,15 +129,14 @@ if __name__ == "__main__":
annotated_frame = frame.copy()
annotated_frame = trace_annotator.annotate(
scene=annotated_frame,
detections=detections)
scene=annotated_frame, detections=detections
)
annotated_frame = bounding_box_annotator.annotate(
scene=annotated_frame,
detections=detections)
scene=annotated_frame, detections=detections
)
annotated_frame = label_annotator.annotate(
scene=annotated_frame,
detections=detections,
labels=labels)
scene=annotated_frame, detections=detections, labels=labels
)
sink.write_frame(annotated_frame)
cv2.imshow("frame", annotated_frame)

73
examples/speed_estimation/yolo_nas_example.py
View File

@ -1,34 +1,29 @@
import argparse
from collections import defaultdict, deque
from super_gradients.training import models
from super_gradients.common.object_names import Models
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]
])
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]
])
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)
@ -63,10 +58,7 @@ def parse_arguments() -> argparse.Namespace:
type=float,
)
parser.add_argument(
"--iou_threshold",
default=0.7,
help="IOU threshold for the model",
type=float
"--iou_threshold", default=0.7, help="IOU threshold for the model", type=float
)
return parser.parse_args()
@ -79,29 +71,30 @@ if __name__ == "__main__":
model = models.get(Models.YOLO_NAS_L, pretrained_weights="coco")
byte_track = sv.ByteTrack(
frame_rate=video_info.fps,
track_thresh=args.confidence_threshold)
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)
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)
text_position=sv.Position.BOTTOM_CENTER,
)
trace_annotator = sv.TraceAnnotator(
thickness=thickness,
trace_length=video_info.fps * 2,
position=sv.Position.BOTTOM_CENTER)
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)
polygon=SOURCE, frame_resolution_wh=video_info.resolution_wh
)
view_transformer = ViewTransformer(source=SOURCE, target=TARGET)
coordinates = defaultdict(lambda: deque(maxlen=video_info.fps))
@ -115,7 +108,8 @@ if __name__ == "__main__":
detections = byte_track.update_with_detections(detections=detections)
points = detections.get_anchors_coordinates(
anchor=sv.Position.BOTTOM_CENTER)
anchor=sv.Position.BOTTOM_CENTER
)
points = view_transformer.transform_points(points=points).astype(int)
for tracker_id, [_, y] in zip(detections.tracker_id, points):
@ -135,18 +129,17 @@ if __name__ == "__main__":
annotated_frame = frame.copy()
annotated_frame = trace_annotator.annotate(
scene=annotated_frame,
detections=detections)
scene=annotated_frame, detections=detections
)
annotated_frame = bounding_box_annotator.annotate(
scene=annotated_frame,
detections=detections)
scene=annotated_frame, detections=detections
)
annotated_frame = label_annotator.annotate(
scene=annotated_frame,
detections=detections,
labels=labels)
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()
cv2.destroyAllWindows()

2
examples/traffic_analysis/.gitignore vendored
View File

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