本文共 4763 字,大约阅读时间需要 15 分钟。
本文来源公众号“OpenCV与AI深度学习”,仅用于学术分享,侵权删,干货满满。
监控摄像头可以有效地用于各种场景下的车辆计数和交通流量统计。先进的计算机视觉技术(例如对象检测和跟踪)可应用于监控录像,以识别和跟踪车辆在摄像机视野中移动。本文将详细介绍如何使用YOLOv9和OpenCV实现车辆跟踪计数的实现步骤和源码。
首先,我们需要安装Ultralytics库,用于直接使用YOLOv9预训练模型。
pip install ultralytics
接下来,我们创建了一个名为tracker.py的Python文件,用于跟踪对象。以下是该文件的实现代码:
import mathclass CustomTracker: def __init__(self): # 存储对象的中心位置 self.custom_center_points = {} # 记录对象ID的数量 self.custom_id_count = 0 def custom_update(self, custom_objects_rect): # 参数:新的对象矩形列表 custom_objects_bbs_ids = [] # 遍历每个新对象矩形 for custom_rect in custom_objects_rect: x, y, w, h = custom_rect # 计算矩形中心点 cx = (x + x + w) // 2 cy = (y + y + h) // 2 # 检查是否是重复检测的对象 same_object_detected = False for custom_id, pt in self.custom_center_points.items(): dist = math.hypot(cx - pt[0], cy - pt[1]) if dist < 35: self.custom_center_points[custom_id] = (cx, cy) custom_objects_bbs_ids.append([x, y, w, h, custom_id]) same_object_detected = True break # 如果是新对象,分配ID if not same_object_detected: self.custom_center_points[self.custom_id_count] = (cx, cy) custom_objects_bbs_ids.append([x, y, w, h, self.custom_id_count]) self.custom_id_count += 1 # 清理未使用的中心点 new_custom_center_points = {} for custom_obj_bb_id in custom_objects_bbs_ids: _, _, _, _, custom_object_id = custom_obj_bb_id center = self.custom_center_points[custom_object_id] new_custom_center_points[custom_object_id] = center self.custom_center_points = new_custom_center_points.copy() return custom_objects_bbs_ids 接下来,我们编写车辆计数的主要代码。以下是完整的代码实现:
import cv2import pandas as pdfrom ultralytics import YOLOfrom tracker import *# 导入模型model = YOLO('yolov9c.pt')# 定义类别列表class_list = [ 'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']# 初始化跟踪器和计数器tracker = CustomTracker()count = 0# 加载视频cap = cv2.VideoCapture('traffictrim.mp4')# 获取视频属性fps = int(cap.get(cv2.CAP_PROP_FPS))width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 创建视频写入器fourcc = cv2.VideoWriter_fourcc(*'mp4v')output_video_path = 'output_video.mp4'out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))# 主循环while True: ret, frame = cap.read() if not ret: break count += 1 # 检测对象 results = model.predict(frame) a = results[0].boxes.data a = a.detach().cpu().numpy() px = pd.DataFrame(a).astype("float") # 过滤车辆 list = [] for index, row in px.iterrows(): x1 = int(row[0]) y1 = int(row[1]) x2 = int(row[2]) y2 = int(row[3]) d = int(row[5]) c = class_list[d] if 'car' in c: list.append([x1, y1, x2, y2]) # 更新跟踪器 bbox_id = tracker.custom_update(list) # 绘制结果 for bbox in bbox_id: x3, y3, x4, y4, id = bbox cx = (x3 + x4) // 2 cy = (y3 + y4) // 2 # 绘制中心点 cv2.circle(frame, (cx, cy), 4, (0, 0, 255), -1) # 绘制边框 cv2.rectangle(frame, (x3, y3), (x4, y4), (0, 255, 0), 2) # 绘制ID cv2.putText(frame, str(id), (cx, cy), cv2.FONT_HERSHEY_COMPLEX, 0.8, (0, 255, 255), 2) # 绘制红线 y = 308 offset = 7 cv2.line(frame, (282, 308), (1004, 308), (0, 0, 255), 3) cv2.putText(frame, 'red line', (280, 308), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA) # 绘制计数器 downwards = len(counter_down) cv2.putText(frame, ('Vehicle Counter - ') + str(downwards), (60, 40), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 1, cv2.LINE_AA) # 写入输出视频 out.write(frame) 以上代码实现了车辆的检测、跟踪和计数功能。通过YOLOv9模型检测车辆,并结合OpenCV图像处理技术,实现了对车辆的实时跟踪和计数。
转载地址:http://gesfk.baihongyu.com/