Spaces:
Running
on
T4
Running
on
T4
from menpo_functions import * | |
from deep_heatmaps_model_fusion_net import DeepHeatmapsModel | |
from scipy.misc import imsave | |
# *************** define parameters and paths *************** | |
data_dir = '~/AF_dataset2/' | |
test_data = 'Fernand_Leger' # subdirectory containing portraits for landmark detection (under data_dir) | |
use_gt_bb = False # use ground truth bounding box to crop images. if False, use face detector bounding box (relevant | |
# for challenging, common, full & training sets only) | |
out_dir = 'out_pred_landmarks' # directory for saving predicted landmarks | |
if not os.path.exists(out_dir): | |
os.mkdir(out_dir) | |
model_path = '~/model_foa/deep_heatmaps-60000' # model for estimation stage | |
pdm_path = 'pdm_clm_models/pdm_models/' # models for correction stage | |
clm_path = 'pdm_clm_models/clm_models/g_t_all' # model for tuning stage | |
outline_tune = False # if true use tuning stage on eyebrows+jaw, else use tuning stage on jaw only | |
# (see paper for details) | |
save_cropped_imgs = False # save input images in their cropped version to out_dir. | |
map_landmarks_to_original_image = True # if True, landmark predictions will be mapped to match original | |
# input image size. otherwise the predicted landmarks will match the cropped version (256x256) of the images | |
# *************** load images and model *************** | |
# load images | |
bb_dir = os.path.join(data_dir, 'Bounding_Boxes') | |
bb_dictionary = load_bb_dictionary(bb_dir, mode='TEST', test_data=test_data) | |
if use_gt_bb: | |
bb_type = 'gt' | |
else: | |
bb_type = 'init' | |
img_list = load_menpo_image_list( | |
img_dir=data_dir, test_data=test_data, train_crop_dir=data_dir, img_dir_ns=data_dir, bb_type=bb_type, | |
bb_dictionary=bb_dictionary, mode='TEST', return_transform=map_landmarks_to_original_image) | |
# load model | |
heatmap_model = DeepHeatmapsModel( | |
mode='TEST', img_path=data_dir, test_model_path=model_path, test_data=test_data, menpo_verbose=False) | |
# *************** predict landmarks *************** | |
print ("\npredicting landmarks for: "+os.path.join(data_dir, test_data)) | |
print ("\nsaving landmarks to: "+out_dir) | |
for i, img in enumerate(img_list): | |
if i == 0: | |
reuse = None | |
else: | |
reuse = True | |
preds = heatmap_model.get_landmark_predictions(img_list=[img], pdm_models_dir=pdm_path, clm_model_path=clm_path, | |
reuse=reuse, map_to_input_size=map_landmarks_to_original_image) | |
if map_landmarks_to_original_image: | |
img = img[0] | |
if outline_tune: | |
pred_lms = preds['ECpTp_out'] | |
else: | |
pred_lms = preds['ECpTp_jaw'] | |
mio.export_landmark_file(PointCloud(pred_lms[0]), os.path.join(out_dir, img.path.stem + '.pts'), | |
overwrite=True) | |
if save_cropped_imgs: | |
imsave(os.path.join(out_dir, img.path.stem + '.png'), img.pixels_with_channels_at_back()) | |
print ("\nDONE!") | |