数据集文件夹的标签可能是一堆名字,那我们当然不希望他们是这样了,如何将其与数字对应起来呢
当然我写了好一段时间,结果发现datasets.ImageFolder已经帮我实现了这个功能(伤心有点),但我还是记录一下,内部实现可能和它的也不同。
首先说明我的每个文件夹的名字都是类的名字,如果不是的话要改改
以下是全部代码
import os
path = 'mth/test'
path_names = []
labelnames= []
for root, dirs, files in os.walk(path, topdown=False):
for name in files:
path_names.append(os.path.join(root, name))
label = root.split('\\')[-1]
labelnames.append(label)
#lab = sorted(list(set(labelnames))) #集合每次都顺序不同,因此我排了一下顺序
#print(len(lab))
#生成标签字典 ,可能比较难懂,大概就是我们这里很多字当然我们需要把他们对应成数字用range(len())
labdict = dict(zip(labelnames,list(range(0,len(labelnames)))))
#这里巧妙就体现出字典的价值,从labelnames里面找到对应位置的类别
labels = [str(labdict[i]) for i in labelnames]
writefile = open('testlabels.txt', 'w',encoding='utf-8')
for label in labels:
writefile.write(label+'\n')
writefile.close()
with open('testpathes.txt','w',encoding='utf-8') as f:
for path_name in path_names:
f.write(path_name+'\n')
通过walk的root来获取同一个文件夹下面的所有file名字,然后把他们的名字保存下来也就是每一类的名字。
关键是如何和数字对应呢?
上面有一行注释里面的是set
是因为之前在别的里面由于读取到的类有重复要去重。
在这里要建立一个和类别数量相对应的range出的字典
然后针对于每次读出的i不同去字典中获取之前建立好的值(这个关键想法也是从别的大神那里借鉴来的。
后面我们当然为了写入文件我加了str()