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