Spaces:
Sleeping
Sleeping
File size: 4,522 Bytes
7e2a2a5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 |
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]
|