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

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

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

本文的訓(xùn)練數(shù)據(jù)使用的是開源數(shù)據(jù)集SHWD。

俗話說，環(huán)境配不對(duì)，學(xué)習(xí)兩行淚，首先我們需要安裝Anaconda（Anaconda安裝非常簡單并且百度上有大量資料），然后創(chuàng)建一個(gè)專門用來訓(xùn)練YOLOv5的虛擬環(huán)境。按win+r打開“運(yùn)行對(duì)話框”，輸入“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或下載到本地，可以直接點(diǎn)擊Download ZIP進(jìn)行下載，

下載地址：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"替換為自己的路徑。

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

python -m pip install -r requirements.txt

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

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

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

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

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

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

我們也可以對(duì)視頻進(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'），檢測結(jié)果同樣放在'./runs/detect'下。

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

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

yolov5的參數(shù)是由argparse包傳入的，我們可以通過命令行傳入?yún)?shù)，也可以直接設(shè)置參數(shù)默認(rèn)值。我們打開yolov5-master文件夾下的detect.py文件，傳參的代碼在主函數(shù)中，幾個(gè)關(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ù)是我們訓(xùn)練好的權(quán)重文件，yolov5共有四種模型：yolov5s、yolov5m、yolov5l、yolov5x，它們的網(wǎng)絡(luò)主干深度和寬度依次遞增，一般情況下檢測效果也遞增，yolov5x的效果最好，yolov5s最差，但是yolov5s網(wǎng)絡(luò)參數(shù)最少（14MB），yolov5x參數(shù)最多（166MB）。官方提供的預(yù)訓(xùn)練模型也有四種，大家可以根據(jù)需要設(shè)置。

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

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

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

標(biāo)注圖片步驟如下：

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

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

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

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

以訓(xùn)練集為例，標(biāo)簽文件名需要與圖片對(duì)應(yīng)：

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

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

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

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

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

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

   修改后的文件如下：

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

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

在修改完這兩個(gè)配置文件后就可以開始訓(xùn)練了！我們?cè)赾md中輸入（記住在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)開始訓(xùn)練了：

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

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

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

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

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

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

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

• img-size：訓(xùn)練圖片的尺寸

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

• resume：恢復(fù)最近保存的模型開始訓(xùn)練

• nosave：僅保存最終checkpoint

• notest：僅測試最后的epoch

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

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

• bucket：gsutil bucket

• cache-images：緩存圖像以加快訓(xùn)練速度

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

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

• multi-scale：多尺度訓(xùn)練，img-size +/- 50%

• single-cls：單類別的訓(xùn)練集

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

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

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

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

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

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

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

• warmup_epochs：預(yù)熱學(xué)習(xí)輪數(shù)

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

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

• giou：GIoU損失收益

• cls：類別損失收益

• cls_pw：類別交叉熵?fù)p失正類權(quán)重

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

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

• iou_t：iou閾值

• anchor_t：多尺度anchor閾值

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

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

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

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

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

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

• scale：圖片縮放

• shear：圖片仿射變換

• perspec：****變換

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

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

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

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