From 351ef464d95b9769d1af7719a5fc801eda00ac8d Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 8 Jan 2024 12:04:29 +0100 Subject: [PATCH] initial version of `yolo_nas_example.py` script --- examples/speed_estimation/.gitignore | 2 +- examples/speed_estimation/README.md | 4 +- examples/speed_estimation/requirements.txt | 9 ++-- .../speed_estimation/ultralytics_example.py | 28 +++++++----- examples/speed_estimation/yolo_nas_example.py | 45 +++++++++++++++++++ 5 files changed, 70 insertions(+), 18 deletions(-) create mode 100644 examples/speed_estimation/yolo_nas_example.py diff --git a/examples/speed_estimation/.gitignore b/examples/speed_estimation/.gitignore index 9435cafc..27cca3bd 100644 --- a/examples/speed_estimation/.gitignore +++ b/examples/speed_estimation/.gitignore @@ -1,2 +1,2 @@ data/ -venv/ \ No newline at end of file +venv*/ \ No newline at end of file diff --git a/examples/speed_estimation/README.md b/examples/speed_estimation/README.md index f7e5b0df..7833e18b 100644 --- a/examples/speed_estimation/README.md +++ b/examples/speed_estimation/README.md @@ -41,7 +41,7 @@ supervision package for multiple tasks such as tracking, annotations, etc. 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` (optional): The path to save the output video with +- `--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 @@ -56,8 +56,8 @@ supervision package for multiple tasks such as tracking, annotations, etc. ```bash python ultralytics_example.py \ --source_weights_path yolov8x.pt \ - --lines_configuration_path data/vehicles-config.json \ --source_video_path data/vehicles.mp4 \ + --target_video_path data/vehicles-result.mp4 \ --confidence_threshold 0.3 \ --iou_threshold 0.5 ``` diff --git a/examples/speed_estimation/requirements.txt b/examples/speed_estimation/requirements.txt index 537df5eb..7365d559 100644 --- a/examples/speed_estimation/requirements.txt +++ b/examples/speed_estimation/requirements.txt @@ -1,5 +1,6 @@ gdown -inference -supervision[asserts] -tqdm -ultralytics \ No newline at end of file +supervision + +ultralytics +super-gradients +inference \ No newline at end of file diff --git a/examples/speed_estimation/ultralytics_example.py b/examples/speed_estimation/ultralytics_example.py index 16e419bb..9b5d84a0 100644 --- a/examples/speed_estimation/ultralytics_example.py +++ b/examples/speed_estimation/ultralytics_example.py @@ -7,7 +7,6 @@ from ultralytics import YOLO import supervision as sv - SOURCE = np.array([ [1252, 787], [2298, 803], @@ -39,7 +38,7 @@ class ViewTransformer: return transformed_points.reshape(-1, 2) -if __name__ == "__main__": +def parse_arguments() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Vehicle Speed Estimation using Supervision Package" ) @@ -57,7 +56,7 @@ if __name__ == "__main__": ) parser.add_argument( "--target_video_path", - default=None, + required=True, help="Path to the target video file (output)", type=str, ) @@ -74,39 +73,46 @@ if __name__ == "__main__": type=float ) - args = parser.parse_args() + 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(args.source_weights_path) + 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) - box_corner_annotator = sv.BoundingBoxAnnotator( thickness=thickness) label_annotator = sv.LabelAnnotator( - text_scale=text_scale, text_thickness=thickness, + 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, + 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) + 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, imgsz=1280)[0] + result = model(frame)[0] detections = sv.Detections.from_ultralytics(result) detections = detections[detections.confidence > args.confidence_threshold] detections = detections[polygon_zone.trigger(detections)] diff --git a/examples/speed_estimation/yolo_nas_example.py b/examples/speed_estimation/yolo_nas_example.py new file mode 100644 index 00000000..79eabe15 --- /dev/null +++ b/examples/speed_estimation/yolo_nas_example.py @@ -0,0 +1,45 @@ +import argparse + +from super_gradients.training import models +from super_gradients.common.object_names import Models + +import supervision as sv + + +def parse_arguments() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Vehicle Speed Estimation using Supervision Package" + ) + 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") \ No newline at end of file