数据集格式转换的要点是从对应的格式中解析出对应的class和bbox
整体说明：
lables_voc_dir： 为存放所有xml的标注文件夹，labels_yolo_dir: 为转换后存放所有xml的标注文件夹images_dir: 为所有图像文件的文件夹。、
注意：代码部分可以不要看，重点关注调用函数的接口
以下代码资源
一. VOC 转YOLO
import os
import random
import shutil
import glob
from PIL import Image
import xml.etree.ElementTree as ETdef convert_xml_to_txt_format(xml_file_path: str, class_list: list) -> list:"""将单个xml文件转换为txt格式。参数:xml_file_path (str): XML文件的路径。class_list (list): 包含所有类别的列表。返回:list: 包含YOLO格式的标签列表。"""# 根据class_list构建class_to_id_mapclass_to_id_map = {class_name: index for index, class_name in enumerate(class_list)}tree = ET.parse(xml_file_path)root = tree.getroot()width = int(root.find('.//size/width').text)height = int(root.find('.//size/height').text)txt_format = []for obj in root.findall('.//object'):class_name = obj.find('name').textif class_name not in class_to_id_map:continueclass_id = class_to_id_map[class_name]bbox = obj.find('bndbox')xmin = float(bbox.find('xmin').text)ymin = float(bbox.find('ymin').text)xmax = float(bbox.find('xmax').text)ymax = float(bbox.find('ymax').text)x_center = (xmin + xmax) / 2 / widthy_center = (ymin + ymax) / 2 / heightbbox_width = (xmax - xmin) / widthbbox_height = (ymax - ymin) / heighttxt_format.append(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")return txt_formatdef write_yolo_format_labels_to_file(yolo_label_path: str, yolo_content: list):"""将YOLO格式的标签写入到文件。参数:yolo_label_path (str): YOLO标签文件的名称。yolo_content (list): 包含YOLO格式标签的列表。"""with open(yolo_label_path, 'w') as f:for line in yolo_content:f.write(line + '\n')def convert_voc2yolo(voc_annotations_dir: str, yolo_labels_dir: str, class_list: list):"""将VOC格式的XML注释转换为YOLO格式的注释。参数:voc_annotations_dir (str): VOC注释目录的路径。yolo_labels_dir (str): YOLO标签目录的路径。class_list (list): 类别名称的list。"""if not os.path.exists(yolo_labels_dir):os.makedirs(yolo_labels_dir)for xml_file in os.listdir(voc_annotations_dir):if xml_file.endswith('.xml'):yolo_label_path = os.path.join(yolo_labels_dir,xml_file.replace('.xml', '.txt'))xml_file_path = os.path.join(voc_annotations_dir, xml_file)yolo_content = convert_xml_to_txt_format(xml_file_path,class_list)write_yolo_format_labels_to_file(yolo_label_path, yolo_content)with open(os.path.join(yolo_labels_dir,"classes.txt"),"w",encoding='utf-8') as f:for c in class_list:f.write(c+'\n')如下代码所示：lables_voc_dir： 为存放所有xml的标注文件夹，labels_yolo_dir: 为转换后存放所有xml的标注文件夹，class_names为类别列表，当调用convert_voc2yolo()后会labels_yolo_dir目录下会生成所有的txt文件。
if __name__ == "__main__":dataset_dir = "./dataset"labels_voc_dir = "labels_voc"labels_yolo_dir = "labels_yolo"class_names = ["橘子", "香蕉", "草莓"]convert_voc2yolo(labels_voc_dir,labels_yolo_dir,class_names)
二. YOLO转VOC
import os
import random
import shutil
import glob
from PIL import Image
import xml.etree.ElementTree as ETdef convert_txt_to_xml_format(txt_file_path: str, voc_file_path: str, class_names: list, images_dir: str):"""将单个txt文件转换为xml格式。参数:txt_file_path (str): YOLO标签文件txt的路径。voc_file_path (str): XML文件的路径。class_names (list): 包含类别名称的列表。images_dir (str): 包含图像文件的目录路径。"""with open(txt_file_path, 'r') as file:lines = file.readlines()# 获取txt文件的文件名，然后去掉后缀txt_filename = os.path.basename(txt_file_path)txt_filename_without_ext = os.path.splitext(txt_filename)[0]# 使用glob库在images_dir下搜索与txt文件同名的图片文件pattern = os.path.join(images_dir, f"{txt_filename_without_ext}.*")matching_files = glob.glob(pattern)if not matching_files:raise FileNotFoundError(f"No matching image found for {txt_filename_without_ext} in {images_dir}")img_file = matching_files[0] # 假设找到的第一个匹配文件就是正确的图片# 从与YOLO标签同名的图像文件中读取宽度和高度img = Image.open(img_file)width, height = img.sizeroot = ET.Element("annotation")folder = ET.SubElement(root, "folder")folder.text = os.path.basename(images_dir)filename = ET.SubElement(root, "filename")filename.text = os.path.basename(img_file)path = ET.SubElement(root, "path")path.text = os.path.abspath(img_file)source = ET.SubElement(root, "source")database = ET.SubElement(source, "database")database.text = "Unknown"size = ET.SubElement(root, "size")ET.SubElement(size, "width").text = str(width)ET.SubElement(size, "height").text = str(height)ET.SubElement(size, "depth").text = "3"segmented = ET.SubElement(root, "segmented")segmented.text = "0"for line in lines:values = line.strip().split()if len(values) < 5:continueclass_id, x_center, y_center, bbox_width, bbox_height = map(float, values)object = ET.SubElement(root, "object")name = ET.SubElement(object, "name")name.text = class_names[int(class_id)]pose = ET.SubElement(object, "pose")pose.text = "Unspecified"truncated = ET.SubElement(object, "truncated")truncated.text = "0"difficult = ET.SubElement(object, "difficult")difficult.text = "0"bndbox = ET.SubElement(object, "bndbox")xmin = int((x_center - bbox_width / 2) * width)ymin = int((y_center - bbox_height / 2) * height)xmax = int((x_center + bbox_width / 2) * width)ymax = int((y_center + bbox_height / 2) * height)ET.SubElement(bndbox, "xmin").text = str(xmin)ET.SubElement(bndbox, "ymin").text = str(ymin)ET.SubElement(bndbox, "xmax").text = str(xmax)ET.SubElement(bndbox, "ymax").text = str(ymax)# 使用prettify函数来格式化输出pretty_xml_as_string = prettify(root)with open(voc_file_path, 'w', encoding='utf-8') as f:f.write(pretty_xml_as_string)def prettify(elem):from xml.dom import minidom"""Return a pretty-printed XML string for the Element."""rough_string = ET.tostring(elem, 'utf-8')reparsed = minidom.parseString(rough_string)return reparsed.toprettyxml(indent=" ")def convert_yolo2voc(yolo_labels_dir: str, voc_labels_dir: str, images_dir: str, class_list: list):"""将YOLO格式的标签文件夹转换为VOC格式的XML注释文件夹。参数:yolo_labels_dir (str): YOLO标签文件夹的路径。voc_labels_dir (str): 转换后的VOC XML文件夹的路径。images_dir (str): 包含图像文件的目录路径。class_names (list): 包含类别名称的列表。"""# 确保voc_labels_dir存在if not os.path.exists(voc_labels_dir):os.makedirs(voc_labels_dir)# 遍历YOLO标签文件夹for label_file in os.listdir(yolo_labels_dir):if label_file.endswith(".txt"):yolo_file_path = os.path.join(yolo_labels_dir, label_file)voc_file_path = os.path.join(voc_labels_dir, label_file.replace(".txt", ".xml"))convert_txt_to_xml_format(yolo_file_path, voc_file_path, class_list, images_dir)如下代码所示：lables_voc_dir： 为转换后存放所有xml的标注文件夹，labels_yolo_dir: 为存放所有xml的标注文件夹，class_names为类别列表，当调用convert_voc2yolo()后会labels_voc_dir目录下会生成所有的txt文件。
if __name__ == "__main__":labels_voc_dir = "labels_voc"labels_yolo_dir = "labels_yolo"imgs_dir="imgs"class_names = ["橘子", "香蕉", "草莓"]convert_yolo2voc(labels_yolo_dir,labels_voc_dir,imgs_dir,class_names)三. Json 转YOLO
提示:可以采用 data_df = pd.read_json(json_path),然后利用pandas的语法进行提取bbox和class_id
四. 划分数据集
def split_dataset(images_origin_dir: str, labels_origin_dir: str, dataset_dir: str, train_ratio: float = 0.8):"""将原始数据集分割为训练集和验证集，并将相应的文件复制到新的文件夹中。参数:images_origin_dir (str): 存放所有图片文件夹路径labels_origin_dir (str): 存放所有标签文件夹路径。dataset_dir (str): 新的数据集根目录。train_ratio (float, optional): 训练集的比例，默认为 0.8。"""images_dir = os.path.join(dataset_dir, "images")labels_dir = os.path.join(dataset_dir, "labels")train_dir = os.path.join(images_dir, "train")val_dir = os.path.join(images_dir, "val")# 创建训练集和验证集的文件夹os.makedirs(train_dir, exist_ok=True)os.makedirs(val_dir, exist_ok=True)os.makedirs(os.path.join(labels_dir, "train"), exist_ok=True)os.makedirs(os.path.join(labels_dir, "val"), exist_ok=True)image_files=[]for file in os.listdir(images_origin_dir):if file.endswith(".jpg") or file.endswith(".png") or file.endswith(".jpeg"):image_files.append(os.path.join(images_origin_dir, file))# 使用 random.sample 函数进行随机划分random.seed(42) # 设置随机种子train_images = random.sample(image_files, k=int(len(image_files) * train_ratio))val_images = [f for f in image_files if f not in train_images]# 复制文件到训练集和验证集文件夹for subset, images in [('train', train_images), ('val', val_images)]:for image_file in images:# 复制图片文件target_image_path = os.path.join(images_dir, subset, os.path.basename(image_file))shutil.copyfile(image_file, target_image_path)# 复制标签文件label_file = os.path.splitext(os.path.basename(image_file))[0] + '.txt'label_path = os.path.join(labels_origin_dir, label_file)if os.path.exists(label_path):target_label_path = os.path.join(labels_dir, subset, label_file)shutil.copyfile(label_path, target_label_path)print(f"训练集大小: {len(train_images)}")print(f"验证集大小: {len(val_images)}")
如下代码所示：dataset_dir:为待生成数据集的根目录，imgs_dir：为存放所有图片的根目录，labels_yolo:为存放所有txt的根目录，调用split_dataset()以后，会生成如下的目录结构：
if __name__ == "__main__":dataset_dir = "./dataset"imgs_dir="imgs"labels_yolo_dir = "labels_yolo"split_dataset(imgs_dir,labels_yolo_dir,dataset_dir)
