本文不生产技术,只做技术的搬运工!!!
前言
在yolo训练时,我们需要对图像进行标注,而使用labelimg标注时如果直接选择输出yolo格式的数据集,则原始数据的很多信息无法被保存,因此一版使用xml格式的标签,这时再去训练时就需要对标签进行转换。
代码
import os
import xml.etree.ElementTree as ET
import cv2
def getFileList(dir, Filelist, ext=None):
"""
获取文件夹及其子文件夹中文件列表
输入 dir:文件夹根目录
输入 ext: 扩展名
返回: 文件路径列表
"""
newDir = dir
if os.path.isfile(dir):
if ext is None:
Filelist.append(dir)
else:
if ext in dir:
Filelist.append(dir)
elif os.path.isdir(dir):
for s in os.listdir(dir):
newDir = os.path.join(dir, s)
getFileList(newDir, Filelist, ext)
return Filelist
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)
def convert_annotation(xml_file,xml_name, image_file,dst, class_names):
tree = ET.parse(xml_file)
root = tree.getroot()
# 读取图像以获取准确的尺寸
image = cv2.imread(image_file)
if image is None:
raise FileNotFoundError(f"Image file not found: {image_file}")
image_size = [image.shape[1], image.shape[0]] # [width, height]
output_path = os.path.join(dst, xml_name.replace(".xml", '.txt'))
with open(output_path, 'w') as out_file:
for obj in root.iter('object'):
difficult = obj.find('difficult').text
class_name = obj.find('name').text
if class_name not in class_names or int(difficult) == 1:
continue
class_id = class_names.index(class_name)
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(image_size, b)
out_file.write(f"{class_id} {bb[0]} {bb[1]} {bb[2]} {bb[3]}\n")
if __name__ == '__main__':
xml_dir = "/data2/dataset/abcd/20250305/xml" #xml标签路径
image_dir = "/data2/dataset/abcd/20250305/images" #图像路径
output_dir = "/data2/dataset/abcd/20250305/labels" #输出的txt路径
class_names = ['aaa', 'bbb', 'ccc', 'ddd'] # 类别名称列表
xml_file_list = []
xml_file_list = getFileList(xml_dir, xml_file_list, '.xml')
lenth = len(xml_file_list)
i = 1
for xml_file in xml_file_list:
xml_name = os.path.basename(xml_file)
image_name = xml_name.replace(".xml", ".jpg")
image_file = os.path.join(image_dir, image_name)
if not os.path.exists(image_file):
print(f"image file not found: {image_file}")
continue
convert_annotation(xml_file, xml_name, image_file, output_dir, class_names)
print('{}/{}'.format(i, lenth))
i += 1