1、环境搭建
https://github.com/ultralytics/yolov5
2、安装需要的软件
pip install -U -r requirements.txt
3、准备数据
在data文件下建立上面三个文件(Annotations、images与ImageSets,labels后续我们脚本生成)其中Annotations存放xml文件,images图像,ImageSets新建Main文件存放train与test文件(脚本生成),labels是标签文件
划分训练集与测试集(为了充分利用数据集我们只划分这两个),生成的在ImageSets / Main文件下
import os
import randomtrainval_percent = 0.2 #可自行进行调节
train_percent = 1
xmlfilepath = 'Annotations'
txtsavepath = 'ImageSets\Main'
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/Main/trainval.txt', 'w')
ftest = open('ImageSets/Main/test.txt', 'w')
ftrain = open('ImageSets/Main/train.txt', 'w')
#fval = open('ImageSets/Main/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_labels文件生成labels标签文件
import xml.etree.ElementTree as ETimport pickleimport osfrom os import listdir, getcwdfrom os.path import joinsets = ['train', 'test']classes = ['apple','orange'] #自己训练的类别def convert(size, box):dw = 1. / size[0]dh = 1. / size[1]x = (box[0] + box[1]) / 2.0y = (box[2] + box[3]) / 2.0w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn (x, y, w, h)def convert_annotation(image_id):in_file = open('data/Annotations/%s.xml' % (image_id))out_file = open('data/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').textcls = obj.find('name').textif cls not in classes or int(difficult) == 1:continuecls_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')wd = getcwd()for image_set in sets:if not os.path.exists('data/labels/'):os.makedirs('data/labels/')image_ids = open('data/ImageSets/Main/%s.txt' % (image_set)).read().strip().split()list_file = open('data/%s.txt' % (image_set), 'w')for image_id in image_ids:list_file.write('data/images/%s.jpg\n' % (image_id))convert_annotation(image_id)list_file.close()
4、配置训练文件
在data目录下新建fruit.yaml,配置训练的数据
# COCO 2017 dataset http://cocodataset.org# Download command: bash yolov5/data/get_coco2017.sh# Train command: python train.py --data ./data/coco.yaml# Dataset should be placed next to yolov5 folder:# /parent_folder# /coco# /yolov5# train and val datasets (image directory or *.txt file with image paths)train: xx/xx/train2017.txt # 上面我们生成的train,根据自己的路径进行更改val: xx/xx/val2017.txt # 上面我们生成的test#test: ../coco/test-dev2017.txt # 20k images for submission to https://competitions.codalab.org/competitions/20794# number of classesnc: 2 #训练的类别# class namesnames: ['apple','orange']# Print classes# with open('data/coco.yaml') as f:# d = yaml.load(f, Loader=yaml.FullLoader) # dict# for i, x in enumerate(d['names']):# print(i, x)
models文件(采用那个yaml我们更改那个),例如采用yolov5s.yaml:
# parametersnc: 2 # number of classes 训练的类别数depth_multiple: 0.33 # model depth multiplewidth_multiple: 0.50 # layer channel multiple# anchorsanchors:- [10,13, 16,30, 33,23] # P3/8- [30,61, 62,45, 59,119] # P4/16- [116,90, 156,198, 373,326] # P5/32# yolov5 backbonebackbone:# [from, number, module, args][[-1, 1, Focus, [64, 3]], # 1-P1/2[-1, 1, Conv, [128, 3, 2]], # 2-P2/4[-1, 3, Bottleneck, [128]],[-1, 1, Conv, [256, 3, 2]], # 4-P3/8[-1, 9, Bottleneck, [256]],[-1, 1, Conv, [512, 3, 2]], # 6-P4/16[-1, 9, Bottleneck, [512]],[-1, 1, Conv, [1024, 3, 2]], # 8-P5/32[-1, 1, SPP, [1024, [5, 9, 13]]],[-1, 3, Bottleneck, [1024]], # 10]# yolov5 headhead:[[-1, 3, Bottleneck, [1024, False]], # 11[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 12 (P5/32-large)[-2, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 6], 1, Concat, [1]], # cat backbone P4[-1, 1, Conv, [512, 1, 1]],[-1, 3, Bottleneck, [512, False]],[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 17 (P4/16-medium)[-2, 1, nn.Upsample, [None, 2, 'nearest']],[[-1, 4], 1, Concat, [1]], # cat backbone P3[-1, 1, Conv, [256, 1, 1]],[-1, 3, Bottleneck, [256, False]],[-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]], # 22 (P3/8-small)[[], 1, Detect, [nc, anchors]], # Detect(P3, P4, P5)]
5、训练
python train.py --data data/fruit.yaml --cfg models/yolov5s.yaml --weights '' --batch-size 16 --epochs 100
这里想加载预训练权重要更改下代码,不然会报错:
train.py中115行(日期2020.6.9)
try:#ckpt['model'] = \#{k: v for k, v in ckpt['model'].state_dict().items() if model.state_dict()[k].numel() == v.numel()}ckpt['model'] = \{k: v for k, v in ckpt['model'].state_dict().items() if k in model.state_dict().keys()\and model.state_dict()[k].numel() == v.numel()
训练命令:
python train.py --data data/fruit.yaml --cfg models/yolov5s.yaml --weights weights/yolov5s.pt --batch-size 16 --epochs 100
6、测试:
python detect.py --source file.jpg # imagefile.mp4 # video./dir # directory0 # webcamrtsp://170.93.143.139/rtplive/470011e600ef003a004ee33696235daa # rtsp streamhttp://112.50.243.8/PLTV/88888888/224/3221225900/1.m3u8 # http stream
跟大家点小建议,自己去尝试。
YOLO V5建议:如果需要训练较小的自定义数据集,Adam是更合适的选择,尽管Adam的学习率通常比SGD低。但是如果训练大型数据集,对于YOLOV5来说SGD效果比Adam好
官方指南,你值得拥有:https://github.com/ultralytics/yolov5