Pytorch：加载灰度图像数据集。

Question

我想加载一组灰度图像数据集。 我使用了 ImageFolder，但是它默认不会加载灰度图像，因为它会将图像转换为RGB格式。

我发现可以使用 ImageFolder 加载图像，然后使用以下方法将图像转换为灰度：

transforms.Grayscale(num_output_channels=1)

或者

ImageOps.grayscale(image)

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

- user3582433

---

这个回答解决了你的问题吗？如何在PyTorch数据加载器中将RGB图像转换为灰度图像？ - Andreas K.

Answer 1

假设数据集存储在如下的“Dataset”文件夹中，请将根目录设置为“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应该是train/test数据加载器大小的因子，否则将会出现错误。

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

- Tanya Jain

Answer 2

是的，那是正确的，据我所知，默认情况下，pillow 加载图像时使用的是 RGB，例如请参见 此问题的答案。因此，转换为 grayscale 是唯一的方法，尽管这需要时间。

纯 PyTorch 解决方案（如果 ImageFolder 不适用）

您可以编写自己的数据加载功能，如果我是您，我不会选择 fastai 路线，因为它非常高级并且会剥夺您的控制权（您可能根本不需要这些功能）。

原则上，您只需创建类似于以下内容的东西：

import pathlib

import torch
from PIL import Image


class ImageDataset(torch.utils.data.Dataset):
    def __init__(self, path: pathlib.Path, images_class: int, regex="*.png"):
        self.files = [file for file in path.glob(regex)]
        self.images_class: int = images_class

    def __getitem__(self, index):
        return Image.open(self.files[index]).convert("LA"), self.images_class


# Assuming you have `png` images, can modify that with regex
final_dataset = (
    ImageDataset(pathlib.Path("/path/to/dogs/images"), 0)
    + ImageDataset(pathlib.Path("/path/to/cats/images"), 1)
    + ImageDataset(pathlib.Path("/path/to/turtles/images"), 2)
)

上面的代码会从提供的路径中获取图像，每个图像都会返回相应的类别。

这样做可以让你拥有更多的灵活性（不同于 torchvision.datasets.ImageFolder 的文件夹设置）需要写几行额外的代码。

当然，你也可以添加更多的内容或使用循环或其他方法。你还可以应用 torchvision.transforms，例如将上述图像转换为张量，请参阅下文。

torchdata 解决方案

免责声明，作者在这里。如果你关心数据加载时间和灰度转换，你可以使用第三方库 torchdata 来处理 pytorch。

使用它，你可以创建与上面相同的东西，但使用缓存或映射（以便轻松使用 torchvision.transforms 或其他转换），以及一些其他的东西，如来自 tensorflow.data 模块的知识，见下文：

import pathlib

from PIL import Image

import torchdata


# Change inheritance
class ImageDataset(torchdata.Dataset):
    def __init__(self, path: pathlib.Path, images_class: int, regex="*.png"):
        super().__init__()  # And add constructor call and that's it
        self.files = [file for file in path.glob(regex)]
        self.images_class: int = images_class

    def __getitem__(self, index):
        return Image.open(self.files[index]), self.images_class


final_dataset = (
    ImageDataset(pathlib.Path("/path/to/dogs/images"), 0)
    + ImageDataset(pathlib.Path("/path/to/cats/images"), 1)
    + ImageDataset(pathlib.Path("/path/to/turtles/images"), 2)
).cache()  # will cache data in-memory after first pass
# You could apply transformations after caching for possible speed-up

torchvision ImageFolder loader

正如评论中@jodag所正确指出的，可以使用具有单个参数path的loader可调用函数来进行自定义数据打开，例如对于灰度图像，它可以是：

from PIL import Image

import torchvision

dataset = torchvision.datasets.ImageFolder(
    "/path/to/images", loader=lambda path: Image.open(path).convert("LA")
)

请注意，您也可以将其用于其他类型的文件，这些文件不一定是图像。

- Szymon Maszke

---

值得一提的是，您可以为 ImageFolder 提供一个 loader 参数来覆盖默认的加载行为。 - jodag

@jodag 谢谢，我甚至不知道这个。 - Szymon Maszke

Answer 3

制作自定义加载器，将其提供给ImageFolder：

import numpy as np
from PIL import Image, ImageOps

def gray_reader(image_path):
    im = Image.open(image_path)
    im2 = ImageOps.grayscale(im)
    im.close()
    return np.array(im2)   # return np array
    # return im2           # return PIL Image

some_dataset = ImageFolder(image_root_path, loader=gray_reader)

编辑：

下面的代码比之前更好，获取彩色图像并在 transform() 中转换为灰度。

def get_transformer(h, w):
    valid_transform = transforms.Compose([            
        transforms.ToPILImage(),                                    
        transforms.Grayscale(num_output_channels=1),
        transforms.Resize((h, w)),                                    
        transforms.ToTensor(),                                    
        transforms.Normalize([0.5], [0.5]) ])
    return valid_transform

- sailfish009