add comments: datasets

Author: junyanz

data/aligned_dataset.py

@@ -14,7 +14,11 @@ class AlignedDataset(BaseDataset):
     """
 
     def __init__(self, opt):
-        """Initialize this dataset class."""
+        """Initialize this dataset class.
+
+        Parameters:
+            opt -- options (needs to be a subclass of BaseOptions)
+        """
         BaseDataset.__init__(self, opt)
         self.dir_AB = os.path.join(opt.dataroot, opt.phase)  # get the image directory
         self.AB_paths = sorted(make_dataset(self.dir_AB, opt.max_dataset_size))  # get image paths
@@ -32,7 +36,7 @@ def __getitem__(self, index):
         Parameters:
             index - - a random integer for data indexing
 
-        Returns a dictionary of A, B, A_paths and B_paths
+        Returns a dictionary that contains A, B, A_paths and B_paths
             A(tensor) - - an image in the input domain
             B(tensor) - - its corresponding image in the target domain
             A_paths(str) - - image paths

data/base_dataset.py

@@ -1,16 +1,43 @@
+"""This module implements an abstract base class (ABC) 'BaseDataset' for datasets.
+
+It also includes common transformation functions (e.g., get_transform, __scale_width), which can be later used in subclasses.
+"""
 import torch.utils.data as data
 from PIL import Image
 import torchvision.transforms as transforms
 from abc import ABC, abstractmethod
 
 
 class BaseDataset(data.Dataset, ABC):
+    """This class is an abstract base class (ABC) for datasets.
+
+    To create a subclass, you need to implement four functions:
+    -- <__init__> (initialize the class, first call BaseDataset.__init__(self, opt))
+    -- <__len__> (return the size of dataset)
+    -- <__getitem__>　(get a data point)
+    -- (optionally) <modify_commandline_options> (add dataset-specific options and set default options).
+    """
+
     def __init__(self, opt):
+        """Initialize the class; save the options in the class
+
+        Parameters:
+            opt -- options (needs to be a subclass of BaseOptions)
+        """
         self.opt = opt
         self.root = opt.dataroot
 
     @staticmethod
     def modify_commandline_options(parser, is_train):
+        """Add new dataset-specific options, and rewrite default values for existing options.
+
+        Parameters:
+            parser -- original option parser
+            is_train -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
+
+        Returns:
+            the modified parser.
+        """
         return parser
 
     @abstractmethod
@@ -20,13 +47,21 @@ def __len__(self):
 
     @abstractmethod
     def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns:
+            a dictionary of data with their names. It ususally contains the data itself and its metadata information.
+        """
         pass
 
 
 def get_transform(opt, grayscale=False, convert=True, crop=True, flip=True):
     """Create a torchvision transformation function
 
-    The type of transformation is defined by option (e.g., [preprocess], [load_size], [crop_size])
+    The type of transformation is defined by option(e.g., [preprocess], [load_size], [crop_size])
     and can be overwritten by arguments such as [convert], [crop], and [flip]
     """
     transform_list = []
@@ -105,7 +140,7 @@ def __scale_width(img, target_width):
 
 
 def __print_size_warning(ow, oh, w, h):
-    """Print warning information about image size (only print once)"""
+    """Print warning information about image size(only print once)"""
     if not hasattr(__print_size_warning, 'has_printed'):
         print("The image size needs to be a multiple of 4. "
               "The loaded image size was (%d, %d), so it was adjusted to "

data/colorization_dataset.py

@@ -8,6 +8,10 @@
 
 
 class ColorizationDataset(BaseDataset):
+    """This dataset class can load a set of nature images in RGB, and convert RGB format into (L, ab) pairs in Lab color space.
+
+    This dataset is required by pix2pix-based colorization model ('--model colorization')
+    """
     @staticmethod
     def modify_commandline_options(parser, is_train):
         parser.set_defaults(input_nc=1, output_nc=2, direction='AtoB')

data/image_folder.py

@@ -1,7 +1,7 @@
-"""Modified Image folder class
-Code from https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py
-Modified the original code so that it also loads images from the current
-directory as well as the subdirectories
+"""A modified image folder class
+
+We modify the official PyTorch image folder (https://github.com/pytorch/vision/blob/master/torchvision/datasets/folder.py)
+so that this class can load images from both current directory and its subdirectories.
 """
 
 import torch.utils.data as data

data/single_dataset.py

@@ -4,21 +4,37 @@
 
 
 class SingleDataset(BaseDataset):
-    @staticmethod
-    def modify_commandline_options(parser, is_train):
-        return parser
+    """This dataset class can load a set of images specified by the path --dataroot /path/to/data.
+
+    It can be used for generating CycleGAN results only for one side with the model option '-model test'.
+    """
 
     def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt -- options (needs to be a subclass of BaseOptions)
+        """
         BaseDataset.__init__(self, opt)
         self.A_paths = sorted(make_dataset(opt.dataroot, opt.max_dataset_size))
         input_nc = self.opt.output_nc if self.opt.direction == 'BtoA' else self.opt.input_nc
         self.transform = get_transform(opt, input_nc == 1)
 
     def __getitem__(self, index):
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns a dictionary that contains A and A_paths
+            A(tensor) - - an image in one domain
+            A_paths(str) - - the path of the image
+        """
         A_path = self.A_paths[index]
         A_img = Image.open(A_path).convert('RGB')
         A = self.transform(A_img)
         return {'A': A, 'A_paths': A_path}
 
     def __len__(self):
+        """Return the total number of images in the dataset."""
         return len(self.A_paths)

data/template_dataset.py

@@ -23,8 +23,8 @@ def modify_commandline_options(parser, is_train):
         """Add new dataset-specific options, and rewrite default values for existing options.
 
         Parameters:
-            parser -- the option parser
-            is_train -- if it is training phase or test phase. You can use this flag to add training-specific or test-specific options.
+            parser -- original option parser
+            is_train -- whether training phase or test phase. You can use this flag to add training-specific or test-specific options.
 
         Returns:
             the modified parser.

data/unaligned_dataset.py

@@ -6,38 +6,65 @@
 
 
 class UnalignedDataset(BaseDataset):
-    @staticmethod
-    def modify_commandline_options(parser, is_train):
-        return parser
+    """
+    This dataset class can load unaligned/unpaired datasets.
+
+    It requires two directories to host training images from domain A '/path/to/data/trainA'
+    and from domain B '/path/to/data/trainB' respectively.
+    You can train the model with the dataset flag '--dataroot /path/to/data'.
+    Similarly, you need to prepare two directories:
+    '/path/to/data/testA' and '/path/to/data/testB' during test time.
+    """
 
     def __init__(self, opt):
+        """Initialize this dataset class.
+
+        Parameters:
+            opt -- options (needs to be a subclass of BaseOptions)
+        """
         BaseDataset.__init__(self, opt)
-        self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')
-        self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')
+        self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')  # create a path '/path/to/data/trainA'
+        self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')  # create a path '/path/to/data/trainB'
 
-        self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size))
-        self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size))
-        self.A_size = len(self.A_paths)
-        self.B_size = len(self.B_paths)
+        self.A_paths = sorted(make_dataset(self.dir_A, opt.max_dataset_size))   # load images from '/path/to/data/trainA'
+        self.B_paths = sorted(make_dataset(self.dir_B, opt.max_dataset_size))    # load images from '/path/to/data/trainB'
+        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'
-        input_nc = self.opt.output_nc if btoA else self.opt.input_nc
-        output_nc = self.opt.input_nc if btoA else self.opt.output_nc
-        self.transform_A = get_transform(opt, input_nc == 1)
-        self.transform_B = get_transform(opt, output_nc == 1)
+        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(opt, grayscale=(input_nc == 1))   # if nc == 1, we convert RGB to grayscale image
+        self.transform_B = get_transform(opt, grayscale=(output_nc == 1))  # if nc == 1, we convert RGB to grayscale image
 
     def __getitem__(self, index):
-        A_path = self.A_paths[index % self.A_size]
-        if self.opt.serial_batches:
+        """Return a data point and its metadata information.
+
+        Parameters:
+            index - - a random integer for data indexing
+
+        Returns a dictionary that contains A, B, A_paths and B_paths
+            A(tensor) - - an image in the input domain
+            B(tensor) - - its corresponding image in the target domain
+            A_paths(str) - - image paths
+            B_paths(str) - - image paths
+        """
+        A_path = self.A_paths[index % self.A_size]  # make sure index is within then range
+        if self.opt.serial_batches:   # make sure index is within then range
             index_B = index % self.B_size
-        else:
+        else:   # randomize the index for domain B to avoid fixed pairs.
             index_B = random.randint(0, self.B_size - 1)
         B_path = self.B_paths[index_B]
         A_img = Image.open(A_path).convert('RGB')
         B_img = Image.open(B_path).convert('RGB')
-
+        # apply image transformation
         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}
 
     def __len__(self):
+        """Return the total number of images in the dataset.
+
+        As we have two datasets with potentially different number of images,
+        we take a maximum of
+        """
         return max(self.A_size, self.B_size)