在視頻中計數對象可能看起來有挑戰性,但借助Python和OpenCV的強大功能,變得令人意外地易于實現。在本文中,我們將探討如何使用YOLO(You Only Look Once)目標檢測模型在視頻流或文件中計數對象。我們將該過程分解為簡單的步驟,使初學者能夠輕松跟隨。
本文將分為以下幾個部分:
-
需求
-
啟發式:汽車計數
-
檢測過濾
-
啟發式:實現
-
結論
需求
在我們深入了解該過程之前,讓我們確保已安裝所需的庫。主要需要:
-
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
一旦我們查看了主要要求,就該了解我們將開發用于從視頻中計數對象的啟發式的時間了。
啟發式:汽車計數
在此示例中,我們將使用一個視頻場景,其中將對汽車進行計數。圖2顯示了一個示例幀。
用于計數汽車的視頻幀
為了計數汽車,我們將使用Yolov5來檢測視頻中的對象。基于檢測到的對象,我們將過濾與汽車、公共汽車和卡車有關的類別。由于檢測基于邊界框(具有坐標xmin、ymin、xmax、ymax的多邊形),我們將需要獲取每個邊界框的中心點(xc, yc),該中心點將是我們對象的參考點。
最后,我們將繪制一個多邊形,該多邊形將是計數對象的參考,也就是說,如果對象的參考點在多邊形內,我們將增加對象計數器,否則計數器不受影響。在下圖中,我們可以看到多邊形和多邊形內的汽車數量的表示。
檢測(綠色點)、多邊形(紅色線)和計數器
到此為止,我們已經知道了需求是什么,以及我們將實施用于計數對象的啟發式的方法。現在可以加載模型: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行,我們生成一個迭代器,只要有要顯示的幀,它就會保持活動狀態。一旦幀完成,與迭代器相關的對象就會被釋放(第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分別計算中心點。
在這一點上,我們已經下載了模型,過濾了預測,并獲得了每個對象的中心點。現在,我們唯一需要的是生成決定啟發式區域的多邊形。因此,讓我們繼續下一節!
啟發式:實現
我們將定義的多邊形可能會因視頻、透視等而有所不同。在這種情況下,例如此示例,我們將使用8個點,如下圖所示:
具有坐標的多邊形
一旦我們定義了多邊形,我們唯一需要做的就是驗證每個對象的參考點是否在多邊形內。如果在多邊形內,我們就會增加一個計數器,如果不在,我們就繼續。
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)
結論
在本文中,我們看到了如何從視頻中實現一個對象計數器。我們開發了一種計算汽車、卡車和公共汽車的實現,基于一個定義的多邊形,即如果對象在多邊形內,計數器就會增加。
-
計數器
+關注
關注
32文章
2256瀏覽量
94562 -
模型
+關注
關注
1文章
3243瀏覽量
48836 -
目標檢測
+關注
關注
0文章
209瀏覽量
15610
原文標題:基于YOLOv5的視頻計數 — 汽車計數實現
文章出處:【微信號:vision263com,微信公眾號:新機器視覺】歡迎添加關注!文章轉載請注明出處。
發布評論請先 登錄
相關推薦
評論