153 lines
7.6 KiB
Plaintext
153 lines
7.6 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# Object Tracking\n",
|
|
"\n",
|
|
"---\n",
|
|
"\n",
|
|
"[](https://colab.research.google.com/github/roboflow/supervision/blob/develop/docs/notebooks/object-tracking.ipynb)\n",
|
|
"\n",
|
|
"In some cases, it's important for us to track objects across multiple frames of a video. For example, we may need to figure out the direction a vehicle is moving, or count objects in a frame. Some Supervision [Annotators](https://supervision.roboflow.com/latest/annotators/) and Tools like [LineZone](https://supervision.roboflow.com/latest/detection/tools/line_zone/) require tracking to be setup. In this cookbook, we'll cover how to get a tracker up and running for use in your computer vision applications.\n",
|
|
"\n",
|
|
"## What is a Tracker?\n",
|
|
"\n",
|
|
"Trackers are a piece of code that identifies objects across frames and assigns them a unique `tracker_id`. There are a few popular trackers at the time of writing this including ByteTrack and Bot-SORT. Supervision makes using trackers a breeze and comes with ByteTrack built-in. But first, let's get our deppendencies installed. \n",
|
|
"\n",
|
|
"## Install Dependencies"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {
|
|
"vscode": {
|
|
"languageId": "powershell"
|
|
}
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"!pip install -q inference \"supervision[assets]\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Download a Video Asset\n",
|
|
"\n",
|
|
"Now that we have our enviornment setup, lets download a video that we can detect objects in. Supervision comes with a great utility to help us hit the ground running. We can use the below snippet to he video is save a video asset in our local directory. It can also be accessed with the variable `path_to_video` for additional application logic."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from supervision.assets import download_assets, VideoAssets\n",
|
|
"\n",
|
|
"# Download a supervision video asset \n",
|
|
"path_to_video = download_assets(VideoAssets.PEOPLE_WALKING)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Tracking Objects\n",
|
|
"\n",
|
|
"Now that we have our video installed, let's get to work on tracking objects. We'll first pull in a model from roboflow Inference to detect people in our video. Then let's create a `byte_tracker` object that we'll pass our detections to. This will give us a `tracker_id`. We'll then utilize that tracker id to label our detections with a `label_annotator` to display the tracker id. Finally, we'll use a utility called `VideoSink` to save the annotated frames to a video. Let's dive in to the code."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"import supervision as sv\n",
|
|
"from supervision.assets import download_assets, VideoAssets\n",
|
|
"from inference.models.utils import get_roboflow_model\n",
|
|
"\n",
|
|
"\n",
|
|
"if __name__ == '__main__':\n",
|
|
"\n",
|
|
" # Install our video from supervision assets.\n",
|
|
" PATH_TO_VIDEO = download_assets(VideoAssets.PEOPLE_WALKING)\n",
|
|
"\n",
|
|
" # Load a pre trained yolov8 nano model from Roboflow Inference.\n",
|
|
" model = get_roboflow_model('yolov8n-640')\n",
|
|
"\n",
|
|
" # Create a video info object from the video path.\n",
|
|
" video_info = sv.VideoInfo.from_video_path(PATH_TO_VIDEO)\n",
|
|
"\n",
|
|
" # Create a label annotator for labeling detections with our tracker_id.\n",
|
|
" label = sv.LabelAnnotator()\n",
|
|
"\n",
|
|
" # Create a ByteTrack object to track detections.\n",
|
|
" byte_tracker = sv.ByteTrack(frame_rate=video_info.fps)\n",
|
|
"\n",
|
|
" # Create a frame generator from video path for iteration of frames.\n",
|
|
" frame_generator = sv.get_video_frames_generator(PATH_TO_VIDEO)\n",
|
|
"\n",
|
|
" # Create a video sink context manager to save resulting video.\n",
|
|
" with sv.VideoSink(target_path=\"output.mp4\", video_info=video_info) as sink:\n",
|
|
"\n",
|
|
" # Iterate over frames of the video.\n",
|
|
" for frame in frame_generator:\n",
|
|
"\n",
|
|
" # Run inference on the frame by passing it to our model.\n",
|
|
" result = model.infer(frame)[0]\n",
|
|
"\n",
|
|
" # Convert model results to a supervision detection object.\n",
|
|
" detections = sv.Detections.from_inference(result)\n",
|
|
"\n",
|
|
" # Update detections with tracker ids fro byte_tracker.\n",
|
|
" tracked_detections = byte_tracker.update_with_detections(detections)\n",
|
|
"\n",
|
|
" # Create labels with tracker_id for label annotator.\n",
|
|
" labels = [ f\"{tracker_id}\" for tracker_id in tracked_detections.tracker_id ]\n",
|
|
"\n",
|
|
" # Apply label annotator to frame.\n",
|
|
" annotated_frame = label.annotate(scene=frame.copy(), detections=tracked_detections, labels=labels)\n",
|
|
"\n",
|
|
" # Save the annotated frame to an output video.\n",
|
|
" sink.write_frame(frame=annotated_frame)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"metadata": {},
|
|
"source": [
|
|
"# View the Video\n",
|
|
"\n",
|
|
"Let's take a look at our resulting video. It will be created in your current directory with the name `output.mp4` Notice how even with a little flicker, we can see the `tracker_id` on the people walking in the video. With trackers under your belt, there are now a wide variety of use cases you can solve for! Happy building! "
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "venv",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.9.7"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 2
|
|
}
|