image - Pytorch:加载灰度图像数据集

Question

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

我找到了使用 ImageFolder 加载图像并在将图像转换为灰度图像后使用以下方法的解决方案:

transforms.Grayscale(num_output_channels=1)

或

ImageOps.grayscale(image)

正确吗？ 如何在不转换的情况下加载灰度图像？我尝试了 ImageDataBunch，但我无法导入 fastai.vision

Answer 1

假设数据集存储在下面给出的“Dataset”文件夹中，将根目录设置为“Dataset”:

数据集

• class_1
  • img1.png
  • img2.png
• class_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)