def __getitem__(self, index): """Return a data point and its metadata information. Parameters: index (int) -- 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 A1_path = self.A1_paths[index % self.A1_size] B_path = self.B_paths[index % self.B_size] #A_img = cv2.imread(A_path, -1) A1_img = exrlib.read_exr_float32( A1_path, list(self.input_names[self.input_channels]), 512, 512) B_img = exrlib.read_exr_float32( B_path, list(self.output_names[self.output_channels]), 512, 512) #A = torch.Tensor(A_img) A1 = self.convert_input(A1_img) B = self.convert_output(B_img) return {'A': A1, 'B': B, 'A_paths': A1_path, 'B_paths': B_path}
def get_val_item(self, index): A1_path = self.A1_test_paths[index % self.A1_test_size] B_path = self.B_test_paths[index % self.B_test_size] A1_img = exrlib.read_exr_float32(A1_path, list(self.input_names[self.input_channels]), 512, 512) B_img = exrlib.read_exr_float32(B_path, list(self.output_names[self.output_channels]), 512, 512) A1 = self.convert_input(A1_img) B = self.convert_output(B_img) return {'A': A1, 'B': B, 'A_orig': torch.Tensor(np.transpose(A1_img, (2, 0, 1))), 'B_orig': torch.Tensor(np.transpose(B_img, (2, 0, 1))), 'A_paths': A1_path, 'B_paths': B_path}
def get_val_item(self, index): dict = super().get_val_item(index) Flowmap_path = self.Flowmap_test_paths[index % self.Flowmap_size] Flowmap_img = exrlib.read_exr_float32(Flowmap_path, list(['0']), 512, 512) dict["Flowmap"] = (torch.Tensor(np.transpose(Flowmap_img, (2, 0, 1))) - 0.5) / 0.5 return dict
def get_val_item(self, index): dict = super().get_val_item(index) Extra_path = self.Extra_test_paths[index % self.Extra_size] Extra_img = exrlib.read_exr_float32(Extra_path, list(['0', '1']), 512, 512) Extra_img[:, :, 0] = (Extra_img[:, :, 0] - 0.5) / 0.5 Extra_img[:, :, 1] = (Extra_img[:, :, 1] - 412) / 996 * 2 dict["Extra"] = torch.Tensor(np.transpose(Extra_img, (2, 0, 1))) return dict
def __getitem__(self, index): """Return a data point and its metadata information. Parameters: index (int) -- 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 """ dict = super().__getitem__(index) Flowmap_path = self.Flowmap_paths[index % self.Flowmap_size] Flowmap_img = exrlib.read_exr_float32(Flowmap_path, list(['0']), 512, 512) dict["Flowmap"] = (torch.Tensor(np.transpose(Flowmap_img, (2, 0, 1))) - 0.5) / 0.5 return dict
def __getitem__(self, index): """Return a data point and its metadata information. Parameters: index (int) -- 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 """ dict = super().__getitem__(index) Extra_path = self.Extra_paths[index % self.Extra_size] Extra_img = exrlib.read_exr_float32(Extra_path, list(['0', '1']), 512, 512) Extra_img[:, :, 0] = (Extra_img[:, :, 0] - 0.5) / 0.5 Extra_img[:, :, 1] = (Extra_img[:, :, 1] - 412) / 996 * 2 dict["Extra"] = torch.Tensor(np.transpose(Extra_img, (2, 0, 1))) return dict
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) assert (opt.image_type == 'exr') #self.A = os.path.join(opt.dataroot, opt.phase + '_input') self.A1 = os.path.join(opt.dataroot, opt.phase + '_input_terraform') self.B = os.path.join(opt.dataroot, opt.phase + '_output') #self.A_paths = sorted(make_dataset(self.A, opt.max_dataset_size)) self.A1_paths = sorted(make_dataset(self.A1, opt.max_dataset_size)) self.B_paths = sorted(make_dataset(self.B, opt.max_dataset_size)) #self.A_size = len(self.A_paths) # get the size of dataset A self.A1_size = len(self.A1_paths) 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.A1_test_paths = sorted( make_dataset(os.path.join(opt.dataroot, 'test_input_terraform'))) self.B_test_paths = sorted( make_dataset(os.path.join(opt.dataroot, 'test_output'))) self.A1_test_size = len(self.A1_test_paths) self.B_test_size = len(self.B_test_paths) self.input_names = np.array([ "RockDetailMask.RockDetailMask", "SoftDetailMask.SoftDetailMask", "cliffs.cliffs", "height.height", "mesa.mesa", "slope.slope", "slopex.slopex", "slopez.slopez" ]) self.output_names = np.array([ "RockDetailMask.RockDetailMask", "SoftDetailMask.SoftDetailMask", "bedrock.bedrock", "cliffs.cliffs", "flow.flow", "flowx.flowx", "flowz.flowz", "height.height", "mesa.mesa", "sediment.sediment", "water.water" ]) self.input_channels = np.array([3, 6, 7]) #height, slopex, slopez self.output_channels = np.array([7]) #height if not self.opt.compute_bounds: self.i_channels_min = np.array([[[0, -400, -400]]]) self.i_channels_max = np.array([[[824, 20, 20]]]) self.o_channels_min = np.array([[[-4]]]) self.o_channels_max = np.array([[[819]]]) return channels_min = np.array([2**16 for _ in self.input_channels]) channels_max = np.array([0 for _ in self.input_channels]) examples = 0 for A1_path in self.A1_paths: A1_img = exrlib.read_exr_float32( A1_path, list(self.input_names[self.input_channels]), 512, 512).transpose(2, 0, 1).reshape(len(self.input_channels), -1) channels_min = np.min( np.concatenate((np.expand_dims( channels_min, 1), np.expand_dims(np.min(A1_img, 1), 1)), 1), 1) channels_max = np.max( np.concatenate((np.expand_dims( channels_min, 1), np.expand_dims(np.max(A1_img, 1), 1)), 1), 1) examples += 1 if examples >= 1000: break print(channels_min) self.i_channels_min = np.expand_dims( np.expand_dims(np.array(channels_min), 1), 2) print(channels_max) self.i_channels_max = np.expand_dims( np.expand_dims(np.array(channels_max), 1), 2) channels_min = np.array([2**16 for _ in self.output_channels]) channels_max = np.array([0 for _ in self.output_channels]) examples = 0 for B_path in self.B_paths: B_img = exrlib.read_exr_float32( B_path, list(self.output_names[self.output_channels]), 512, 512).transpose(2, 0, 1).reshape(len(self.output_channels), -1) channels_min = np.min( np.concatenate((np.expand_dims( channels_min, 1), np.expand_dims(np.min(B_img, 1), 1)), 1), 1) channels_max = np.max( np.concatenate((np.expand_dims( channels_min, 1), np.expand_dims(np.max(B_img, 1), 1)), 1), 1) examples += 1 if examples >= 1000: break print(channels_min) self.o_channels_min = np.expand_dims(np.expand_dims(channels_min, 1), 2) print(channels_max) self.o_channels_max = np.expand_dims(np.expand_dims(channels_max, 1), 2)