基于YOLOv5的視頻計數(shù) — 汽車計數(shù)實(shí)現(xiàn)

在視頻中計數(shù)對象可能看起來有挑戰(zhàn)性，但借助Python和OpenCV的強(qiáng)大功能，變得令人意外地易于實(shí)現(xiàn)。在本文中，我們將探討如何使用YOLO（You Only Look Once）目標(biāo)檢測模型在視頻流或文件中計數(shù)對象。我們將該過程分解為簡單的步驟，使初學(xué)者能夠輕松跟隨。

本文將分為以下幾個部分：

• 需求

• 啟發(fā)式：汽車計數(shù)

• 檢測過濾

• 啟發(fā)式：實(shí)現(xiàn)

• 結(jié)論

需求

在我們深入了解該過程之前，讓我們確保已安裝所需的庫。主要需要：

• PyTorch：通過PyTorch Hub，我們將訪問Ultralytics存儲庫以下載Yolov5模型。

• OpenCV：用于加載、操作和顯示視頻的所有實(shí)用程序。

• Matplotlib（可選）：我們將使用此實(shí)用程序在多邊形內(nèi)進(jìn)行點(diǎn)驗(yàn)證。

如代碼片段1所示，requirements.txt文件中列出了這些要求。

opencv-python==4.8.1.78
torch==2.1.0
matplotlib==3.8.0
ultralytics==8.0.203
pandas==2.1.2
requests==2.31.0

一旦我們查看了主要要求，就該了解我們將開發(fā)用于從視頻中計數(shù)對象的啟發(fā)式的時間了。

啟發(fā)式：汽車計數(shù)

在此示例中，我們將使用一個視頻場景，其中將對汽車進(jìn)行計數(shù)。圖2顯示了一個示例幀。

用于計數(shù)汽車的視頻幀

為了計數(shù)汽車，我們將使用Yolov5來檢測視頻中的對象?；跈z測到的對象，我們將過濾與汽車、公共汽車和卡車有關(guān)的類別。由于檢測基于邊界框（具有坐標(biāo)xmin、ymin、xmax、ymax的多邊形），我們將需要獲取每個邊界框的中心點(diǎn)（xc, yc），該中心點(diǎn)將是我們對象的參考點(diǎn)。

最后，我們將繪制一個多邊形，該多邊形將是計數(shù)對象的參考，也就是說，如果對象的參考點(diǎn)在多邊形內(nèi)，我們將增加對象計數(shù)器，否則計數(shù)器不受影響。在下圖中，我們可以看到多邊形和多邊形內(nèi)的汽車數(shù)量的表示。

檢測（綠色點(diǎn)）、多邊形（紅色線）和計數(shù)器

到此為止，我們已經(jīng)知道了需求是什么，以及我們將實(shí)施用于計數(shù)對象的啟發(fā)式的方法?，F(xiàn)在可以加載模型：Yolov5 Nano

在本例中，我們將使用Yolov5的nano版本（即yolov5n），我們將通過PyTorch Hub從Ultralytics存儲庫中擴(kuò)展它。同樣，為了加載和在每一幀上生成迭代器，我們將使用OpenCV（即cv2），下述代碼是具體的實(shí)現(xiàn)方式：

import cv2
import torch


VIDEO_PATH="data/traffic.mp4"
HUB="ultralytics/yolov5"
YOLO="yolov5n"


def count_cars(cap: cv2.VideoCapture):


    model = torch.hub.load(HUB, model=YOLO, pretrained=True)


    while cap.isOpened():
        status, frame = cap.read()


        if not status:
            break


        # Detection filtering and heuristic
        # will be implemented here.


        cv2.imshow("frame", frame)


        if cv2.waitKey(10) & 0xFF == ord('q'):
             break


    cap.release()




if __name__ == '__main__':


    cap = cv2.VideoCapture(VIDEO_PATH)
    count_cars(cap)

正如我們所看到的，我們已經(jīng)定義了count_cars()函數(shù)，我們將在整個項(xiàng)目中對其進(jìn)行更新。在4-6行，我們定義了視頻所在路徑、hub和模型名稱的變量。從那里，讓我們迅速跳到31-32行，在那里通過初始化cap對象加載視頻，然后將其傳遞給count_cars()函數(shù)。

返回到第10行，通過PyTorch Hub，我們下載并初始化了yolov5n模型。隨后，在第12行，我們生成一個迭代器，只要有要顯示的幀，它就會保持活動狀態(tài)。一旦幀完成，與迭代器相關(guān)的對象就會被釋放（第26行）。

在第13行，我們讀取幀，驗(yàn)證是否成功讀取，并顯示它們（第21行）。在這一部分，將出現(xiàn)一個窗口，用于查看從此迭代器顯示的視頻。最后，第23行是在按q鍵時刪除彈出窗口。

檢測過濾

過濾檢測是指從Yolo預(yù)測中提取感興趣的類別的過程。在這種情況下，我們將過濾掉分?jǐn)?shù)大于0.5且類別為汽車、公共汽車或卡車的檢測。同樣，我們將需要找到邊界框的中心點(diǎn)，我們將其用作對象的參考點(diǎn)。下面代碼顯示了這兩個函數(shù)的實(shí)現(xiàn)。

import pandas as pd


def get_bboxes(preds: object):
    df = preds.pandas().xyxy[0]
    df = df[df["confidence"] >= 0.5]
    df = df[df["name"].isin(["car", "bus", "truck"])]


    return df[["xmin", "ymin", "xmax", "ymax"]].values.astype(int)


def get_center(bbox):
    center = ((bbox[0] + bbox[2]) // 2, (bbox[1] + bbox[3]) // 2)
    return center

正如我們所看到的，我們定義了兩個函數(shù)get_bboxes()和get_center()。get_bboxes()函數(shù)（第3行）旨在提取所有分?jǐn)?shù)大于0.5并過濾掉已經(jīng)提到的類別的預(yù)測，返回一個坐標(biāo)形式的邊界框的numpy數(shù)組[xmin, ymin, xmax, ymax]。

get_center()函數(shù)（第10行）接收一個帶有邊界框坐標(biāo)的numpy數(shù)組，并使用方程xc, yc = (xmin + xmax) // 2, (ymin + ymax) // 2分別計算中心點(diǎn)。

在這一點(diǎn)上，我們已經(jīng)下載了模型，過濾了預(yù)測，并獲得了每個對象的中心點(diǎn)?，F(xiàn)在，我們唯一需要的是生成決定啟發(fā)式區(qū)域的多邊形。因此，讓我們繼續(xù)下一節(jié)！

啟發(fā)式：實(shí)現(xiàn)

我們將定義的多邊形可能會因視頻、透視等而有所不同。在這種情況下，例如此示例，我們將使用8個點(diǎn)，如下圖所示：

具有坐標(biāo)的多邊形

一旦我們定義了多邊形，我們唯一需要做的就是驗(yàn)證每個對象的參考點(diǎn)是否在多邊形內(nèi)。如果在多邊形內(nèi)，我們就會增加一個計數(shù)器，如果不在，我們就繼續(xù)。

import cv2
import numpy as np
import matplotlib.path as mplPath


POLYGON = np.array([
    [333, 374],
    [403, 470],
    [476, 655],
    [498, 710],
    [1237, 714],
    [1217, 523],
    [1139, 469],
    [1009, 393],
])




def is_valid_detection(xc, yc):
    return mplPath.Path(POLYGON).contains_point((xc, yc))


def count_cars(cap: object):


    model = torch.hub.load(HUB, model=YOLO, pretrained=True)


    while cap.isOpened():
        status, frame = cap.read()


        if not status:
            break


        preds = model(frame)
        bboxes = get_bboxes(preds)


        detections = 0
        for box in bboxes:
            xc, yc = get_center(box)


            if is_valid_detection(xc, yc):
                detections += 1

讓我們注意到在第5行，我們定義了多邊形。在第17行，我們定義了關(guān)鍵函數(shù)：is_valid_detection()，它旨在驗(yàn)證參考點(diǎn)（xc, yc）是否在多邊形內(nèi)。這個函數(shù)在第37行調(diào)用，如果為真，它會增加有效檢測計數(shù)器，否則什么也不做。

最后，為了可視化，我們將添加一些OpenCV行來顯示計數(shù)器、每輛檢測到的汽車的參考點(diǎn)和多邊形。

def count_cars(cap: object):


    model = torch.hub.load(HUB, model=YOLO, pretrained=True)
    
    while cap.isOpened():
        status, frame = cap.read()


        if not status:
            break


        preds = model(frame)
        bboxes = get_bboxes(preds)


        detections = 0
        for box in bboxes:
            xc, yc = get_center(box)
            
            if is_valid_detection(xc, yc):
                detections += 1
                
            # Draw poit of reference for each detection
            cv2.circle(img=frame, center=(xc, yc), radius=5, color=(0,255,0), thickness=-1)
            # Draw bounding boxes for each detection
            cv2.rectangle(img=frame, pt1=(box[0], box[1]), pt2=(box[2], box[3]), color=(255, 0, 0), thickness=1)
        # Draw the counter
        cv2.putText(img=frame, text=f"Cars: {detections}", org=(100, 100), fontFace=cv2.FONT_HERSHEY_PLAIN, fontScale=3, color=(0,0,0), thickness=3)
        # Draw the polygon
        cv2.polylines(img=frame, pts=[POLYGON], isClosed=True, color=(0,0,255), thickness=4)
        # Display frame
        cv2.imshow("frame", frame)

結(jié)論

在本文中，我們看到了如何從視頻中實(shí)現(xiàn)一個對象計數(shù)器。我們開發(fā)了一種計算汽車、卡車和公共汽車的實(shí)現(xiàn)，基于一個定義的多邊形，即如果對象在多邊形內(nèi)，計數(shù)器就會增加。