YOLO标注文件清洗案例代码-学习篇

背景简介

YOLO标注文件清洗
训练一个人工智能算法需要一个庞大的数据集，这个数据集需要进行人为标注
但由于出现意外，造成部分数据丢失，使得标注文件和图片文件的文件名前缀不能一一对应
需要写一段代码将可以文件名前缀一一对应的文件保存到一个新的文件夹中，已完成数据的清洗

问题背景

待清洗的文件目录，images中是图片，labels中是txt标注文件，标注后images和labels中的文件名时相同的，现在数据丢失后，部分标注文件images和labels对不上了，
需要筛选出images和labels文件中img和txt文件名字相同的文件，

比对完导出的文件目录

代码文件如下：

#抽取的公共方法，拼接返回某个文件夹下所有的文件的路径集合
def get_info(root_from):
    file_full_path_list = []
    for root, dirs, files in os.walk(root_from):
        for file in files:
            file_full_path = os.path.join(root, file)
            file_full_path1 = file_full_path.split('\\')
            file_short_path  = os.path.join(file_full_path1[-2],file_full_path1[-1])#以\\分隔路径后，获取倒数第一第二个元素 拼接起来
            file_i_name = file_short_path.split(".")[0] #以.分隔字符串，去掉后缀
            file_full_path_list.append(file_i_name)
    return file_full_path_list

root_path_from = r'./fire_yolo_format' #待清洗的目录
root_path_save = r'./clean_data'  #清洗后保存的目录
#拼接出待清洗的文件夹下的子文件夹
root_images_from = os.path.join(root_path_from,'images')
root_labels_from = os.path.join(root_path_from,'labels')
#拼接出清洗后存储的子文件夹
root_images_save = os.path.join(root_path_save,'images')
root_labels_save = os.path.join(root_path_save,'labels')

print(root_images_from)
print(root_labels_from)
print(root_images_save)
print(root_labels_save)

#循环遍历创建清洗后文件存储文件夹
"""dir_list_1 = ['images','labels']
dir_name_list = ['train','test','val']
for dir1 in dir_list_1:
    for dir2 in dir_name_list:
        dir_full_path = os.path.join(root_path_save,dir1,dir2)
        if not os.path.exists(dir_full_path):
            os.makedirs(dir_full_path)"""

#遍历待清洗的图片文件夹下的文件
image_file_full_path_list = get_info(root_images_from)
#print(image_file_full_path_list)
labels_file_full_path_list = get_info(root_labels_from)
print(labels_file_full_path_list)
#labels下的txt和images下的图片名称比较，只比较\\val\\small (28).txt 后面的这部分
#使用集合做交集比较
image_set = set(image_file_full_path_list)
label_set = set(labels_file_full_path_list)
intersection_set = image_set & label_set
print(len(image_set))
print(len(label_set))
print(len(intersection_set))

print(intersection_set)#这些是两个集合里都有的，是能一一对上的。
#把这些路径拼接成完整路径再复制到clean_data文件夹
for intersection_i in intersection_set:
    #file_yolo_format中的能一一对上的文件
    intersection_i_image_full_path_from = os.path.join(root_images_from,intersection_i) + '.jpg'
    intersection_i_label_full_path_from = os.path.join(root_labels_from,intersection_i) + '.txt'
    #要保存到的clean_data的文件完整路径
    intersection_i_image_full_path_save = os.path.join(root_images_save, intersection_i) + '.jpg'
    intersection_i_label_full_path_save = os.path.join(root_labels_save, intersection_i) + '.txt'

    shutil.copy(intersection_i_image_full_path_from, intersection_i_image_full_path_save)
    shutil.copy(intersection_i_label_full_path_from, intersection_i_label_full_path_save)