如何在 PyTorch 数据加载器中将 RGB 图像转换为灰度图像？

Question

我已经从 MNIST 数据集中以 .jpg 格式下载了一些示例图像。现在我正在加载这些图像以测试我的预训练模型。

 # transforms to apply to the data
trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])

# MNIST dataset
test_dataset = dataset.ImageFolder(root=DATA_PATH, transform=trans)

# Data loader
test_loader = DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

这里 DATA_PATH 包含一个带有示例图像的子文件夹。

这是我的网络定义

# Convolutional neural network (two convolutional layers)
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        self.network2D = nn.Sequential(
           nn.Conv2d(1, 32, kernel_size=5, stride=1, padding=2),
           nn.ReLU(),
           nn.MaxPool2d(kernel_size=2, stride=2),
           nn.Conv2d(32, 64, kernel_size=5, stride=1, padding=2),
           nn.ReLU(),
           nn.MaxPool2d(kernel_size=2, stride=2))
        self.network1D = nn.Sequential(
           nn.Dropout(),
           nn.Linear(7 * 7 * 64, 1000),
           nn.Linear(1000, 10))

    def forward(self, x):
        out = self.network2D(x)
        out = out.reshape(out.size(0), -1)
        out = self.network1D(out)
        return out

这是我的推理部分

# Test the model
model = torch.load("mnist_weights_5.pth.tar")
model.eval()

for images, labels in test_loader:
   outputs = model(images.cuda())

当我运行此代码时，出现以下错误：

 RuntimeError: Given groups=1, weight of size [32, 1, 5, 5], expected input[1, 3, 28, 28] to have 1 channels, but got 3 channels instead

我知道图像是作为 3 通道 (RGB) 加载的。那么如何将它们转换为 dataloader 中的单通道？

更新：我改变了 transforms 包括 Grayscale 选项

trans = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,)), transforms.Grayscale(num_output_channels=1)])

但是现在我得到了这个错误

TypeError: img should be PIL Image. Got <class 'torch.Tensor'>

原文由 Harsh Wardhan 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 1

我找到了一个非常简单的解决这个问题的方法。张量的所需尺寸为 [1,1,28,28] 而输入张量的形式为 [1,3,28,28] 。所以我只需要从中读取 1 个频道

images = images[:,0,:,:]

这给了我 [1,28,28] 形式的张量。现在我需要将其转换为 [1,1,28,28] 形式的张量。可以这样做

images = images.unsqueeze(0)

所以把上面两行放在一起，代码的预测部分可以这样写

for images, labels in test_loader:
   images = images[:,0,:,:].unsqueeze(0) ## Extract single channel and reshape the tensor
   outputs = model(images.cuda())

原文由 Harsh Wardhan 发布，翻译遵循 CC BY-SA 4.0 许可协议

Answer 2

当使用 ImageFolder 类并且没有自定义加载器时，pytorch 使用 PIL 加载图像并将其转换为 RGB。如果 torchvision 图像后端是 PIL，则默认加载程序：

 def pil_loader(path):
    with open(path, 'rb') as f:
        img = Image.open(f)
        return img.convert('RGB')

您可以在转换中使用 torchvision 的灰度 函数。它将 3 通道 RGB 图像转换为 1 通道灰度。在 此处 了解更多信息

示例代码如下，

 import torchvision as tv
import numpy as np
import torch.utils.data as data
dataDir         = 'D:\\general\\ML_DL\\datasets\\CIFAR'
trainTransform  = tv.transforms.Compose([tv.transforms.Grayscale(num_output_channels=1),
                                    tv.transforms.ToTensor(),
                                    tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
trainSet        = tv.datasets.CIFAR10(dataDir, train=True, download=False, transform=trainTransform)
dataloader      = data.DataLoader(trainSet, batch_size=1, shuffle=False, num_workers=0)
images, labels  = iter(dataloader).next()
print (images.size())

原文由 Thiagarajan 发布，翻译遵循 CC BY-SA 4.0 许可协议