[人工智能]pytorch 数据分类 数据集制作 txt

?标签一个一个写太麻烦了，就写了个脚本

先展示图片 然后输入图片的标签 之后自动生成txt格式 用于一些分类模型

import os
import matplotlib.pyplot as plt
import matplotlib.image as mpimg


files = os.listdir(".\\train-image\\")

for filename in files:
    portion = os.path.splitext(filename)

    if portion[1] == ".jpg":
        print(filename)
        l = mpimg.imread(".\\train-image\\"+filename)
        plt.imshow(l)
        plt.show()
        print('Input grade:')
        label=input()

        f=open('train.txt','a')
        f.write(filename+' ')
        f.write(label+'\n')

f.close()