为了方便,我将代码封装成了桌面程序,GUI部分我就不介绍了,泛泛而谈到时啥都没学会。
一、yolo标签格式
我们看一下yolo标签的格式:
<class index> <x center> <y center> <width> <height> \text{<class index> <x center> <y center> <width> <height>} <class index> <x center> <y center> <width> <height>
上述格式【class_index】是某类别对应的标签索引,为整型;后面4个参数均为比例,而不是实际参数,我们举个例子
1 0.759623 0.749454 0.079443 0.081878
我假设“汽车”这个类别对应的标签索引为1。那么上述例子表示我们标了一个框,这个框的类别索引为1,表示为类别“汽车”。
前两个浮点参数表示这个框的中心坐标分别相对于整个图像宽度和高度的比例,用比例来表示中心点的位置;后两者表示框的宽度和高度分别相对于整个图像的比例,用比例来表达框的大小。
二、载入XML文件
首先我们使用python自带模块【os】来将指定目录(XML文件所在目录)【self.source_path】的所有文件名载入到文件名列表【Files】中。
然后我们使用for训练,遍历文件列表的每一个文件名,如果文件名的扩展名是【.XML】,则表示是标签文件,此时调用Process进行转换
def LoadVocTag(self):Files = os.listdir(self.source_path)# 处理每一个xml文件for File in Files:if File.endswith('.xml'):self.Process(File)
三、获取类别标签对应类别索引字典
代码中【self.classes_file】表示含有对应关系的txt文件路径,现在我们要把它载入,做成字典。我们看一个类别标签和类别索引的对应关系的txt文件内容:
纸类 1
塑料类 2
金属类 3
衣品类 4
其他 5
上述内容,类别标签和类别索引之间用空格隔开了,因此我们打开文件后,逐行读取内容,然后使用空格进行分割成两段。前段是类别标签,后段时类别索引,类别索引读出来是文本,我们转换成整型后添加到字典就好了。
我相信现在看下面的代码就一目了然了:
def GetTag(self):# 类别字典class_dict = {}# 打开含有号码和类别对应关系的txt文件with open(self.classes_file, 'r', encoding='utf-8') as CF:# 获取txt文件的每一行信息for line in CF:parts = line.strip().split(' ')# 每行的第一部分是类别的中文描述,第二部分是相应的数字class_dict[parts[0]] = int(parts[1])# 返回类别字典return class_dict
四、格式转换
这一部分需要解析XML文件,解析XML文件,我呢使用轻量级的【xml.etree.ElementTree】模块,在小文件中,这个模块用来解析是非常方便的。
先获取对应关系的字典,便于后面解析文件得到类别后转换成类别索引:
# 获取类别字典,方便后续得到类别对应的号码
Class_Dict = self.GetTag()
for循环用tree.iterfind(‘object’)自动解析XML文件,循环查找 <object></object> \text{<object></object>} <object></object>标记,因为一个object表示我们锚的一个框。
# 查找object标签,每一个object标签,代表一个框
for object in tree.iterfind('object'):
VOC标签文档中object标记下第一行是 <name></name> \text{<name></name>} <name></name>标记,表示类别名字,也就是object[0]表示 <name></name> \text{<name></name>} <name></name>标记,其属性text就是标记的内容,该内容就是类别名字。
# 获取当前框所对应的类别名字
name = object[0].text
从字典中得到类别索引:
# 从类别字典中获取该类别对应的号码
tag = Class_Dict[name]
接下来就是计算转换部分了,下面代码中,【self.image_width】和【self.image_height】是图片尺寸,这个是需要输入的,其他的部分计算原理来自yolo标签的原理,可见第一部分的讲解
# 获取边框盒,该标记包含了目标框的4个角的坐标
box = object.find('bndbox')
# 获取4个角坐标(相对于图像左上角)
xmin = float(box.find('xmin').text)
ymin = float(box.find('ymin').text)
xmax = float(box.find('xmax').text)
ymax = float(box.find('ymax').text)# 计算中心点坐标
center_x = (xmin + xmax) / 2
center_y = (ymin + ymax) / 2
# 计算宽度和高度
width = xmax - xmin
height = ymax - ymin
# 图像的宽度和高度(实际应根据图像尺寸来调整)# 归一化为相对于图像宽度和高度的比例,YOLO标签以比例表示
normalized_center_x = center_x / self.image_width
normalized_center_y = center_y / self.image_height
normalized_width = width / self.image_width
normalized_height = height / self.image_height
接着我们将计算得到的结果严格按照yolo标签的格式组合在一起,每个参数之间用空格隔开,每个比例保留6位小数。content是一个列表,存储一张图片中包含的所有锚框标签信息。
# 格式化将写入txt文件的标签信息
info = f'{tag} {normalized_center_x:.6f} {normalized_center_y:.6f} {normalized_width:.6f} {normalized_height:.6f}'
# 给当前txt文件添加一行,表示处理完一个框
content.append(info)
五、保存txt标签文件
这部分需要将content列表的内容,逐个写入,每个占一行
def SaveToYoloTag(self, FileName, content):name = os.path.join(self.object_path, FileName.split('.')[0] + '.txt')with open(name, 'w') as yolo:# 按行写入for item in content:yolo.write(f'{item}\n')# 在滚动窗口中显示输出信息showname=name.split("\\")[-1]show=f'{showname} has been saved'self.output_textedit.append(show)
六、完整程序
import os
import xml.etree.ElementTree as ET
import sys
from PyQt5.QtWidgets import QApplication, QWidget, QVBoxLayout, QLabel, QPushButton, QFileDialog, QLineEdit, QTextEdit,QDesktopWidgetclass YoloConverterApp(QWidget):def __init__(self):super().__init__()self.source_path = Noneself.object_path = Noneself.classes_file = Noneself.image_width = Noneself.image_height = Noneself.init_ui()def init_ui(self):layout = QVBoxLayout()self.source_label = QLabel('VOC标签文件路径:')layout.addWidget(self.source_label)self.object_label = QLabel('YOLO标签文件路径:')layout.addWidget(self.object_label)self.classes_label = QLabel('类别指示文件路径:')layout.addWidget(self.classes_label)self.width_label = QLabel('图像宽度:')layout.addWidget(self.width_label)self.width_input = QLineEdit()layout.addWidget(self.width_input)self.height_label = QLabel('图像高度:')layout.addWidget(self.height_label)self.height_input = QLineEdit()layout.addWidget(self.height_input)source_button = QPushButton('选择VOC标签文件路径')source_button.clicked.connect(self.choose_source_path)layout.addWidget(source_button)object_button = QPushButton('选择YOLO标签文件路径')object_button.clicked.connect(self.choose_object_path)layout.addWidget(object_button)classes_button = QPushButton('选择类别指示文件路径')classes_button.clicked.connect(self.choose_classes_file)layout.addWidget(classes_button)convert_button = QPushButton('转换')convert_button.clicked.connect(self.convert)layout.addWidget(convert_button)# 添加滚动窗口用于显示输出信息self.output_textedit = QTextEdit()self.output_textedit.setReadOnly(True)layout.addWidget(self.output_textedit)self.setLayout(layout)self.setWindowTitle('Yolo Converter')# 获取桌面尺寸desktop = QDesktopWidget()screen_rect = desktop.screenGeometry()width = screen_rect.width() * 0.5height = screen_rect.height() * 0.5# 设置窗口大小为桌面大小的50%self.setGeometry(100, 100, int(width), int(height))def open_file_dialog(self, label):# 创建文件对话框的选项options = QFileDialog.Options()options |= QFileDialog.DontUseNativeDialog# 创建 QFileDialog 实例file_dialog = QFileDialog()file_dialog.setOptions(options)# 将文件模式设置为选择目录file_dialog.setFileMode(QFileDialog.Directory)# 如果用户选择了目录并点击了确定if file_dialog.exec_():# 获取所选路径selected_path = file_dialog.selectedFiles()[0]# 将提供的标签的文本设置为所选路径label.setText(selected_path)def convert(self):self.image_width = int(self.width_input.text())self.image_height = int(self.height_input.text())self.LoadVocTag()def choose_source_path(self):self.open_file_dialog(self.source_label)self.source_path = self.source_label.text()def choose_object_path(self):self.open_file_dialog(self.object_label)self.object_path = self.object_label.text()def choose_classes_file(self):# 创建文件对话框的选项options = QFileDialog.Options()options |= QFileDialog.DontUseNativeDialog# 创建 QFileDialog 实例file_dialog = QFileDialog()file_dialog.setOptions(options)# 如果用户选择了文件并点击了确定if file_dialog.exec_():# 获取所选文件selected_file = file_dialog.selectedFiles()[0]# 将类别标签的文本设置为所选文件self.classes_label.setText(selected_file)# 将类别文件属性设置为所选文件路径self.classes_file = selected_filedef LoadVocTag(self):Files = os.listdir(self.source_path)# 处理每一个xml文件for File in Files:if File.endswith('.xml'):self.Process(File)# 处理单个XML文件def Process(self, File):# 获取解析树tree = ET.parse(os.path.join(self.source_path, File))# 获取类别字典,方便后续得到类别对应的号码Class_Dict = self.GetTag()# 存储要写入txt文件的内容,content的元素个数就是该图片含有的目标框数content = []# 查找object标签,每一个object标签,代表一个框for object in tree.iterfind('object'):# 获取当前框所对应的类别名字name = object[0].text# 从类别字典中获取该类别对应的号码tag = Class_Dict[name]# 获取边框盒,该标记包含了目标框的4个角的坐标box = object.find('bndbox')# 获取4个角坐标(相对于图像左上角)xmin = float(box.find('xmin').text)ymin = float(box.find('ymin').text)xmax = float(box.find('xmax').text)ymax = float(box.find('ymax').text)# 计算中心点坐标center_x = (xmin + xmax) / 2center_y = (ymin + ymax) / 2# 计算宽度和高度width = xmax - xminheight = ymax - ymin# 图像的宽度和高度(实际应根据图像尺寸来调整)# 归一化为相对于图像宽度和高度的比例,YOLO标签以比例表示normalized_center_x = center_x / self.image_widthnormalized_center_y = center_y / self.image_heightnormalized_width = width / self.image_widthnormalized_height = height / self.image_height# 格式化将写入txt文件的标签信息info = f'{tag} {normalized_center_x:.6f} {normalized_center_y:.6f} {normalized_width:.6f} {normalized_height:.6f}'# 给当前txt文件添加一行,表示处理完一个框content.append(info)# 保存为yolo格式的标签self.SaveToYoloTag(File, content)# 获取类别标签对应的字典def GetTag(self):# 类别字典class_dict = {}# 打开含有号码和类别对应关系的txt文件with open(self.classes_file, 'r', encoding='utf-8') as CF:# 获取txt文件的每一行信息for line in CF:parts = line.strip().split(' ')# 每行的第一部分是类别的中文描述,第二部分是相应的数字class_dict[parts[0]] = int(parts[1])# 返回类别字典return class_dict# 将处理后的标签信息保存为YOLO格式的文件def SaveToYoloTag(self, FileName, content):name = os.path.join(self.object_path, FileName.split('.')[0] + '.txt')with open(name, 'w') as yolo:# 按行写入for item in content:yolo.write(f'{item}\n')# 在滚动窗口中显示输出信息showname=name.split("\\")[-1]show=f'{showname} has been saved'self.output_textedit.append(show)if __name__ == '__main__':app = QApplication(sys.argv)converter_app = YoloConverterApp()converter_app.show()sys.exit(app.exec_())