目录
相比yolov3版本,整理了生成train、val txt文档的方式,可以将数据集放在别的目录下了
1.下载源码
https://github.com/ultralytics/yolov5,下载源码并解压
https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data,官方教程
2.制作数据集
仍由VOC格式为例
JPEGImages改名为images,目录下只含有两个文件夹即可:
datasetname
-images
-Annotations
3.生成文件
代码1:split_train_val.py
import os
import random
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--xml_path', type=str, help='input xml label path')
parser.add_argument('--txt_path', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 1.0
train_percent = 0.8
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
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(txtsavepath + '/trainval.txt', 'w')
ftest = open(txtsavepath + '/test.txt', 'w')
ftrain = open(txtsavepath + '/train.txt', 'w')
fval = open(txtsavepath + '/val.txt', 'w')
for i in list:
name=total_xml[i][:-4]+'\n'
if i in trainval:
ftrainval.write(name)
if i in train:
ftrain.write(name)
else:
fval.write(name)
else:
ftest.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()
把这个文件随便放在哪,执行python split_train_val.py --xml_path xx --txt_path xx
指定xml的label存放位置和希望输出的训练验证集划分txt位置,代码会从xml_path读取label文件,然后进行划分,将文件名保存在txt_path的txt文件中,txt_path 一般为 .../ImageSets/Main
代码2:voc_label.py
# -*- coding: utf-8 -*-
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets=['train', 'val', 'test'] #替æ¢ä¸ºè‡ªå·±çš„æ•°æ®é›?
classes = ['一次性快餐盒', '书籍纸张', '充电宝', '剩饭剩菜', '包', '垃圾桶', '塑料器皿', '塑料玩具', '塑料衣架', '大骨头', '干电池', '快递纸袋',
'插头电线', '旧衣服', '易拉罐', '枕头', '果皮果肉', '毛绒玩具', '污损塑料', '污损用纸', '洗护用品', '烟蒂', '牙签', '玻璃器皿', '砧板', '筷子',
'纸盒纸箱', '花盆', '茶叶渣', '菜帮菜叶', '蛋壳', '调料瓶', '软膏', '过期药物', '酒瓶', '金属厨具', '金属器皿', '金属食品罐', '锅', '陶瓷器皿', '鞋', '食用油桶', '饮料瓶', '鱼骨']
#修改为自己的类别
#classes = ["eye", "nose"]
abs_path = os.getcwd()
def convert(size, box):
dw = 1./(size[0])
dh = 1./(size[1])
x = (box[0] + box[1])/2.0 - 1
y = (box[2] + box[3])/2.0 - 1
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 = open('Annotations/%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')
wd = getcwd()
for image_set in sets:
if not os.path.exists('labels/'):
os.makedirs('labels/')
image_ids = open('ImageSets/Main/%s.txt'%(image_set)).read().strip().split()
list_file = open('%s.txt'%(image_set), 'w')
for image_id in image_ids:
list_file.write(abs_path+'/images/%s.jpg\n'%(image_id))
convert_annotation(image_id)
list_file.close()
#os.system("cat 2007_train.txt 2007_val.txt > train.txt")
把这个文件放在数据集根目录下,修改其中的classes列表,然后运行 python voc_label.py
代码会根据上一步划分的txt文件读取xml格式的label文件,转换成txt格式的label保存在labels目录下,并在当前目录下生成train、val、test.txt文件,里面保存的是完整的文件名,是绝对路径
两个代码运行结束后目录如下:
datasetname
-Annotations
-xxx.xml
-images
-xxx.jpg
-ImageSets
-Main
-train.txt
-val.txt
-labels
-xxx.txt
4.两个yaml文件
创建./data/xx.yaml
train: ./data/train.txt # 128 images
val: ./data/val.txt # 128 images
# number of classes
nc: 44
# class names
names: ['一次性快餐盒', '书籍纸张', '充电宝', '剩饭剩菜', '包', '垃圾桶', '塑料器皿', '塑料玩具', '塑料衣架', '大骨头', '干电池', '快递纸袋',
'插头电线', '旧衣服', '易拉罐', '枕头', '果皮果肉', '毛绒玩具', '污损塑料', '污损用纸']
其中train和val是voc_label.py生成的txt文件,nc是类别数量, names中的内容与顺序要和voc_label.py文件中保持一致
修改./models/xx.yaml
在./models里选一个模型进行修改,只修改其中的nc为自己的类别数量即可。
5.训练
cd 到yolov5目录下
python train.py --img 640 --batch 8 --epoch 300 --data ./data/HuaWeigarbage.yaml --cfg ./models/yolov5s.yaml --weights weights/yolov5s.pt --logdir runs/temp --device '1'
日志保存目录默认保存在./runs/exp0下,想要修改目录可以指定--logdir参数
6.测试
python detect.py --weights runs/exp0/weights/best.pt --source data/val.txt --device 0
对原始代码进行了修改,可以指定存放待检测图片路径的txt文件
7.python test.py
python test.py --weights runs/mask/exp2/weights/best.pt --data data/mask.yaml --batch-size 16 --conf-thres 0.4 --device '0'
一些说明:
默认保存模型的标准:
实际应用中可以根据需要修改这个比重。
保存模型的内容:中间过程保存的是模型、优化器、epoch等信息,所以会比单一模型占用空间大很多。
with open(results_file, 'r') as f: # create checkpoint
ckpt = {'epoch': epoch,
'best_fitness': best_fitness,
'training_results': f.read(),
'model': ema.ema,
'optimizer': None if final_epoch else optimizer.state_dict()}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
在训练结束后进行了strip optimizer操作,只保存模型,所以占用空间会小很多:
def strip_optimizer(f='weights/best.pt', s=''): # from utils.general import *; strip_optimizer()
# Strip optimizer from 'f' to finalize training, optionally save as 's'
x = torch.load(f, map_location=torch.device('cpu'))
x['optimizer'] = None
x['training_results'] = None
x['epoch'] = -1
x['model'].half() # to FP16
for p in x['model'].parameters():
p.requires_grad = False
torch.save(x, s or f)
mb = os.path.getsize(s or f) / 1E6 # filesize
print('Optimizer stripped from %s,%s %.1fMB' % (f, (' saved as %s,' % s) if s else '', mb))