使用 YOLO 和計算機視覺進行目標跟蹤
YOLO 是一種能夠實時進行目標檢測的深度學習算法。您可以使用預訓練的 YOLO 模型,如 YOLOv8 或 YOLOv9,或者在需要時在自定義數據集上訓練自己的模型。在本文中,我將帶您了解如何使用預訓練的 YOLO 模型進行目標跟蹤。這是最簡單的教程,我們只處理簡單的目標檢測。
讓我們看看目標檢測所涉及的步驟。
- 我們將加載模型
- 我們將加載視頻
- 我們將逐幀提取視頻
- 對每一幀進行預測并使用 OpenCV 顯示結果
下面的圖像總結了這些步驟:
目標檢測步驟 讓我們開始編碼。
第 0 步:環(huán)境設置
我們需要兩個庫 ultralytics 和 opencv,您可以使用以下代碼安裝這些庫。
pip install ultralytics opencv-python
導入所有必要的庫。
import cv2 as cv
from ultralytics import YOLO
第 1 步:加載模型
#load the model
model = YOLO('yolov8n.pt')
在 YOLO 中加載模型非常容易。我們可以直接使用 ultralytics 提供的 YOLO() 函數,并指定模型名稱。一旦運行,它會自動在代碼所在的文件夾中下載指定的 YOLO 模型。
第 2 步:加載視頻
我們使用 OpenCV 的 VideoCapture() 函數來加載視頻。這個函數的輸入是視頻的路徑。
path = "../videos/football_test.mp4"
#read video
vs = cv.VideoCapture(path)
第 3 步:提取幀
一旦我們加載了視頻,在下一步我們將逐個提取視頻中的幀。視頻是由連續(xù)拼接在一起的幀組成的。
- vs = cv.VideoCapture(path):包含指向視頻的指針。
- vs.read():允許我們逐幀讀取視頻。一旦我們調用vs上的read函數,它將以迭代的方式返回當前的視頻幀。
- cv.imshow('image', frame):顯示每個提取出的幀。
path = "../videos/football_test.mp4"
#read video
vs = cv.VideoCapture(path)
#load the model
model = YOLO('yolov8n.pt')
while True:
(grabbed,frame) = vs.read()
if not grabbed:
break
cv.imshow('image', frame)
cv.waitKey(500)
第 4 步:應用預測
模型的predict()方法是用于對每一幀進行預測的方法。預測結果存儲在results變量中。在while循環(huán)中,我們從視頻中提取每一幀,并在每一幀上進行預測。簡單來說,我正在嘗試在視頻中的每一幀中檢測對象。
results = model.predict(frame,stream=False)
path = "../videos/football_test.mp4"
#read video
vs = cv.VideoCapture(path)
#load the model
model = YOLO('yolov8n.pt')
while True:
(grabbed,frame) = vs.read()
if not grabbed:
break
results = model.predict(frame,stream=False)
cv.imshow('image', frame)
cv.waitKey(500)
第 5 步:繪制邊界框
為了繪制邊界框,我創(chuàng)建了一個函數,該函數將輸入圖像、數據(包含預測結果)、檢測到的類別名稱,并返回帶有檢測到的對象邊界框的圖像。drawBox(data, image, name):data是預測詳情,image是幀,name是檢測到的類別名稱,例如:人、汽車等。
def drawBox(data,image,name):
x1, y1, x2, y2, conf, id = data
p1 = (int(x1), int(y1))
p2 = (int(x2), int(y2))
cv.rectangle(image, p1, p2, (0, 0, 255), 3)
cv.putText(image, name, p1, cv.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 3)
return image
results = model.predict(frame,stream=False)
for result in results:
for data in result.boxes.data.tolist():
#print(data)
id = data[5]
drawBox(data, frame,detection_classes[id])
由于結果包含了幀中所有檢測到的對象的詳細信息,我們在上述代碼中對每個檢測到的對象進行解碼。對于每個檢測,我們調用`drawBox(data, frame, detection_classes[id])`來在檢測周圍繪制邊界框。
第 6 步:完整代碼
您也可以在此處訪問完整代碼:https://github.com/sunnykumar1516/YOLO/blob/main/ReadVideoApplyYoloV8.py。
import cv2 as cv
from ultralytics import YOLO
def drawBox(data,image,name):
x1, y1, x2, y2, conf, id = data
p1 = (int(x1), int(y1))
p2 = (int(x2), int(y2))
cv.rectangle(image, p1, p2, (0, 0, 255), 3)
cv.putText(image, name, p1, cv.FONT_HERSHEY_SIMPLEX, 3, (0, 0, 255), 3)
return image
detection_classes= []
path = "../videos/football_test.mp4"
#read video
vs = cv.VideoCapture(path)
#load the model
model = YOLO('yolov8n.pt')
while True:
(grabbed,frame) = vs.read()
if not grabbed:
break
results = model.predict(frame,stream=False)
detection_classes = results[0].names
for result in results:
for data in result.boxes.data.tolist():
#print(data)
id = data[5]
drawBox(data, frame,detection_classes[id])
完整代碼參考:https://github.com/sunnykumar1516/YOLO/blob/main