技術博客丨使用YOLOv5模型進行目標檢測！

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

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

本文的訓練數據使用的是開源數據集SHWD。

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

conda create -n course_yolov5 python==3.8

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

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

接下來激活剛剛創建的虛擬環境并解壓剛下好的壓縮文件，將工作路徑切換到解壓好的文件夾下：

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默認安裝pytorch-cpu版，如果有gpu可以安裝pytorch-gpu版。

在正確配置好環境后就可以檢測自己的圖片或視頻了。YOLOv5已經在COCO數據集上訓練好，COCO數據集一共有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']

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

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

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

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

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'），檢測結果同樣放在'./runs/detect'下。

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

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

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

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

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

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

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

標注圖片步驟如下：

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

我們隨便打開一個xml文件，xml文件是DOM樹結構，python的xml模塊可以解析它。

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

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

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

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

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

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

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

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

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

   修改后的文件如下：

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

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

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

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

如果出現下面界面，說明已經開始訓練了：

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

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

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

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

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

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

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

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

• rect：進行矩形訓練

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

• nosave：僅保存最終checkpoint

• notest：僅測試最后的epoch

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

• evolve：進化超參數

• bucket：gsutil bucket

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

• image-weights：給圖片加上權重進行訓練

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

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

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

• adam：使用adam優化器

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

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

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

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

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

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

• warmup_epochs：預熱學習輪數

• warmup_momentum：預熱學習初始動量

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

• giou：GIoU損失收益

• cls：類別損失收益

• cls_pw：類別交叉熵損失正類權重

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

• obj_pw：是否有物體交叉熵正類權重

• 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數據增強

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

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