simplify dataset loader code

Author: junyanz

data/base_dataset.py

@@ -78,25 +78,30 @@ def get_params(opt, size):
     return {'crop_pos': (x, y), 'flip': flip}
 
 
-def get_transform(opt, params, grayscale=False, method=Image.BICUBIC, convert=True):
+def get_transform(opt, params=None, grayscale=False, method=Image.BICUBIC, convert=True):
     transform_list = []
     if grayscale:
         transform_list.append(transforms.Grayscale(1))
     if 'resize' in opt.preprocess:
         osize = [opt.load_size, opt.load_size]
-        transform_list.append(transforms.Scale(osize, method))
+        transform_list.append(transforms.Resize(osize, method))
     elif 'scale_width' in opt.preprocess:
         transform_list.append(transforms.Lambda(lambda img: __scale_width(img, opt.load_size, method)))
 
     if 'crop' in opt.preprocess:
-        transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size)))
+        if params is None:
+            transform_list.append(transforms.RandomCrop(opt.crop_size))
+        else:
+            transform_list.append(transforms.Lambda(lambda img: __crop(img, params['crop_pos'], opt.crop_size)))
 
     if opt.preprocess == 'none':
-        base = 4
-        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base, method)))
+        transform_list.append(transforms.Lambda(lambda img: __make_power_2(img, base=4, method=method)))
 
-    if not opt.no_flip and params['flip']:
-        transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
+    if not opt.no_flip:
+        if params is None:
+            transform_list.append(transforms.RandomHorizontalFlip())
+        elif params['flip']:
+            transform_list.append(transforms.Lambda(lambda img: __flip(img, params['flip'])))
 
     if convert:
         transform_list += [transforms.ToTensor(),

data/colorization_dataset.py

@@ -1,5 +1,5 @@
 import os.path
-from data.base_dataset import BaseDataset, get_params, get_transform
+from data.base_dataset import BaseDataset, get_transform
 from data.image_folder import make_dataset
 from skimage import color  # require skimage
 from PIL import Image
@@ -39,6 +39,7 @@ def __init__(self, opt):
         self.dir = os.path.join(opt.dataroot)
         self.AB_paths = sorted(make_dataset(self.dir, opt.max_dataset_size))
         assert(opt.input_nc == 1 and opt.output_nc == 2 and opt.direction == 'AtoB')
+        self.transform = get_transform(self.opt, convert=False)
 
     def __getitem__(self, index):
         """Return a data point and its metadata information.
@@ -54,9 +55,7 @@ def __getitem__(self, index):
         """
         path = self.AB_paths[index]
         im = Image.open(path).convert('RGB')
-        transform_params = get_params(self.opt, im.size)
-        transform = get_transform(self.opt, transform_params, convert=False)
-        im = transform(im)
+        im = self.transform(im)
         im = np.array(im)
         lab = color.rgb2lab(im).astype(np.float32)
         lab_t = transforms.ToTensor()(lab)

data/single_dataset.py

@@ -1,4 +1,4 @@
-from data.base_dataset import BaseDataset, get_params, get_transform
+from data.base_dataset import BaseDataset, get_transform
 from data.image_folder import make_dataset
 from PIL import Image
 
@@ -17,8 +17,8 @@ def __init__(self, opt):
         """
         BaseDataset.__init__(self, opt)
         self.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size))
-        self.input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
-        # self.transform = get_transform(opt, input_nc == 1)
+        input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
+        self.transform = get_transform(opt, grayscale=(input_nc == 1))
 
     def __getitem__(self, index):
         """Return a data point and its metadata information.
@@ -32,9 +32,7 @@ def __getitem__(self, index):
         """
         A_path = self.A_paths[index]
         A_img = Image.open(A_path).convert('RGB')
-        transform_params = get_params(self.opt, A_img.size)
-        transform = get_transform(self.opt, transform_params, grayscale=(self.input_nc == 1))
-        A = transform(A_img)
+        A = self.transform(A_img)
         return {'A': A, 'A_paths': A_path}
 
     def __len__(self):

data/unaligned_dataset.py

@@ -1,5 +1,5 @@
 import os.path
-from data.base_dataset import BaseDataset, get_params, get_transform
+from data.base_dataset import BaseDataset, get_transform
 from data.image_folder import make_dataset
 from PIL import Image
 import random
@@ -31,8 +31,10 @@ def __init__(self, opt):
         self.A_size = len(self.A_paths)  # get the size of dataset A
         self.B_size = len(self.B_paths)  # get the size of dataset B
         btoA = self.opt.direction == 'BtoA'
-        self.input_nc = self.opt.output_nc if btoA else self.opt.input_nc       # get the number of channels of input image
-        self.output_nc = self.opt.input_nc if btoA else self.opt.output_nc      # get the number of channels of output image
+        input_nc = self.opt.output_nc if btoA else self.opt.input_nc       # get the number of channels of input image
+        output_nc = self.opt.input_nc if btoA else self.opt.output_nc      # get the number of channels of output image
+        self.transform_A = get_transform(self.opt, grayscale=(input_nc == 1))
+        self.transform_B = get_transform(self.opt, grayscale=(output_nc == 1))
 
     def __getitem__(self, index):
         """Return a data point and its metadata information.
@@ -55,13 +57,8 @@ def __getitem__(self, index):
         A_img = Image.open(A_path).convert('RGB')
         B_img = Image.open(B_path).convert('RGB')
         # apply image transformation
-        A_transform_params = get_params(self.opt, A_img.size)
-        A_transform = get_transform(self.opt, A_transform_params, grayscale=(self.input_nc == 1))
-        A = A_transform(A_img)
-
-        B_transform_params = get_params(self.opt, B_img.size)
-        B_transform = get_transform(self.opt, B_transform_params, grayscale=(self.output_nc == 1))
-        B = B_transform(B_img)
+        A = self.transform_A(A_img)
+        B = self.transform_B(B_img)
 
         return {'A': A, 'B': B, 'A_paths': A_path, 'B_paths': B_path}
 

models/base_model.py

@@ -41,7 +41,7 @@ def __init__(self, opt):
         self.visual_names = []
         self.optimizers = []
         self.image_paths = []
-        self.metric = 0 #  used for learning rate policy 'plateau'
+        self.metric = 0  # used for learning rate policy 'plateau'
 
     @staticmethod
     def modify_commandline_options(parser, is_train):