基于python的图片CSV标签文档绘制

背景：

        最近在做计算机视觉方面图片瑕疵检测的工作，考虑要用到分类模型，于是打算使用Xception模型，无奈该原始模型要求针对训练图片绘制CSV文档，CSV文档中存放相对应的图片名称及OK和NG标签(其中NG标签对应之类不良)。考虑针对该想法完整代码的编写和验证。

数据集格式：

要求images/train文件夹下面存放一系列图片，图片后缀为JPG格式，图片命名要有一定的规则。

比如：

 文件名称：000000000000000001_1_2_TA07_02_20211116155646412_03.jpg

核心代码展示：

# -*- coding: utf-8 -*-
# @Time    : 2021-12-29 15:31
# @Author  : Pegasus
'''
功能说明：
根据train和test中的图片生成相对应的标签
'''

import os
import csv

#要读取的文件的根目录
root_path=r'./images/train' 

#将所有目录下的文件信息放到列表中
def get_Write_file_infos(path):
    # 文件信息列表
    file_infos_list=[]

    # 遍历并写入文件信息
    for root, dirnames, filenames in os.walk(path):
        for filename in filenames:
            file_infos = {}
            dirname=root
            filename1 = filename.split('.jpg')[0]
            flag = filename1[-1]
            file_infos["ImageId"] = filename1
            if flag=='0':
                file_infos["ClassId"]='0'
            else:
                file_infos["ClassId"]='1'
            file_infos["Flag"] = flag
            #将数据追加字典到列表中
            file_infos_list.append(file_infos)
            
    return file_infos_list


#写入csv文件
def write_csv(file_infos_list):
    with open('train_label.csv','a+',newline='') as csv_file_train:
        csv_writer = csv.DictWriter(csv_file_train,fieldnames=['ImageId','ClassId','Flag'])
        csv_writer.writeheader()
        for each in file_infos_list:
            print(each)
            csv_writer.writerow(each)
            
#主函数
def main():
#调用获取文件信息的函数
    file_infos_list =get_Write_file_infos(root_path)
    write_csv(file_infos_list)

#主程序入口

if __name__ == '__main__':
    main()
    print('处理结束！')

运行完毕如下：

同时运行结束后会在代码同级目录下生成对应的CSV标签文档：