图像与标注文件分离

本步骤是针对使用 labelimg 时，默认图像文件和标注文件存放到同一目录，因此需要分开处理，最终结果为图像在 images文件夹下，标注文件在 labels文件夹下。如果指定的是不同目录，可以省略此步骤

separate.py

# 数据集图像与标注文件分离脚本

import os
import shutil

import argparse

# 创建 ArgumentParser 对象
parser = argparse.ArgumentParser(description='数据集划分工具')

# 添加参数
parser.add_argument('--dir', help='目录位置', required=True)

# 解析命令行参数
args = parser.parse_args()

# 当前文件夹路径
folder_path = args.dir

# print(folder_path)

# 删除已存在的 images 文件夹和 labels 文件夹及其内容
shutil.rmtree(os.path.join(folder_path, 'images'), ignore_errors=True)
shutil.rmtree(os.path.join(folder_path, 'labels'), ignore_errors=True)

# 创建 images 文件夹和 labels 文件夹
os.makedirs(os.path.join(folder_path, 'images'))
os.makedirs(os.path.join(folder_path, 'labels'))

# 递归查找当前文件夹下的所有文件
for root, dirs, files in os.walk(folder_path):
        # 排除 label 和 images 文件夹
    if 'labels' in dirs:
        dirs.remove('labels')
    if 'images' in dirs:
        dirs.remove('images')
    print(root, dirs, files)
    for file in files:
        file_path = os.path.join(root, file)
        if file.endswith('.jpg'):
            # 复制 jpg 文件到 images 文件夹下（copy2 同名强制覆盖）
            shutil.copy(file_path, os.path.join(folder_path, 'images', file))
        elif file.endswith('.txt'):
            # 复制 txt 文件到 labels 文件夹下（copy2 同名强制覆盖）
            shutil.copy(file_path, os.path.join(folder_path, 'labels', file))

运行脚本

python separate.py --dir /home/hualiujie/baoxinshagnchuan/ultralytics-main-cgh/datasets/TVDSDataSet

划分数据集

将数据集以 7:2:1的比例划分为训练集，验证集、测试集

divide_dataset.py

# 划分数据集（7:2:1）

import os
import numpy as np
import math
import shutil
import tkinter as tk
from tkinter import filedialog

# 使用时请选择数据集根目录，根目录必须包含images与labels两个子目录
# images（原图）与labels（标签）对应
# train：训练集（训练模型）
# val：验证集（选择模型）
# test：测试集（评估模型）

# 遍历images、labels目录中的文件树
# 得到单个文件对象
# 利用随机数将文件名打乱放入数组对象中
# 对数组进行切割，分成三份元组，按照比例 train(7),val(2),test(1)
# 遍历单个元组，按照元组中的文本查找文件并分类

# walk 函数，一个目录，返回三个元组，第一个是目录路径，第二个参数是子目录树，第三个参数是该目录下的文件树


# 数据提取函数，将文件夹中的名称提取出来
# 解释下面代码含义
def data_extraction(paths):
    file_name_list = np.array([])
    for dirs, subdir, fileList in os.walk(paths):
        for file in fileList:
            # 文件名与文件后辍分割
            file_name = os.path.splitext(file)[0]
            # 将文件名放入数组中
            if file_name != 'classes':
                file_name_list = np.append(file_name_list, file_name)
    return file_name_list


# 分配资源（随机）
def allocation(file_list: np.array([]), percentage: list):
    # 随机为 train 分配资源
    train_num = math.floor(np.size(file_list) * percentage[0])
    train_list = np.random.choice(file_list, train_num, replace=False)

    # 得到剩余资源
    leftover_list = np.setdiff1d(file_list, train_list)

    # 随机为 val 分配资源
    val_num = math.floor(np.size(file_list) * percentage[1])
    val_list = np.random.choice(leftover_list, val_num, replace=False)

    # test 拾取剩下资源
    test_list = np.setdiff1d(leftover_list, val_list)
    return np.array([train_list, val_list, test_list], dtype=object)


# 检测文件夹是否存在
def dir_make(path):
    if not os.path.exists(path):
        os.makedirs(path)


# 根据文件名移动文件
def file_move(path, path2, __name__):
    # print(path,path2,__name__)
    # 遍历老目录
    for root, dirs, files in os.walk(path):
        # 查找文件
        for filename in files:
            # print(filename,name)
            # print(name in filename)
            if name in filename:
                f1 = os.path.join(path, filename)
                f2 = os.path.join(path2, filename)
                shutil.copy(f1, f2)
                print(f1 + ' 已经被复制划分到 ' + f2)
                break
        break


# 主函数
if __name__ == '__main__':

    import argparse
    # 创建 ArgumentParser 对象
    parser = argparse.ArgumentParser(description='数据集划分脚本工具')
    # 添加参数
    parser.add_argument('--dir', help='目录位置', required=True)

    # 解析命令行参数
    args = parser.parse_args()

    PATH = args.dir

    # 图片路径
    ImgPath = os.path.join(PATH, 'images')
    # 标注路径
    LabelPath = os.path.join(PATH, 'labels')
    # 训练占比，分别对应 train、val，分配前两个参数时，请不要超过0.9，因为test会自动得到剩下的图片
    PERCENTAGE = [0.7, 0.2]
    print(ImgPath,LabelPath)
    # 检测 train 文件夹
    ImgTrainPath = os.path.join(ImgPath, 'train')
    dir_make(ImgTrainPath)
    LabTrainPath = os.path.join(LabelPath, 'train')
    dir_make(LabTrainPath)

    # 检测 val 文件夹
    ImgValPath = os.path.join(ImgPath, 'val')
    dir_make(ImgValPath)
    LabValPath = os.path.join(LabelPath, 'val')
    dir_make(LabValPath)

    # 检测 test 文件夹
    ImgTestPath = os.path.join(ImgPath, 'test')
    dir_make(ImgTestPath)
    LabTestPath = os.path.join(LabelPath, 'test')
    dir_make(LabTestPath)

    # 获得资源分配
    res = allocation(data_extraction(ImgPath), PERCENTAGE)
    # 遍历三个数组
    for res_item in res:
        # 遍历当前数组
        for name in res_item:
            # train 分类
            if name in res[0]:
                file_move(ImgPath, ImgTrainPath, name)
                file_move(LabelPath, LabTrainPath, name)
            # val 分类
            elif name in res[1]:
                file_move(ImgPath, ImgValPath, name)
                file_move(LabelPath, LabValPath, name)
            # test 分类
            else:
                file_move(ImgPath, ImgTestPath, name)
                file_move(LabelPath, LabTestPath, name)

运行脚本

python divide_dataset.py --dir /home/hualiujie/baoxinshagnchuan/ultralytics-main-cgh/datasets/TVDSDataSet