Spaces:
Sleeping
Sleeping
| import os | |
| from utils.util import check_path_is_static_data | |
| from utils.data_utils import Transforms | |
| from utils.augmentation import ImagePathToImage, NumpyToTensor | |
| def add_test_data(data, transforms, config): | |
| A_paths = [] | |
| B_paths = [] | |
| if not config['testing']['test_img'] is None: | |
| A_paths.append(config['testing']['test_img']) | |
| B_paths.append(config['testing']['test_img']) | |
| else: | |
| files = os.listdir(config['testing']['test_folder']) | |
| for fn in files: | |
| if not check_path_is_static_data(fn): | |
| continue | |
| full_path = os.path.join(config['testing']['test_folder'], fn) | |
| A_paths.append(full_path) | |
| B_paths.append(full_path) | |
| btoA = config['dataset']['direction'] == 'BtoA' | |
| # get the number of channels of input image | |
| input_nc = config['model']['output_nc'] if btoA else config['model']['input_nc'] | |
| output_nc = config['model']['input_nc'] if btoA else config['model']['output_nc'] | |
| transform = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1)) | |
| transform.create_transforms_from_list(config['testing']['preprocess']) | |
| transform.get_transforms().insert(0, ImagePathToImage()) | |
| transform = transform.compose_transforms() | |
| transform_np = Transforms(config, input_grayscale_flag=(input_nc == 1), output_grayscale_flag=(output_nc == 1)) | |
| transform_np.transform_list.append(NumpyToTensor()) | |
| transform_np = transform_np.compose_transforms() | |
| data['test_A_path'] = A_paths | |
| data['test_B_path'] = B_paths | |
| transforms['test'] = transform | |
| transforms['test_np'] = transform_np | |
| def apply_test_transforms(index, data, transforms, return_dict): | |
| if len(data['test_A_path']) > 0: | |
| ext_name = os.path.splitext(data['test_A_path'][index])[1] | |
| if not ext_name.lower() in ['.npy', '.npz']: | |
| return_dict['test_A'], return_dict['test_B'] = transforms['test'] \ | |
| (data['test_A_path'][index], data['test_B_path'][index]) | |
| else: | |
| return_dict['test_A'], return_dict['test_B'] = transforms['test_np'] \ | |
| (data['test_A_path'][index], data['test_B_path'][index]) | |
| return_dict['test_A_path'] = data['test_A_path'][index] | |
| return_dict['test_B_path'] = data['test_B_path'][index] | |