由于课题需要,需要自己制作数据集进行训练,目前是自己制作的第二个数据集,发现有某些细节已经忘记,记录备忘,同时为后来者提供借鉴。文章以car-tank数据集做为例子介绍
整体流程:
1.准备数据:从各种途径获取原始数据,博主的car-tank原始数据是在网上搜索然后批量保存的
2.标注数据:使用labelimg标注软件标注数据,最终得到images、label文件
3.分配数据转换格式:label标注的标签文件为xml文件,需要划分数据集,并转换到txt格式
4.修改yolo-v5代码,使用car-tank数据集训练、测试
1.准备数据
使用Fatkun浏览器插件批量下载图片(博主试验多种插件后推荐此款),可同时下载界面中的所有图片。若只为学习,推荐分开下载分开保存,标注数据会省劲。如果你还没尝试过自己训练数据集,推荐以此方法做个小的数据集学习此过程,博主用了一两天标注测试了200多张图片就做好了自己的数据集。
注意,图片最好为jpg格式,标注过程中不知为啥png的标注有点问题,时间久远记不清现象了,特意转换格式后才可以。----好像是标注完后xml文件中没有图片的宽度和高度
2.标注数据
labelimg自行下载,是常用的一款图片标注软件
(1)打开图片保存目录
(2)选定label文件(xml文件)保存目录
(3)view-Auto saving 设置自动保存
开始标注吧,关键记住快捷键
w:框选目标
a: 上一张
d: 下一张
博主习惯先按一类标注,右上角使用默认label,按下w键后,点击鼠标左键框选目标,点击a进行下一张,框选完继续下一张~~~,注意文件目录最好全部为英文路径
3.分配数据转换格式
(1)分配数据
将所有的图片(不同种类的)放到一个文件夹image,将所有的label放到一个文件夹xmllabel,注意不同来源的图片不要重名。至此数据标注完整理到了一块。
分配数据,将所有标注的数据分为train,test,trainval,val等,makeTXT.py将数据划分到几个txt文件,这几个文件保存的都是图片的名字(不带后缀的)。
makeTXT.py
# 划分数据集,划分到4个TXT文件
import os
import random
# 初始化划分比例和xml文件路径
trainval_percent = 0.1
train_percent = 0.9
xmlfilepath = 'xmllabel' # xml 标签目录
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('SplitTXT/trainval.txt', 'w')
ftest = open('SplitTXT/test.txt', 'w')
ftrain = open('SplitTXT/train.txt', 'w')
fval = open('SplitTXT/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()
至此分割数据集完成
(2)转换格式
先上代码
voc_label2.py
import xml.etree.ElementTree as ET
import pickle
import os
from os import listdir, getcwd
from os.path import join
import shutil
# 1.makeTXT生成四个txt,记录切分情况(只记录图片名字,不记录路径),test、train 、val、trainval
# 2.创建label文件夹,再创建test、train 、val的txt文件,用于保存各部分图片的路径
# 3.将所有图片的xml文件从一个文件夹读取,转存到另一个文件夹(label)变成多个txt文件(yolov5格式的,每个图片一个txt,每一行标注有1个目标的类别及位置)
# 4.将各个部分的图片按照刚才保存的路径,拷贝到目标路径
sets = [('2012', 'train'), ('2012', 'val'), ('2007', 'train'), ('2007', 'val'), ('2007', 'test')]
# classes = ["aeroplane", "bicycle", "bird", "boat", "bottle", "bus", "car", "cat", "chair", "cow", "diningtable", "dog",
# "horse", "motorbike", "person", "pottedplant", "sheep", "sofa", "train", "tvmonitor"]
classes = ["1", "2"]
# 转换box位置标签
def convert(size, box):
dw = 1. / size[0]
dh = 1. / size[1]
x = (box[0] + box[1]) / 2.0
y = (box[2] + box[3]) / 2.0
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)
# 从xml标签转换到txt标签
def convert_annotation(image_id):
in_file = open('xmllabel/%s.xml' % (image_id)) # 打开对应的xml文件
out_file = open('labels/%s.txt' % (image_id), 'w') # 输出文件,yolov5格式的label
tree = ET.parse(in_file)
root = tree.getroot()
print("image_id", image_id)
size = root.find('size')
w = int(size.find('width').text)
h = int(size.find('height').text)
print("size", size)
print("w", w)
print("h", h)
for obj in root.iter('object'):
print("")
print("hhhhhhhhhh")
if obj.find('difficult'):
difficult = obj.find('difficult').text
else:
difficult = 0
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)
print("bbbbbb", str(cls_id), bb)
out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + '\n') # 将数据写入输出文件
for year, image_set in sets:
# 2.创建label文件夹,再创建test、train 、val的txt文件,用于保存各部分图片的路径
if not os.path.exists('labels/'): # 创建标签文件夹
os.makedirs('labels/')
# 读取TXT文件,保存图片ID
image_ids = open('SplitTXT/%s.txt' % (image_set)).read().strip().split() # 读取已经分好的标签
list_file = open('%s.txt' % (image_set), 'w') # 保存各个部分数据(test、val、train)图片的路径
print("image_set", image_set)
for image_id in image_ids:
list_file.write('image/%s.jpg\n' % image_id) # 将每个图片的路径暂时保存到txt------成功
# 3.将所有图片的xml文件从一个文件夹读取,转存到另一个文件夹(label)变成多个txt文件(yolov5格式的,每个图片一个txt,每一行标注有1个目标的类别及位置)
convert_annotation(image_id) # 将xml文件的标签转存到label下的每个TXT
list_file.close()
# 合并多个文本文件
def mergeTxt(file_list, outfile):
with open(outfile, 'w') as wfd:
for f in file_list:
with open(f, 'r') as fd:
shutil.copyfileobj(fd, wfd)
# 创建VOC文件夹及子文件夹
wd = os.getcwd()
data_base_dir = os.path.join(wd, "cartank_mlw/")
if not os.path.isdir(data_base_dir):
os.mkdir(data_base_dir)
img_dir = os.path.join(data_base_dir, "images/")
if not os.path.isdir(img_dir):
os.mkdir(img_dir)
img_train_dir = os.path.join(img_dir, "train/")
if not os.path.isdir(img_train_dir):
os.mkdir(img_train_dir)
img_val_dir = os.path.join(img_dir, "val/")
if not os.path.isdir(img_val_dir):
os.mkdir(img_val_dir)
label_dir = os.path.join(data_base_dir, "labels/")
if not os.path.isdir(label_dir):
os.mkdir(label_dir)
label_train_dir = os.path.join(label_dir, "train/")
if not os.path.isdir(label_train_dir):
os.mkdir(label_train_dir)
label_val_dir = os.path.join(label_dir, "val/")
if not os.path.isdir(label_val_dir):
os.mkdir(label_val_dir)
# 使用train.txt中的图片作为yolov5的训练集
print(os.path.exists('train.txt'))
f = open('train.txt', 'r')
lines = f.readlines()
# 4.将各个部分的图片按照刚才保存的路径,拷贝到目标路径
for line in lines:
line = line.replace('\n', '')
if (os.path.exists(line)):
shutil.copy(line, "cartank_mlw/images/train") # 复制图片到voc_mlw/images/train
print('coping train img file %s' % line + '\n')
print("line", line)
line = line.replace('image', 'labels') # 复制label,将图片路径名字换成对应TXT的路径名字
line = line.replace('jpg', 'txt')
print("line", line)
if (os.path.exists(line)):
shutil.copy(line, "cartank_mlw/labels/train")
print('copying train label file %s' % line + '\n')
# 使用test.txt中的图片作为yolov5验证集
print(os.path.exists('test.txt'))
f = open('test.txt', 'r')
lines = f.readlines()
for line in lines:
line = line.replace('\n', '')
if (os.path.exists(line)):
shutil.copy(line, "cartank_mlw/images/val") #line是图片路径,复制图片到 voc_mlw/images/val
print('coping val img file %s' % line + '\n')
line = line.replace('image', 'labels') # 复制label,将图片路径名字换成对应TXT的路径名字
line = line.replace('jpg', 'txt')
if (os.path.exists(line)):
shutil.copy(line, "cartank_mlw/labels/val")
print('copying val img label %s' % line + '\n')
这部分比较主要也稍复杂,加上makeTXT一块考虑: 1.makeTXT生成四个txt,记录切分情况,test、train 、val、trainval,作为后期文件转移用 2.创建labels文件夹,用于保存所有转化完的txt 再创建test、train 、val的txt文件,用于暂时保存各部分图片的路径,便于后期拷贝图片 3.将所有图片的xml文件从一个文件夹读取,转换成txt文件(yolov5格式的,每个图片一个txt,每一行标注有1个目标的类别及位置),转存到另一个文件夹(labels) 4.将各个部分的图片按照刚才保存的路径,拷贝到目标路径(train or val目录),将各个图片的txt 标签拷贝到对应目录。
至此数据集转换格式完成,具体细节可研究voc_label2.py
4.修改yolo-v5部分代码,使用car-tank数据集训练、测试
https://github.com/ultralytics/yolov5https://github.com/ultralytics/yolov5博主使用上面的yolov5代码
(1)仿照已有的yaml文件,在 data目录下创建car-tank数据集的yaml文件-VOC_cartank_mlw.yaml文件,修改其中的数据集目录
(2)修改models,yolo-v5s.yaml
其他的研究清楚可以自己改,如anchors大小,关键改nc, number of class,car-tank共2种目标,应该为2
(3)修改train.py,
数据变量改为准备用的数据集的yaml~如VOC_cartank_mlw.yaml ,batchsize对应自己的GPU能力,然后run即可