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',
'{}albedo.exr'.format(prefix))
B_path = os.path.join(self.root, 'Displacement', 'high',
'{}displacement_high.exr'.format(prefix))
geo_path = os.path.join(self.root, 'Pointcloud',
'{}pointcloud.exr'.format(prefix))
# load 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=4))
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):
geometry_tensor = 0
norm = normalize()
#### Before get_transform: convert to png
A_name = self.names_label[index]
### input A (label maps)
A_path = self.A_paths[index]
exr = cv2.imread(A_path, -1)
img = np.clip(exr, 0, 1)
img = img.astype(np.float32)
A_tensor = norm(toTensor(np.transpose(img, (2, 0, 1))))
if not self.opt.no_geometry:
geo_path = os.path.join(
self.dir_geo,
re.sub(r'diffuse_albedo(_\d\d)*', 'pointcloud', A_name))
geometry = cv2.imread(geo_path, cv2.IMREAD_UNCHANGED)
geo = rescale(np.copy(geometry), type=0)
geometry_tensor = norm(geo)
B_tensor = inst_tensor = feat_tensor = 0
### input B (real images)
if self.opt.isTrain or self.opt.use_encoded_image:
B_path = os.path.join(
self.dir_specular,
re.sub(r'diffuse_albedo(_\d\d)*', 'specular_unlit', A_name))
exr = cv2.imread(B_path, cv2.IMREAD_UNCHANGED)
specular = rescale(np.copy(exr), type=1)
B_tensor = norm(specular)
if not self.opt.isTrain and self.opt.loadGT:
B_path = os.path.join(
self.dir_specular,
re.sub(r'diffuse_albedo(_\d\d)*', 'specular_unlit', A_name))
exr = cv2.imread(B_path, cv2.IMREAD_UNCHANGED)
specular = rescale(np.copy(exr), type=1)
B_tensor = norm(specular)
input_dict = {
'label': A_tensor,
'inst': inst_tensor,
'image': B_tensor,
'feat': feat_tensor,
'path': A_path,
'geometry': geometry_tensor
}
return input_dict
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
