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

我了解图片以三通道（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

Answer 1

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

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

您可以在transforms中使用torchvision的Grayscale函数。它将把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

---

我在使用你的代码时遇到了以下错误： RuntimeError: 输出形状为[1, 32, 32]与广播形状[3, 32, 32]不匹配。 - ma3oun

这个错误是通过实施以下答案得到解决的： https://dev59.com/DFMI5IYBdhLWcg3w4Pk_ - salRad

我认为这并没有将图像转换为灰度，因为您要么使用图像的红色、绿色或蓝色部分，而没有一个是灰度。 - Manjit Ullal

Answer 2

你可以实现Dataloader不是从ImageFolder，而是从Datagenerator中直接在__getitem__函数中加载图像。使用PIL.Image.open("..")打开图像，然后将其转换为灰度图，再转换为numpy数组和Tensor。

另一个选择是通过公式Y = 0.299 R + 0.587 G + 0.114 B.从RGB计算出灰度(Y)通道，然后切片数组并转换为一个通道。

但是如何训练您的模型？通常训练和测试数据以相同的方式加载。

- chro

Answer 3

我找到了一个非常简单的解决方案来解决这个问题。所需张量的尺寸为[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

---

你只从一个红色通道中获取信息，这是不精确的。 - chro

不，如果R = G = B，则这可能实际上可以起作用。在这种情况下，没有理由进行转换，您只需切割一个通道即可。但是，我担心pytorch如何将灰度图像转换为RGB，这可能不是简单地将灰度通道复制到R和G和B。默认的pytorch加载器是：def pil_loader(path): with open(path, 'rb') as f: img = Image.open(f) return img.convert('RGB') - saurabheights