Implement DataLoader to provide two copies of image

 Changes to be committed:
	modified:   .gitignore
	new file:   dataloader.py
	new file:   main.py

Author: dheeraj-coding

.gitignore

@@ -1 +1,4 @@
 env/
+.idea/
+data/
+__pycache__

dataloader.py

@@ -0,0 +1,45 @@
+import os
+
+import torch
+import torchvision.transforms as transforms
+from PIL import Image
+import numpy as np
+
+
+class ImageDataset(torch.utils.data.Dataset):
+    def __init__(self, path, transform=None, target_transform=None):
+        folders = os.listdir(path)
+        self.img_names = [os.path.join(path, folder, name) for folder in folders
+                          for name in os.listdir(os.path.join(path, folder))]
+        self.transform = transform
+        self.target_transform = target_transform
+
+    def __len__(self):
+        return len(self.img_names)
+
+    def __getitem__(self, index):
+        x = Image.open(self.img_names[index])
+        xo = x.copy()
+
+        if self.transform:
+            x = self.transform(x)
+        if self.target_transform:
+            xo = self.target_transform(xo)
+        xo = np.array(xo)
+        return x, xo
+
+
+def get_dataloader(path, scale_size=256, crop_size=224, batch_size=4, shuffle=True, num_workers=4):
+    transformer = transforms.Compose([
+        transforms.Resize(scale_size),
+        transforms.CenterCrop(crop_size),
+        transforms.ToTensor()
+    ])
+    target_transformer = transforms.Compose([
+        transforms.Resize(scale_size),
+        transforms.CenterCrop(crop_size),
+    ])
+
+    dataset = ImageDataset(path, transformer, target_transformer)
+
+    return torch.utils.data.DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers)

main.py

@@ -0,0 +1,18 @@
+import matplotlib.pyplot as plt
+
+import dataloader
+
+
+def main():
+    datagen = dataloader.get_dataloader('data')
+    for x, xo in datagen:
+        print(xo.shape)
+        axes = plt.gca()
+        axes.axis('off')
+        axes.imshow(xo[0].numpy())
+        plt.show()
+        break
+
+
+if __name__ == '__main__':
+    main()