第一個 YOLOv5 的辨識模型

John Hu

Published in

No silver bullet

12 min readFeb 17, 2022

前陣子做完 Lunar Lander 後，小弟沒有花很多時間在 RL 學習上，主要懶得一直研究數學演算法。就在茶餘飯後的聊天中發現一個需求：物件辨識。所以，小弟就開始學習最近很紅的 YOLO 了。

YOLO 的歷史，就不在這裡贅述了，網路上已經很多文章在討論。至於選擇 YOLOv5 的原因，單純認爲它是最新的版本。雖然，有人說 YOLOv4 還沒有結束，YOLOv5 也還沒有被原作者認可。就還是覺得用一下 v5 好了。

準備資料

每次訓練之前，我們都要準備訓練用的資料。我們可以在 Kaggle 上搜尋 YOLO 這個關鍵字來找到一些可以訓練用的現有資料集。

爲了簡化整個訓練的時間跟模型，我們選擇口罩偵測，也就是讓電腦檢查照片裡面有沒有「有戴口罩的臉」、或是「沒戴口罩的臉」這兩種情境。所以，我們的資料集就變成這個搜尋結果。意外發現，上面居然也有不少臺灣的照片，例如：

照官方的說法，我們最好能夠找到符合下面規則的資料集：

夠多的圖片，例如：1,500 張以上的圖片。
每一個辨識類別需要有 10,000 個標示才會比較好。以我們這個案例來看，辨識類別有兩種：有口罩、跟沒口罩；所以，我們需要有 20,000 個標示。
照片的多樣性也是很重要的一件事，所以我們選用的資料集必須要包含每個國家的情境、每種口罩（素色跟花紋）。
每個標示要很一致，例如：戴一半的口罩算不算有戴口罩。
正確的標示出每一個物品，例如：有的資料集只有在標示的時候，只有框出口罩的位置，而沒有框出人臉加口罩的位置。
需要有 0 ~ 10% 的單純背景圖片，以我們這個例子就是，資料集中要有同時沒有「有戴口罩的臉」跟「沒戴口罩的臉」的圖片。

為了簡化資料集處理流程，我們先忽略標示一致性、標示正確性、跟背景圖片這三部分，單純把網路上找到的 YOLO 口罩資料集合併成一個，以下就是總數：

Training set：8,414 張圖片、20,951 個標示
Validation set：1,665 張圖片、3,519 個標示

下載 YOLOv5

YOLOv5 可以直接在 GitHub 下載。我們只需要依照 installation 的指示執行，就可以安裝完成：

git clone https://github.com/ultralytics/yolov5
cd yolov5
pip install -r requirements.txt

這邊的一個條件就是，它需要 Python 3.7 以上的版本，我們這篇使用 Python 3.9。

下載完成後，我們可以拿 YOLOv5 官方預先訓練好的模型來做辨識測試：

python detect.py --source data/images/bus.jpg

辨識結果會被儲存在 runs/detect/exp1/bus.jpg 。如果執行很多次，其中的 exp1 就會一路累加上去。辨識前與辨識後的圖片是：

開始訓練

根據官方文件，預設的模型大小有五種，模型越大執行的速度就越慢，但是準確性就越高，反之亦然；這個是一個 trade-off。由於，我們希望可以把模型放到 GitHub Pages 上面直接用手機來執行，所以，我們選擇最小的模型：YOLOv5n (Nano)。

YOLOv5 的訓練方式很簡單，大概是小弟用過最簡單的方式。不過，會這麼簡單的原因是，YOLOv5 把全部的 config 跟 hyperparameters 都設定上了預設值，我們可以用這個指令來執行預設的訓練：

python train.py --data mask/obj.data\
                --weights yolov5n.pt

其中的 data 是我們的 YOLO dataset format 的入口：

path: mask
train: images/train
val: images/valid
nc: 2
names: ['mask', 'no_mask']

path: 資料集目錄，相對於程式執行目錄。
train: 訓練用的資料集，相對於資料集目錄。
val: 驗證用的資料集，相對於資料集目錄。
nc: number of classes 的縮寫，我們要偵測的類別數量
names: 顯示在辨識圖片上面的名字

為了方便起見，我們給了這個專案一個名字（mask-detector-n），同時我們也明確指定設定檔，所以，訓練指令就變成：

python train.py --data mask/obj.data\
                --cfg models/yolov5n.yaml\
                --weights yolov5n.pt\
                --project mask-detector-n\
                --img-size 320

cfg：我們直接給 yolov5n 的模型。基本上，我們可以不用給這個參數。如果，大家要全新訓練一個，不要繼承官方預先訓練的模型，就一定要給這個參數。
weights：我們直接繼承官方的 yolov5n.pt 的模型接續著訓練。
project：是等下訓練結果放置的資料夾，如果沒指定的話會是 runs 資料夾。
img-size：是預計輸入圖片最長邊的大小，由於我們要放在網站上執行，必須要把輸入的檔案縮小才能有比較好的 FPS。

如果，大家只要測試的話，記得加上 --epochs 10 先跑個 10 輪來看看。接著要繼續訓練下去的話，記得把 weights 改成跑出來的結果，以我們這個例子來看就是要改用 mask-detector-n/exp/weights/best.pt 。YOLOv5 就會用之前訓練的結果繼續訓練下去了。

解讀結果

訓練時，YOLOv5 會給我們下面的訊息：

Epoch   gpu_mem       box       obj       cls    labels  img_size                                                                                                                                   0/9       0.65G   0.02388   0.01097 0.0004858        60       320

訓練的項目說明說明如下：

Epoch: 第幾次的訓練。
gpu_mem: 目前使用多少 GPU 的記憶體。
box: 偵測出來外框位置的 loss，越小越好。
obj: 偵測出來物件數量的 loss，越小越好。
cls: 物件判讀類別的 loss，越小越好。
labels: 這一個批次執行中標註了幾個物件。
img_size: 單一測試的圖片大小。

每次 Epoch 執行完成後，驗證階段會有下面的訊息：

Class Images     Labels        P          R     mAP@.5  mAP@.5:.95:   
all      451       1033    0.886       0.76      0.842       0.568

驗證的項目說明如下：

Images: 驗證圖片數量，我們這裡用一個小的資料集來做解說。
Labels: 模型標註出來的物件數量，我們這裡用一個小的資料集來做解說。
P: Precision，標註出來的結果有多少是正確的，越大越好。
R: Recall，全部正確的結果中，有多少被標註出來，越大越好。
mAP@.5: 標註出來的面積與正確面積有 50% 以上重疊正確性的平均值，越大越好。
mAP@.5:.95: 標註出來的面積與正確面積有 50% 以上到 95% 中間取 10 個點的到重疊正確性的平均值，越大越好。

由於 mAP 的計算較為複雜一些，大家如果想知道多一點可以參考這篇文章。

下面是我們用正式的資料集訓練 300 輪後的最佳結果：

P: 0.93033
R: 0.88909
mAP@.5: 0.93557
mAP@.5:.95: 0.61845

訓練時的 loss 如下：

box_loss: 0.02412
obj_loss: 0.01234
cls_loss: 0.00101

這個的訓練是用一台 ASUS Zephyrus G14 筆電來執行，規格如下：

CPU：AMD Ryzen 9 5900HS 3.1GHz
RAM：32 GB DDR4–3200
獨立顯示卡： NVIDIA GeForce RTX 3060 6GB GDDR6

全部訓練的時間是約 7 個小時。

轉換模型

訓練完成後，我們需要決定要將 YOLOv5 的模型轉換成什麼格式的模型。一開始，我們決定把這個模型放到網站上面執行，所以就只有 TensorFlow.js 的這個選項。

依照官方文件的講法，YOLOv5 自帶了一個 export.py 可以把 YOLOv5 的模型轉換成許多格式：torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs 。我們這邊有兩種選項，一個是轉 saved_model 在轉 tfjs、另一個就是直接轉 tfjs。為了方便起見，我們就直接轉 tfjs，不經過 saved_model 這個格式：

python export.py --weights mask/model/best.pt\
                 --include tfjs\
                 --img-size 320\
                 --data mask/obj.data

weights：最後訓練出來最好的 model。
include：輸出格式；如果要一次輸出多種格式，可以一次給多個。
img-size：圖片大小。訓練時我們使用 320x320 當成輸入的圖片，這邊最好也使用一樣的設定。我們測試過，如果給不同大小的圖片，也是可以運作，不過，效果會沒原本的這麼好。也因為在手機上運作，圖片越小，運作的速度越快。
data：資料集格式的設定。

完成後，我們會在 mask/model/best.pt 旁邊找到兩個資料夾：best_saved_model、best_web_model 。沒想到，export.py 也是先轉 saved_model 後再轉 tfjs。依照名字來看，我們可以知道 best_web_model，就是 tfjs 的模型。

部署到網站

老實說，我們剛開始看 YOLOv5 的時候，只知道它的輸入是一張圖片：[320, 320, 3] 的大小，但是不知道圖片的像素值是不是要經過任何轉換。後來，我們發現了這個 GitHub repo，它解答了我們全部問題：

tf.image
  .resizeBilinear(tf.browser.fromPixels(canvas), MODEL_SIZE)
  .div(255.0)
  .expandDims(0)

所以，我們修改了一下程式碼做出這個 GitHub repo：https://github.com/john-hu/mask。最後，部署網址是：https://john.hu/mask。

這個 WebApp 使用 React + WebRTC + TensorFlow.js 做出來的，它會在啟動時，打開電腦或是手機的攝影機（預設是背後鏡頭），同時載入 TensorFlow.js 的模型。接著，它會做一次 Warm up 的辨識，把大部分的資源都載入到 WebGL 中。全部完成後，就可以拿著這個 WebApp 來拍看看是不是要有人沒戴口罩了。

如果，大家訓練的模型不同，你們只需要修改一下這個檔案，把裡面的常數改掉就可以直接使用自己的模型：

export const WEIGHTS_PATH = "/mask/modeln/model.json"; // model 路徑
export const CLASS_NAMES = ["mask", "no mask"]; // 不同類別的名稱
export const CLASS_COLORS = ["#70FF70", "#FF7070"]; // 不同類別的顏色
export const MODEL_SIZE = [320, 320]; // 輸入的圖片大小
export const DETECTION_THRESHOLD = 0.6; // 判讀的最低分數

最後別忘了把轉出的 best_web_model 的內容放到 public/modeln 裡面。