Загрузка оригинальных изображений помимо преобразованных в оттенки серого с помощью ImageFroder - PullRequest
Новая
       4

Загрузка оригинальных изображений помимо преобразованных в оттенки серого с помощью ImageFroder

1 голос
/ 09 июня 2019

Я пытаюсь настроить GAN для раскрашивания изображений, чтобы я использовал ImageFolder torchvision для загрузки его в виде изображений в градациях серого, но мне также нужны исходные данные вместе с преобразованными.

Я хочу, чтобы это было как можно быстрее, поскольку данные велики, поэтому я хочу, чтобы ImageFolder загружал оба файла одновременно, чтобы сэкономить больше времени. 

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

1 Ответ

1 голос
/ 10 июня 2019

Я предполагаю, что у вас есть 2 переменные набора данных, то есть dataset_bw и dataset_color, которые вы можете загрузить, как вы упомянули, используя ImageFolder. Тогда вы можете сделать следующее: 

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, вы получите два изображения по вашему желанию.

...