DataLoader for a VDSR network

username: Fabrice_Auzanneau
cooked: <p>Hello all<br>
I’m trying to run a VDSR network in Pytorch and I wonder how I should use the DataLoader.</p>
<p>I have made 4 directories:</p>
<ul>
<li>HR: the ground truth images</li>
<li>LR_div2: each image from the HR folder is downsampled by a factor 2 then upsampled again using bilinear interpolation. They have the same name as the original image</li>
<li>LR_div3 : the same with factor 3</li>
<li>LR_div4 : the same with factor 4</li>
</ul>
<p>Each folder has subfolders ‘train’ and ‘test’. Now I’m wondering how I can write the DataLoader.</p>
<p>I have found this piece of code:</p>
<pre><code class="lang-auto">class AutoEncoderDataSet(Dataset):
    def __init__(self, dir_lr, dir_gt, traintest, transform=None):
        self.dir_lr = self.load_dir_single(join(dir_lr, traintest))
        self.dir_gt = self.load_dir_single(join(dir_gt, traintest))
        self.transform = transform

    def is_image_file(self, filename):
        return any(filename.endswith(extension) for extension in [".png", ".PNG", ".jpg", ".JPG", ".jpeg", ".JPEG"])

    def load_img(self, filename):
        img = Image.open(filename)
        return img

    def load_dir_single(self, directory):
        return [join(directory, x) for x in listdir(directory) if self.is_image_file(x)]

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

    def __getitem__(self, index):
        img_lr = self.load_img(self.dir_lr[index])
        img_gt = self.load_img(self.dir_gt[index])
        sample = {'img_in': img_lr, 'img_gt': img_gt}

        if self.transform:
            sample = self.transform(sample)

        return sample
</code></pre>
<p>together with</p>
<pre><code class="lang-auto">ps = {
    'DIR_LR2': PATH + 'LR_div2',
    'DIR_LR3': PATH + 'LR_div3',
    'DIR_LR4': PATH + 'LR_div4',
    'DIR_HR':  PATH + 'HR'
}
train_set = AutoEncoderDataSet(ps['DIR_LR2'], ps['DIR_HR'], 'train', composed)
train_loader = DataLoader(train_set, batch_size=BATCH_SIZE_TRAIN, shuffle=True, num_workers=4)
</code></pre>
<p>But how can I take into account all 3 LR folders?<br>
Thanks for your help.</p>

username: ptrblck
cooked: <p>If you want to concatenate the samples from all 3 directories, you could add another class method which is able to accept multiple folders, could call <code>load_dir_single</code> on each one of them, and concatenate the samples in the end.</p>

username: Fabrice_Auzanneau
cooked: <p>Thanks for your answer. I’m not sure to understand what you mean, as I’m quite new to pytorch.</p>
<p>I thought that the first argument of DataLoader should be a list of couples, made of an input image and its associated ground truth.<br>
Do you suggest that it could be a list of 4 images (3 LR and the ground truth) ?</p>

username: ptrblck
cooked: <p>The first argument to a <a href="https://pytorch.org/docs/stable/data.html?highlight=dataloader#torch.utils.data.DataLoader"><code>DataLoader</code></a> is a <code>dataset</code>.<br>
In your current code snippet you are already working with a custom <code>Dataset</code> and you can adapt it as needed. My suggestion was to pass all folders into <code>AutoEncoderDataset</code> and to load all directories into a single list.</p>