7.6 KiB
Vendored
7.6 KiB
Vendored
In [ ]:
!pip install -q inference "supervision[assets]"In [ ]:
from supervision.assets import download_assets, VideoAssets
# Download a supervision video asset
path_to_video = download_assets(VideoAssets.PEOPLE_WALKING)In [ ]:
import supervision as sv
from supervision.assets import download_assets, VideoAssets
from inference.models.utils import get_roboflow_model
if __name__ == '__main__':
# Install our video from supervision assets.
PATH_TO_VIDEO = download_assets(VideoAssets.PEOPLE_WALKING)
# Load a pre trained yolov8 nano model from Roboflow Inference.
model = get_roboflow_model('yolov8n-640')
# Create a video info object from the video path.
video_info = sv.VideoInfo.from_video_path(PATH_TO_VIDEO)
# Create a label annotator for labeling detections with our tracker_id.
label = sv.LabelAnnotator()
# Create a ByteTrack object to track detections.
byte_tracker = sv.ByteTrack(frame_rate=video_info.fps)
# Create a frame generator from video path for iteration of frames.
frame_generator = sv.get_video_frames_generator(PATH_TO_VIDEO)
# Create a video sink context manager to save resulting video.
with sv.VideoSink(target_path="output.mp4", video_info=video_info) as sink:
# Iterate over frames of the video.
for frame in frame_generator:
# Run inference on the frame by passing it to our model.
result = model.infer(frame)[0]
# Convert model results to a supervision detection object.
detections = sv.Detections.from_inference(result)
# Update detections with tracker ids fro byte_tracker.
tracked_detections = byte_tracker.update_with_detections(detections)
# Create labels with tracker_id for label annotator.
labels = [ f"{tracker_id}" for tracker_id in tracked_detections.tracker_id ]
# Apply label annotator to frame.
annotated_frame = label.annotate(scene=frame.copy(), detections=tracked_detections, labels=labels)
# Save the annotated frame to an output video.
sink.write_frame(frame=annotated_frame)