简化版代码如下:
import cv2
import supervision as sv
from ultralytics import YOLO
from supervision.draw.color import Color
model = YOLO('yolov8m.pt')
# 越线检测位置
line_counter = sv.LineZone(start=sv.Point(0, 400), end=sv.Point(1280, 400))
# 可视化配置
line_annotator = sv.LineZoneAnnotator(thickness=2, text_thickness=2, text_scale=2, color=Color(r=224, g=57, b=151))
box_annotator = sv.BoxAnnotator(thickness=1, text_thickness=1, text_scale=1)
input_path ='test.mp4'
output_path = "out-" + input_path.split('/')[-1]
cap = cv2.VideoCapture(input_path)
frame_size = (cap.get(cv2.CAP_PROP_FRAME_WIDTH), cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
fps = cap.get(cv2.CAP_PROP_FPS)
out = cv2.VideoWriter(output_path, fourcc, fps, (int(frame_size[0]), int(frame_size[1])))
# 视频逐帧追踪
for result in model.track(source=input_path, show=False, stream=True, verbose=False, device='cuda:0'):
frame = result.orig_img
detections = sv.Detections.from_yolov8(result) # 用supervision解析预测结果
detections.tracker_id = result.boxes.id.cpu().numpy().astype(int) # 解析追踪ID
# 获取每个目标的:追踪ID、类别名称、置信度
class_ids = detections.class_id # 类别ID
confidences = detections.confidence # 置信度
tracker_ids = detections.tracker_id # 多目标追踪ID
labels = ['#{} {} {:.1f}'.format(tracker_ids[i], model.names[class_ids[i]], confidences[i]) for i in range(len(class_ids))]
frame = box_annotator.annotate(scene=frame, detections=detections, labels=labels) # 绘制目标检测可视化结果
# 越线检测
line_counter.trigger(detections=detections)
line_annotator.annotate(frame=frame, line_counter=line_counter)
out.write(frame)
cv2.destroyAllWindows()
out.release()
cap.release()
print('跨线进入车辆数:', line_counter.in_count)
print('跨线离开车辆数:', line_counter.out_count)
使用的视频同两百行C++代码实现yolov5车辆计数部署(通俗易懂版)中的视频。经过试验,LZ发现使用yolov8n(检出47辆)、yolov8s(检出51辆)等小模型时会发生车辆漏检,而yolov8m及更大的模型能检测出所有通过的52辆车。