整理YOLO数据集（图片及XML标签同时重命名）

最新推荐文章于 2024-09-25 22:46:17 发布

阿甘的笔记本

最新推荐文章于 2024-09-25 22:46:17 发布

阅读量559

点赞数 4

分类专栏：目标检测YOLO系列文章标签： YOLO python 图像处理

本文链接：https://blog.csdn.net/weixin_58404533/article/details/137566395

版权

目标检测YOLO系列专栏收录该内容

3 篇文章 0 订阅

订阅专栏

一、重命名数据集的必要性

不同类别数据集整合到一个文件夹下，可能会变得杂乱无章。若不同数据集的命名相同则需进行覆盖操作，由此上传的数据集将混乱，进而影响模型训练的精度。因此有必要对数据集进行重命名，对数据集进行规范化处理。

二、复现重命名程序

此例重命名文件目录如下图：

复现重命名代码时，保证安装好对应的python包，并且按照注释中的4个序号（如“1、2、3、4、”）步骤进行更改。

import os
import xml.etree.ElementTree as ET
import glob


# 重命名xml与图片文件函数
def rename(xmldir, imgdir, outdir):
    # 获取图片文件列表
    os.chdir(imgdir)
    imgs = os.listdir('.')
    imgs = glob.glob(str(imgs) + '*.jpg')

    # 获取xml文件列表
    os.chdir(xmldir)
    annotations = os.listdir('.')
    annotations = glob.glob(str(annotations) + '*.xml')

    labels = ['nect']  # 1、  改为数据集得标签名
    nums = [1, 1, 1, 1, 1]  # 每一类标签对应的图片数量，最后一个代表混合标签mix，即一张图片中有多种标签对应的目标

    for i, file in enumerate(annotations):  # 遍历xml文件列表
        # actual parsing
        in_file = open(file, encoding='utf-8')  # 打开xml文件
        tree = ET.parse(in_file)
        root = tree.getroot()

        lastCla = -1  # 遍历标签时的前一个标签
        isSingle = 1  # 文件所含标签是否单一，初始化为1（代表标签单一）

        for obj in root.iter('object'):  # 遍历当前xml文件中的所有标签
            name = obj.find('name').text  # 获得标签名称
            if name == "nect":   #2、  改为数据集得标签名
                cla = 0
            # elif name == "face":
            #     cla = 1
            # elif name == "mask_chin":
            #     cla = 2
            # elif name == 'mask_mouth_chin':
            #     cla = 3

            if lastCla == -1:
                lastCla = cla  # 如果是第一次遍历，则将cla赋值给lastCla
            # elif lastCla != cla:  # 如果上一次遍历的标签与这次遍历的标签不同
                isSingle = 1  # 说明该文件包含的标签种类不是单一的
                break

        in_file.close()  # 关闭xml文件

        name = ""
        if isSingle == 1:  # 如果文件标签单一
            # name = str(lastCla) + "_" + labels[lastCla] + "_" + str(nums[lastCla])
            name = "1" + "_" + labels[lastCla] + "_" + str(nums[lastCla])  #3、  “1”表示第几个类别
            nums[lastCla] = nums[lastCla] + 1  # 标签对应的文件数量+1
        else:  # 如果文件标签不单一
            name = "4_mix_" + str(nums[4])
            nums[4] = nums[4] + 1  # mix标签的文件数量+1

        # 重命名xml文件
        old_xml = os.path.join(os.path.abspath(xmldir), file)  # 原xml路径 + 文件名
        new_xml = os.path.join(os.path.abspath(outdir), name + ".xml")  # 新xml路径 + 文件名
        os.rename(old_xml, new_xml)

        # 重命名jpg文件
        old_img = os.path.join(os.path.abspath(imgdir), file.split('.')[0] + '.jpg')  # 原jpg路径 + 文件名
        new_img = os.path.join(os.path.abspath(outdir), name + ".jpg")  # 新jpg路径 + 文件名
        os.rename(old_img, new_img)

#4、  改对应的文件夹绝对路径
xmldir = 'E:/Project_Practice/Transmission_Line_Defect_Detection/1_Data/Compression/Origin_Data_XML/'  # xml文件所在目录
imgdir = 'E:/Project_Practice/Transmission_Line_Defect_Detection/1_Data/Compression/Origin_Data/'  # 图片文件所在目录
outdir = 'E:/Project_Practice/Transmission_Line_Defect_Detection/1_Data/Compression/Rename_Data/'  # 保存目录（本代码将重命名后的xml文件与图片文件保存在同一文件夹下，亦可修改代码保存在不同的文件夹下）

rename(xmldir, imgdir, outdir)  # 调用函数进行重命名