1

我有一个数据集，其中包含作为输入的图像和作为图像的标签/目标。文件夹中的结构如下：

> DATASET/
> ---TRAIN/
> ------image_xx.png
> ------label_xx.png
> ---TEST/
> ------image_xx.png
> ------label_xx.png

我目前尝试使用 torchvisions 数据集中的“ImageFolder”来加载图像，如下所示：

TRAIN_PATH = '/path/to/dataset/DATASET'
train_data = datasets.ImageFolder(root=TRAIN_PATH, transform=transforms.ToTensor())
train_loader = DataLoader(train_data, batch_size=16, shuffle=True)

不过如下图：

for img, label in train_loader:
    print(img.shape)
    print(label.shape)
    break


torch.Size([16, 3, 128, 128])
torch.Size([16])

标签不是图像，而是指示或类似的东西。有没有一种方便的方法来导入具有上述结构的数据集？

4

1 回答 1

1

ImageFolder当您对每个图像都有离散的标量类时，该数据集是合适的。它期望目录结构使得每个子目录都包含一个特定的类。

对于您的情况，您可以简单地定义自己的torch.nn.Dataset. 本教程可能会有所帮助。

一个简单的例子（我没有尝试运行它来查看它是否正常工作）：

import torch
import os
import cv2

class MyDataset(torch.utils.data.Dataset):

    def __init__(self, root_path, transform=None):
        self.data_paths = [f for f in sorted(os.listdir(root_path)) if f.startswith("image")]
        self.label_paths = [f for f in sorted(os.listdir(root_path)) if f.startswith("label")]
        self.transform = transform

    def __getitem__(self, idx):
        img = cv2.imread(self.data_paths[idx])
        label = cv2.imread(self.label_paths[idx])
        if self.transform:
            img = self.transform(img)
        return img, label

    def __len__(self):
        return len(self.data_paths)

TRAIN_PATH = '/path/to/dataset/DATASET/TRAIN/'
train_data = MyDataset(root_path=TRAIN_PATH, transform=transforms.ToTensor())
train_loader = torch.utils.data.DataLoader(train_data, batch_size=16, shuffle=True)
于 2021-05-05T19:28:11.873 回答