darknet训练自己的数据集(pytorch版本)
2019/7/24 11:08 如云漂泊
话不多说,直接开始训练步骤。
1.文件目录
从github下载darknet-master文件夹,在文件夹根目录新建myData文件夹,在myData文件夹下再新建三个文件夹,分别为annotations、ImageSets、JPEGImages,其中annotations文件下存放xml文件,ImageSets下新建main文件夹,其中建一个train.txt,其中放图片名称(不要带.jpg,只是单独的图片名称),JPEGImages下放要训练的图片。
2.制作voc数据集
推荐使用labelimg标注软件对图像进行标注,可以得到xml文件,其中包含了标签,坐标等内容,在后面需要用到这些信息。labelimg下载直接在网上搜索即可。
3.将voc数据集转换成yolo数据集
推荐脚本如下:
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
sets=[('myData', 'train')] #元组数组
classes = ['screen'] # each category's name
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(year, image_id):
in_file = open('myData/Annotations/%s.xml'%(image_id))
out_file = open('myData/labels/%s.txt'%(image_id), 'w')
tree=ET.parse(in_file)
root = tree.getroot() #这两行表示从xml文件中读取内容
size = root.find('size') #从读取到的xml内容中发现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
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))
try:
bb = convert((w,h), b)
except:
print(image_id)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n')
if __name__ == '__main__':
wd = getcwd() #返回当前工作目录
for year, image_set in sets:
if not os.path.exists('myData/labels/'): #os.path.exists("")表示如果路径 path 存在,返回 True;如果路径 path 不存在,返回 False。
os.makedirs('myData/labels/') #创建目录
image_ids = open('myData/ImageSets/Main/%s.txt'%(image_set)).read().strip().split() #打开该文件,并一次性读取该文件所有内容,再去掉文本中句子开头与结尾的符号的
list_file = open('myData/%s_%s.txt'%(year, image_set), 'w') #打开myData_train该文件
for image_id in image_ids:
list_file.write('%s/myData/JPEGImages/%s.jpg\n'%(wd, image_id)) #在myData_train.txt里面写入图片地址+名称
convert_annotation(year, image_id)
list_file.close()
其中代码已经详细注释,针对自己的文件目录进行修改。在darknet-master目录下新建py文件,输入以上代码,运行该脚本后会在myData文件夹下面生成labels文件夹,其中包含每个图片标注后的坐标信息,以txt文件形式存在。
4.修改类别名称及调用gpu
打开cfg文件夹下面的voc.data,进行如下修改:
第一行为类别数,第二行为训练数据集名称,第三行位测试数据集名称,第四行为权重保存路径。
打开makefile文件,如下:
如果电脑安装了cuda,则将第一行和第二行参数改为1。
打开cfg文件夹下面的yolov3-voc.cfg,主要内容如下:
在进行训练的时候将testing下面两行注释,training下面两行取消注释,测试的时候再改回来。classes为类别数,根据自己的情况修改,filters=(classes+5)*3,根据自己的情况进行修改(要修改的地方有三处,classes和filters都要修改)。修改好了之后在命令行进入darknet文件夹,输入make进行编译。
5.训练数据集
做好以上的准备工作后,在darknet官网下载darknet53.conv.74,放在darknet根目录下面,同时新建backup文件夹,用来存放训练生成的权重文件,做好上面的步骤后在终端输入:
./darknet detector train cfg/my_data.data cfg/my_yolov3.cfg darknet53.conv.74
即可进行训练。推荐使用gpu训练,但是需要在系统装cuda(pytorch自带的cuda不行),挺麻烦的。训练成功后会看到backup文件夹下面有生成的权重文件。
6.测试数据集
测试时如图所示,将training下面两行注释,testing下面两行取消注释,然后在darknet-master目录下打开终端,输入
./darknet detector test cfg/voc.data cfg/yolov3-voc.cfg cfg/yolov3-voc_10000.weights
(yolov3-voc_10000.weights为自己训练得到的权重)即可进行测试。