南通网站推广公司哪家好360优化大师app
本文不生产技术,只做技术的搬运工!!!
前言
在yolo训练时,我们需要对图像进行标注,而使用labelimg标注时如果直接选择输出yolo格式的数据集,则原始数据的很多信息无法被保存,因此一版使用xml格式的标签,这时再去训练时就需要对标签进行转换。
代码
import os
import xml.etree.ElementTree as ET
import cv2def getFileList(dir, Filelist, ext=None):"""获取文件夹及其子文件夹中文件列表输入 dir:文件夹根目录输入 ext: 扩展名返回: 文件路径列表"""newDir = dirif 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 Filelistdef convert(size, box):dw = 1. / size[0]dh = 1. / size[1]x = (box[0] + box[1]) / 2.0y = (box[2] + box[3]) / 2.0w = box[1] - box[0]h = box[3] - box[2]x = x * dww = w * dwy = y * dhh = h * dhreturn (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').textclass_name = obj.find('name').textif class_name not in class_names or int(difficult) == 1:continueclass_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 = 1for 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}")continueconvert_annotation(xml_file, xml_name, image_file, output_dir, class_names)print('{}/{}'.format(i, lenth))i += 1