diff --git a/examples/README.md b/examples/README.md index f3f04c71..936a7402 100644 --- a/examples/README.md +++ b/examples/README.md @@ -7,6 +7,7 @@ interfaces with diverse applications. - [tracking](./tracking) by [@SkalskiP](https://github.com/SkalskiP) - [count people in zone](./count_people_in_zone) by [@SkalskiP](https://github.com/SkalskiP) - [traffic analysis](./traffic_analysis) by [@SkalskiP](https://github.com/SkalskiP) +- [speed estimation](./speed_estimation) by [@SkalskiP](https://github.com/SkalskiP) - [heatmap and track](./heatmap_and_track/) by [@HinePo](https://github.com/HinePo) ## contributing diff --git a/examples/speed_estimation/.gitignore b/examples/speed_estimation/.gitignore new file mode 100644 index 00000000..5b855697 --- /dev/null +++ b/examples/speed_estimation/.gitignore @@ -0,0 +1,2 @@ +data/ +venv*/ diff --git a/examples/speed_estimation/README.md b/examples/speed_estimation/README.md new file mode 100644 index 00000000..35685938 --- /dev/null +++ b/examples/speed_estimation/README.md @@ -0,0 +1,111 @@ +# speed estimation + +## 👋 hello + +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. + +https://github.com/roboflow/supervision/assets/26109316/0542fd3c-bb5f-475e-b96c-793560abeb18 + +## 💻 install + +> [!NOTE] +> YOLO-NAS is compatible with Python versions up to and including Python 3.10. + +- clone repository and navigate to example directory + + ```bash + git clone https://github.com/roboflow/supervision.git + cd supervision/examples/speed_estimation + ``` + +- setup python environment and activate it [optional] + + ```bash + python3.10 -m venv venv + source venv/bin/activate + ``` + +- install required dependencies + + ```bash + pip install -r requirements.txt + ``` + +- download `vehicles.mp4` file + + ```bash + python3.10 video_downloader.py + ``` + +## 🛠️ script arguments + +- `--roboflow_api_key` (optional): The API key for Roboflow services. If not provided + directly, the script tries to fetch it from the `ROBOFLOW_API_KEY` environment + variable. Follow [this guide](https://docs.roboflow.com/api-reference/authentication#retrieve-an-api-key) + to acquire your `API KEY`. +- `--model_id` (optional): Designates the Roboflow model ID to be used. The default + value is `"yolov8x-1280"`. + +- `--source_weights_path`: Required. Specifies the path to the YOLO model's weights + file, which is essential for the object detection process. This file contains the + 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`: 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 + model to filter detections. Default is `0.3`. This determines how confident the + model should be to recognize an object in the video. +- `--iou_threshold` (optional): Specifies the IOU (Intersection Over Union) threshold + for the model. Default is 0.7. This value is used to manage object detection + accuracy, particularly in distinguishing between different objects. + +## ⚙️ run + +- 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 + ``` + +- inference + + ```bash + python inference_example.py \ + --roboflow_api_key \ + --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 + +This demo integrates two main components, each with its own licensing: + +- ultralytics: The object detection model used in this demo, YOLOv8, is distributed + under the [AGPL-3.0 license](https://github.com/ultralytics/ultralytics/blob/main/LICENSE). + You can find more details about this license here. + +- supervision: The analytics code that powers the zone-based analysis in this demo is + 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. diff --git a/examples/speed_estimation/inference_example.py b/examples/speed_estimation/inference_example.py new file mode 100644 index 00000000..1043b46b --- /dev/null +++ b/examples/speed_estimation/inference_example.py @@ -0,0 +1,167 @@ +import argparse +import os +from collections import defaultdict, deque + +import cv2 +import numpy as np +from inference.models.utils import get_roboflow_model + +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 Inference and Supervision" + ) + parser.add_argument( + "--model_id", + default="yolov8x-640", + help="Roboflow model ID", + type=str, + ) + parser.add_argument( + "--roboflow_api_key", + default=None, + help="Roboflow API KEY", + type=str, + ) + 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() + + api_key = args.roboflow_api_key + api_key = os.environ.get("ROBOFLOW_API_KEY", api_key) + if api_key is None: + raise ValueError( + "Roboflow API key is missing. Please provide it as an argument or set the " + "ROBOFLOW_API_KEY environment variable." + ) + args.roboflow_api_key = api_key + + video_info = sv.VideoInfo.from_video_path(video_path=args.source_video_path) + 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 + ) + + 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: + results = model.infer(frame)[0] + detections = sv.Detections.from_inference(results) + detections = detections[detections.confidence > args.confidence_threshold] + 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() diff --git a/examples/speed_estimation/requirements.txt b/examples/speed_estimation/requirements.txt new file mode 100644 index 00000000..63f20432 --- /dev/null +++ b/examples/speed_estimation/requirements.txt @@ -0,0 +1,6 @@ +supervision==0.18.0rc1 +tqdm +requests +ultralytics +super-gradients +inference diff --git a/examples/speed_estimation/ultralytics_example.py b/examples/speed_estimation/ultralytics_example.py new file mode 100644 index 00000000..6bb2b2f6 --- /dev/null +++ b/examples/speed_estimation/ultralytics_example.py @@ -0,0 +1,145 @@ +import argparse +from collections import defaultdict, deque + +import cv2 +import numpy as np +from ultralytics import YOLO + +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 Ultralytics and Supervision" + ) + 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 = YOLO("yolov8x.pt") + + 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(frame)[0] + detections = sv.Detections.from_ultralytics(result) + detections = detections[detections.confidence > args.confidence_threshold] + 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() diff --git a/examples/speed_estimation/video_downloader.py b/examples/speed_estimation/video_downloader.py new file mode 100644 index 00000000..4272120b --- /dev/null +++ b/examples/speed_estimation/video_downloader.py @@ -0,0 +1,8 @@ +import os + +from supervision.assets import VideoAssets, download_assets + +if not os.path.exists("data"): + os.makedirs("data") +os.chdir("data") +download_assets(VideoAssets.VEHICLES) diff --git a/examples/speed_estimation/yolo_nas_example.py b/examples/speed_estimation/yolo_nas_example.py new file mode 100644 index 00000000..5847d759 --- /dev/null +++ b/examples/speed_estimation/yolo_nas_example.py @@ -0,0 +1,145 @@ +import argparse +from collections import defaultdict, deque + +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]]) + +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 YOLO-NAS and Supervision" + ) + 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") + + 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() diff --git a/examples/traffic_analysis/.gitignore b/examples/traffic_analysis/.gitignore index 8fce6030..cea2fc37 100644 --- a/examples/traffic_analysis/.gitignore +++ b/examples/traffic_analysis/.gitignore @@ -1 +1,2 @@ data/ +venv/