基于Yolo v5目標檢測代碼實戰(zhàn)教程

目標檢測是計算機視覺領域的一大任務，大致分為一階段目標檢測與兩階段目標檢測。其中一階段目標檢測模型以YOLO系列為代表。最新的YOLOv5在各個數(shù)據(jù)集上體現(xiàn)出收斂速度快、模型可定制性強的特點，值得關(guān)注。本文主要講解如何從零訓練自己的YOLOv5模型與一些重要參數(shù)的含義。

本文的訓練數(shù)據(jù)使用的是開源數(shù)據(jù)集SHWD，已上傳開源數(shù)據(jù)平臺Graviti，在文末可下載。

在學習或研究目標檢測的同學，后臺回復“210702”可進群一起交流。

一、配置環(huán)境

1.1 創(chuàng)建虛擬環(huán)境

俗話說，環(huán)境配不對，學習兩行淚，首先我們需要安裝Anaconda（Anaconda安裝非常簡單并且百度上有大量資料），然后創(chuàng)建一個專門用來訓練YOLOv5的虛擬環(huán)境。按win+r打開“運行對話框”，輸入“cmd”打開cmd。輸入下面代碼創(chuàng)建虛擬環(huán)境：

conda create -n course_yolov5 python==3.8

其中“course_yolov5”是虛擬環(huán)境的名稱，“python==3.8”是虛擬環(huán)境的python版本。然后我們需要將Ultralytics開源的YOLOv5代碼Clone或下載到本地，可以直接點擊Download ZIP進行下載，

下載地址：https://github.com/ultralytics/yolov5

接下來激活剛剛創(chuàng)建的虛擬環(huán)境并解壓剛下好的壓縮文件，將工作路徑切換到解壓好的文件夾下：

conda activate course_yolov5
cd D:\Study\PyCharm20\PycharmProjects\course_yolov5\yolov5-master
d:

注意：這里需要將" D:\Study\PyCharm20\PycharmProjects\course_yolov5"替換為自己的路徑。

1.2 安裝模塊：

在安裝模塊之前，最好先更換pip源為阿里源或國科大源，然后安裝yolov5需要的模塊，記住工作路徑要在yolov5文件夾下：

python -m pip install -r requirements.txt

如果沒有安裝cuda默認安裝pytorch-cpu版，如果有g(shù)pu可以安裝pytorch-gpu版。

二、檢測

2.1 COCO數(shù)據(jù)集

在正確配置好環(huán)境后就可以檢測自己的圖片或視頻了。YOLOv5已經(jīng)在COCO數(shù)據(jù)集上訓練好，COCO數(shù)據(jù)集一共有80個類別，如果您需要的類別也在其中的話，可以直接用訓練好的模型進行檢測。這80個類分別是：

['person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']

2.2 用預訓練模型進行測試

下面我們先演示如何檢測圖片中的目標。我們一般將要檢測的數(shù)據(jù)放在'./data/images'路徑下，其中 '.' 代表當前路徑即解壓好的yolov5-master文件夾，然后我們在cmd中輸入下面代碼：

python detect.py --source ./data/images/example.jpg --weights weights/yolov5s.pt --conf-thres 0.25

如果沒有下載預訓練模型需要等預訓練模型下好，沒有報錯就說明我們的圖像檢測成功了！檢測好的圖片會放在'./runs/detect'下，各個參數(shù)的含義下面會有具體的介紹。

我們也可以對視頻進行檢測：

python detect.py --source ./data/images/happysheep.mp4 --weights weights/yolov5s.pt --conf-thres 0.25

或一個文件夾中的所有圖片和視頻（圖片支持的格式：'bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'，視頻支持的格式：'mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'），檢測結(jié)果同樣放在'./runs/detect'下。

python detect.py --source ./data/images/ --weights weights/yolov5s.pt --conf-thres 0.25

仔細觀察圖中，我們可以發(fā)現(xiàn)，官方訓練好的模型雖然能將物體框出來，但是對物體的分類存在問題。一會兒將那只邊牧預測為奶牛，一會兒預測為綿羊。

這是目標檢測現(xiàn)階段的難點之一，即不容易區(qū)分圖像中與目標物體外形相似的非目標物體，對于這個問題，我們可以在檢測出目標區(qū)域后再接一個分類器對物體進行分類。另外，地上的樹枝被誤認為是小鳥，我們可以通過調(diào)整置信度閾值來解決這一問題。

2.3 檢測文件參數(shù)說明

yolov5的參數(shù)是由argparse包傳入的，我們可以通過命令行傳入?yún)?shù)，也可以直接設置參數(shù)默認值。我們打開yolov5-master文件加下的detect.py文件，傳參的代碼在主函數(shù)中，幾個關(guān)鍵參數(shù)如下：

parser.add_argument('--weights', nargs='+', type=str, default='yolov5x.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='data/images/happysheep.mp4', help='source')
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')

weights參數(shù)是我們訓練好的權(quán)重文件，yolov5共有四種模型：yolov5s、yolov5m、yolov5l、yolov5x，它們的網(wǎng)絡主干深度和寬度依次遞增，一般情況下檢測效果也遞增，yolov5x的效果最好，yolov5s最差，但是yolov5s網(wǎng)絡參數(shù)最少（14MB），yolov5x參數(shù)最多（166MB）。官方提供的預訓練模型也有四種，大家可以根據(jù)需要設置。

source參數(shù)為檢測數(shù)據(jù)路徑。img-size參數(shù)為檢測時圖像大小，最好與訓練時相同，默認為640。conf-thres為檢測置信度閾值，預測出的置信度高于這個閾值的物體就會顯示在圖中。iou-thres是NMS的IOU閾值，一般為0.3~0.5。

3 訓練自己的數(shù)據(jù)集

3.1 使用labelimg標注圖片

我們訓練模型的第一步就是獲取數(shù)據(jù)集，數(shù)據(jù)集一般通過拍照、爬蟲或直接下載等方式獲得，直接下載的數(shù)據(jù)集如比賽數(shù)據(jù)集都會有標注，我們可以直接使用所給的數(shù)據(jù)進行訓練，但是通過拍照和爬蟲獲得的數(shù)據(jù)需要我們自己進行數(shù)據(jù)標注。目標檢測標注工具有很多，今天主要講解labelimg的標注步驟（文末附labeling下載地址）。

下載并解壓完后打開data文件夾中的predefined_classes.txt文件，可以在文件中寫入自己要定義的類名，如安全帽檢測中有兩類：安全帽和人，我們就在predefined_classes.txt文件中的第一行寫helmet，第二行寫person。

標注圖片步驟如下：

1. 把要標注的數(shù)據(jù)放在img_whole/datasets文件夾中，打開labelimg.exe，
2. 點擊“Open Dir”打開img_whole/datasets文件夾，屏幕中就會顯示圖片，
3. 點擊“Create RectBox”將需要標注的物體框出來并注明其類別，
4. 在標注完所有物體后點擊“Save”保存標注文件至img_whole/xml文件夾中，標注文件格式為xml，
5. 點擊“Next Image”標注下一張圖片，然后繼續(xù)進行步驟3直至標注完全部圖片。

3.2 將xml文件轉(zhuǎn)換為YOLO標注文件

我們隨便打開一個xml文件，xml文件是DOM樹結(jié)構(gòu)，python的xml模塊可以解析它。

我們需要的信息是圖像寬度width、圖像高度height、檢測框左上角坐標xmin和ymin以及檢測框右下角坐標xmax、ymax。我們運行voc_label.py便可在labels文件夾中生成YOLOv5標簽文件，標簽文件中每一行的數(shù)據(jù)為class, x, y, w, h，class是該物體的類別，x,y是檢測框中心坐標，w,h是檢測框的寬和高。voc_label.py代碼：

import xml.etree.ElementTree as ET
import os
from os import getcwd
from tqdm import tqdm

classes = ["helmet", "person"]

def convert(size, box):
    dw = 1. / size[0]
    dh = 1. / size[1]
    x = (box[0] + box[1]) / 2.0
    y = (box[2] + box[3]) / 2.0
    w = box[1] - box[0]
    h = box[3] - box[2]
    x = x * dw
    w = w * dw
    y = y * dh
    h = h * dh
    return (x, y, w, h)

def convert_annotation(image_id):
    in_file = './xml/%s.xml' % (image_id)
    out_file = open('./labels/%s.txt' % (image_id), 'w')
    tree = ET.parse(in_file)
    root = tree.getroot()
    size = root.find('size')
    w = int(size.find('width').text)
    h = int(size.find('height').text)

    for obj in root.iter('object'):
        difficult = obj.find('Difficult').text
        cls = obj.find('name').text
        if cls not in classes or int(difficult) == 1:
            continue
        cls_id = classes.index(cls)
        xmlbox = obj.find('bndbox')
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        bb = convert((w, h), b)
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')

if __name__ == "__main__":
    wd = getcwd()
    print(wd)
    if not os.path.exists('./labels/'):
        os.makedirs('./labels/')
    image_ids = os.listdir('./datasets')
    for image_id in tqdm(image_ids):
        convert_annotation(image_id.split('.')[0])

3.3 訓練

本文的訓練數(shù)據(jù)使用的是開源數(shù)據(jù)集SHWD，數(shù)據(jù)已上傳開源數(shù)據(jù)平臺Graviti，在文末可下載。當我們數(shù)據(jù)很多的時候我們可以將數(shù)據(jù)集劃分為訓練集、驗證集和測試集，比例大致為98:1:1，數(shù)據(jù)較少劃分時比例大概為6:2:2。我們新建一個文件夾datasets，將劃分好的圖片和標簽放入其中，文件夾結(jié)構(gòu)如下：

以訓練集為例，標簽文件名需要與圖片對應：

然后將datasets文件夾放入yolov5-master中的data文件夾中，數(shù)據(jù)準備階段就算完成了。

下面需要修改YOLOv5的配置文件，需要修改的配置文件有兩個，我們先復制一份data/coco.yaml，這里將其重命名為helmet.yaml，修改圖中橫線中的參數(shù)：

1. 在 download前加上一個#注釋掉這段代碼

2. 將train、val、test修改為自己的路徑，以train為例   ./data/helmet/images/train

3. 將nc修改為數(shù)據(jù)的類別數(shù)，如安全帽檢測只檢測人和安全帽，故修改為2

4. 將names修改自己數(shù)據(jù)的類名，如['helmet', 'person']

   修改后的文件如下：

下一個需要修改的文件為模型配置文件，在models文件夾中有四個模型的配置文件：yolov5s.yaml、yolov5m.yaml、yolov5l.yaml和yolov5x.yaml，可以根據(jù)需要選擇相應的模型，這里以yolovx.yaml為例，打開文件，修改文件中的nc為自己的類別數(shù)即可。

yolov5加入了自動判斷是否需要重新k-means聚類anchors的功能，所以該文件中的anchors參數(shù)可以不做修改。如果有需要只需將這里anchors修改為自己計算出的即可。

在修改完這兩個配置文件后就可以開始訓練了！我們在cmd中輸入（記住在yolov5-master路徑下）：

python train.py --weights weights/yolov5x.pt --cfg models/yolov5x.yaml --data data/helmet.yaml --epoch 50 --batch-size 32

如果出現(xiàn)下面界面，說明已經(jīng)開始訓練了：

3.4 訓練參數(shù)解釋

我們打開train.py文件，滑到主函數(shù)部分就可以看到需要傳入的參數(shù)，下面一一介紹一下這些參數(shù)：

• weights：權(quán)重文件路徑，如果是''則重頭訓練參數(shù)，如果不為空則做遷移學習，權(quán)重文件的模型需與cfg參數(shù)中的模型對應

• cfg：存儲模型結(jié)構(gòu)的配置文件

• data：訓練、驗證數(shù)據(jù)配置文件

• hyp：超參數(shù)配置文件，其中的參數(shù)意義下面會解釋

• epochs：指的就是訓練過程中整個數(shù)據(jù)集將被迭代多少次

• batch-size：每次梯度更新的批量數(shù)，太大會導致顯存不足

• img-size：訓練圖片的尺寸

• rect：進行矩形訓練

• resume：恢復最近保存的模型開始訓練

• nosave：僅保存最終checkpoint

• notest：僅測試最后的epoch

• noautoanchor：不進行anchors的k-means聚類

• evolve：進化超參數(shù)

• bucket：gsutil bucket

• cache-images：緩存圖像以加快訓練速度

• image-weights：給圖片加上權(quán)重進行訓練

• device：cuda device, i.e. 0 or 0,1,2,3 or cpu

• multi-scale：多尺度訓練，img-size +/- 50%

• single-cls：單類別的訓練集

• adam：使用adam優(yōu)化器

• name：重命名results.txt to results_name.txt

超參數(shù)配置文件./data/hyp.scratch.yaml參數(shù)解釋：

• lr0：學習率，可以理解為模型的學習速度

• lrf：OneCycleLR學習率變化策略的最終學習率系數(shù)

• momentum：動量，梯度下降法中一種常用的加速技術(shù)，加快收斂

• weight_decay：權(quán)值衰減，防止過擬合。在損失函數(shù)中，weight decay是正則項（regularization）前的一個系數(shù)

• warmup_epochs：預熱學習輪數(shù)

• warmup_momentum：預熱學習初始動量

• warmup_bias_lr：預熱學習初始偏差學習率

• giou：GIoU損失收益

• cls：類別損失收益

• cls_pw：類別交叉熵損失正類權(quán)重

• obj：是否有物體損失收益

• obj_pw：是否有物體交叉熵正類權(quán)重

• iou_t：iou閾值

• anchor_t：多尺度anchor閾值

• fl_gamma：focal loss gamma系數(shù)

• hsv_h：色調(diào)Hue，增強系數(shù)

• hsv_s：飽和度Saturation，增強系數(shù)

• hsv_v：明度Value，增強系數(shù)

• degrees：圖片旋轉(zhuǎn)角度

• translate：圖片轉(zhuǎn)換

• scale：圖片縮放

• shear：圖片仿射變換

• perspec：透視變換

• mosaic：mosaic數(shù)據(jù)增強

• mixup：mixup數(shù)據(jù)增強

  由于時間和能力有限，上面的解釋沒有包含所有參數(shù)，讀者可以通過閱讀源代碼進行進一步理解。

4 訓練結(jié)果與測試

訓練結(jié)束后，在runs/train文件夾下會自動生成訓練結(jié)果，其中包括模型權(quán)重、混淆矩陣、PR曲線等。測試代碼與2.2中的類似，只需將權(quán)重文件路徑修改為訓練好的權(quán)重文件路徑即可，安全帽檢測結(jié)果如下：