2021-03-26

最新推荐文章于 2024-07-30 21:09:58 发布

bhhabc

最新推荐文章于 2024-07-30 21:09:58 发布

阅读量108

点赞数 1

分类专栏： Python 文章标签： python

本文链接：https://blog.csdn.net/bhhabc/article/details/115253795

版权

Python 专栏收录该内容

3 篇文章 0 订阅

订阅专栏

采坑记录

我想通过torchvision.datasets.ImageFolder这个类来读取一个文件夹，其子文件夹如下图所示：我以为读进来会是{‘1’: 0, ‘2’: 1, ‘3’: 2, ‘4’: 3, ‘5’: 4,…，‘20’: 19}这样的格式，即我希望读进来的顺序是子文件夹1，2，3…，20，结果却是子文件夹1,10,11,12,13,14…。
大文件夹下的子文件夹
在这里插入图片描述
实际上，我们只需要把上面那些子文件夹命名成同一长度的数字，即可实现上述效果。

最后贴上全部的代码：

import os
import torch
import torchvision

path = 'THU_Class/'

transform = torchvision.transforms.Compose([
    torchvision.transforms.Resize((224, 224)),
    torchvision.transforms.ToTensor(),
])
dataset = torchvision.datasets.ImageFolder(root=path, transform=transform)
# print(dataset)
train_size = int(0.9 * (len(dataset)))
test_size = len(dataset) - train_size
train_set, test_set = torch.utils.data.random_split(dataset, [train_size, test_size],
                                                    generator=torch.Generator().manual_seed(42))
print(dataset.class_to_idx)

for file in os.listdir(path):
    file_name = file.zfill(5)
    os.rename(path+file, path+file_name)