当前位置：首页 > news >正文

【图像标签转换】XML转为TXT图像数据集标签

news 来源：原创 2024/9/19 14:47:25

引言

该脚本用于将包含对象标注的 XML 文件转换为 YOLO（You Only Look Once）对象检测格式的 TXT 文件。脚本读取 XML 文件，提取对象信息，规范化边界框坐标，并将数据写入相应的 TXT 文件。此外，它还生成一个 classes.txt 文件，列出所有对象类别。

效果：

环境和依赖

Python 3.x
库：os, xml.etree.ElementTree

目录结构

输入目录：包含对象标注的 XML 文件。
输出目录：将包含转换后的 TXT 文件和 classes.txt 文件。

1. 引入必要的库

2. 定义输入和输出目录，以及类别列表

3. 获取所有 XML 文件的文件名

4. 获取所有分类

5. 转换坐标到 YOLO 格式

6. 读取 XML 文件并转换为 TXT 文件

7. 主函数

完整程序

1. 引入必要的库

os 用于处理文件和目录操作。
xml.etree.ElementTree 用于解析 XML 文件

import os
import xml.etree.ElementTree as ET

2. 定义输入和输出目录，以及类别列表

input_dir = r'D:/lenovo/Archie/shujukuochongv1.0/xmlall1'
out_dir = r'D:/lenovo/Archie/shujukuochongv1.0/txtall1'
class_list = ['1','2','3','4','5','6','7','8','9']

input_dir 是存放 XML 文件的目录。
out_dir 是保存转换后 TXT 文件的目录。
class_list 是初始定义的类别名列表。

3. 获取所有 XML 文件的文件名

遍历输入目录，获取所有 XML 文件的文件名（不带扩展名）。

def file_name(input_dir):F = []for root, dirs, files in os.walk(input_dir):for file in files:if os.path.splitext(file)[1] == '.xml':t = os.path.splitext(file)[0]F.append(t)  # 将所有的文件名添加到 F 列表中return F  # 返回 F 列表

4. 获取所有分类

解析每个 XML 文件，获取对象的类别，并添加到类别列表中（跳过 difficult 属性为 1 的对象）。

def get_class(filelist):for i in filelist:f_dir = input_dir + "\\" + i + ".xml"in_file = open(f_dir, encoding='UTF-8')filetree = ET.parse(in_file)in_file.close()root = filetree.getroot()for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in class_list and int(difficult) == 0:class_list.append(cls)

5. 转换坐标到 YOLO 格式

将 XML 文件中的像素坐标转换为 YOLO 格式的归一化坐标。

def ConverCoordinate(imgshape, bbox):xmin, xmax, ymin, ymax = bboxwidth = imgshape[0]height = imgshape[1]dw = 1. / widthdh = 1. / heightx = (xmin + xmax) / 2.0y = (ymin + ymax) / 2.0w = xmax - xminh = ymax - yminx = x * dwy = y * dhw = w * dwh = h * dhreturn x, y, w, h

6. 读取 XML 文件并转换为 TXT 文件

读取 XML 文件，提取目标信息，将其转换为 YOLO 格式，并写入 TXT 文件。

def readxml(i):f_dir = input_dir + "\\" + i + ".xml"txtresult = ''outfile = open(f_dir, encoding='UTF-8')filetree = ET.parse(outfile)outfile.close()root = filetree.getroot()size = root.find('size')width = int(size.find('width').text)height = int(size.find('height').text)imgshape = (width, height)for obj in root.findall('object'):obj_name = obj.find('name').textobj_id = class_list.index(obj_name)bbox = obj.find('bndbox')xmin = float(bbox.find('xmin').text)xmax = float(bbox.find('xmax').text)ymin = float(bbox.find('ymin').text)ymax = float(bbox.find('ymax').text)bbox_coor = (xmin, xmax, ymin, ymax)x, y, w, h = ConverCoordinate(imgshape, bbox_coor)txt = '{} {} {} {} {}\n'.format(obj_id, x, y, w, h)txtresult += txtwith open(out_dir + "\\" + i + ".txt", 'w+') as f:f.write(txtresult)

7. 主函数

获取 XML 文件列表。
提取所有分类并打印。
逐个将 XML 文件转换为 YOLO 格式的 TXT 文件。
生成一个 classes.txt 文件，列出所有类别。

filelist = file_name(input_dir)
get_class(filelist)
print(class_list)for i in filelist:readxml(i)with open(out_dir + "\\classes.txt", 'a') as f:classresult = '\n'.join(class_list) + '\n'f.write(classresult)

完整程序

该脚本用于将 XML 文件中的对象标注转换为 YOLO 格式的 TXT 文件，并生成一个包含所有类别的 classes.txt 文件。通过这些步骤，可以方便地将标注数据用于 YOLO 模型的训练。

import os
import xml.etree.ElementTree as ET# xml文件存放目录（不要以\结尾）
input_dir = r'D:/lenovo/Archie/shujukuochongv1.0/xmlall1'# 输出txt文件目录（不要以\结尾）
out_dir = r'D:/lenovo/Archie/shujukuochongv1.0/txtall1'
#类别名
class_list = ['1','2','3','4','5','6','7','8','9']# 获取目录所有xml文件
def file_name(input_dir):F = []for root, dirs, files in os.walk(input_dir):for file in files:# print file.decode('gbk')    #文件名中有中文字符时转码if os.path.splitext(file)[1] == '.xml':t = os.path.splitext(file)[0]F.append(t)  # 将所有的文件名添加到L列表中return F  # 返回L列表# 获取所有分类
def get_class(filelist):for i in filelist:f_dir = input_dir + "\\" + i + ".xml"in_file = open(f_dir, encoding='UTF-8')filetree = ET.parse(in_file)in_file.close()root = filetree.getroot()for obj in root.iter('object'):difficult = obj.find('difficult').textcls = obj.find('name').textif cls not in class_list or int(difficult) == 1:class_list.append(cls)def ConverCoordinate(imgshape, bbox):# 将xml像素坐标转换为txt归一化后的坐标xmin, xmax, ymin, ymax = bboxwidth = imgshape[0]height = imgshape[1]dw = 1. / widthdh = 1. / heightx = (xmin + xmax) / 2.0y = (ymin + ymax) / 2.0w = xmax - xminh = ymax - ymin# 归一化x = x * dwy = y * dhw = w * dwh = h * dhreturn x, y, w, hdef readxml(i):f_dir = input_dir + "\\" + i + ".xml"txtresult = ''outfile = open(f_dir, encoding='UTF-8')filetree = ET.parse(outfile)outfile.close()root = filetree.getroot()# 获取图片大小size = root.find('size')width = int(size.find('width').text)height = int(size.find('height').text)imgshape = (width, height)# 转化为yolov的格式for obj in root.findall('object'):# 获取类别名obj_name = obj.find('name').textobj_id = class_list.index(obj_name)# 获取每个obj的bbox框的左上和右下坐标bbox = obj.find('bndbox')xmin = float(bbox.find('xmin').text)xmax = float(bbox.find('xmax').text)ymin = float(bbox.find('ymin').text)ymax = float(bbox.find('ymax').text)bbox_coor = (xmin, xmax, ymin, ymax)x, y, w, h = ConverCoordinate(imgshape, bbox_coor)txt = '{} {} {} {} {}\n'.format(obj_id, x, y, w, h)txtresult = txtresult + txt# print(txtresult)f = open(out_dir + "\\" + i + ".txt", 'w+')f.write(txtresult)f.close()# 获取文件夹下的所有文件
filelist = file_name(input_dir)# 获取所有分类
get_class(filelist)# 打印class
print(class_list)# xml转txt
for i in filelist:readxml(i)# 在out_dir下生成一个class文件
f = open(out_dir + "\\classes.txt", 'a')
classresult = ''
for i in class_list:classresult = classresult + i + "\n"
f.write(classresult)
f.close()