零基礎 從 yolo8 入門計算機視覺超簡單:物體識別、影像分類、軌跡追蹤、姿勢識別

痴者工良發表於2023-12-14


Ultralytics YOLOv8 是備受好評的實時目標檢測和影像分割模型,主要功能是物體識別、分割圖片物體、分類、姿態識別和跟蹤等。Ultralytics 支援使用 CPU、GPU 進行訓練,支援 x64、arm64 等 CPU 架構,支援蘋果的 M1/M2 晶片,支援在邊緣裝置中訓練和使用。

Ultralytics 對於個人免費,使用 【AGPL-3.0 許可】 開源協議,對於企業則需要付費。

Ultralytics 是框架的名稱,YOLOv8 是該框架預設攜帶的模型的版本號,框架預設會自帶多種用途模型,我們一般會選擇在官方的模型上做訓練。

image-20231213082807854

Ultralytics YOLOv8 的物體識別功能很強大,如下圖所示,能夠準確識別圖片中化了妝的妹子為 person(人類),甚至把床也識別出來了。

2

使用 Ultralytics,一般會經歷以下過程:

  • 訓練(Train)模式:在自定義或預載入的資料集上微調您的模型。
  • 驗證(Val)模式:訓練後進行校驗,以驗證模型效能。
  • 預測(Predict)模式:在真實世界資料上釋放模型的預測能力。
  • 匯出(Export)模式:以各種格式使模型準備就緒,部署至生產環境。
  • 跟蹤(Track)模式:將您的目標檢測模型擴充套件到實時跟蹤應用中。
  • 基準(Benchmark)模式:在不同部署環境中分析模型的速度和準確性。

所以,本文也會按照該順序,逐步講解。

安裝 Ultralytics

一般來說,直接使用 pip 安裝即可。

# 從PyPI安裝ultralytics包
pip install ultralytics

官方文件的安裝方式是最詳細的,其它安裝方式直接參考官方文件即可。

官方安裝文件:https://docs.ultralytics.com/zh/quickstart/#ultralytics

安裝 ultralytics 之後,在專案中使用 YOLO 名稱引入 ultralytics:

from ultralytics import YOLO

ultralytics 預設是從執行目錄中讀寫資料集、訓練模型的,這樣可能會比較亂。

image-20231213083616208

可以在程式啟動時,修改配置執行配置:

# 模型和訓練模組,核心
from ultralytics import YOLO
# 設定模組
from ultralytics import settings

# 更新設定
settings.update({'runs_dir': './' , 'tensorboard': False})

settings 也可以匯出配置,或者從檔案中載入配置:

    def load(self):
        """Loads settings from the YAML file."""
        super().update(yaml_load(self.file))

    def save(self):
        """Saves the current settings to the YAML file."""
        yaml_save(self.file, dict(self))

    def update(self, *args, **kwargs):
        """Updates a setting value in the current settings."""
        super().update(*args, **kwargs)
        self.save()

    def reset(self):
        """Resets the settings to default and saves them."""
        self.clear()
        self.update(self.defaults)
        self.save()

settings 可以設定的全部配置如下:

名稱 示例值 資料型別 描述
settings_version '0.0.4' str Ultralytics settings 版本
datasets_dir '/path/to/datasets' str 儲存資料集的目錄
weights_dir '/path/to/weights' str 儲存模型權重的目錄
runs_dir '/path/to/runs' str 儲存實驗執行的目錄
uuid 'a1b2c3d4' str 當前設定的唯一識別符號
sync True bool 是否將分析和崩潰同步到HUB
api_key '' str Ultralytics HUB API Key
clearml True bool 是否使用ClearML記錄
comet True bool 是否使用Comet ML進行實驗跟蹤和視覺化
dvc True bool 是否使用DVC進行實驗跟蹤和版本控制
hub True bool 是否使用Ultralytics HUB整合
mlflow True bool 是否使用MLFlow進行實驗跟蹤
neptune True bool 是否使用Neptune進行實驗跟蹤
raytune True bool 是否使用Ray Tune進行超引數調整
tensorboard True bool 是否使用TensorBoard進行視覺化
wandb True bool 是否使用Weights & Biases記錄

訓練

ultralytics 中常見的檔案格式有兩種,模型以 .pt 結尾,模型或資料集以 .yaml 結尾。

可以基於官方模型進行訓練,或者從已有模型進行訓練,甚至在沒有模型的情況下訓練出自己的模型。

如果已經有模型,則無需再次訓練。

官方文件的 demo 如下:

下面的程式碼不能直接啟動。

from ultralytics import YOLO

# 載入一個模型
model = YOLO('yolov8n.yaml')  # 從YAML建立一個新模型
model = YOLO('yolov8n.pt')  # 載入預訓練模型(推薦用於訓練)
model = YOLO('yolov8n.yaml').load('yolov8n.pt')  # 從YAML建立並轉移權重

# 訓練模型
results = model.train(data='coco128.yaml', epochs=100, imgsz=640)

官方的 yolov8n.pt 模型和 coco128.yaml 資料集,主要是人、動物、常見物體,可用於物體識別。

如果選擇使用官方模型,則第一次使用時,會自動下載 yolov8n.pt 到程式碼所在目錄。

我們先來看看 .train() 訓練模型時的引數。epochs 表示訓練多少輪,一般 10 輪即可,100輪需要非常久的!imgsz 表示圖片大小,當資料集的圖片大小太大時,可以適當降低畫素,官方的資料集圖片大小是 640*480 ,已經提前處理好大小了。

coco128.yaml 是官方提供的資料集合,有 128 張圖片,在程式首次執行時,會自動從官方倉庫中拉取資料集儲存到 datasets/coco128 下面,其中裡面包含了一些圖片和標註。當然,我們也可以從開源社群中獲取更多的資料集,以及自己製作資料集,用來訓練模型。

image-20231213182507127

然後再回到載入模型這裡。

我們透過官方配置檔案,從零訓練出一個模型:

from ultralytics import YOLO

# 載入一個模型
model = YOLO('yolov8n.yaml')  # 從YAML建立一個新模型

# 訓練模型
results = model.train(data='coco128.yaml', epochs=10, imgsz=640)

image-20231213183814026

由於筆者只有 AMD 5600G,是整合顯示卡,因此跑起來特別慢。

如果你有指定的多個,可以在訓練時指定:

results = model.train(data='coco128.yaml', epochs=10, imgsz=640, device=0)

如果不指定,框架會自動選擇 GPU。

可能筆者是 AMD 的 CPU,也有可能是因為不支援整合顯示卡,所以筆者是使用 CPU 訓練的,超級慢。

image-20231213184913056

訓練出的模型會儲存在 detect 目錄下。其它兩種方式訓練的模型也是如此,都會在 detect 目錄下。

1702464360893

我們也可以匯入已有的模型:

# 載入官方預訓練模型
# model = YOLO('yolov8n.pt')
# 訓練過的模型
model = YOLO('detect/train/weights/best.pt') 

在使用 .train() 訓練模型時,可以傳遞很多引數,全部引數說明如下:

描述
model None 模型檔案路徑,例如 yolov8n.pt, yolov8n.yaml
data None 資料檔案路徑,例如 coco128.yaml
epochs 100 訓練的輪次數量
patience 50 早停訓練的等待輪次
batch 16 每批影像數量(-1為自動批大小)
imgsz 640 輸入影像的大小,以整數表示
save True 儲存訓練檢查點和預測結果
save_period -1 每x輪次儲存檢查點(如果<1則禁用)
cache False True/ram, disk 或 False。使用快取載入資料
device None 執行裝置,例如 cuda device=0 或 device=0,1,2,3 或 device=cpu
workers 8 資料載入的工作執行緒數(如果DDP則為每個RANK)
project None 專案名稱
name None 實驗名稱
exist_ok False 是否覆蓋現有實驗
pretrained True (bool 或 str) 是否使用預訓練模型(bool)或從中載入權重的模型(str)
optimizer 'auto' 使用的最佳化器,選擇範圍=[SGD, Adam, Adamax, AdamW, NAdam, RAdam, RMSProp, auto]
verbose False 是否列印詳細輸出
seed 0 隨機種子,用於可重複性
deterministic True 是否啟用確定性模式
single_cls False 將多類資料作為單類訓練
rect False 矩形訓練,每批為最小填充整合
cos_lr False 使用餘弦學習率排程器
close_mosaic 10 (int) 最後輪次禁用馬賽克增強(0為禁用)
resume False 從最後檢查點恢復訓練
amp True 自動混合精度(AMP)訓練,選擇範圍=[True, False]
fraction 1.0 訓練的資料集比例(預設為1.0,即訓練集中的所有影像)
profile False 在訓練期間為記錄器分析ONNX和TensorRT速度
freeze None (int 或 list, 可選) 在訓練期間凍結前n層,或凍結層索引列表
lr0 0.01 初始學習率(例如 SGD=1E-2, Adam=1E-3)
lrf 0.01 最終學習率 (lr0 * lrf)
momentum 0.937 SGD動量/Adam beta1
weight_decay 0.0005 最佳化器權重衰減5e-4
warmup_epochs 3.0 熱身輪次(小數ok)
warmup_momentum 0.8 熱身初始動量
warmup_bias_lr 0.1 熱身初始偏差lr
box 7.5 框損失增益
cls 0.5 cls損失增益(根據畫素縮放)
dfl 1.5 dfl損失增益
pose 12.0 姿態損失增益(僅限姿態)
kobj 2.0 關鍵點obj損失增益(僅限姿態)
label_smoothing 0.0 標籤平滑(小數)
nbs 64 標稱批大小
overlap_mask True 訓練期間掩碼應重疊(僅限分割訓練)
mask_ratio 4 掩碼降取樣比率(僅限分割訓練)
dropout 0.0 使用dropout正則化(僅限分類訓練)
val True 訓練期間驗證/測試

模型驗證

前一小節中,我們使用了官方帶有 128 張圖片的 coco128.yaml 資料集進行訓練,訓練完成後,需要進行模型驗證,以便在 coco128.yaml 資料集上驗證訓練過的模型的準確性。

此外,資料集也分為多種:

  • 訓練集(Training Set):用於訓練和調整模型引數。
  • 驗證集(Validation Set):用於驗證模型精度和調整模型引數。
  • 測試集(Test Set):用於驗證模型的泛化能力。

對於模型驗證這一部分,官方文件的描述不多,同時本文作為入門文章,不會深入講解這些細節,所以我們只需要記得自行訓練的模型,都加上模型驗證的程式碼模板即可。

在已經訓練的模型中驗證的模板如下:

from ultralytics import YOLO

# 載入訓練後的模型
model = YOLO('detect/train3/weights/best.pt')

# 驗證模型
metrics = model.val()  # 無需引數,資料集和設定記憶
metrics.box.map    # map50-95
metrics.box.map50  # map50
metrics.box.map75  # map75
metrics.box.maps   # 包含每個類別的map50-95列表

一般,在訓練之後立即驗證:

from ultralytics import YOLO

model = YOLO('yolov8n.pt')  # 載入預訓練模型(推薦用於訓練)

# 訓練模型
results = model.train(data='coco128.yaml', epochs=10, imgsz=640)

# 驗證模型
metrics = model.val()  # 無需引數,資料集和設定記憶
metrics.box.map    # map50-95
metrics.box.map50  # map50
metrics.box.map75  # map75
metrics.box.maps   # 包含每個類別的map50-95列表

驗證結果:

前面提到過,coco128.yaml 資料集主要是人、動物和常見物體的圖片,因此控制檯列出了每種標記(person 人類、car 汽車等) 的識別驗證結果。

當然,我們是入門階段,這些資訊對我們作用不大。

image-20231213191057501

預測 & 識別

來到了激動人心的階段,我們可以使用現有的模型或官方模型對圖片進行識別。

你可以將下面兩張圖片下載儲存:

img1

img2

下面兩張圖片,是程式碼識別後的結果,會自動將圖片中識別到的物體標記名稱以及使用矩形框選出來。

1

2

from ultralytics import YOLO
from PIL import Image
# 載入訓練後的模型
model = YOLO('detect/train3/weights/best.pt')


# 匯入並識別圖片,results 是處理後的結果
results = model(['img1.jpg', 'img2.jpg'])

# 儲存處理後的圖片到目錄中
i = 0
for r in results:
    i += 1
    im_array = r.plot()  # 繪製包含預測結果的BGR numpy陣列
    im = Image.fromarray(im_array[..., ::-1])  # RGB PIL影像
    im.save('results/' + str(i) + '.jpg')  # 儲存影像

當然,我們也可以在處理圖片之前預覽圖片:

im.show()  # 彈出視窗,預覽影像
im.save('results/' + str(i) + '.jpg')  # 儲存影像

如果只需要獲取識別資訊,不需要處理圖片,那麼可以這樣列印資訊:

from ultralytics import YOLO
from PIL import Image

# 載入訓練後的模型
model = YOLO('detect/train3/weights/best.pt')

# 匯入並識別圖片,results 是處理後的結果
results = model(['img1.jpg', 'img2.jpg'])

# 處理結果列表
for result in results:
    boxes = result.boxes  # 邊界框輸出的 Boxes 物件
    masks = result.masks  # 分割掩碼輸出的 Masks 物件
    keypoints = result.keypoints  # 姿態輸出的 Keypoints 物件
    probs = result.probs  # 分類輸出的 Probs 物件
    print("圖片邊框")
    print(boxes)
    print("分割掩碼")
    print(masks)
    print("姿勢輸出")
    print(keypoints)
    print("分類輸出")
    print(probs)

注意,上面的結果除了邊框有值,其它都是 None,因為我們只識別圖片,並沒有做分類、切割、姿勢識別等。

前面提到,results 是處理之後的結果,results 是一個 Results 物件,Results 物件具有的屬性如下:

屬性 型別 描述
orig_img numpy.ndarray 原始影像的numpy陣列。
orig_shape tuple 原始影像的形狀,格式為(高度,寬度)。
boxes Boxes, 可選 包含檢測邊界框的Boxes物件。
masks Masks, 可選 包含檢測掩碼的Masks物件。
probs Probs, 可選 包含每個類別的機率的Probs物件,用於分類任務。
keypoints Keypoints, 可選 包含每個物件檢測到的關鍵點的Keypoints物件。
speed dict 以毫秒為單位的每張圖片的預處理、推理和後處理速度的字典。
names dict 類別名稱的字典。
path str 影像檔案的路徑。

Results具有以下方法:

方法 返回型別 描述
__getitem__() Results 返回指定索引的Results物件。
__len__() int 返回Results物件中的檢測數量。
update() None 更新Results物件的boxes, masks和probs屬性。
cpu() Results 將所有張量移動到CPU記憶體上的Results物件的副本。
numpy() Results 將所有張量轉換為numpy陣列的Results物件的副本。
cuda() Results 將所有張量移動到GPU記憶體上的Results物件的副本。
to() Results 返回將張量移動到指定裝置和dtype的Results物件的副本。
new() Results 返回一個帶有相同影像、路徑和名稱的新Results物件。
keys() List[str] 返回非空屬性名稱的列表。
plot() numpy.ndarray 繪製檢測結果。返回帶有註釋的影像的numpy陣列。
verbose() str 返回每個任務的日誌字串。
save_txt() None 將預測儲存到txt檔案中。
save_crop() None 將裁剪的預測儲存到save_dir/cls/file_name.jpg
tojson() None 將物件轉換為JSON格式。

除了本地圖片外,ultralytics 還支援影片流,可以從本地或網路中匯入要識別的資源。對於影片流的處理,本文的後續章節再介紹。

ultralytics 支援識別的其它資源:

來源 引數 型別 備註
image 'image.jpg' strPath 單個影像檔案。
URL 'https://ultralytics.com/images/bus.jpg' str 影像的 URL 地址。
screenshot 'screen' str 擷取螢幕影像。
PIL Image.open('im.jpg') PIL.Image RGB 通道的 HWC 格式影像。
OpenCV cv2.imread('im.jpg') np.ndarray BGR 通道的 HWC 格式影像 uint8 (0-255)
numpy np.zeros((640,1280,3)) np.ndarray BGR 通道的 HWC 格式影像 uint8 (0-255)
torch torch.zeros(16,3,320,640) torch.Tensor RGB 通道的 BCHW 格式影像 float32 (0.0-1.0)
CSV 'sources.csv' strPath 包含影像、影片或目錄路徑的 CSV 檔案。
video✅ 'video.mp4' strPath 如 MP4, AVI 等格式的影片檔案。
directory✅ 'path/' strPath 包含影像或影片檔案的目錄路徑。
glob ✅ 'path/*.jpg' str 匹配多個檔案的萬用字元模式。使用 * 字元作為萬用字元。
YouTube ✅ 'https://youtu.be/LNwODJXcvt4' str YouTube 影片的 URL 地址。
stream ✅ 'rtsp://example.com/media.mp4' str RTSP, RTMP, TCP 或 IP 地址等流協議的 URL 地址。
multi-stream✅ 'list.streams' strPath 一個流 URL 每行的 *.streams 文字檔案,例如 8 個流將以 8 的批處理大小執行。

如果要了解每種資源的匯入方式,可以開啟 https://docs.ultralytics.com/zh/modes/predict/#_4

本文不再贅述所有匯入方式,不過本文後續會介紹影片流等資源識別。

image-20231213193519004

匯出

ultralytics 支援將訓練後的模型匯出為 ONNX、 TensorRT、 CoreML 等格式,然後我們可以使用其它框架或其它語言進行呼叫,例如 C# 的 ML.NET 框架。

from ultralytics import YOLO

# 載入模型 .pt 檔案模型
model = YOLO('path/to/best.pt')

# 匯出為其它型別的模型
model.export(format='onnx')

匯出為其它模型時,需要安裝對應支援包,如果本地沒有安裝過,則第一次會自動安裝。

image-20231213194310259

程式碼執行後,控制檯會列印出模型檔案被匯出到哪個位置:

image-20231213194354063

ultralytics 支援將模型匯出為以下格式:

格式 format 引數 模型 後設資料 引數
PyTorch - yolov8n.pt -
TorchScript torchscript yolov8n.torchscript imgsz, optimize
ONNX onnx yolov8n.onnx imgsz, half, dynamic, simplify, opset
OpenVINO openvino yolov8n_openvino_model/ imgsz, half
TensorRT engine yolov8n.engine imgsz, half, dynamic, simplify, workspace
CoreML coreml yolov8n.mlpackage imgsz, half, int8, nms
TF SavedModel saved_model yolov8n_saved_model/ imgsz, keras
TF GraphDef pb yolov8n.pb imgsz
TF Lite tflite yolov8n.tflite imgsz, half, int8
TF Edge TPU edgetpu yolov8n_edgetpu.tflite imgsz
TF.js tfjs yolov8n_web_model/ imgsz
PaddlePaddle paddle yolov8n_paddle_model/ imgsz
ncnn ncnn yolov8n_ncnn_model/ imgsz, half

追蹤

Ultralytics 的追蹤可以處理影片中的物體。前面筆者已經介紹了透過圖片進行物體識別,其實還支援對圖片進行姿勢識別、圖片分類、將物體從圖片中切割出來,這些等後面的章節再介紹。Ultralytics 的追蹤也支援物體識別、姿勢識別等,在本節中,筆者將介紹如何在影片中識別和追蹤物體,在後面的章節中會介紹更多的示例。

Ultralytics YOLO 擴充套件了其物體檢測功能,以提供強大且多功能的物體追蹤:

  • 實時追蹤: 在高幀率影片中無縫追蹤物體。
  • 支援多個追蹤器: 從多種成熟的追蹤演算法中選擇。
  • 自定義追蹤器配置: 透過調整各種引數來定製追蹤演算法,以滿足特定需求。

Ultralytics 預設有兩種追蹤器。

  • BoT-SORT - 模型檔案為 botsort.yaml ,預設使用,不需要配置。
  • ByteTrack - 模型檔案為bytetrack.yaml

以下程式碼演示了使用 cv2 載入影片檔案,並逐幀識別圖片中的物體,然後每識別一幀,就顯示到桌面中。

from ultralytics import YOLO
import cv2

# 載入自己訓練的物體識別模型
model = YOLO('detect/train3/weights/best.pt')
# 或者使用官方訓練的物體識別模型
# model = YOLO('yolov8n.pt')

# 使用 cv2 載入影片檔案
video_path = "1.mp4"
cap = cv2.VideoCapture(video_path)

# 迴圈遍歷影片幀
while cap.isOpened():
    # 從影片讀取一幀
    success, frame = cap.read()

    if success:
        # 在幀上執行YOLOv8追蹤,持續追蹤幀間的物體
        results = model.track(frame, persist=True)

        # 在幀上展示結果
        annotated_frame = results[0].plot()

        # 使用 cv2 彈出視窗,並展示帶註釋的幀
        # 也就是一邊識別,一邊播放影片
        cv2.imshow("YOLOv8 Tracking", annotated_frame)

        # 如果按下'q'則退出迴圈
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        # 如果影片結束則退出迴圈
        break

# 釋放影片捕獲物件並關閉顯示視窗
cap.release()
cv2.destroyAllWindows()

狗追兔子

當然,由於筆者下載的影片比較模糊,以及訓練的資料集不大,識別出來的物體名稱不是很理想,不過其追蹤器確實確實牛批。

影像分割提取

影像分割用於在識別圖片中的物體後,將物體從圖片中提取出來。

示例程式碼如下:

from ultralytics import YOLO

# 載入 yolov8n-seg 模型
model = YOLO('yolov8n-seg.pt')

# 訓練模型
# results = model.train(data='coco128-seg.yaml', epochs=100, imgsz=640)

# 驗證
# ... ...

# 對影像進行處理
results = model(['img1.jpg', 'img2.jpg'])

# 儲存處理後的圖片到目錄中
i = 0
for r in results:
    i += 1
    im_array = r.save_crop(save_dir="results")

提取到的物體圖片會被儲存到:

image-20231213202642747

如果不需要提取物體,由於 ultralytics 會先建立圖片蒙版,因此我們可以匯出帶有物體蒙版的圖片。

from ultralytics import YOLO
from PIL import Image

# 載入 yolov8n-seg 模型
model = YOLO('yolov8n-seg.pt')

# 訓練模型
# results = model.train(data='coco128-seg.yaml', epochs=100, imgsz=640)

# 驗證
# ... ...

# 對影像進行處理
results = model(['img1.jpg', 'img2.jpg'])

# 儲存處理後的圖片到目錄中
i = 0
for r in results:
    i += 1
    im_array = r.plot()  # 繪製包含預測結果的BGR numpy陣列
    im = Image.fromarray(im_array[..., ::-1])  # RGB PIL影像
    im.save('results/' + str(i) + '.jpg')  # 儲存影像

image-20231213203023876

分類

很明顯,用於分類圖片。

分類圖片需要使用官方的 yolov8n-cls.pt 模型。

官方提供了一個 mnist160 資料集,該資料集是從 0-9 的手寫圖片,因此我們訓練模型之後,也是用來做手寫數字識別的。

這裡隨便寫三個數字:

1702471024196

由於資料量不大,因此我們可以直接訓練然後使用訓練後的模型提取圖片中的文字:

from ultralytics import YOLO

# 載入 yolov8n-cls 模型
model = YOLO('yolov8n-cls.pt')

# 訓練模型
results = model.train(data='mnist160', epochs=100, imgsz=64)

# 驗證
# ... ...

# 對影像進行處理
results = model(['666.png'])

# 儲存處理後的圖片到目錄中
i = 0
for r in results:
    r.save_txt(txt_file="results/cls"+ str(i) + '.txt')
    i += 1

因為訓練的資料集是單數字的,因此對多數字的支援有點問題。

image-20231213204204767

改成大大的 6 ,再次識別:

1702471564364

image-20231213204702260

姿勢識別

姿勢識別用於在圖片或影片中給人體進行打點、劃線,然後追蹤人體的姿勢變化。

姿勢識別

示例程式碼如下:

import cv2
from ultralytics import YOLO

# 載入 yolov8n-pose 模型,並進行訓練
model = YOLO('yolov8n-pose.pt')
results = model.train(data='coco8-pose.yaml', epochs=10, imgsz=640)

# 開啟影片檔案
video_path = "1.mp4"
cap = cv2.VideoCapture(video_path)

# 迴圈遍歷影片幀
while cap.isOpened():
    # 從影片讀取一幀
    success, frame = cap.read()

    if success:
        # 在幀上執行YOLOv8追蹤,持續追蹤幀間的物體
        results = model.track(frame, persist=True)

        # 在幀上展示結果
        annotated_frame = results[0].plot()

        # 展示帶註釋的幀
        cv2.imshow("YOLOv8 Tracking", annotated_frame)

        # 如果按下'q'則退出迴圈
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        # 如果影片結束則退出迴圈
        break

# 釋放影片捕獲物件並關閉顯示視窗
cap.release()
cv2.destroyAllWindows()

軌跡生成

在影片中識別追蹤物體時,還可以對物體的執行軌跡進行追蹤。

手頭上沒有更多影片演示,將就一下。

軌跡生成

示例程式碼如下:

from collections import defaultdict

import cv2
import numpy as np

from ultralytics import YOLO

# 載入 yolov8n.pt 模型,也可以使用其它模型
model = YOLO('yolov8n.pt')

# 開啟影片
video_path = "0.mp4"
cap = cv2.VideoCapture(video_path)

# 儲存歷史記錄
track_history = defaultdict(lambda: [])

# 處理影片的每一幀
while cap.isOpened():
    success, frame = cap.read()

    if success:
        # 追蹤
        results = model.track(frame, persist=True)

        # 讀取當前幀中物體的框
        boxes = results[0].boxes.xywh.cpu()
        # 如果未識別到影像有框時,則忽略
        if results[0].boxes.id is None:
            continue
        # 獲取每個被追蹤的物體 id
        track_ids = results[0].boxes.id.int().cpu().tolist()

        annotated_frame = results[0].plot()

        # 繪製
        for box, track_id in zip(boxes, track_ids):
            x, y, w, h = box
            track = track_history[track_id]
            track.append((float(x), float(y))) 
            if len(track) > 30:
                track.pop(0)

            points = np.hstack(track).astype(np.int32).reshape((-1, 1, 2))
            cv2.polylines(annotated_frame, [points], isClosed=False, color=(230, 230, 230), thickness=10)

        cv2.imshow("YOLOv8 Tracking", annotated_frame)

        if cv2.waitKey(1) & 0xFF == ord("q"):
            break
    else:
        break

cap.release()
cv2.destroyAllWindows()

相關文章