当前位置：首页 > news >正文

Supervision 计算机视觉工具

news 来源：原创 2024/9/28 1:21:36

Supervision 是一个强大的计算机视觉工具库，它提供了丰富的功能来处理各种视觉任务，如目标检测、语义分割、目标跟踪等。这个库的最大特点是它能够简化代码的编写，使得视觉任务的处理变得更加高效和直观。

安装

Supervision 库可以在 Python 3.8 及以上版本中安装。如果你需要完整的 GUI 支持，可以使用以下命令安装：

pip install "supervision[desktop]" -i https://pypi.tuna.tsinghua.edu.cn/simple

如果不需要 GUI 支持，可以简化安装：

pip install supervision -i https://pypi.tuna.tsinghua.edu.cn/simple

目标检测与跟踪

Supervision 库结合 YOLO V8 可以轻松实现目标检测和跟踪。以下是使用 Supervision 和 YOLO V8 进行目标检测和跟踪的代码示例：

import cv2
import supervision as sv
from ultralytics import YOLO# 加载模型
model = YOLO("yolov8n.pt")# 初始化追踪器
tracker = sv.ByteTrack()# 读取视频
video_name = sv.download_assets(sv.VideoAssets.PEOPLE_WALKING)
cap = cv2.VideoCapture(video_name)while True:ret, frame = cap.read()if not ret:breakresult = model(frame)[0]detections = sv.Detections.from_ultralytics(result)detections = tracker.update_with_detections(detections)# 绘制边框和标签annotated_frame = sv.BoundingBoxAnnotator().annotate(frame, detections)labels = [f"{d.class_id} {model.names[d.class_id]}" for d in detections]annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)cv2.imshow('video', annotated_frame)if cv2.waitKey(1) & 0xFF == ord('q'):breakcap.release()
cv2.destroyAllWindows()

轨迹平滑

轨迹平滑是目标跟踪中的一个重要环节，Supervision 提供了 DetectionsSmoother 类来实现轨迹平滑：

import cv2
import supervision as sv
from ultralytics import YOLO# 加载模型
model = YOLO("yolov8n.pt")# 初始化追踪器和平滑器
tracker = sv.ByteTrack()
smoother = sv.DetectionsSmoother(length=4)# 读取视频
video_name = sv.download_assets(sv.VideoAssets.PEOPLE_WALKING)
cap = cv2.VideoCapture(video_name)while True:ret, frame = cap.read()if not ret:breakresult = model(frame)[0]detections = sv.Detections.from_ultralytics(result)detections = tracker.update_with_detections(detections)detections = smoother.update_with_detections(detections)# 绘制边框和标签annotated_frame = sv.BoundingBoxAnnotator().annotate(frame, detections)labels = [f"{d.class_id} {model.names[d.class_id]}" for d in detections]annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)cv2.imshow('video', annotated_frame)if cv2.waitKey(1) & 0xFF == ord('q'):breakcap.release()
cv2.destroyAllWindows()

越线数量统计

Supervision 还提供了 LineZone 类来实现越线数量统计功能：

import supervision as sv
from ultralytics import YOLO# 加载模型
model = YOLO("yolov8n.pt")# 初始化追踪器和越线检测器
tracker = sv.ByteTrack()
line_zone = sv.LineZone(start=sv.Point(0, 400), end=sv.Point(1280, 400))# 读取视频
video_name = sv.download_assets(sv.VideoAssets.VEHICLES)
cap = cv2.VideoCapture(video_name)while True:ret, frame = cap.read()if not ret:breakresult = model(frame)[0]detections = sv.Detections.from_ultralytics(result)detections = tracker.update_with_detections(detections)crossed_in, crossed_out = line_zone.trigger(detections)# 绘制边框和标签annotated_frame = sv.BoundingBoxAnnotator().annotate(frame, detections)labels = [f"{d.class_id} {model.names[d.class_id]}" for d in detections]annotated_frame = sv.LabelAnnotator().annotate(annotated_frame, detections, labels)cv2.imshow('video', annotated_frame)if cv2.waitKey(1) & 0xFF == ord('q'):breakcap.release()
cv2.destroyAllWindows()

切片推理

对于小目标的检测，Supervision 提供了 InferenceSlicer 类来进行切片推理：

import cv2
import supervision as sv
from ultralytics import YOLO
import numpy as np# 加载模型
model = YOLO("yolov8n.pt")# 切片回调函数
def callback(image_slice: np.ndarray) -> sv.Detections:result = model(image_slice, verbose=False)[0]return sv.Detections.from_ultralytics(result)# 初始化切片器
slicer = sv.InferenceSlicer(callback=callback, slice_wh=(320, 320), overlap_ratio_wh=(0.3, 0.3), iou_threshold=0.5, thread_workers=4)# 读取图像
image = cv2.imread("img/3.png")
detections = slicer(image)# 绘制边框和标签
annotated_image = sv.BoundingBoxAnnotator().annotate(image, detections)
labels = [f"{d.class_id} {model.names[d.class_id]}" for d in detections]
annotated_image = sv.LabelAnnotator().annotate(annotated_image, detections, labels)cv2.imshow('img', annotated_image)
cv2.waitKey(0)

辅助展现任务

Supervision 提供了多种方式来展现检测结果，包括边框、角点边框、遮罩等：

import cv2
import supervision as sv
from ultralytics import YOLO# 加载模型
model = YOLO("yolov8n.pt")# 读取图像
image = cv2.imread("img/img.png")
result = model(image)[0]
detections = sv.Detections.from_ultralytics(result)# 角点边框展现
corner_annotator = sv.BoxCornerAnnotator(corner_length=15, thickness=2, color=sv.Color(r=255, g=0, b=0))
image = corner_annotator.annotate(image, detections)# 绘制标签
labels = [f"{d.class_id} {model.names[d.class_id]}" for d in detections]
image = sv.LabelAnnotator().annotate(image, detections, labels)cv2.imshow('img', image)
cv2.waitKey(0)