GoodWin/Deep-Multi-scale

1import os.path2from data.base_dataset import BaseDataset, get_transform3from data.image_folder import make_dataset4from PIL import Image5import random6 7 8class UnalignedDataset(BaseDataset):9    @staticmethod10    def modify_commandline_options(parser, is_train):11        return parser12 13    def initialize(self, opt):14        self.opt = opt15        self.root = opt.dataroot16        self.dir_A = os.path.join(opt.dataroot, opt.phase + 'A')17        self.dir_B = os.path.join(opt.dataroot, opt.phase + 'B')18 19        self.A_paths = make_dataset(self.dir_A)20        self.B_paths = make_dataset(self.dir_B)21 22        self.A_paths = sorted(self.A_paths)23        self.B_paths = sorted(self.B_paths)24        self.A_size = len(self.A_paths)25        self.B_size = len(self.B_paths)26        self.transform = get_transform(opt)27 28    def __getitem__(self, index):29        A_path = self.A_paths[index % self.A_size]30        if self.opt.serial_batches:31            index_B = index % self.B_size32        else:33            index_B = random.randint(0, self.B_size - 1)34        B_path = self.B_paths[index_B]35        # print('(A, B) = (%d, %d)' % (index_A, index_B))36        A_img = Image.open(A_path).convert('RGB')37        B_img = Image.open(B_path).convert('RGB')38 39        A = self.transform(A_img)40        B = self.transform(B_img)41        if self.opt.which_direction == 'BtoA':42            input_nc = self.opt.output_nc43            output_nc = self.opt.input_nc44        else:45            input_nc = self.opt.input_nc46            output_nc = self.opt.output_nc47 48        if input_nc == 1:  # RGB to gray49            tmp = A[0, ...] * 0.299 + A[1, ...] * 0.587 + A[2, ...] * 0.11450            A = tmp.unsqueeze(0)51 52        if output_nc == 1:  # RGB to gray53            tmp = B[0, ...] * 0.299 + B[1, ...] * 0.587 + B[2, ...] * 0.11454            B = tmp.unsqueeze(0)55        return {'A': A, 'B': B,56                'A_paths': A_path, 'B_paths': B_path}57 58    def __len__(self):59        return max(self.A_size, self.B_size)60 61    def name(self):62        return 'UnalignedDataset'63