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要根據自己視訊記憶體調節,不然會報錯。