制作YOLOv5数据集
1 使用labelImg标注数据集
labelImg是一个可视化的图像标定工具。Faster R-CNN,YOLO,SSD等目标检测网络所需要的数据集,均需要借此工具标定图像中的目标。可以标注两种格式:VOC标签格式(.xml)和YOLO标签格式(.txt)。
1.1 下载labelImg
可通过两种方式下载:
1.前往下载链接下载labelImg。
2.使用命令行安装labelImg
pip install labelimg
1.2 使用labelImg标注数据
1.2.1 数据存放
预先准备好数据,以以下格式新建目录存放数据:
├── VOC2012 根目录
│├── JPEGImages 存放需要标注的图片
│├── Annotations 存放标签
│├── classes.txt 存放类别
先将图片放入JPEGImages文件夹中,再新建classes.txt,在txt中逐行输入类别,如下:
1.2.2 标注数据
如果是下载的labelImg,打开labelImg.exe即可。
如果是命令行安装的labelImg,使用命令:
labelimg JPEGImages classes.txt
打开后如下:
主要功能介绍:
使用快捷键A、D可进行图片切换。使用快捷键W进行标注,输入标签即可。
完成标注后,可在设置的标签保存路径查看标签。
2 将voc标签转换到yolo标签格式
如果标签是voc格式的,需要转换成yolo标签格式。如果标签是yolo格式的,可直接进行数据集划分。
首先准备voc_classes.json文件:
{
"aeroplane": 1,
"bicycle": 2,
"bird": 3,
"boat": 4,
"bottle": 5,
"bus": 6,
"car": 7,
"cat": 8,
"chair": 9,
"cow": 10,
"diningtable": 11,
"dog": 12,
"horse": 13,
"motorbike": 14,
"person": 15,
"pottedplant": 16,
"sheep": 17,
"sofa": 18,
"train": 19,
"tvmonitor": 20
}
第二步准备trans_voc2yolo.py文件进行转换,更改相应目录即可:
"""
将voc数据集标注信息(.xml)转为yolo标注格式(.txt)
"""
import os
from tqdm import tqdm
from lxml import etree
import json
import shutil
# voc数据集根目录可改成自己的数据集路径
voc_root = "VOC2012"
# 转换后的文件保存目录
save_file_root = os.path.join(voc_root, "Yolo_labels")"
# label标签对应json文件
label_json_path = os.path.join(voc_root, "voc_classes.json")
# voc的xml标签路径
voc_xml_path = os.path.join(voc_root, "Annotations")
# 检查文件/文件夹都是否存在
assert os.path.exists(voc_xml_path), "VOC xml path not exist..."
assert os.path.exists(label_json_path), "label_json_path does not exist..."
if os.path.exists(save_file_root) is False:
os.makedirs(save_file_root)
def parse_xml_to_dict(xml):
"""
将xml文件解析成字典形式,参考tensorflow的recursive_parse_xml_to_dict
Args:
xml: xml tree obtained by parsing XML file contents using lxml.etree
Returns:
Python dictionary holding XML contents.
"""
if len(xml) == 0: # 遍历到底层,直接返回tag对应的信息
return {xml.tag: xml.text}
result = {}
for child in xml:
child_result = parse_xml_to_dict(child) # 递归遍历标签信息
if child.tag != 'object':
result[child.tag] = child_result[child.tag]
else:
if child.tag not in result: # 因为object可能有多个,所以需要放入列表里
result[child.tag] = []
result[child.tag].append(child_result[child.tag])
return {xml.tag: result}
def translate_info(file_names: list ,class_dict: dict):
"""
将对应xml文件信息转为yolo中使用的txt文件信息
:param file_names:
:param save_root:
:param class_dict:
:return:
"""
for file in tqdm(file_names, desc="translate xml file..."):
# 检查xml文件是否存在
xml_path = os.path.join(voc_xml_path, file)
assert os.path.exists(xml_path), "file:{} not exist...".format(xml_path)
# read xml
with open(xml_path) as fid:
xml_str = fid.read()
xml = etree.fromstring(xml_str)
data = parse_xml_to_dict(xml)["annotation"]
img_height = int(data["size"]["height"])
img_width = int(data["size"]["width"])
# write object info into txt
assert "object" in data.keys(), "file: '{}' lack of object key.".format(xml_path)
if len(data["object"]) == 0:
# 如果xml文件中没有目标就直接忽略该样本
print("Warning: in '{}' xml, there are no objects.".format(xml_path))
continue
with open(os.path.join(save_file_root, file[:-4] + ".txt"), "w") as f:
for index, obj in enumerate(data["object"]):
# 获取每个object的box信息
xmin = float(obj["bndbox"]["xmin"])
xmax = float(obj["bndbox"]["xmax"])
ymin = float(obj["bndbox"]["ymin"])
ymax = float(obj["bndbox"]["ymax"])
class_name = obj["name"]
class_index = class_dict[class_name] - 1 # 目标id从0开始
# 进一步检查数据,有的标注信息中可能有w或h为0的情况,这样的数据会导致计算回归loss为nan
if xmax <= xmin or ymax <= ymin:
print("Warning: in '{}' xml, there are some bbox w/h <=0".format(xml_path))
continue
# 将box信息转换到yolo格式
xcenter = xmin + (xmax - xmin) / 2
ycenter = ymin + (ymax - ymin) / 2
w = xmax - xmin
h = ymax - ymin
# 绝对坐标转相对坐标,保存6位小数
xcenter = round(xcenter / img_width, 6)
ycenter = round(ycenter / img_height, 6)
w = round(w / img_width, 6)
h = round(h / img_height, 6)
info = [str(i) for i in [class_index, xcenter, ycenter, w, h]]
if index == 0:
f.write(" ".join(info))
else:
f.write("\n" + " ".join(info))
def main():
# read class_indict
json_file = open(label_json_path, 'r')
class_dict = json.load(json_file)
# voc信息转yolo,并将图像文件复制到相应文件夹
file_names = os.listdir(voc_xml_path)
translate_info(file_names, class_dict)
if __name__ == "__main__":
main()
转换后的标签格式如下:
14 0.538066 0.452 0.360082 0.5
文件目录结构:
├── VOC2012 根目录
│├── JPEGImages 存放需要标注的图片(.jpg,.bmp...)
│├── Annotations 存放voc标签(.xml)
│├── Yolo_labels 存放yolo标签(.txt)
│├── voc_classes.json 存放分类json文件
3 数据集划分训练集、验证集和测试集
import os
import random
import shutil
# 改成自己的路径
root_path = 'VOC2012'
txtfilepath = root_path + '\Yolo_labels\\'
imgfilepath = root_path + '\JPEGImages\\'
txtsavepath = root_path + '\ImageSets\MainTest'
yolo_root_path = root_path + '\YOLO'
yolo_images_train_path = yolo_root_path + '\images\\train\\'
yolo_images_val_path = yolo_root_path + '\images\\val\\'
yolo_images_test_path = yolo_root_path + '\images\\test\\'
yolo_labels_train_path = yolo_root_path + '\labels\\train\\'
yolo_labels_val_path = yolo_root_path + '\labels\\val\\'
yolo_labels_test_path = yolo_root_path + '\labels\\test\\'
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
if not os.path.exists(yolo_root_path):
os.makedirs(yolo_root_path)
if not os.path.exists(yolo_images_train_path):
os.makedirs(yolo_images_train_path)
if not os.path.exists(yolo_images_val_path):
os.makedirs(yolo_images_val_path)
if not os.path.exists(yolo_images_test_path):
os.makedirs(yolo_images_test_path)
if not os.path.exists(yolo_labels_train_path):
os.makedirs(yolo_labels_train_path)
if not os.path.exists(yolo_labels_val_path):
os.makedirs(yolo_labels_val_path)
if not os.path.exists(yolo_labels_test_path):
os.makedirs(yolo_labels_test_path)
def copyfile(imgname,name,imgpath,labelpath):
image_copy_to = os.path.join(imgpath,imgname[:-1])
if os.path.exists(image_copy_to) is False:
shutil.copyfile(imgfilepath +imgname[:-1] , image_copy_to)
label_copy_to = os.path.join(labelpath,name[:-1])
if os.path.exists(label_copy_to) is False:
shutil.copyfile(txtfilepath + name[:-1] , label_copy_to)
def main():
train_test_percent = 0.9 # (训练集+验证集)/(训练集+验证集+测试集)
train_valid_percent = 0.9 # 训练集/(训练集+验证集)
total_xml = os.listdir(txtfilepath)
num = len(total_xml)
list = range(num)
tv = int(num * train_test_percent) # 训练集+验证集数量
ts = int(num-tv) # 测试集数量
tr = int(tv * train_valid_percent) # 验证集数量
tz = int(tv-tr) # 训练集数量
trainval = random.sample(list, tv)
train = random.sample(trainval, tr)
print("train and valid size:", tv)
print("train size:", tz)
print("test size:", ts)
print("valid size:", tr)
ftest = open(txtsavepath + '/test.txt', 'w')
ftrain = open(txtsavepath + '/train.txt', 'w')
fvalid = open(txtsavepath + '/valid.txt', 'w')
ftestimg = open(txtsavepath + '/img_test.txt', 'w')
ftrainimg = open(txtsavepath + '/img_train.txt', 'w')
fvalidimg = open(txtsavepath + '/img_valid.txt', 'w')
for i in list:
name = total_xml[i][:-4] + '.txt' + '\n'
imgname = total_xml[i][:-4] + '.jpg' + '\n'
if i in trainval:
if i in train:
ftrain.write(name)
ftrainimg.write(imgname)
copyfile(imgname,name,yolo_images_train_path,yolo_labels_train_path)
else:
fvalid.write(name)
fvalidimg.write(imgname)
copyfile(imgname,name,yolo_images_val_path,yolo_labels_val_path)
else:
ftest.write(name)
ftestimg.write(imgname)
copyfile(imgname,name,yolo_images_test_path,yolo_labels_test_path)
ftrain.close()
fvalid.close()
ftest.close()
ftrainimg.close()
fvalidimg.close()
ftestimg.close()
print("finished!")
if __name__ == "__main__":
main()
文件目录结构:
├── VOC2012 根目录
│├── JPEGImages 存放需要标注的图片(.jpg,.bmp...)
│├── Annotations 存放voc标签(.xml)
│├── Yolo_labels 存放yolo标签(.txt)
│├── ImageSet
││├──Main
│││├──test.txt 存放测试集yolo标签
│││├──train.txt 存放训练集yolo标签
│││├──valid.txt 存放验证集集yolo标签
│││├──img_test.txt 存放测试集图片
│││├──img_train.txt 存放训练集集图片
│││├──img_valid.txt 存放验证集图片
│├── YOLO
││├──images
│││├──test 测试集图片
│││├──train 训练集图片
│││├──val 验证集图片
││├──labels
│││├──test 测试集标签
│││├──train 训练集标签
│││├──val 验证集标签
│├── voc_classes.json 存放分类json文件
后面会继续介绍如何在YOLOv5中使用自己制作的数据集。
已更新博客:
深度学习笔记:使用YOLOv5训练自己的数据集模型