Ejemplo n.º 1
0
    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseDataset.__init__(self, opt)
        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))  # 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  # 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))
Ejemplo n.º 2
0
    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, 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 (same as A_paths)
        """
        # read a image given a random integer index
        AB_path = self.AB_paths[index]
        AB = Image.open(AB_path).convert('RGB')
        # split AB image into A and B
        w, h = AB.size
        w2 = int(w / (self.targets_no + 1))
        A = AB.crop((0, 0, w2, h))

        # apply the same transform to both A and B
        transform_params = get_params(self.preprocess, self.load_size,
                                      self.crop_size, A.size)
        A_transform = get_transform(self.preprocess,
                                    self.load_size,
                                    self.crop_size,
                                    self.no_flip,
                                    transform_params,
                                    grayscale=(self.input_nc == 1))
        B_transform = get_transform(self.preprocess,
                                    self.load_size,
                                    self.crop_size,
                                    self.no_flip,
                                    transform_params,
                                    grayscale=(self.output_nc == 1))

        A = A_transform(A)
        B_Array = []
        for i in range(1, self.targets_no + 1):
            B = AB.crop((w2 * i, 0, w2 * (i + 1), h))
            B = B_transform(B)
            B_Array.append(B)

        return {'A': A, 'B': B_Array, 'A_paths': AB_path, 'B_paths': AB_path}
Ejemplo n.º 3
0
    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; 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, grayscale=(input_nc == 1))
Ejemplo n.º 4
0
    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions
        """
        BaseDataset.__init__(self, opt)
        self.dir = os.path.join(opt.dataroot, opt.phase)
        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)
Ejemplo n.º 5
0
    def __init__(self, opt):
        """Initialize this dataset class.

        Parameters:
            opt (Option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions

        A few things can be done here.
        - save the options (have been done in BaseDataset)
        - get image paths and meta information of the dataset.
        - define the image transformation.
        """
        # save the option and dataset root
        BaseDataset.__init__(self, opt)
        # get the image paths of your dataset;
        self.image_paths = [
        ]  # You can call sorted(make_dataset(self.root, opt.max_dataset_size)) to get all the image paths under the directory self.root
        # define the default transform function. You can use <base_dataset.get_transform>; You can also define your custom transform function
        self.transform = get_transform(opt)