YOLO数据集自动生成脚本
通过labelImage生成的xml转换为YOLO需要的数据集。
脚本可以自动创建YOLO的目录，并且移动文件、生成YOLO格式数据。
"""
labelImg标注的图片自动生成YOLO需要的数据集、生成训练集、验证集、测试集。
"""
import os
import math
import shutil
import xmltodict
from loguru import logger
class YoloData:
"""
将labelImg生成的xml格式转换为YOLO数据集
"""
classes = []
name_map = {}
@staticmethod
def load_xml_file(filename: str) -> dict:
with open(filename, mode="r", encoding='utf8') as f:
xml = f.read()
return xmltodict.parse(xml)['annotation']
@classmethod
def to_yolo_data(cls, annotation: dict) -> str:
filename = annotation['filename']
width, height = (float(annotation['size']["width"]), float(annotation['size']["height"]))
classes = annotation['object']
# 遍历所有的分类
lines = []
for c in classes:
name = c['name']
box = c['bndbox']
x1, y1, x2, y2 = (float(box['xmin']), float(box['ymin']), float(box["xmax"]), float(box["ymax"]))
x = f"{(x1 + x2) / (2 * width):.6f}"
y = f"{(y1 + y2) / (2 * height):.6f}"
w = f"{(x2 - x1) / width:.6f}"
h = f"{(y2 - y1) / height:.6f}"
line = ' '.join([str(len(cls.classes)), x, y, w, h])
lines.append(line)
if name not in cls.classes:
cls.classes.append(name)
return "\n".join(lines)
@classmethod
def save_classes(cls, path: str):
with open(path, mode='w', encoding='utf8') as f:
f.write("\n".join(cls.classes))
@classmethod
def main(cls, img_path: str, xml_path: str, folder: str = "coco", per="7:2:1"):
"""
构建图片与xml文件路径关系
:param img_path: 所有标注图片的路径
:param xml_path: 所有标注好图片的xml路径
:param folder: 在当前目录下自动生成的目录，如果目录存在会被级联删除（特别注意该目录会被删除，不要保存重要数据）
:param per: 训练集、验证集、测试集 图片占比
:return:
"""
# 提前删除文件目录
if os.path.exists(folder):
shutil.rmtree(folder)
# 创建目录结构
paths = ["label/train", "label/val", "label/test", "images/train", "images/val", "images/test"]
for path in paths:
path = os.path.abspath(os.path.join(folder, path))
if not os.path.exists(path):
os.makedirs(path)
# 关联图片与xml文件
img_map = {".".join(i.split('.')[:-1]): os.path.join(img_path, i) for i in os.listdir(img_path)}
xml_map = {".".join(i.split('.')[:-1]): os.path.join(xml_path, i) for i in os.listdir(xml_path)}
path_map = {img_map[name]: xml_map[name] for name in img_map if name in xml_map}
# 计算训练、验证、测试集占比
pers = [float(i) for i in per.split(":")]
train = math.ceil(len(path_map) * pers[0] / sum(pers))
val = math.ceil(len(path_map) * pers[1] / sum(pers))
logger.debug(f"{train}, {val}, {train + val},{len(path_map)}")
# 生成YOLO文件、并且复制图片
for i, (img_path, xml_path) in enumerate(path_map.items()):
logger.debug(f"正在转换第【{i}】个xml, {xml_path}")
annotation = cls.load_xml_file(xml_path)
text = cls.to_yolo_data(annotation)
if i < train:
mid_folder = "train"
elif train <= i < train + val:
mid_folder = "val"
else:
mid_folder = "test"
# 构建目录移动文件
label_path = os.path.abspath(os.path.join(folder, "label", mid_folder, f"{i}.txt"))
img_name = img_path.split(".")[-1]
img_dst_path = os.path.abspath(os.path.join(folder, "images", mid_folder, f"{i}.{img_name}"))
logger.debug(img_dst_path)
# 生成文件
shutil.copy(img_path, img_dst_path)
with open(label_path, mode='w', encoding='utf8') as f:
f.write(text)
# 保存训练、验证、测试图片路径
img_path = os.path.join(*img_dst_path.split(os.path.sep)[-4:])
img_txt_path = os.path.join(folder, "label", f"{mid_folder}_list.txt")
with open(img_txt_path, mode='a', encoding='utf8') as f:
f.write(f"{img_path}\n")
# 保存classes文件
classes_path = os.path.abspath(os.path.join(folder, "classes.txt"))
cls.save_classes(classes_path)
logger.debug(f"训练集【{train}】张,验证集【{val}】张,测试集【{len(path_map) - train - val}】张,共有【{len(cls.classes)}】个分类")
if __name__ == '__main__':
YoloData.main(r"D:\dataset\ctrip\JPEGImages", r"D:\dataset\ctrip\Annotations", folder="coco", per="8:1.5:0.5")
第二版
这一版本抽取出了一些公共方法，
把yolo数据转换为labelImg数据、把labelImg数据转换为yolo数据，labelImg的xml数据提取，把标注数据保存为json、csv、Excel、pickle等格式。通过保存的数据加载出图片与xml的关系，然后自动生成文件目录。
目录格式为：
Annotations：保存所有标注的xml数据
images: 保存YOLO训练是需要的图片数据，分别为训练集、验证集、测试集
JPEGImages: 保存所有的图片数据，汇总，可以用labelImg打开该目录标注，但是生成的folder与这个目录最好不在同一个目录。
labels: 保存所有的标签数据
classes.txt: 保存着所有的分类名称
其他3个txt保存着图片路径，默认在dataset目录下，所以需要在YOLO目录下新建一个dataset目录，然后把数据复制到该目录下。
代码：
import io
import os
import re
import math
import json
import base64
import shutil
import pickle
import hashlib
import zipfile
import dicttoxml
import xmltodict
import pandas as pd
from PIL import Image
from lxml import etree
from loguru import logger
from typing import Tuple, Iterable, List, Union
def yolo2box(
width: float,
height: float,
box_list: Iterable[Tuple[str, float, float, float, float]],
sep: int = 0
) -> List[Tuple[str, float, float, float, float]]:
"""
将YOLO数据转换为坐标数据
:param width:图片宽度
:param height:图片高度
:param box_list:盒子列表
:param sep: 保留小数N位
:return: 转换后的盒子数据
"""
items = list()
for label, x, y, w, h in box_list:
x1 = round((x * 2 * width - w * width) / 2, sep)
x2 = round((w * width + x * 2 * width) / 2, sep)
y1 = round((y * 2 * height - h * height) / 2, sep)
y2 = round((y * 2 * height + h * height) / 2, sep)
if sep == 0:
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
items.append((label, x1, y1, x2, y2))
return items
def box2yolo(
width: float,
height: float,
box_list: Iterable[Tuple[str, float, float, float, float]],
sep: int = 6
) -> list[tuple[str, float, float, float, float]]:
"""
将box转换为YOLO数据
:param width: 图片的宽度
:param height: 图片的高度
:param box_list: 标注盒子的可迭代数据（label, x1, y1, x2, y2）
:param sep: 保留小数位数
:return: 盒子的YOLO数据
"""
items = list()
for label, x1, y1, x2, y2 in box_list:
x = round((x1 + x2) / (2 * width), sep)
y = round((y1 + y2) / (2 * height), sep)
w = round((x2 - x1) / width, sep)
h = round((y2 - y1) / height, sep)
items.append((label, x, y, w, h))
return items
def create_xml(
img_fp: Union[str, bytes],
box_list: Iterable[Tuple[str, int, int, int, int]],
img_size: Tuple[int, int, int] = None,
is_format: bool = True,
save_path: str = None
) -> str:
"""
创建图像标注的xml文件
:param img_fp: 图片地址或字节码
:param box_list: 框选的盒子列表
:param img_size: 图片的大小
:param is_format: 是否格式化生成的xml
:param save_path: 保存文件路径
:return: 格式化的xml对象
"""
# 构建盒子数据
objects = list()
for label, x1, y1, x2, y2 in box_list:
item = {
"name": label, "pose": "Unspecified", "truncated": "0", "difficult": "0",
"bndbox": {"xmin": x1, "ymin": y1, "xmax": x2, "ymax": y2}
}
objects.append(item)
# 判断传入是图片还是图片字节码
if isinstance(img_fp, str):
img_fp = path = os.path.abspath(img_fp)
filename = os.path.basename(path)
else:
path = None
filename = None
img_fp = io.BytesIO(img_fp)
# 没有图片尺寸加载图片，计算大小
if not img_size:
img = Image.open(img_fp)
img_size = (img.width, img.height, 3)
# 生成dict数据
xml_dict = {
"annotation": {
"folder": "JPEGImages", "filename": filename, "path": path,
"source": {"database": "Unknown"},
"size": {"width": img_size[0], "height": img_size[1], "depth": img_size[2]},
"segmented": "0", "object": objects
}
}
xml_txt = dicttoxml.dicttoxml(xml_dict, root=False, attr_type=False, return_bytes=False)
xml_txt = re.sub("</?object>", "", xml_txt).replace('item', "object")
# 格式化
if is_format:
xml_tree = etree.fromstring(xml_txt)
etree.indent(xml_tree, space="\t")
xml_txt = etree.tostring(xml_tree, encoding="utf8").decode('utf8')
if save_path:
with open(save_path, mode='wb', encoding='utf8') as f:
f.write(xml_txt)
return xml_txt
def load_xml(xml_path: str) -> dict:
"""
加载labelImg标注的xml, 转换为字典格式
:param xml_path: 文件路径
:return:
"""
with open(xml_path, mode="r", encoding='utf8') as f:
xml = f.read()
annotation = xmltodict.parse(xml)['annotation']
filename = annotation['filename']
folder = annotation['folder']
path = annotation.get('path')
size = annotation['size']
size = [int(size['width']), int(size['height']), int(size['depth'])]
classes = annotation['object']
classes = classes if isinstance(classes, list) else [classes]
box_list = list()
for c in classes:
name = c['name']
box = c['bndbox']
x1, y1, x2, y2 = (float(box['xmin']), float(box['ymin']), float(box["xmax"]), float(box["ymax"]))
box_list.append((name, x1, y1, x2, y2))
return {"filename": filename, "folder": folder, "path": path, "size": size, "box_list": box_list}
def _create_folders(folder="coco"):
"""创建训练的目录解构"""
# 提前删除文件目录
if os.path.exists(folder):
shutil.rmtree(folder)
# 创建目录结构
paths = [
"labels/train", "labels/val", "labels/test",
"images/train", "images/val", "images/test", "JPEGImages", "Annotations"
]
for path in paths:
path = os.path.abspath(os.path.join(folder, path))
if not os.path.exists(path):
os.makedirs(path)
return True
def relation_xml_img(img_paths: str, xml_paths: str, save_filename: str = None) -> List[dict]:
"""
关联标注的xml与img数据
:param img_paths: 图片的目录
:param xml_paths: 标注后的xml目录
:param save_filename: 保存的CSV文件路径
:return:
"""
# 关联图片与xml文件
img_map = {".".join(i.split('.')[:-1]): os.path.join(img_paths, i) for i in os.listdir(img_paths)}
xml_map = {".".join(i.split('.')[:-1]): os.path.join(xml_paths, i) for i in os.listdir(xml_paths)}
path_map = {img_map[name]: xml_map[name] for name in img_map if name in xml_map}
# 提取xml中的数据
items = list()
for i, (img_path, xml_path) in enumerate(path_map.items()):
logger.debug(f'正在加载第【{i}】张图片，共有【{len(path_map)}】张！{img_path}')
xml_dict = load_xml(xml_path)
with open(img_path, mode='rb') as f:
b64img = base64.b64encode(f.read()).decode('utf8')
xml_dict['b64img'] = b64img
xml_dict['img_id'] = hashlib.md5(b64img.encode('utf8')).hexdigest()
items.append(xml_dict)
# 没有文件名直接返回
if not save_filename:
return items
# 保存文件
if save_filename.endswith('.json'):
with open(save_filename, mode='w', encoding="utf8") as f:
json.dump(items, f, ensure_ascii=False)
elif save_filename.endswith('.pkl'):
with open(save_filename, mode='wb') as f:
pickle.dump(items, f)
elif save_filename.endswith('.csv'):
df = pd.DataFrame(items)
df.to_csv(save_filename, index=False)
elif save_filename.split('.')[-1] in ['xls', 'xlsx']:
df = pd.DataFrame(items)
df.to_excel(save_filename, index=False)
else:
raise Exception("文件名错误！")
return items
def load_data(data_path: str):
"""加载数据"""
if data_path.endswith('.json'):
with open(data_path, mode='r', encoding="utf8") as f:
items = json.load(f)
elif data_path.endswith('.pkl'):
with open(data_path, mode='rb') as f:
pickle.load(f)
elif data_path.endswith('.csv'):
items = pd.read_csv(data_path).to_dict(orient="records")
elif data_path.split('.')[-1] in ['xls', 'xlsx']:
items = pd.read_excel(data_path).to_dict(orient="records")
else:
raise Exception("不支持该类型的数据格式！")
return items
def save_file(path: str, content: bytes):
"""保存文件"""
with open(path, mode="wb") as f:
f.write(content)
def create_yolo_folder(data: Union[str, list], folder: str = "coco", per="7:2:1", classes: dict = None, is_zip=False):
"""
创建YOLO格式数据目录
:param data: 数据或者数据所在目录
:param folder: 保存目录
:param per: 训练、验证、测试图片占比
:param classes: 分类字典{分类标签名：下标}
:param is_zip: 是否压缩文件夹
:return:
"""
# 提前删除文件目录
if os.path.exists(folder):
shutil.rmtree(folder)
# 创建目录结构
paths = [
"labels/train", "labels/val", "labels/test",
"images/train", "images/val", "images/test", "JPEGImages", "Annotations"
]
for path in paths:
path = os.path.abspath(os.path.join(folder, path))
if not os.path.exists(path):
os.makedirs(path)
# 2.加载数据、转换为YOLO格式并保存
images = load_data(data) if isinstance(data, str) else data
# 计算训练、验证、测试集占比
pers = [float(i) for i in per.split(":")]
train = math.ceil(len(images) * pers[0] / sum(pers))
val = math.ceil(len(images) * pers[1] / sum(pers))
logger.debug(f"{train}, {val}, {train + val},{len(images)}")
classes = classes if classes else {}
length = len(str(len(images)))
for i, image in enumerate(images):
# 拆分图片，计算中间目录名称
if i < train:
mid_folder = "train"
elif train <= i <= train + val:
mid_folder = "val"
else:
mid_folder = "test"
# 构建需要保存的文件名称
img_filename = os.path.join(folder, "images", mid_folder, f"{i}.jpg".zfill(length + 4))
label_filename = os.path.join(folder, "labels", mid_folder, f"{i}.txt".zfill(length + 4))
jpg_filename = os.path.join(folder, "JPEGImages", f"{i}.jpg".zfill(length + 4))
annotation_filename = os.path.join(folder, "Annotations", f"{i}.xml".zfill(length + 4))
img_txt_filename = os.path.join(folder, f"{mid_folder}_list.txt")
# 提取数据创建xml文件，构建YOLO格式
size = image['size']
box_list = image['box_list']
b64img = image['b64img']
img_bytes = base64.b64decode(b64img.encode('utf8'))
xml_txt = create_xml(img_bytes, box_list=box_list, img_size=size)
# 生成YOLO格式数据
yolo_box = box2yolo(size[0], size[1], box_list)
# 生成分类字典映射
for box in yolo_box:
if box[0].strip().upper() in [chr(i) for i in range(0, 97 + 40)]:
print("PPP", img_filename)
if box[0] in classes:
continue
classes[box[0]] = len(classes)
label_txt = "\n".join([f"{classes[label]} {x1} {y1} {x2} {y2}" for label, x1, y1, x2, y2 in yolo_box])
# 保存文件
save_file(img_filename, img_bytes)
save_file(jpg_filename, img_bytes)
save_file(annotation_filename, xml_txt.encode('utf8'))
save_file(label_filename, label_txt.encode('utf8'))
with open(img_txt_filename, mode='a', encoding='utf8') as f:
f.write(f"dataset\{img_filename}\n")
# 保存分类
cls_filename = f"{folder}\classes.txt"
with open(cls_filename, mode='w', encoding='utf8') as f:
f.write("\n".join(classes))
if is_zip:
with zipfile.ZipFile(f"{folder}.zip", 'w', zipfile.ZIP_DEFLATED) as zip:
for root, paths, files in os.walk(folder):
for file in files:
zip.write(os.path.join(root, file))
logger.debug(f'所有分类：{classes}')
logger.debug(f'所有分类：{list(classes)}')
logger.debug(f"训练集【{train}】张,验证集【{val}】张,测试集【{len(images) - train - val}】张,共有【{len(classes)}】个分类")
if __name__ == '__main__':
jpeg_path = r"D:\github\yolov7\dataset\mt_color\JPEGImages"
annotation_path = r"D:\github\yolov7\dataset\mt_color\Annotations"
# 关联xml与img数据
folder = "mt_color"
filename = f"{folder}.json"
relation_xml_img(jpeg_path, annotation_path, save_filename=filename)
# 生成YOLO格式文件夹
create_yolo_folder(filename, folder=folder, is_zip=True)
只需要改变floder名称和两个路径即可把xml文件和图片数据关联起来。
扫一扫
206
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
