diff --git a/examples/speed_estimation/README.md b/examples/speed_estimation/README.md index 7833e18b..4dea9e88 100644 --- a/examples/speed_estimation/README.md +++ b/examples/speed_estimation/README.md @@ -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 diff --git a/examples/speed_estimation/ultralytics_example.py b/examples/speed_estimation/ultralytics_example.py index 9b5d84a0..b397a527 100644 --- a/examples/speed_estimation/ultralytics_example.py +++ b/examples/speed_estimation/ultralytics_example.py @@ -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( diff --git a/examples/speed_estimation/yolo_nas_example.py b/examples/speed_estimation/yolo_nas_example.py index 79eabe15..9862dfbf 100644 --- a/examples/speed_estimation/yolo_nas_example.py +++ b/examples/speed_estimation/yolo_nas_example.py @@ -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") \ No newline at end of file + 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() \ No newline at end of file