import torchvision
torchvision.datasets.ImageFolder # 从分类文件夹中创建dataset数据
分目录的方法
torchvision.datasets.ImageFolder # 从分类文件夹中创建dataset数据
base_dir = r'./dataa/4weather'
if not os.path.isdir(base_dir):
os.mkdir(base_dir)
train_dir = os.path.join(base_dir, 'train')
test_dir = os.path.join(base_dir, 'test')
os.mkdir(train_dir)
os.mkdir(test_dir)
specises = ['cloudy', 'rain', 'shine', 'sunrise']
for train_or_test in ['train', 'test']:
for spec in specises:
os.mkdir(os.path.join(base_dir, train_or_test, spec))
创建的结果
Python的enumerate函数_Python 学习者的博客-CSDN博客_enumerate python
# 把图片分别放入文件夹
image_dir = r'./dataa/dataset2'
for i, img in enumerate(os.listdir(image_dir)): # enumerate() 遍历列表输出对应的内容和索引
for spec in specises:
if spec in img:
s = os.path.join(image_dir, img)
if i % 5 == 0:
d = os.path.join(base_dir, 'train', spec, img)
else:
d = os.path.join(base_dir, 'test', spec, img)
shutil.copy(s, d)
# 查看各文件夹数量
for train_or_test in ['train', 'test']:
for spec in specises:
print(train_or_test, spec, len(os.listdir(os.path.join(base_dir, train_or_test, spec))))
torchvision.datasets.ImageFolder使用详解_马鹏森的博客-CSDN博客_datasets.imagefolder