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)
