基于OpenCV+YOLOv5实现车辆跟踪与计数(附源码)

2023-11-27 16:19:46 浏览数 (1)

导 读

本文主要介绍基于OpenCV YOLOv5实现车辆跟踪与计数的应用,并给出源码。

资源下载

基础代码和视频下载地址:

代码语言:javascript复制
https://github.com/freedomwebtech/win11vehiclecount

main.py代码:

代码语言:javascript复制
import cv2
import torch
import numpy as np
from tracker import *


model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)

cap=cv2.VideoCapture('highway.mp4')

count=0
tracker = Tracker()




def POINTS(event, x, y, flags, param):
    if event == cv2.EVENT_MOUSEMOVE :  
        colorsBGR = [x, y]
        print(colorsBGR)
        

cv2.namedWindow('FRAME')
cv2.setMouseCallback('FRAME', POINTS)

while True:
    ret,frame=cap.read()
    if not ret:
        break
    count  = 1
    if count % 3 != 0:
        continue
    frame=cv2.resize(frame,(1020,600))
    results=model(frame)
    results.pandas().xyxy[0]
       
        
    
    cv2.imshow("FRAME",frame)
    if cv2.waitKey(0)&0xFF==27:
        break
cap.release()
cv2.destroyAllWindows()

tracker.py代码:

代码语言:javascript复制
import math


class Tracker:
    def __init__(self):
        # Store the center positions of the objects
        self.center_points = {}
        # Keep the count of the IDs
        # each time a new object id detected, the count will increase by one
        self.id_count = 0


    def update(self, objects_rect):
        # Objects boxes and ids
        objects_bbs_ids = []

        # Get center point of new object
        for rect in objects_rect:
            x, y, w, h = rect
            cx = (x   x   w) // 2
            cy = (y   y   h) // 2

            # Find out if that object was detected already
            same_object_detected = False
            for id, pt in self.center_points.items():
                dist = math.hypot(cx - pt[0], cy - pt[1])

                if dist < 35:
                    self.center_points[id] = (cx, cy)
#                    print(self.center_points)
                    objects_bbs_ids.append([x, y, w, h, id])
                    same_object_detected = True
                    break

            # New object is detected we assign the ID to that object
            if same_object_detected is False:
                self.center_points[self.id_count] = (cx, cy)
                objects_bbs_ids.append([x, y, w, h, self.id_count])
                self.id_count  = 1

        # Clean the dictionary by center points to remove IDS not used anymore
        new_center_points = {}
        for obj_bb_id in objects_bbs_ids:
            _, _, _, _, object_id = obj_bb_id
            center = self.center_points[object_id]
            new_center_points[object_id] = center

        # Update dictionary with IDs not used removed
        self.center_points = new_center_points.copy()
        return objects_bbs_ids

下载测试视频highway.mp4(download.txt中有链接):

安装ultralytics:

代码语言:javascript复制
pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple

实现车辆检测

添加代码查看YoloV5模型检测输出信息:

添加车辆检测结果绘制:

实现车辆跟踪

将车辆矩形坐标保存到list中:

添加对象跟踪:

车辆跟踪效果:

实现车辆计数

【1】划定一个多边形区域(黄色),车辆经过此区域则计数;

代码语言:javascript复制
cv2.polylines(frame,[np.array(areal,np.int32)],True,(0,255,255),3)

【2】绿色点标注车辆矩形框右下角坐标点,如下图:

代码语言:javascript复制
cv2.circle(frame,(x3,y3),4,(0,255,0),-1)

【3】计数原理:判断绿色点是否在黄色四边形内,如果在内计数 1

计算点到多边形的距离使用的是OpenCV中pointPolygonTest函数:

当result > 0表示点在轮廓内,也就是车辆在多边形内。此时将车辆id编号添加到集合area_1中,然后通过len(area_1)作为车辆计数值。

注意:这里为什么不直接用count = 1来计算数量?如果用count = 1会造成很多重复计数,而代码中利用了python集合的特性,add函数添加元素时,如果元素已经存在,则不重复添加,有效的避免了重复计数问题。

最终效果如下:

视频演示:https://mpvideo.qpic.cn/0bc3uqabwaaakqaflodda5svbjgddosaagya.f10002.mp4?

0 人点赞