数据集格式转化

数据集格式转换的要点是从对应的格式中解析出对应的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 ET

def 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_map
    class_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').text
        if class_name not in class_to_id_map:
            continue
        class_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 / width
        y_center = (ymin + ymax) / 2 / height
        bbox_width = (xmax - xmin) / width
        bbox_height = (ymax - ymin) / height

        txt_format.append(f"{class_id} {x_center} {y_center} {bbox_width} {bbox_height}")

    return txt_format



def 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 ET

def 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.size

    root = 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:
            continue

        class_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)