1. 程式人生 > 其它 >yolov5訓練自己資料集

yolov5訓練自己資料集

技術標籤:深度學習深度學習

yolov5訓練自己資料集

專案地址:https://github.com/ultralytics/yolov5
環境:Windows pytorch-gpu==1.6.0

** 一、資料集準備**

├── data
│   ├── Annotations  進行 detection 任務時的標籤檔案,xml 形式,檔名與圖片名一一對應
│   ├── images  存放 .jpg 格式的圖片檔案
│   ├── ImageSets  存放的是分類和檢測的資料集分割檔案,包含train.txt, val.txt,trainval.txt,test.txt
│   ├── labels  存放label標註資訊的txt檔案,與圖片一一對應

├── ImageSets(
train,val,test建議按照8:1:1比例劃分) │ ├── train.txt 寫著用於訓練的圖片名稱 │ ├── val.txt 寫著用於驗證的圖片名稱 │ ├── trainval.txt train與val的合集 │ ├── test.txt 寫著用於測試的圖片名稱

1.ImageSets 下檔案的生成
執行我的data資料夾下的makeTxt.py檔案。生成如圖所示data/ImageSets下四個檔案。

檔案內容如圖所示:
在這裡插入圖片描述
makeTxt.py完整程式碼如下:

import os
import random

trainval_percent =
0.1# 訓練集和驗證集所佔比例,剩下的0.2是測試集比例 train_percent = 0.9# 訓練集所佔比例 xmlfilepath = './Annotations' txtsavepath = './ImageSets' total_xml = os.listdir(xmlfilepath) num = len(total_xml) list = range(num) tv = int(num * trainval_percent) tr = int(tv * train_percent) trainval = random.sample(list, tv) train = random.sample(
trainval, tr) ftrainval = open('./ImageSets/trainval.txt', 'w') #只包含名稱,沒有後綴和路徑 ftest = open('./ImageSets/test.txt', 'w') ftrain = open('./ImageSets/train.txt', 'w') fval = open('./ImageSets/val.txt', 'w') for i in list: name = total_xml[i][:-4] + '\n' if i in trainval: ftrainval.write(name) if i in train: ftest.write(name) else: fval.write(name) else: ftrain.write(name) ftrainval.close() ftrain.close() fval.close() ftest.close()

執行voc_label.py檔案,將圖片資料集標註後的xml檔案中的標註資訊讀取出來並寫入txt檔案,執行後在labels資料夾中出現所有圖片資料集的標註資訊,並在data目錄下生成三個txt檔案。
在這裡插入圖片描述

# xml解析包 在labels中存放txt格式的標籤資訊
import xml.etree.ElementTree as ET
import pickle
import os
# os.listdir() 方法用於返回指定的資料夾包含的檔案或資料夾的名字的列表
from os import listdir, getcwd
from os.path import join


sets = ['train', 'test', 'val']
classes = ['jyz', 'xcxj', 'fzc', 'nc','jyz_gz','fzc_gz']


# 進行歸一化操作
def convert(size, box): # size:(原圖w,原圖h) , box:(xmin,xmax,ymin,ymax)
    dw = 1./size[0]     # 1/w
    dh = 1./size[1]     # 1/h
    x = (box[0] + box[1])/2.0   # 物體在圖中的中心點x座標
    y = (box[2] + box[3])/2.0   # 物體在圖中的中心點y座標
    w = box[1] - box[0]         # 物體實際畫素寬度
    h = box[3] - box[2]         # 物體實際畫素高度
    x = x*dw    # 物體中心點x的座標比(相當於 x/原圖w)
    w = w*dw    # 物體寬度的寬度比(相當於 w/原圖w)
    y = y*dh    # 物體中心點y的座標比(相當於 y/原圖h)
    h = h*dh    # 物體寬度的寬度比(相當於 h/原圖h)
    return (x, y, w, h)    # 返回 相對於原圖的物體中心點的x座標比,y座標比,寬度比,高度比,取值範圍[0-1]


# year ='2012', 對應圖片的id(檔名)
def convert_annotation(image_id):
    '''
    將對應檔名的xml檔案轉化為label檔案,xml檔案包含了對應的bunding框以及圖片長款大小等資訊,
    通過對其解析,然後進行歸一化最終讀到label檔案中去,也就是說
    一張圖片檔案對應一個xml檔案,然後通過解析和歸一化,能夠將對應的資訊儲存到唯一一個label檔案中去
    labal檔案中的格式:calss x y w h  同時,一張圖片對應的類別有多個,所以對應的bunding的資訊也有多個
    '''
    # 對應的通過year 找到相應的資料夾,並且開啟相應image_id的xml檔案,其對應bund檔案
    in_file = open('./Annotations/%s.xml' % (image_id), encoding='utf-8')
    # 準備在對應的image_id 中寫入對應的label,分別為
    # <object-class> <x> <y> <width> <height>
    out_file = open('./labels/%s.txt' % (image_id), 'w', encoding='utf-8')
    # 解析xml檔案
    tree = ET.parse(in_file)
    # 獲得對應的鍵值對
    root = tree.getroot()
    # 獲得圖片的尺寸大小
    size = root.find('size')
    # 獲得寬
    w = int(size.find('width').text)
    # 獲得高
    h = int(size.find('height').text)
    # 遍歷目標obj
    for obj in root.iter('object'):
        # 獲得difficult ??
        difficult = obj.find('difficult').text
        # 獲得類別 =string 型別
        cls = obj.find('name').text
        # 如果類別不是對應在我們預定好的class檔案中,或difficult==1則跳過
        if cls not in classes or int(difficult) == 1:
            continue
        # 通過類別名稱找到id
        cls_id = classes.index(cls)
        # 找到bndbox 物件
        xmlbox = obj.find('bndbox')
        # 獲取對應的bndbox的陣列 = ['xmin','xmax','ymin','ymax']
        b = (float(xmlbox.find('xmin').text), float(xmlbox.find('xmax').text), float(xmlbox.find('ymin').text),
             float(xmlbox.find('ymax').text))
        print(image_id, cls, b)
        # 帶入進行歸一化操作
        # w = 寬, h = 高, b= bndbox的陣列 = ['xmin','xmax','ymin','ymax']
        bb = convert((w, h), b)
        # bb 對應的是歸一化後的(x,y,w,h)
        # 生成 calss x y w h 在label檔案中
        out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')


# 返回當前工作目錄
wd = getcwd()
print(wd)


for image_set in sets:
    '''
    對所有的檔案資料集進行遍歷
    做了兩個工作:
    1.講所有圖片檔案都遍歷一遍,並且將其所有的全路徑都寫在對應的txt檔案中去,方便定位
    2.同時對所有的圖片檔案進行解析和轉化,將其對應的bundingbox 以及類別的資訊全部解析寫到label 檔案中去
         最後再通過直接讀取檔案,就能找到對應的label 資訊
    '''
    # 先找labels資料夾如果不存在則建立
    if not os.path.exists('./labels/'):
        os.makedirs('./labels/')
    # 讀取在ImageSets/Main 中的train、test..等檔案的內容
    # 包含對應的檔名稱
    image_ids = open('./ImageSets/%s.txt' % (image_set)).read().strip().split()
    # 開啟對應的2012_train.txt 檔案對其進行寫入準備
    list_file = open('./%s.txt' % (image_set), 'w')
    # 將對應的檔案_id以及全路徑寫進去並換行
    for image_id in image_ids:
        list_file.write('./images/%s.jpg\n' % (image_id))
        # 呼叫  year = 年份  image_id = 對應的檔名_id
        convert_annotation(image_id)
    # 關閉檔案`在這裡插入程式碼片`
    list_file.close()

二、模型訓練和檢測
1.在data目錄下,複製一份coco.yaml檔案並將其重新命名為my.yaml,放在data目錄下,並對cat.yaml中的引數進行配置。

#修改訓練集、驗證集和測試集txt路徑
train: data/train.txt  # 118287 images
val: data/val.txt  # 5000 images
test: data/test.txt  # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794
# 修改類別數
nc: 6

# 修改類名
names: ['jyz', 'xcxj', 'fzc', 'nc','jyz_gz','fzc_gz']

2.修改models目錄下的yolov5s.yaml檔案。這裡可以根據自己選擇的模型大小修改對應檔案。

nc: 6   #nc:類別數,你的類別有多少就填寫多少。從1開始算起,不是0-14這樣算。
#順便一提, 下面兩個引數可控制網路結構,不需修改,與模型大小已經對應。
depth_multiple: 0.33  # 控制模型的深度,假設yolov5l中有三個Bottleneck,那yolov5s中就只有一個Bottleneck。
width_multiple: 0.50  # 控制卷積核的個數設定為0.5,Focus就變成[32, 3],Conv就變成[64, 3, 2]。以此類推,卷積核的個數都變成了設定的一半。

3.訓練引數修改 train.py,也可在終端通過命令進行指定引數。注意修改以下幾個引數。執行train.py

parser.add_argument('--weights', type=str, default='yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
#自己的yaml檔案
parser.add_argument('--data', type=str, default='data/my.yaml', help='data.yaml path')
parser.add_argument('--epochs', type=int, default=300)
#爆視訊記憶體的話調小一點
parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
# 訓練的裝置,cpu;0(表示一個gpu裝置cuda:0);0,1,2,3(多個gpu裝置)。值為空時,訓練時預設使用計算機自帶的顯示卡或CPU
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# dataloader的最大worker數量。這裡我記得報錯,可以將workers修改為0.
parser.add_argument('--workers', type=int, default=4, help='maximum number of dataloader workers')

視覺化:進入環境
tensorboard --logdir runs
根據提示開啟瀏覽器http://localhost:6006/
4.檢測需要修改的引數,一般修改的以下幾個引數,然後執行detect.py。

# 選用訓練的權重
parser.add_argument('--weights', nargs='+', type=str, default='runs/train/exp/weights/best.pt', help='model.pt path(s)')
# source=0表示攝像頭實時檢測,也可輸入視訊路徑
parser.add_argument('--source', type=str, default='inference/videos/jyz.mp4', help='source') 
# 網路輸入圖片大小
parser.add_argument('--view-img', action='store_true', help='display results')
# 是否將檢測的框座標以txt檔案形式儲存,預設False,當需要報存時,可修改路徑default="inference/output"
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
# 設定只保留某一部分類別,新增語句default=[0,1,2]等id。
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')

三、pytorch轉onnx
1.官方自帶的export.py,選擇模型

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

2.根據錯誤提示pip install coremltools、packaging。然後繼續執行1的命令
3.此時 weights下出現三個檔案 :onnx、mlmodel、torchscript
4.使用神經網路Netron,出現網路結構圖。

import netron
netron.start('yolov5s.onnx')

5 . pt模型轉為了onnx模型,可以說已經脫離了pytorch了,但是轉ncnn還需要對模型做一個Simplifier操作,因為轉出的onnx模型還有許多冗餘,這在ncnn裡也是不支援的,避免轉ncnn時各種報錯安裝簡化 開啟weights簡化

python -m onnxsim yolov5s.onnx yolov5s_sim.onnx 

四、 遇到的問題
1.路徑一定不要包含中文,不讓opencv報錯。如果必須要讀中文的話,讀取路徑那裡新增一句話:

img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), -1)

2.batch-size要根據自己視訊記憶體調節,不然會報錯。