def __getitem__(self, index):
toTensor_3ch = get_transform_3ch(normalize=True)
toTensor_1ch = get_transform_1ch(normalize=True)
# transform func
toTensor_3ch = get_transform_3ch(normalize=True)
toTensor_1ch = get_transform_1ch(normalize=True)
# index the current file prefix
prefix = self.prefixes[index]
# paths
A_path = os.path.join(self.root, 'Albedo',
'{}albedo.exr'.format(prefix))
B_path = os.path.join(self.root, 'Specular',
'{}specular.exr'.format(prefix))
geo_path = os.path.join(self.root, 'Pointcloud',
'{}pointcloud.exr'.format(prefix))
# load 4K files
A = cv2.imread(A_path, cv2.IMREAD_UNCHANGED)
B = cv2.imread(B_path, cv2.IMREAD_UNCHANGED)
geo = cv2.imread(geo_path, cv2.IMREAD_UNCHANGED)
# input tensor
A_tensor = toTensor_3ch(rescale(A, type=2))
B_tensor = toTensor_1ch(rescale(B, type=1))
geo_tensor = toTensor_3ch(rescale(geo, type=0))
return {
'A': A_tensor,
'B': B_tensor,
'path': A_path,
'geometry': geo_tensor,
}
def __getitem__(self, index):
# transform func
toTensor_3ch = get_transform_3ch(normalize=True)
toTensor_1ch = get_transform_1ch(normalize=True)
# index the current file prefix
prefix = self.prefixes[index]
# paths
A_path = os.path.join(self.root, 'Albedo', '{}.png'.format(prefix))
# B_path = os.path.join(self.root, 'Specular', '{}specular.exr'.format(prefix))
geo_path = os.path.join(self.root, 'Pointcloud',
'{}.exr'.format(prefix))
# load 512 x 512 files
A = (cv2.imread(A_path) / 255.0).astype('float32')
# B = cv2.imread(B_path, cv2.IMREAD_UNCHANGED)
geo = cv2.imread(geo_path, cv2.IMREAD_UNCHANGED)
# input tensor
A_tensor = toTensor_3ch(rescale(A, type=2))
B_tensor = 0
geo_tensor = toTensor_3ch(rescale(geo, type=0))
res = {
'A': A_tensor,
'B': B_tensor,
'name': prefix,
'geometry': geo_tensor,
}
if self.task_type == 'specular':
res['path'] = os.path.join(self.root, 'Specular')
elif self.task_type == 'low':
res['path'] = os.path.join(self.root, 'Displacement', 'low')
elif self.task_type == 'high':
res['path'] = os.path.join(self.root, 'Displacement', 'high')
return res