Как преобразовать изображения RGB в оттенки серого в загрузчике данных PyTorch?

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'>

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())

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-канальный оттенки серого.Узнайте больше об этом на https://pytorch.org/docs/stable/torchvision/transforms.html#torchvision.transforms.Grayscale

Пример кода ниже,

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())

Answer 3

Вы можете реализовать Dataloader не из ImageFolder, а из Datagenerator, напрямую загружать изображения в функции __getitem__.PIL.Image.open ("..") затем в градациях серого, в numy и в Tensor.

Другой вариант - рассчитать канал шкалы серого (Y) из RGB по формуле Y = 0.299 R + 0.587 G + 0.114 B. Массив слайсов и преобразовать в один канал..

Но как ты тренируешь свою модель?обычно тренируют и тестируют данные одинаково.