Spaces:
Sleeping
Sleeping
| import os | |
| from utils.util import check_path_is_img | |
| from utils.data_utils import Transforms | |
| from utils.augmentation import ImagePathToImage | |
| import random | |
| def add_unpaired_data(data, transforms, config, shuffle=False): | |
| A_paths = [] | |
| B_paths = [] | |
| if config['dataset']['unpaired_' + config['common']['phase'] + 'A_filelist'] != '': | |
| unpaired_data_file1 = open(config['dataset']['unpaired_' + config['common']['phase'] + 'A_filelist'], 'r') | |
| Lines = unpaired_data_file1.readlines() | |
| if shuffle: | |
| random.shuffle(Lines) | |
| for line in Lines: | |
| if not config['dataset']['use_absolute_datafile']: | |
| file = os.path.join(config['dataset']['dataroot'], line.strip()) | |
| else: | |
| file = line.strip() | |
| if os.path.exists(file): | |
| A_paths.append(file) | |
| unpaired_data_file1.close() | |
| unpaired_data_file2 = open(config['dataset']['unpaired_' + config['common']['phase'] + 'B_filelist'], 'r') | |
| Lines = unpaired_data_file2.readlines() | |
| if shuffle: | |
| random.shuffle(Lines) | |
| for line in Lines: | |
| if not config['dataset']['use_absolute_datafile']: | |
| file = os.path.join(config['dataset']['dataroot'], line.strip()) | |
| else: | |
| file = line.strip() | |
| if os.path.exists(file): | |
| B_paths.append(file) | |
| unpaired_data_file2.close() | |
| elif config['dataset']['unpaired_' + config['common']['phase'] + 'A_folder'] != '' and \ | |
| config['dataset']['unpaired_' + config['common']['phase'] + 'B_folder'] != '': | |
| dir_A = config['dataset']['unpaired_' + config['common']['phase'] + 'A_folder'] | |
| filenames = os.listdir(dir_A) | |
| if shuffle: | |
| random.shuffle(filenames) | |
| for filename in filenames: | |
| if not check_path_is_img(filename): | |
| continue | |
| A_path = os.path.join(dir_A, filename) | |
| if os.path.exists(A_path): | |
| A_paths.append(A_path) | |
| dir_B = config['dataset']['unpaired_' + config['common']['phase'] + 'B_folder'] | |
| filenames = os.listdir(dir_B) | |
| if shuffle: | |
| random.shuffle(filenames) | |
| for filename in filenames: | |
| if not check_path_is_img(filename): | |
| continue | |
| B_path = os.path.join(dir_B, filename) | |
| if os.path.exists(B_path): | |
| B_paths.append(B_path) | |
| else: | |
| dir_A = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'unpairedA') | |
| dir_B = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'unpairedB') | |
| if os.path.exists(dir_A) and os.path.exists(dir_B): | |
| filenames = os.listdir(dir_A) | |
| if shuffle: | |
| random.shuffle(filenames) | |
| for filename in filenames: | |
| if not check_path_is_img(filename): | |
| continue | |
| A_path = os.path.join(dir_A, filename) | |
| A_paths.append(A_path) | |
| filenames = os.listdir(dir_B) | |
| if shuffle: | |
| random.shuffle(filenames) | |
| for filename in filenames: | |
| if not check_path_is_img(filename): | |
| continue | |
| B_path = os.path.join(dir_B, filename) | |
| B_paths.append(B_path) | |
| btoA = config['dataset']['direction'] == 'BtoA' | |
| 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.get_transform_from_config() | |
| transform.get_transforms().insert(0, ImagePathToImage()) | |
| transform = transform.compose_transforms() | |
| data['unpaired_A_path'] = A_paths | |
| data['unpaired_B_path'] = B_paths | |
| transforms['unpaired'] = transform | |
| def apply_unpaired_transforms(index, data, transforms, return_dict): | |
| if len(data['unpaired_A_path']) > 0 and len(data['unpaired_B_path']) > 0: | |
| index_B = random.randint(0, len(data['unpaired_B_path']) - 1) | |
| return_dict['unpaired_A'], return_dict['unpaired_B'] = transforms['unpaired'] \ | |
| (data['unpaired_A_path'][index], data['unpaired_B_path'][index_B]) | |
| return_dict['unpaired_A_path'] = data['unpaired_A_path'][index] | |
| return_dict['unpaired_B_path'] = data['unpaired_B_path'][index_B] | |