initial version of yolo_nas_example.py script

examples/speed_estimation/.gitignore

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

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

examples/speed_estimation/requirements.txt

@@ -1,5 +1,6 @@
 gdown
-inference
-supervision[asserts]
-tqdm
-ultralytics
+supervision
+
+ultralytics
+super-gradients
+inference

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

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