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

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

• 需求

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

• 檢測過濾

• 啟發(fā)式：實現

• 結論

需求

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

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

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

• Matplotlib（可選）：我們將使用此實用程序在多邊形內進行點驗證。

如代碼片段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ā)用于從視頻中計數對象的啟發(fā)式的時間了。

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

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

用于計數汽車的視頻幀

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

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

檢測（綠色點）、多邊形（紅色線）和計數器

到此為止，我們已經知道了需求是什么，以及我們將實施用于計數對象的啟發(fā)式的方法。現在可以加載模型：Yolov5 Nano

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

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)

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

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

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

檢測過濾

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

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

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

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

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

啟發(fā)式：實現

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

具有坐標的多邊形

一旦我們定義了多邊形，我們唯一需要做的就是驗證每個對象的參考點是否在多邊形內。如果在多邊形內，我們就會增加一個計數器，如果不在，我們就繼續(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行，我們定義了關鍵函數：is_valid_detection()，它旨在驗證參考點（xc, yc）是否在多邊形內。這個函數在第37行調用，如果為真，它會增加有效檢測計數器，否則什么也不做。

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

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)

結論

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