1.确认你的目录格式
这里仅展示训练集目录,根据个人需求可设置val,test目录(问就是我懒~)
2.确认你的xml格式
<annotation>
<folder>images</folder>
<filename>vid_000021_frame0000007.jpg</filename>
<size>
<width>480</width>
<height>270</height>
<depth>3</depth>
</size>
<object>
<name>1</name>
<bndbox>
<xmin>373</xmin>
<ymin>241</ymin>
<xmax>408</xmax>
<ymax>270</ymax>
</bndbox>
</object>
<object>
<name>1</name>
<bndbox>
<xmin>444</xmin>
<ymin>222</ymin>
<xmax>480</xmax>
<ymax>269</ymax>
</bndbox>
</object>
</annotation>
3.运行xml2txt.py
需要修改的位置:main()方法
(1)修改root_path路径为你的文件目录
我这里是D:\StudyAPP\PyCharm2023\Projects\PreData\data2---->改成你的文件目录咯
(2)修改name_list数组为你的classes信息
[‘1’, ‘2’, ‘3’, ‘4’, ‘5’, ‘6’]—>该成你的标签哈,[‘car’, ‘person’] 等等
(3)如果还需要转val和test的同学
改成 flags = [‘train’,‘val’ ] 就行~
运行完会在目录下生成新的labels目录,里面就是yolo格式的txt!!!
如下所示:
import os
from tqdm import tqdm
from lxml import etree
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, save_root: str, class_dict: dict, flag: str):
"""
将对应xml文件信息转为yolo中使用的txt文件信息
:param flag:
:param file_names:
:param save_root:
:param class_dict:
:return:
"""
save_txt_path = os.path.join(save_root, "labels")
if os.path.exists(save_txt_path) is False:
os.makedirs(save_txt_path)
voc_images_path = os.path.join(save_root, "images")
voc_xml_path = os.path.join(save_root, "xmls")
for file in tqdm(file_names, desc="translate {} file...".format(flag)):
# 检查下图像文件是否存在
img_path = os.path.join(voc_images_path, file + ".jpg")
assert os.path.exists(img_path), "file:{} not exist...".format(img_path)
# 检查xml文件是否存在
xml_path = os.path.join(voc_xml_path, file + ".xml")
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_txt_path, file + ".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] # 目标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():
"""
修改root_path和name_list即可
"""
root_path = r"D:\StudyAPP\PyCharm2023\Projects\PreData\data2"
flags = ['train']
name_list = ['1', '2', '3', '4', '5', '6', '7',
'8', '9', '10', '11', '12', '13', '14',
'15', '16']
for flag in flags:
voc_root = os.path.join(root_path, flag)
class_dict = dict(zip(name_list, range(0, len(name_list))))
xml_path = os.path.join(voc_root, "xmls")
file_names = [line.split('.')[0] for line in os.listdir(xml_path)]
translate_info(file_names, voc_root, class_dict, flag)
if __name__ == "__main__":
main()