`yolo_nas_example.py` script is ready
This commit is contained in:
SkalskiP
parent
351ef464d9
commit
e009402934
|
|
@ -53,14 +53,25 @@ supervision package for multiple tasks such as tracking, annotations, etc.
|
|||
|
||||
## ⚙️ run
|
||||
|
||||
```bash
|
||||
python ultralytics_example.py \
|
||||
--source_weights_path yolov8x.pt \
|
||||
--source_video_path data/vehicles.mp4 \
|
||||
--target_video_path data/vehicles-result.mp4 \
|
||||
--confidence_threshold 0.3 \
|
||||
--iou_threshold 0.5
|
||||
```
|
||||
- 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
|
||||
```
|
||||
|
||||
- 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
|
||||
|
||||
|
|
|
|||
|
|
@ -40,13 +40,7 @@ class ViewTransformer:
|
|||
|
||||
def parse_arguments() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Vehicle Speed Estimation using Supervision Package"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_weights_path",
|
||||
required=True,
|
||||
help="Path to the source weights file",
|
||||
type=str,
|
||||
description="Vehicle Speed Estimation using Ultralytics and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
|
|
@ -80,7 +74,7 @@ if __name__ == "__main__":
|
|||
args = parse_arguments()
|
||||
|
||||
video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path)
|
||||
model = YOLO(args.source_weights_path)
|
||||
model = YOLO('yolov8x.pt')
|
||||
|
||||
byte_track = sv.ByteTrack(
|
||||
frame_rate=video_info.fps,
|
||||
|
|
@ -90,7 +84,7 @@ if __name__ == "__main__":
|
|||
resolution_wh=video_info.resolution_wh)
|
||||
text_scale = sv.calculate_dynamic_text_scale(
|
||||
resolution_wh=video_info.resolution_wh)
|
||||
box_corner_annotator = sv.BoundingBoxAnnotator(
|
||||
bounding_box_annotator = sv.BoundingBoxAnnotator(
|
||||
thickness=thickness)
|
||||
label_annotator = sv.LabelAnnotator(
|
||||
text_scale=text_scale,
|
||||
|
|
@ -141,7 +135,7 @@ if __name__ == "__main__":
|
|||
annotated_frame = trace_annotator.annotate(
|
||||
scene=annotated_frame,
|
||||
detections=detections)
|
||||
annotated_frame = box_corner_annotator.annotate(
|
||||
annotated_frame = bounding_box_annotator.annotate(
|
||||
scene=annotated_frame,
|
||||
detections=detections)
|
||||
annotated_frame = label_annotator.annotate(
|
||||
|
|
|
|||
|
|
@ -1,14 +1,48 @@
|
|||
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
|
||||
|
||||
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 Supervision Package"
|
||||
description="Vehicle Speed Estimation using YOLO-NAS and Supervision"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--source_video_path",
|
||||
|
|
@ -42,4 +76,77 @@ 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")
|
||||
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()
|
||||
Loading…
Reference in New Issue