Yolo V8 模型训练初学者笔记（番外篇）-- 一键转移labalimg标签文件

前言

• 上一篇文章中，我用 labelimg 工具为每张训练所需图片打标签(标注)后，会保存输出对应到每一张图片的 xxx.txt 文本和classess.txt（记录标签名的文本），这些 xxx.txt 文本记录着对应每一张图片上的标注索引和标注坐标信息

• 并且在执行训练脚本之前，我们需要把打标签的原始图片和标签文本整理到 data/xxx 目录下，还需要在 data/xxx.yaml 中写好配置信息

• 那每次训练模型之前都要手动这么做就太费时间了，于是我写下来这个一劳永逸的 copy_labelimgfiles_to_datadir.py 脚本

• 以下是说明案例截图

• 其中自动写入的 xxx.yaml 案例截图如下

脚本源码在此！！！

• copy_labelimgfiles_to_datadir.py

import os
import shutil
import random

cur_dir = os.path.dirname(__file__).replace('\\', '/')
print(cur_dir)

source_img_dir = cur_dir + '/source_img'
labelimg_output_dir = cur_dir + '/labelimg_output'

data_dir = cur_dir + '/data'
data_images_dir = data_dir + '/images'
data_labels_dir = data_dir + '/labels'
data_train_dir = data_dir + '/train'
data_train_images_dir = data_train_dir + '/images'
data_train_labels_dir = data_train_dir + '/labels'
data_val_dir = data_dir + '/val'
data_val_images_dir = data_val_dir + '/images'
data_val_labels_dir = data_val_dir + '/labels'
data_test_dir = data_dir + '/test'
data_test_images_dir = data_test_dir + '/images'
data_test_labels_dir = data_test_dir + '/labels'
data_yoloyaml_file_path = data_dir + '/yolov8nconfig.yaml'

def read_full_path_list(dir_path):
    full_path_list = []
    files = os.listdir(dir_path)
    for f in files:
        full_path_list.append(dir_path + '/' + f)
    return full_path_list

def create_dir(dir_path):
    if not os.path.exists(dir_path):
        os.mkdir(dir_path)

def read_labelimg_files():
    # 判断读取的目录是否存在
    if not os.path.exists(source_img_dir):
        print(f'【labelimg_output_dir】: {source_img_dir} not exists')
        return
    
    if not os.path.exists(labelimg_output_dir):
        print(f'【labelimg_output_dir】: {labelimg_output_dir} not exists')
        return

    # 读取 source_img_dir 目录下的所有文件
    source_img_files_fullpath_list = read_full_path_list(source_img_dir)
    # 读取 labelimg_output_dir 目录下的所有文件
    labelimg_output_files_fullpath_list = read_full_path_list(labelimg_output_dir)
    labelimg_txt_files_fullpath_list = []
    labelimg_classestxt_file_fullpath = ''
    for f in labelimg_output_files_fullpath_list:
        # 提取出 classes.txt 的文件
        if f.endswith('classes.txt'):
            labelimg_classestxt_file_fullpath = f
        # 提取出 txt 文件
        elif f.endswith('.txt'):
            labelimg_txt_files_fullpath_list.append(f)

    # print(f'source_img_files_fullpath_list: {source_img_files_fullpath_list}')
    # print(f'labelimg_txt_files_fullpath_list: {labelimg_txt_files_fullpath_list}')
    # print(f'labelimg_classestxt_file_fullpath: {labelimg_classestxt_file_fullpath}')
    
    return source_img_files_fullpath_list, labelimg_txt_files_fullpath_list, labelimg_classestxt_file_fullpath
    
def copy_labelimg_files_to_data_dir(source_img_files_fullpath_list, labelimg_txt_files_fullpath_list, labelimg_classestxt_file_fullpath):
    # 从 labelimg_classestxt_file_fullpath 文件中逐行读取出类名
    label_classes = []
    with open(labelimg_classestxt_file_fullpath, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            line = line.strip()
            if line != '':
                label_classes.append(line)
    # 写入 yolo yaml 配置文件
    yoloymal_config_contents = [
        'train: ' + data_train_dir,
        '\n',
        'val: ' + data_val_dir,
        '\n',
        'test: ' + data_test_dir,
        '\n',
        '\n',
        'nc: ' + str(len(label_classes)),
        '\n',
        '\n',
        'names: ' + str(label_classes)
    ]
    # 根据 yoloymal_config_contents 数组中的内容，逐行写入到 data_yoloyaml_file_path 文件中
    with open(data_yoloyaml_file_path, 'w', encoding='utf-8') as f:
        f.writelines(yoloymal_config_contents)


    # 把参数中的文件复制到 data_dir 目录下对应的目录中
    for simg in source_img_files_fullpath_list:
        shutil.copy2(simg, data_images_dir)
        shutil.copy2(simg, data_train_images_dir)

    for ltxt in labelimg_txt_files_fullpath_list:
        shutil.copy2(ltxt, data_labels_dir)
        shutil.copy2(ltxt, data_train_labels_dir)

    # data/val 和 data/test 目录下不需要放所有的文件，只需随机取一部分的文件即可
    simg_files_pre_parts = []
    ltxt_files_pre_parts = []
    simg_files_end_parts = []
    ltxt_files_end_parts = []
    # 取文件列表一部分的几个文件，放到 data/val 和 data/test 目录下
    files_pre_parts_i_list = []
    files_end_parts_i_list = []
    files_len = len(source_img_files_fullpath_list)
    half_files_len = int(files_len/2)
    for i in range(half_files_len):
        files_pre_parts_i_list.append(random.randint(0, half_files_len))
        files_end_parts_i_list.append(random.randint(half_files_len+1, files_len-1))
    for i in files_pre_parts_i_list:
        simg_files_pre_parts.append(source_img_files_fullpath_list[i])
        ltxt_files_pre_parts.append(labelimg_txt_files_fullpath_list[i])
    for i in files_end_parts_i_list:
        simg_files_end_parts.append(source_img_files_fullpath_list[i])
        ltxt_files_end_parts.append(labelimg_txt_files_fullpath_list[i])
    # 把前面随机提取出来的文件，分别复制到 data/val 和 data/test 目录下
    for f in simg_files_pre_parts:
            shutil.copy2(f, data_val_images_dir)
    for f in ltxt_files_pre_parts:
            shutil.copy2(f, data_val_labels_dir)
    for f in simg_files_end_parts:
            shutil.copy2(f, data_test_images_dir)
    for f in ltxt_files_end_parts:
            shutil.copy2(f, data_test_labels_dir)
           
def remove_data_dir():
    if os.path.exists(data_dir):
        # 删除 data_dir 目录及其下面的所有子目录和文件
        shutil.rmtree(data_dir)

def create_data_dir():
    remove_data_dir()

    create_dir(data_dir)
    create_dir(data_images_dir)
    create_dir(data_labels_dir)
    create_dir(data_train_dir)
    create_dir(data_train_images_dir)    
    create_dir(data_train_labels_dir)
    create_dir(data_val_dir)
    create_dir(data_val_images_dir)
    create_dir(data_val_labels_dir)
    create_dir(data_test_dir)
    create_dir(data_test_images_dir)
    create_dir(data_test_labels_dir)

if __name__ == '__main__':
    source_img_files_fullpath_list, labelimg_txt_files_fullpath_list, labelimg_classestxt_file_fullpath = read_labelimg_files()
    create_data_dir()
    copy_labelimg_files_to_data_dir(source_img_files_fullpath_list, labelimg_txt_files_fullpath_list, labelimg_classestxt_file_fullpath)