PyTorch数据集:ImageFolder和子文件夹过滤

时间:2019-03-30 21:29:57

标签: python pytorch

我想使用ImageFolder创建一个图像数据集。

我当前的图像目录结构如下:

/root
-- train/
---- 001.jpg
---- 002.jpg
---- ....
-- test/
---- 001.jpg
---- 002.jpg
---- ....

我想拥有一个专门用于训练数据的数据集和一个专门用于测试数据的数据集。

据我了解,这样做:

dataset = ImageFolder(root='root/train')

找不到任何图像。

dataset = ImageFolder(root='root')

找到图像,但是traintest图像只是一起被加扰。

ImageFolder具有参数loader,但我没有找到任何用例。

如何根据根文件夹中的子文件夹来区分它们?

3 个答案:

答案 0 :(得分:1)

ImageFolder期望数据文件夹(您作为root参数传递)包含表示其图像所属类的子文件夹。像这样:

data/
├── train/
|   ├── class_0/
|   |   ├── 001.jpg
|   |   ├── 002.jpg
|   |   └── 003.jpg
|   └── class_1/
|       ├── 004.jpg
|       └── 005.jpg
└── test/
    ├── class_0/
    |   ├── 006.jpg
    |   └── 007.jpg
    └── class_1/
        ├── 008.jpg
        └── 009.jpg

为简单起见,我假设只有两个类。具有以上正确的文件夹结构,您可以执行以下操作:

train_dataset = ImageFolder(root='data/train')
test_dataset  = ImageFolder(root='data/test')

由于您没有这种结构,因此一个显而易见的选择是创建类子文件夹并将图像放入其中。另一个选择是创建自定义数据集,而不是使用ImageFolder-如何执行取决于类 存储图像的标签,例如您可能会有一个带有“文件名”和“类”的.csv文件,请参见here

答案 1 :(得分:0)

我也遇到错误。我尝试了不同的方法,但是只能通过创建一个子文件夹来解决它。

答案 2 :(得分:0)

我找到了创建每个类的子文件夹的方法,分别针对火车/验证/测试,如ImageFolder预期的那样可以很好地工作。这是我为自己的用例创建的脚本,您可以针对自己的用例进行修改

data_dir = '/content/data/oxford-102-flowers/'
files = ['train.txt','test.txt','valid.txt']

for i in files:
  with open(data_dir + i) as myfile:
    for line in myfile:
      curr = i.split('.')[0]
      l = line.split()
      src = os.path.join(data_dir + l[0])

      dir = os.path.join(data_dir + curr)
      if not os.path.isdir(dir):
        os.mkdir(dir)
      
      sub_dir = os.path.join(dir + '/' + l[1])
      if not os.path.isdir(sub_dir):
        os.mkdir(sub_dir)

      os.system('cp "%s" "%s"' % (src, sub_dir))
print("All files copied to the subfolders")

我正在Oxford-102-Dataset上工作,每个训练,验证和测试集都有三个.txt文件。 txt文件包含图像的位置和名称(例如:jpg / image_05038.jpg 58,其中58表示实际类的基本真值,“ jpg”是存储所有图像的源文件夹)< / p>