Pytorch：加载灰度图像数据集

Question

我想加载灰度图像的数据集。我使用过ImageFolder，但默认情况下不会加载灰色图像，因为它将图像转换为 RGB。

我找到了使用 ImageFolder 加载图像并在灰度转换图像后使用的解决方案：

transforms.Grayscale(num_output_channels=1)

或者

ImageOps.grayscale(image)

这是对的吗？如何在不进行转换的情况下加载灰度图像？我尝试了ImageDataBunch，但导入时遇到问题fastai.vision

Answer 1

假设数据集存储在“Dataset”文件夹中，如下所示，将根目录设置为“Dataset”：

数据集

• 类_1
  • img1.png
  • img2.png
• 类_2
  • img1.png
  • img2.png

from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader, random_split
from torchvision import transforms

root = 'Dataset/'

data_transform = transforms.Compose([transforms.Grayscale(num_output_channels=1),
                                     transforms.ToTensor()])
dataset = ImageFolder(root, transform=data_transform)

作为参考，训练和测试数据集分别分为 70% 和 30%。

# Split test and train dataset 
train_size = int(0.7 * len(dataset))
test_size = len(dataset) - train_size
train_data, test_data = random_split(dataset, [train_size, test_size])

该数据集可以进一步分为训练数据加载器和测试数据加载器，如下所示以批量执行操作。

通常您会看到数据集被分配一次batch_size，用于训练和测试加载器。但是，我尝试单独定义它。这个想法是给出batch_size，使其成为训练/测试数据加载器大小的一个因素，否则会给出错误。

# Set batch size of train data loader
batch_size_train = 20

# Set batch size of test data loader
batch_size_test = 22

# load the split train and test data into batches via DataLoader()
train_loader = DataLoader(train_data, batch_size=batch_size_train, shuffle=True)
test_loader = DataLoader(test_data, batch_size=batch_size_test, shuffle=True)