Pytorch从Dataloader加载两个图像

时间:2018-11-30 00:50:38

标签: python image-processing machine-learning pytorch

我正在尝试制作可拍摄低分辨率图像的GAN,并尝试从中创建高分辨率图像。为此,我需要使用一个同时装有高分辨率和低分辨率训练图像的数据加载器。

     data_transform = transforms.Compose([transforms.Resize(imageSize),
                                         transforms.Grayscale(num_output_channels=1),
                                         transforms.ToTensor()])

    dataset_hi = "./hi-res-train"
    dataset_lo = "./low-res-train"

    img_data_hi = dset.ImageFolder(root=dataset_hi,transform=data_transform)
    img_data_lo = dset.ImageFolder(root=dataset_lo,transform=data_transform)

    dataloader_hi = torch.utils.data.DataLoader(img_data_hi, batch_size = batchSize, shuffle = True, num_workers = 2) 
    dataloader_lo = torch.utils.data.DataLoader(img_data_lo, batch_size = batchSize, shuffle = True, num_workers = 2) 

我曾尝试使用两个单独的数据加载器(如上所示),但在对其进行混洗时,由于高分辨图像和低分辨率图像不匹配,因此无法枚举它们。我该如何做才能用pytorch枚举和混洗?

1 个答案:

答案 0 :(得分:1)

假设您为高分辨率图像和低分辨率图像使用相似的名称(例如img01_hi和img01_low),一种选择是创建一个自定义Dataloader,该数据加载器通过覆盖__getitem__方法来返回两个图像。

当两个图像在一次调用中返回时,您可以通过在文件名后附加_hi和_low来确保它们匹配。

您可能需要创建一个“提示”文本文件,其中包含所有图像文件名的列表,以确保每个图像文件仅处理一次。