diff --git a/examples/speed_estimation/.gitignore b/examples/speed_estimation/.gitignore index 27cca3bd..5b855697 100644 --- a/examples/speed_estimation/.gitignore +++ b/examples/speed_estimation/.gitignore @@ -1,2 +1,2 @@ data/ -venv*/ \ No newline at end of file +venv*/ diff --git a/examples/speed_estimation/README.md b/examples/speed_estimation/README.md index a982f9f8..0500ab0d 100644 --- a/examples/speed_estimation/README.md +++ b/examples/speed_estimation/README.md @@ -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. \ No newline at end of file + projects. diff --git a/examples/speed_estimation/inference_example.py b/examples/speed_estimation/inference_example.py index 28441ddb..d7762c5e 100644 --- a/examples/speed_estimation/inference_example.py +++ b/examples/speed_estimation/inference_example.py @@ -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) diff --git a/examples/speed_estimation/requirements.txt b/examples/speed_estimation/requirements.txt index 7365d559..8a43a381 100644 --- a/examples/speed_estimation/requirements.txt +++ b/examples/speed_estimation/requirements.txt @@ -3,4 +3,4 @@ supervision ultralytics super-gradients -inference \ No newline at end of file +inference diff --git a/examples/speed_estimation/ultralytics_example.py b/examples/speed_estimation/ultralytics_example.py index b397a527..6bb2b2f6 100644 --- a/examples/speed_estimation/ultralytics_example.py +++ b/examples/speed_estimation/ultralytics_example.py @@ -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) diff --git a/examples/speed_estimation/yolo_nas_example.py b/examples/speed_estimation/yolo_nas_example.py index 9862dfbf..5847d759 100644 --- a/examples/speed_estimation/yolo_nas_example.py +++ b/examples/speed_estimation/yolo_nas_example.py @@ -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() \ No newline at end of file + cv2.destroyAllWindows() diff --git a/examples/traffic_analysis/.gitignore b/examples/traffic_analysis/.gitignore index 9435cafc..cea2fc37 100644 --- a/examples/traffic_analysis/.gitignore +++ b/examples/traffic_analysis/.gitignore @@ -1,2 +1,2 @@ data/ -venv/ \ No newline at end of file +venv/