技術(shù)博客丨使用YOLOv5模型進(jìn)行目標檢測！

目標檢測是計算機視覺(jué)領(lǐng)域的一大任務(wù)，大致分為一階段目標檢測與兩階段目標檢測。其中一階段目標檢測模型以YOLO系列為代表。

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

本文的訓練數據使用的是開(kāi)源數據集SHWD。

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

conda create -n course_yolov5 python==3.8

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

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

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

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

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

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

python -m pip install -r requirements.txt

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

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

['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']

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

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

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

我們也可以對視頻進(jìn)行檢測：

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

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

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

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

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

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

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

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

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

標注圖片步驟如下：

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

我們隨便打開(kāi)一個(gè)xml文件，xml文件是DOM樹(shù)結構，python的xml模塊可以解析它。

我們需要的信息是圖像寬度width、圖像高度height、檢測框左上角坐標xmin和ymin以及檢測框右下角坐標xmax、ymax。我們運行voc_label.py便可在labels文件夾中生成YOLOv5標簽文件，標簽文件中每一行的數據為class, x, y, w, h，class是該物體的類(lèi)別，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])

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

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

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

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

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

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

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

4. 將names修改自己數據的類(lèi)名，如['helmet', 'person']

   修改后的文件如下：

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

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

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

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

如果出現下面界面，說(shuō)明已經(jīng)開(kāi)始訓練了：

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

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

• cfg：存儲模型結構的配置文件

• data：訓練、驗證數據配置文件

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

• epochs：指的就是訓練過(guò)程中整個(gè)數據集將被迭代多少次

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

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

• rect：進(jìn)行矩形訓練

• resume：恢復最近保存的模型開(kāi)始訓練

• nosave：僅保存最終checkpoint

• notest：僅測試最后的epoch

• noautoanchor：不進(jìn)行anchors的k-means聚類(lèi)

• evolve：進(jìn)化超參數

• bucket：gsutil bucket

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

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

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

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

• single-cls：?jiǎn)晤?lèi)別的訓練集

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

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

超參數配置文件./data/hyp.scratch.yaml參數解釋?zhuān)?/p>

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

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

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

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

• warmup_epochs：預熱學(xué)習輪數

• warmup_momentum：預熱學(xué)習初始動(dòng)量

• warmup_bias_lr：預熱學(xué)習初始偏差學(xué)習率

• giou：GIoU損失收益

• cls：類(lèi)別損失收益

• cls_pw：類(lèi)別交叉熵損失正類(lèi)權重

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

• obj_pw：是否有物體交叉熵正類(lèi)權重

• iou_t：iou閾值

• anchor_t：多尺度anchor閾值

• fl_gamma：focal loss gamma系數

• hsv_h：色調Hue，增強系數

• hsv_s：飽和度Saturation，增強系數

• hsv_v：明度Value，增強系數

• degrees：圖片旋轉角度

• translate：圖片轉換

• scale：圖片縮放

• shear：圖片仿射變換

• perspec：****變換

• mosaic：mosaic數據增強

• mixup：mixup數據增強

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

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