YOLOv5實(shí)現(xiàn)自定義對(duì)象訓(xùn)練與OpenVINO部署全解析

點(diǎn)擊上方“小白學(xué)視覺(jué)”，選擇加"星標(biāo)"或“置頂”

重磅干貨，第一時(shí)間送達(dá)

大家好，前面寫了一個(gè)OpenVINO部署YOLOv5推理的教程，收到很多反饋！這里就再寫了一篇如何使用YOLOv5訓(xùn)練自定義對(duì)象檢測(cè)，從數(shù)據(jù)準(zhǔn)備到推理整個(gè)環(huán)節(jié)，幫助大家更好的使用YOLOv5來(lái)解決實(shí)際問(wèn)題。整個(gè)文章主要分為三個(gè)部分，分別是數(shù)據(jù)準(zhǔn)備與YOLO格式數(shù)據(jù)轉(zhuǎn)換，模型訓(xùn)練與推理測(cè)試，模型轉(zhuǎn)換為ONNX與部署。

關(guān)于環(huán)境搭建與測(cè)試請(qǐng)看這里：

YOLOv5在最新OpenVINO 2021R02版本的部署與代碼演示詳解

https://github.com/EscVM/OIDv4_ToolKit

用這個(gè)工具可以很方便的下載想要的數(shù)據(jù)，Open Image數(shù)據(jù)集包括包括100W+張圖像，對(duì)象檢測(cè)支持600個(gè)類別，所以很多圖像都可以通過(guò)這個(gè)來(lái)下載，避免自己收集圖像數(shù)據(jù)與標(biāo)注數(shù)據(jù)的時(shí)間。所以首先需要執(zhí)行

git clone https://github.com/EscVM/OIDv4_ToolKit

然后運(yùn)行

pip3 install -r requirements.txt

安裝好所有依賴，之后再下載數(shù)據(jù)集

python3 main.py downloader --classes Apple Orange --type_csv validation

意思是下載蘋果跟橘子兩個(gè)類別的圖像標(biāo)注數(shù)據(jù)。下載好的數(shù)據(jù)集文件夾結(jié)構(gòu)如下：

分為三個(gè)文件夾，分別對(duì)應(yīng)訓(xùn)練集、測(cè)試集、驗(yàn)證集。其中jpg文件是每個(gè)圖像，labels里面的txt同名文件是每個(gè)圖像中對(duì)象標(biāo)注信息，格式如下：

name_of_the_class left top right bottomname_of_the_class 對(duì)象類別名稱left top標(biāo)注框左上角坐標(biāo)right bottom 標(biāo)注框右下角坐標(biāo)

轉(zhuǎn)換為YOLO標(biāo)注格式數(shù)據(jù)，首先看一下標(biāo)注文本文件與圖像的文件夾結(jié)構(gòu)

其中images是所有圖像的集合，labels是所有標(biāo)簽信息集合。train表示訓(xùn)練集、valid表示測(cè)試集。這里需要注意的是，labels信息中的標(biāo)簽信息YOLO格式標(biāo)注框需要把原始標(biāo)注信息的標(biāo)注框格式從：

Left top right bottom

轉(zhuǎn)換為

Center_x, center_y, width, height

并歸一化到0~1之間，這部分我寫了一個(gè)腳本來(lái)完成label標(biāo)簽的生成。代碼如下：

for f in files:
    if os.path.isfile(os.path.join(current_dir, f)):
        image = cv.imread(os.path.join(current_dir, f))
        label_file = os.path.join(current_dir, "label", f.replace(".jpg", ".txt"))
        yolo_label = f.replace(".jpg", ".txt")
        data_label_text_f = os.path.join(valid_label_dir, yolo_label)
        file_write_obj = open(data_label_text_f, 'w')

        with open(label_file) as f:
            boxes = [line.strip() for line in f.readlines()]
        clazz_index = -1
        # create new file
        for box in boxes:
            anno_info = box.split(" ")
            if anno_info[0] == "Cricket":
                print("class name: ", anno_info[0] + " ball")
                x1 = float(anno_info[2])
                y1 = float(anno_info[3])
                x2 = float(anno_info[4])
                y2 = float(anno_info[5])
                clazz_index = 0
            else:
                print("class name: ", anno_info[0])
                x1 = float(anno_info[1])
                y1 = float(anno_info[2])
                x2 = float(anno_info[3])
                y2 = float(anno_info[4])
                clazz_index = 1
            h, w, c = image.shape
            cx = (x1 + (x2 - x1) / 2) / w
            cy = (y1 + (y2 - y1) / 2) / h
            sw = (x2 - x1) / w
            sh = (y2 - y1) / h
            file_write_obj.write("%d %f %f %f %f\n"%(clazz_index, cx, cy, sw, sh))
            # cv.rectangle(image, (int(x1), int(y1)), (int(x2), int(y2)),(0, 0, 255), 2, 8)
        file_write_obj.close()

最后需要?jiǎng)?chuàng)建一個(gè)dataset.ymal文件，放在與data文件夾同一層，它的內(nèi)容如下：

# train andval datasets (image directory or *.txt file with image paths)train:football_training/data/images/train/val: football_training/data/images/valid/ # number ofclassesnc: 2 # class namesnames:['Cricketball', 'Football']

其中

Cricketball與Football是我從Open Image中下載的兩個(gè)類別數(shù)據(jù)。nc:2 表示識(shí)別兩個(gè)類別。

這樣就完成了整個(gè)數(shù)據(jù)集準(zhǔn)備部分。

這里我分別基于yolov5l.ymal與yolov5s.ymal完成了模型訓(xùn)練，需要修改的只有一個(gè)地方，就是把類別數(shù)目從80改為2。圖示如下：

然后執(zhí)行訓(xùn)練腳本的命令行如下：

python train.py --data --epochs 300 football_training/dataset.yaml --cfg football_traing/yolov5l.yaml --weights '' --batch-size 1

推薦的參數(shù)中沒(méi)有這么小的batchi-size，我是用筆記本訓(xùn)練，發(fā)現(xiàn)2就爆內(nèi)存了，人窮！有條件的把這個(gè)參數(shù)調(diào)大點(diǎn)，效果會(huì)好！訓(xùn)練時(shí)可以通過(guò)tensorboard查看訓(xùn)練實(shí)時(shí)變換，圖示如下：

最終每個(gè)類別的AP得分

最終訓(xùn)練完成得到best.pt文件

運(yùn)行測(cè)試視頻

python detect.py --source football_training/test.mp4 --weights football_training/best.pt --conf 0.5

通過(guò)下面的腳本轉(zhuǎn)換為ONNX文件

python models/export.py --weights football_training/best.pt --img 640 --batch 1

然后再轉(zhuǎn)換為OpenVINO的IR中間文件格式，C++推理與部署代碼實(shí)現(xiàn)參見

YOLOv5在最新OpenVINO 2021R02版本的部署與代碼演示詳解

運(yùn)行效果如下：