Spaces:
Sleeping
Sleeping
File size: 4,818 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 |
import os
from PIL import Image
import numpy as np
from utils.data_utils import check_create_shuffled_order, check_equal_length
def landmark_path_to_numpy(lmk_path, image_path, image_tensor):
"""Convert an landmark path to the actual landmarks in a numpy array. Also applies scaling to the landmarks
according to final image' size.
Parameters:
lmk_path -- the landmark file path.
image_path -- the original image file path.
image_tensor -- the image tensor after all transformations.
"""
lmk = np.load(lmk_path)
ow, oh = Image.open(image_path).size
h, w = image_tensor.size()[1:]
lmk[:, 0] *= w / ow
lmk[:, 1] *= h / oh
return lmk
def add_landmark_data(data, config, paired_data_order):
A_lmk_paths = []
B_lmk_paths = []
if config['dataset']['paired_' + config['common']['phase'] + '_filelist'] != '':
paired_data_file = open(config['dataset']['paired_' + config['common']['phase'] + '_filelist'], 'r')
Lines = paired_data_file.readlines()
paired_data_order = check_create_shuffled_order(Lines, paired_data_order)
check_equal_length(Lines, paired_data_order, data)
for i in paired_data_order:
line = Lines[i]
if not config['dataset']['use_absolute_datafile']:
file3 = os.path.join(config['dataset']['dataroot'], line.split(" ")[2]).strip()
file4 = os.path.join(config['dataset']['dataroot'], line.split(" ")[3]).strip()
else:
file3 = line.split(" ")[2].strip()
file4 = line.split(" ")[3].strip()
if os.path.exists(file3) and os.path.exists(file4):
A_lmk_paths.append(file3)
B_lmk_paths.append(file4)
paired_data_file.close()
elif config['dataset']['paired_' + config['common']['phase'] + 'A_folder'] != '' and \
config['dataset']['paired_' + config['common']['phase'] + 'B_folder'] != '' and \
os.path.exists(config['dataset']['paired_' + config['common']['phase'] + 'A_lmk_folder']) and \
os.path.exists(config['dataset']['paired_' + config['common']['phase'] + 'B_lmk_folder']):
dir_A = config['dataset']['paired_' + config['common']['phase'] + 'A_folder']
dir_A_lmk = config['dataset']['paired_' + config['common']['phase'] + 'A_lmk_folder']
dir_B_lmk = config['dataset']['paired_' + config['common']['phase'] + 'B_lmk_folder']
filenames = os.listdir(dir_A)
paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
check_equal_length(filenames, paired_data_order, data)
for i in paired_data_order:
filename = filenames[i]
A_lmk_path = os.path.join(dir_A_lmk, os.path.splitext(filename)[0] + '.npy')
B_lmk_path = os.path.join(dir_B_lmk, os.path.splitext(filename)[0] + '.npy')
if os.path.exists(A_lmk_path) and os.path.exists(B_lmk_path):
A_lmk_paths.append(A_lmk_path)
B_lmk_paths.append(B_lmk_path)
else:
dir_A = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedA')
dir_A_lmk = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedA_lmk')
dir_B_lmk = os.path.join(config['dataset']['dataroot'], config['common']['phase'] + 'pairedB_lmk')
if os.path.exists(dir_A_lmk) and os.path.exists(dir_B_lmk):
filenames = os.listdir(dir_A)
paired_data_order = check_create_shuffled_order(filenames, paired_data_order)
check_equal_length(filenames, paired_data_order, data)
for i in paired_data_order:
filename = filenames[i]
A_lmk_path = os.path.join(dir_A_lmk, os.path.splitext(filename)[0] + '.npy')
B_lmk_path = os.path.join(dir_B_lmk, os.path.splitext(filename)[0] + '.npy')
if os.path.exists(A_lmk_path) and os.path.exists(B_lmk_path):
A_lmk_paths.append(A_lmk_path)
B_lmk_paths.append(B_lmk_path)
else:
print(dir_A_lmk + " or " + dir_B_lmk + " doesn't exist. Skipping landmark data.")
data['A_lmk_path'] = A_lmk_paths
data['B_lmk_path'] = B_lmk_paths
return paired_data_order
def apply_landmark_transforms(index, data, return_dict):
if len(data['A_lmk_path']) > 0:
return_dict['A_lmk'] = landmark_path_to_numpy(data['A_lmk_path'][index], data['paired_A_path'][index], return_dict['paired_A'])
return_dict['B_lmk'] = landmark_path_to_numpy(data['B_lmk_path'][index], data['paired_B_path'][index], return_dict['paired_B'])
return_dict['A_lmk_path'] = data['A_lmk_path'][index]
return_dict['B_lmk_path'] = data['B_lmk_path'][index]
|