Yolov5的label格式要求的是txt格式,如果遇到xml格式的代码,可以使用下列代码进行转换。
import os
import xml.etree.ElementTree as ET
dirpath = ''#存放xml的目录
newdir = ''#存放txt的目录
if not os.path.exists(newdir):
os.makedirs(newdir)
nameList = ['a','b']#数据有多少类别
for fp in os.listdir(dirpath):
root = ET.parse(os.path.join(dirpath,fp)).getroot()
xmin, ymin, xmax, ymax = 0,0,0,0
sz = root.find('size')
width = float(sz[0].text)
height = float(sz[1].text)
filename = root.find('filename').text
for child in root.findall('object'):
sub = child.find('bndbox')
xmin = float(sub[0].text)
ymin = float(sub[1].text)
xmax = float(sub[2].text)
ymax = float(sub[3].text)
name = child.find('name').text
if name not in nameList:
nameList.append(name)
idx = nameList.index(name)
try:
x_center = (xmin + xmax) / (2 * width)
y_center = (ymin + ymax) / (2 * height)
w = (xmax - xmin) / width
h = (ymax - ymin) / height
except ZeroDivisionError:
print(filename,'的 width有问题')
with open(os.path.join(newdir, fp.split('.')[0]+'.txt'), 'a+') as f:
f.write(' '.join([str(idx), str(x_center), str(y_center), str(w), str(h) + '\n']))
可以创建一个新的文件夹,里面需要包含三部分文件。分别为train文件夹、valid文件夹和一个yaml格式的文件。trian和valid文件下面都包含有名为images的文件夹和名为labels的文件夹,用以储存图片和标签。
三个文件中train为训练集,valid为验证集,yaml用来存放数据的位置信息,具体需要包含以下四个信息,训练集图片的位置,验证机图片的位置,类别的数量,以及数据的类别以及名称。
示例:
数据存放在voc文件夹,voc文件夹与train.py文件同级别,voc文件夹里面有train、valid和名为data.yaml。
假设数据仅有a和b两个类别,那么mydata.yaml里需要包含以下四个信息,书写方式如下:
train: ./voc/train/images
val: ./voc/valid/images
nc: 2
names: ['a', 'b']
yaml文件里的路径不对,按照报错修改路径即可。
yaml里的类别对不上,可能是yaml里n或names有问题。
准备完数据后打开train.py文件,然后对一些参数进行修改。
首先需要修改参数位置,需要修改的代码如下:
parser.add_argument('--data', type=str, default='修改为前面写好的yaml的相对路径', help='data.yaml path')
接着就可以运行了。不过在第一次运行时,建议先将epoch和权重改小一点,运行成功了再改大参数。
打开test.py文件,把里面的库和定义test的函数全部复制到train.py文件。
打开detect.py,里面需要修改的参数有三个,weight,source和output。
需要需改的代码如下:
parser.add_argument('--weights', nargs='+', type=str, default='修改为runs文件下最新的一次迭代里weights文件夹里的best.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='需要检测的图片位置,为相对路径', help='source')
parser.add_argument('--output', type=str, default='检测结果输出的位置,需要创建一个空的文件夹,不是空的会被清空', help='output folder')
这里参考这个解决方案
将utils/datasets.py文件中的
p = str(Path(path).absolute())
改为p = str(Path(path)),重新运行代码即可