五十行python代码实现yolov8车辆计数

参考:https://www.bilibili.com/video/BV1Dg4y1N7nz/?spm_id_from=333.1007.tianma.2-2-5.click&vd_source=d897272a73b9faef6629bccb52b418e3

简化版代码如下:

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辆车。

猜你喜欢

转载自blog.csdn.net/taifyang/article/details/131276674