使用ImageFolder加载经过转换的图像以及原始图像 [英] Loading original images besides the transformed ones using ImageFolder

问题描述

我正在尝试训练GAN使图像着色。为此，我正在使用 torchvision 的 ImageFolder 来加载灰度图像，但是我还需要原始数据以及转换后的

I'm trying to train a GAN to colorize images. For that, I'm using ImageFolder of torchvision to load grayscale images but I also need the original data alongwith the transformed ones.

我想要最快的方式，因为数据量很大。我想同时加载 ImageFolder 来减少时间复杂度。

I want it in the fastest way as the data is large. I want to make ImageFolder load both at the same time to reduce the time complexity.

def load_data_bw(opt):
    datapath = '/content/gdrive/My Drive/faces/2003'

    dataset = torchvision.datasets.ImageFolder(datapath,
                       transform=transforms.Compose([
                                                    transforms.Grayscale(num_output_channels=3), #load images as grayscale with three channels
                                                    transforms.RandomChoice(
                                                       [transforms.Resize(opt['loadSize'], interpolation=1),
                                                        transforms.Resize(opt['loadSize'], interpolation=2),
                                                        transforms.Resize(opt['loadSize'], interpolation=3),
                                                        transforms.Resize((opt['loadSize'], opt['loadSize']),
                                                                          interpolation=1),
                                                        transforms.Resize((opt['loadSize'], opt['loadSize']),
                                                                          interpolation=2),
                                                        transforms.Resize((opt['loadSize'], opt['loadSize']),
                                                                          interpolation=3)]
                                                    ),
                                                    transforms.RandomChoice(
                                                       [transforms.RandomResizedCrop(opt['fineSize'], interpolation=1),
                                                        transforms.RandomResizedCrop(opt['fineSize'], interpolation=2),
                                                        transforms.RandomResizedCrop(opt['fineSize'], interpolation=3)]
                                                    ),
                                                    transforms.ColorJitter(brightness=0.1, contrast=0.1),
                                                    transforms.RandomHorizontalFlip(),
                                                    transforms.ToTensor()
                                               ]))
    return dataset

我期望得到：

for iteration, orig_data, gray_data in enumerate(training_data_loader, 1):
    # code..

推荐答案

我假设您有2个数据集变量，即 dataset_bw 和 dataset_color ，您可以使用 ImageFolder 加载它们。然后您可以执行以下操作：

I assume you have 2 dataset variables i.e. dataset_bw and dataset_color that you can load as you mention using ImageFolder. Then you could do the following :

class GAN_dataset(Dataset):
    def __init__(self, dataset_bw, dataset_color):
        self.dataset1 = dataset_bw
        self.dataset2 = dataset_color

    def __getitem__(self, index):
        x1 = self.dataset1[index]
        x2 = self.dataset2[index]

        return x1, x2

    def __len__(self):
        return len(self.dataset1)

dataset = GAN_dataset(dataset_bw, dataset_color)
loader = DataLoader(dataset, batch_size = ...)

这样，当您遍历 loader 时，将根据需要获得两个图像。

This way you when you iterate through loader, you will get two images as you require.

这篇关于使用ImageFolder加载经过转换的图像以及原始图像的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！