Spaces:
Running
on
T4
Running
on
T4
| import scipy.io | |
| import scipy.misc | |
| from glob import glob | |
| import os | |
| import numpy as np | |
| from thirdparty.face_of_art.ops import * | |
| import tensorflow as tf | |
| from tensorflow import contrib | |
| from thirdparty.face_of_art.menpo_functions import * | |
| from thirdparty.face_of_art.logging_functions import * | |
| from thirdparty.face_of_art.data_loading_functions import * | |
| class DeepHeatmapsModel(object): | |
| """facial landmark localization Network""" | |
| def __init__(self, mode='TRAIN', train_iter=100000, batch_size=10, learning_rate=1e-3, l_weight_primary=1., | |
| l_weight_fusion=1.,l_weight_upsample=3.,adam_optimizer=True,momentum=0.95,step=100000, gamma=0.1,reg=0, | |
| weight_initializer='xavier', weight_initializer_std=0.01, bias_initializer=0.0, image_size=256,c_dim=3, | |
| num_landmarks=68, sigma=1.5, scale=1, margin=0.25, bb_type='gt', win_mult=3.33335, | |
| augment_basic=True,augment_texture=False, p_texture=0., augment_geom=False, p_geom=0., | |
| output_dir='output', save_model_path='model', | |
| save_sample_path='sample', save_log_path='logs', test_model_path='model/deep_heatmaps-50000', | |
| pre_train_path='model/deep_heatmaps-50000', load_pretrain=False, load_primary_only=False, | |
| img_path='data', test_data='full', valid_data='full', valid_size=0, log_valid_every=5, | |
| train_crop_dir='crop_gt_margin_0.25', img_dir_ns='crop_gt_margin_0.25_ns', | |
| print_every=100, save_every=5000, sample_every=5000, sample_grid=9, sample_to_log=True, | |
| debug_data_size=20, debug=False, epoch_data_dir='epoch_data', use_epoch_data=False, menpo_verbose=True): | |
| # define some extra parameters | |
| self.log_histograms = False # save weight + gradient histogram to log | |
| self.save_valid_images = True # sample heat maps of validation images | |
| self.sample_per_channel = False # sample heatmaps separately for each landmark | |
| # for fine-tuning, choose reset_training_op==True. when resuming training, reset_training_op==False | |
| self.reset_training_op = False | |
| self.fast_img_gen = True | |
| self.compute_nme = True # compute normalized mean error | |
| self.config = tf.ConfigProto() | |
| self.config.gpu_options.allow_growth = True | |
| # sampling and logging parameters | |
| self.print_every = print_every # print losses to screen + log | |
| self.save_every = save_every # save model | |
| self.sample_every = sample_every # save images of gen heat maps compared to GT | |
| self.sample_grid = sample_grid # number of training images in sample | |
| self.sample_to_log = sample_to_log # sample images to log instead of disk | |
| self.log_valid_every = log_valid_every # log validation loss (in epochs) | |
| self.debug = debug | |
| self.debug_data_size = debug_data_size | |
| self.use_epoch_data = use_epoch_data | |
| self.epoch_data_dir = epoch_data_dir | |
| self.load_pretrain = load_pretrain | |
| self.load_primary_only = load_primary_only | |
| self.pre_train_path = pre_train_path | |
| self.mode = mode | |
| self.train_iter = train_iter | |
| self.learning_rate = learning_rate | |
| self.image_size = image_size | |
| self.c_dim = c_dim | |
| self.batch_size = batch_size | |
| self.num_landmarks = num_landmarks | |
| self.save_log_path = save_log_path | |
| self.save_sample_path = save_sample_path | |
| self.save_model_path = save_model_path | |
| self.test_model_path = test_model_path | |
| self.img_path=img_path | |
| self.momentum = momentum | |
| self.step = step # for lr decay | |
| self.gamma = gamma # for lr decay | |
| self.reg = reg # weight decay scale | |
| self.l_weight_primary = l_weight_primary # primary loss weight | |
| self.l_weight_fusion = l_weight_fusion # fusion loss weight | |
| self.l_weight_upsample = l_weight_upsample # upsample loss weight | |
| self.weight_initializer = weight_initializer # random_normal or xavier | |
| self.weight_initializer_std = weight_initializer_std | |
| self.bias_initializer = bias_initializer | |
| self.adam_optimizer = adam_optimizer | |
| self.sigma = sigma # sigma for heatmap generation | |
| self.scale = scale # scale for image normalization 255 / 1 / 0 | |
| self.win_mult = win_mult # gaussian filter size for cpu/gpu approximation: 2 * sigma * win_mult + 1 | |
| self.test_data = test_data # if mode is TEST, this choose the set to use full/common/challenging/test/art | |
| self.train_crop_dir = train_crop_dir | |
| self.img_dir_ns = os.path.join(img_path,img_dir_ns) | |
| self.augment_basic = augment_basic # perform basic augmentation (rotation,flip,crop) | |
| self.augment_texture = augment_texture # perform artistic texture augmentation (NS) | |
| self.p_texture = p_texture # initial probability of artistic texture augmentation | |
| self.augment_geom = augment_geom # perform artistic geometric augmentation | |
| self.p_geom = p_geom # initial probability of artistic geometric augmentation | |
| self.valid_size = valid_size | |
| self.valid_data = valid_data | |
| # load image, bb and landmark data using menpo | |
| self.bb_dir = os.path.join(img_path, 'Bounding_Boxes') | |
| self.bb_dictionary = load_bb_dictionary(self.bb_dir, mode, test_data=self.test_data) | |
| # use pre-augmented data, to save time during training | |
| if self.use_epoch_data: | |
| epoch_0 = os.path.join(self.epoch_data_dir, '0') | |
| self.img_menpo_list = load_menpo_image_list( | |
| img_path, train_crop_dir=epoch_0, img_dir_ns=None, mode=mode, bb_dictionary=self.bb_dictionary, | |
| image_size=self.image_size, test_data=self.test_data, augment_basic=False, augment_texture=False, | |
| augment_geom=False, verbose=menpo_verbose) | |
| else: | |
| self.img_menpo_list = load_menpo_image_list( | |
| img_path, train_crop_dir, self.img_dir_ns, mode, bb_dictionary=self.bb_dictionary, | |
| image_size=self.image_size, margin=margin, bb_type=bb_type, test_data=self.test_data, | |
| augment_basic=augment_basic, augment_texture=augment_texture, p_texture=p_texture, | |
| augment_geom=augment_geom, p_geom=p_geom, verbose=menpo_verbose) | |
| if mode == 'TRAIN': | |
| train_params = locals() | |
| print_training_params_to_file(train_params) # save init parameters | |
| self.train_inds = np.arange(len(self.img_menpo_list)) | |
| if self.debug: | |
| self.train_inds = self.train_inds[:self.debug_data_size] | |
| self.img_menpo_list = self.img_menpo_list[self.train_inds] | |
| if valid_size > 0: | |
| self.valid_bb_dictionary = load_bb_dictionary(self.bb_dir, 'TEST', test_data=self.valid_data) | |
| self.valid_img_menpo_list = load_menpo_image_list( | |
| img_path, train_crop_dir, self.img_dir_ns, 'TEST', bb_dictionary=self.valid_bb_dictionary, | |
| image_size=self.image_size, margin=margin, bb_type=bb_type, test_data=self.valid_data, | |
| verbose=menpo_verbose) | |
| np.random.seed(0) | |
| self.val_inds = np.arange(len(self.valid_img_menpo_list)) | |
| np.random.shuffle(self.val_inds) | |
| self.val_inds = self.val_inds[:self.valid_size] | |
| self.valid_img_menpo_list = self.valid_img_menpo_list[self.val_inds] | |
| self.valid_images_loaded =\ | |
| np.zeros([self.valid_size, self.image_size, self.image_size, self.c_dim]).astype('float32') | |
| self.valid_gt_maps_small_loaded =\ | |
| np.zeros([self.valid_size, self.image_size / 4, self.image_size / 4, | |
| self.num_landmarks]).astype('float32') | |
| self.valid_gt_maps_loaded =\ | |
| np.zeros([self.valid_size, self.image_size, self.image_size, self.num_landmarks] | |
| ).astype('float32') | |
| self.valid_landmarks_loaded = np.zeros([self.valid_size, num_landmarks, 2]).astype('float32') | |
| self.valid_landmarks_pred = np.zeros([self.valid_size, self.num_landmarks, 2]).astype('float32') | |
| load_images_landmarks_approx_maps_alloc_once( | |
| self.valid_img_menpo_list, np.arange(self.valid_size), images=self.valid_images_loaded, | |
| maps_small=self.valid_gt_maps_small_loaded, maps=self.valid_gt_maps_loaded, | |
| landmarks=self.valid_landmarks_loaded, image_size=self.image_size, | |
| num_landmarks=self.num_landmarks, scale=self.scale, win_mult=self.win_mult, sigma=self.sigma, | |
| save_landmarks=self.compute_nme) | |
| if self.valid_size > self.sample_grid: | |
| self.valid_gt_maps_loaded = self.valid_gt_maps_loaded[:self.sample_grid] | |
| self.valid_gt_maps_small_loaded = self.valid_gt_maps_small_loaded[:self.sample_grid] | |
| else: | |
| self.val_inds = None | |
| self.epoch_inds_shuffle = train_val_shuffle_inds_per_epoch( | |
| self.val_inds, self.train_inds, train_iter, batch_size, save_log_path) | |
| def add_placeholders(self): | |
| if self.mode == 'TEST': | |
| self.images = tf.placeholder( | |
| tf.float32, [None, self.image_size, self.image_size, self.c_dim], 'images') | |
| self.heatmaps = tf.placeholder( | |
| tf.float32, [None, self.image_size, self.image_size, self.num_landmarks], 'heatmaps') | |
| self.heatmaps_small = tf.placeholder( | |
| tf.float32, [None, int(self.image_size/4), int(self.image_size/4), self.num_landmarks], 'heatmaps_small') | |
| self.lms = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'lms') | |
| self.pred_lms = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'pred_lms') | |
| elif self.mode == 'TRAIN': | |
| self.images = tf.placeholder( | |
| tf.float32, [None, self.image_size, self.image_size, self.c_dim], 'train_images') | |
| self.heatmaps = tf.placeholder( | |
| tf.float32, [None, self.image_size, self.image_size, self.num_landmarks], 'train_heatmaps') | |
| self.heatmaps_small = tf.placeholder( | |
| tf.float32, [None, int(self.image_size/4), int(self.image_size/4), self.num_landmarks], 'train_heatmaps_small') | |
| self.train_lms = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'train_lms') | |
| self.train_pred_lms = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'train_pred_lms') | |
| self.valid_lms = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'valid_lms') | |
| self.valid_pred_lms = tf.placeholder(tf.float32, [None, self.num_landmarks, 2], 'valid_pred_lms') | |
| # self.p_texture_log = tf.placeholder(tf.float32, []) | |
| # self.p_geom_log = tf.placeholder(tf.float32, []) | |
| # self.sparse_hm_small = tf.placeholder(tf.float32, [None, int(self.image_size/4), int(self.image_size/4), 1]) | |
| # self.sparse_hm = tf.placeholder(tf.float32, [None, self.image_size, self.image_size, 1]) | |
| if self.sample_to_log: | |
| row = int(np.sqrt(self.sample_grid)) | |
| self.log_image_map_small = tf.placeholder( | |
| tf.uint8, [None, row * int(self.image_size/4), 3 * row * int(self.image_size/4), self.c_dim], | |
| 'sample_img_map_small') | |
| self.log_image_map = tf.placeholder( | |
| tf.uint8, [None, row * self.image_size, 3 * row * self.image_size, self.c_dim], | |
| 'sample_img_map') | |
| if self.sample_per_channel: | |
| row = np.ceil(np.sqrt(self.num_landmarks)).astype(np.int64) | |
| self.log_map_channels_small = tf.placeholder( | |
| tf.uint8, [None, row * int(self.image_size/4), 2 * row * int(self.image_size/4), self.c_dim], | |
| 'sample_map_channels_small') | |
| self.log_map_channels = tf.placeholder( | |
| tf.uint8, [None, row * self.image_size, 2 * row * self.image_size, self.c_dim], | |
| 'sample_map_channels') | |
| def heatmaps_network(self, input_images, reuse=None, name='pred_heatmaps'): | |
| with tf.name_scope(name): | |
| if self.weight_initializer == 'xavier': | |
| weight_initializer = contrib.layers.xavier_initializer() | |
| else: | |
| weight_initializer = tf.random_normal_initializer(stddev=self.weight_initializer_std) | |
| bias_init = tf.constant_initializer(self.bias_initializer) | |
| with tf.variable_scope('heatmaps_network'): | |
| with tf.name_scope('primary_net'): | |
| l1 = conv_relu_pool(input_images, 5, 128, conv_ker_init=weight_initializer, conv_bias_init=bias_init, | |
| reuse=reuse, var_scope='conv_1') | |
| l2 = conv_relu_pool(l1, 5, 128, conv_ker_init=weight_initializer, conv_bias_init=bias_init, | |
| reuse=reuse, var_scope='conv_2') | |
| l3 = conv_relu(l2, 5, 128, conv_ker_init=weight_initializer, conv_bias_init=bias_init, | |
| reuse=reuse, var_scope='conv_3') | |
| l4_1 = conv_relu(l3, 3, 128, conv_dilation=1, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_1') | |
| l4_2 = conv_relu(l3, 3, 128, conv_dilation=2, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_2') | |
| l4_3 = conv_relu(l3, 3, 128, conv_dilation=3, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_3') | |
| l4_4 = conv_relu(l3, 3, 128, conv_dilation=4, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_4_4') | |
| l4 = tf.concat([l4_1, l4_2, l4_3, l4_4], 3, name='conv_4') | |
| l5_1 = conv_relu(l4, 3, 256, conv_dilation=1, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_1') | |
| l5_2 = conv_relu(l4, 3, 256, conv_dilation=2, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_2') | |
| l5_3 = conv_relu(l4, 3, 256, conv_dilation=3, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_3') | |
| l5_4 = conv_relu(l4, 3, 256, conv_dilation=4, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_5_4') | |
| l5 = tf.concat([l5_1, l5_2, l5_3, l5_4], 3, name='conv_5') | |
| l6 = conv_relu(l5, 1, 512, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_6') | |
| l7 = conv_relu(l6, 1, 256, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_7') | |
| primary_out = conv(l7, 1, self.num_landmarks, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_8') | |
| with tf.name_scope('fusion_net'): | |
| l_fsn_0 = tf.concat([l3, l7], 3, name='conv_3_7_fsn') | |
| l_fsn_1_1 = conv_relu(l_fsn_0, 3, 64, conv_dilation=1, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_1_1') | |
| l_fsn_1_2 = conv_relu(l_fsn_0, 3, 64, conv_dilation=2, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_1_2') | |
| l_fsn_1_3 = conv_relu(l_fsn_0, 3, 64, conv_dilation=3, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_1_3') | |
| l_fsn_1 = tf.concat([l_fsn_1_1, l_fsn_1_2, l_fsn_1_3], 3, name='conv_fsn_1') | |
| l_fsn_2_1 = conv_relu(l_fsn_1, 3, 64, conv_dilation=1, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_2_1') | |
| l_fsn_2_2 = conv_relu(l_fsn_1, 3, 64, conv_dilation=2, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_2_2') | |
| l_fsn_2_3 = conv_relu(l_fsn_1, 3, 64, conv_dilation=4, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_2_3') | |
| l_fsn_2_4 = conv_relu(l_fsn_1, 5, 64, conv_dilation=3, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_2_4') | |
| l_fsn_2 = tf.concat([l_fsn_2_1, l_fsn_2_2, l_fsn_2_3, l_fsn_2_4], 3, name='conv_fsn_2') | |
| l_fsn_3_1 = conv_relu(l_fsn_2, 3, 128, conv_dilation=1, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_3_1') | |
| l_fsn_3_2 = conv_relu(l_fsn_2, 3, 128, conv_dilation=2, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_3_2') | |
| l_fsn_3_3 = conv_relu(l_fsn_2, 3, 128, conv_dilation=4, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_3_3') | |
| l_fsn_3_4 = conv_relu(l_fsn_2, 5, 128, conv_dilation=3, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_3_4') | |
| l_fsn_3 = tf.concat([l_fsn_3_1, l_fsn_3_2, l_fsn_3_3, l_fsn_3_4], 3, name='conv_fsn_3') | |
| l_fsn_4 = conv_relu(l_fsn_3, 1, 256, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_4') | |
| fusion_out = conv(l_fsn_4, 1, self.num_landmarks, conv_ker_init=weight_initializer, | |
| conv_bias_init=bias_init, reuse=reuse, var_scope='conv_fsn_5') | |
| with tf.name_scope('upsample_net'): | |
| out = deconv(fusion_out, 8, self.num_landmarks, conv_stride=4, | |
| conv_ker_init=deconv2d_bilinear_upsampling_initializer( | |
| [8, 8, self.num_landmarks, self.num_landmarks]), conv_bias_init=bias_init, | |
| reuse=reuse, var_scope='deconv_1') | |
| self.all_layers = [l1, l2, l3, l4, l5, l6, l7, primary_out, l_fsn_1, l_fsn_2, l_fsn_3, l_fsn_4, | |
| fusion_out, out] | |
| return primary_out, fusion_out, out | |
| def build_model(self): | |
| self.pred_hm_p, self.pred_hm_f, self.pred_hm_u = self.heatmaps_network(self.images,name='heatmaps_prediction') | |
| def create_loss_ops(self): | |
| def nme_norm_eyes(pred_landmarks, real_landmarks, normalize=True, name='NME'): | |
| """calculate normalized mean error on landmarks - normalize with inter pupil distance""" | |
| with tf.name_scope(name): | |
| with tf.name_scope('real_pred_landmarks_rmse'): | |
| # calculate RMS ERROR between GT and predicted lms | |
| landmarks_rms_err = tf.reduce_mean( | |
| tf.sqrt(tf.reduce_sum(tf.square(pred_landmarks - real_landmarks), axis=2)), axis=1) | |
| if normalize: | |
| # normalize RMS ERROR with inter-pupil distance of GT lms | |
| with tf.name_scope('inter_pupil_dist'): | |
| with tf.name_scope('left_eye_center'): | |
| p1 = tf.reduce_mean(tf.slice(real_landmarks, [0, 42, 0], [-1, 6, 2]), axis=1) | |
| with tf.name_scope('right_eye_center'): | |
| p2 = tf.reduce_mean(tf.slice(real_landmarks, [0, 36, 0], [-1, 6, 2]), axis=1) | |
| eye_dist = tf.sqrt(tf.reduce_sum(tf.square(p1 - p2), axis=1)) | |
| return landmarks_rms_err / eye_dist | |
| else: | |
| return landmarks_rms_err | |
| if self.mode is 'TRAIN': | |
| # calculate L2 loss between ideal and predicted heatmaps | |
| primary_maps_diff = self.pred_hm_p - self.heatmaps_small | |
| fusion_maps_diff = self.pred_hm_f - self.heatmaps_small | |
| upsample_maps_diff = self.pred_hm_u - self.heatmaps | |
| self.l2_primary = tf.reduce_mean(tf.square(primary_maps_diff)) | |
| self.l2_fusion = tf.reduce_mean(tf.square(fusion_maps_diff)) | |
| self.l2_upsample = tf.reduce_mean(tf.square(upsample_maps_diff)) | |
| self.total_loss = 1000.*(self.l_weight_primary * self.l2_primary + self.l_weight_fusion * self.l2_fusion + | |
| self.l_weight_upsample * self.l2_upsample) | |
| # add weight decay | |
| self.total_loss += self.reg * tf.add_n( | |
| [tf.nn.l2_loss(v) for v in tf.trainable_variables() if 'bias' not in v.name]) | |
| # compute normalized mean error on gt vs. predicted landmarks (for validation) | |
| if self.compute_nme: | |
| self.nme_loss = tf.reduce_mean(nme_norm_eyes(self.train_pred_lms, self.train_lms)) | |
| if self.valid_size > 0 and self.compute_nme: | |
| self.valid_nme_loss = tf.reduce_mean(nme_norm_eyes(self.valid_pred_lms, self.valid_lms)) | |
| elif self.mode == 'TEST' and self.compute_nme: | |
| self.nme_per_image = nme_norm_eyes(self.pred_lms, self.lms) | |
| self.nme_loss = tf.reduce_mean(self.nme_per_image) | |
| def predict_valid_landmarks_in_batches(self, images, session): | |
| num_images=int(images.shape[0]) | |
| num_batches = int(1.*num_images/self.batch_size) | |
| if num_batches == 0: | |
| batch_size = num_images | |
| num_batches = 1 | |
| else: | |
| batch_size = self.batch_size | |
| for j in range(num_batches): | |
| batch_images = images[j * batch_size:(j + 1) * batch_size,:,:,:] | |
| batch_maps_pred = session.run(self.pred_hm_u, {self.images: batch_images}) | |
| batch_heat_maps_to_landmarks_alloc_once( | |
| batch_maps=batch_maps_pred, batch_landmarks=self.valid_landmarks_pred[j * batch_size:(j + 1) * batch_size, :, :], | |
| batch_size=batch_size,image_size=self.image_size,num_landmarks=self.num_landmarks) | |
| reminder = num_images-num_batches*batch_size | |
| if reminder > 0: | |
| batch_images = images[-reminder:, :, :, :] | |
| batch_maps_pred = session.run(self.pred_hm_u, {self.images: batch_images}) | |
| batch_heat_maps_to_landmarks_alloc_once( | |
| batch_maps=batch_maps_pred, | |
| batch_landmarks=self.valid_landmarks_pred[-reminder:, :, :], | |
| batch_size=reminder, image_size=self.image_size, num_landmarks=self.num_landmarks) | |
| def create_summary_ops(self): | |
| """create summary ops for logging""" | |
| # loss summary | |
| l2_primary = tf.summary.scalar('l2_primary', self.l2_primary) | |
| l2_fusion = tf.summary.scalar('l2_fusion', self.l2_fusion) | |
| l2_upsample = tf.summary.scalar('l2_upsample', self.l2_upsample) | |
| l_total = tf.summary.scalar('l_total', self.total_loss) | |
| self.batch_summary_op = tf.summary.merge([l2_primary,l2_fusion,l2_upsample,l_total]) | |
| if self.compute_nme: | |
| nme = tf.summary.scalar('nme', self.nme_loss) | |
| self.batch_summary_op = tf.summary.merge([self.batch_summary_op, nme]) | |
| if self.log_histograms: | |
| var_summary = [tf.summary.histogram(var.name,var) for var in tf.trainable_variables()] | |
| grads = tf.gradients(self.total_loss, tf.trainable_variables()) | |
| grads = list(zip(grads, tf.trainable_variables())) | |
| grad_summary = [tf.summary.histogram(var.name+'/grads',grad) for grad,var in grads] | |
| activ_summary = [tf.summary.histogram(layer.name, layer) for layer in self.all_layers] | |
| self.batch_summary_op = tf.summary.merge([self.batch_summary_op, var_summary, grad_summary, activ_summary]) | |
| if self.valid_size > 0 and self.compute_nme: | |
| self.valid_summary = tf.summary.scalar('valid_nme', self.valid_nme_loss) | |
| if self.sample_to_log: | |
| img_map_summary_small = tf.summary.image('compare_map_to_gt_small', self.log_image_map_small) | |
| img_map_summary = tf.summary.image('compare_map_to_gt', self.log_image_map) | |
| if self.sample_per_channel: | |
| map_channels_summary = tf.summary.image('compare_map_channels_to_gt', self.log_map_channels) | |
| map_channels_summary_small = tf.summary.image('compare_map_channels_to_gt_small', | |
| self.log_map_channels_small) | |
| self.img_summary = tf.summary.merge( | |
| [img_map_summary, img_map_summary_small,map_channels_summary,map_channels_summary_small]) | |
| else: | |
| self.img_summary = tf.summary.merge([img_map_summary, img_map_summary_small]) | |
| if self.valid_size >= self.sample_grid: | |
| img_map_summary_valid_small = tf.summary.image('compare_map_to_gt_small_valid', self.log_image_map_small) | |
| img_map_summary_valid = tf.summary.image('compare_map_to_gt_valid', self.log_image_map) | |
| if self.sample_per_channel: | |
| map_channels_summary_valid_small = tf.summary.image('compare_map_channels_to_gt_small_valid', | |
| self.log_map_channels_small) | |
| map_channels_summary_valid = tf.summary.image('compare_map_channels_to_gt_valid', | |
| self.log_map_channels) | |
| self.img_summary_valid = tf.summary.merge( | |
| [img_map_summary_valid,img_map_summary_valid_small,map_channels_summary_valid, | |
| map_channels_summary_valid_small]) | |
| else: | |
| self.img_summary_valid = tf.summary.merge([img_map_summary_valid, img_map_summary_valid_small]) | |
| def train(self): | |
| # set random seed | |
| tf.set_random_seed(1234) | |
| np.random.seed(1234) | |
| # build a graph | |
| # add placeholders | |
| self.add_placeholders() | |
| # build model | |
| self.build_model() | |
| # create loss ops | |
| self.create_loss_ops() | |
| # create summary ops | |
| self.create_summary_ops() | |
| # create optimizer and training op | |
| global_step = tf.Variable(0, trainable=False) | |
| lr = tf.train.exponential_decay(self.learning_rate,global_step, self.step, self.gamma, staircase=True) | |
| if self.adam_optimizer: | |
| optimizer = tf.train.AdamOptimizer(lr) | |
| else: | |
| optimizer = tf.train.MomentumOptimizer(lr, self.momentum) | |
| train_op = optimizer.minimize(self.total_loss,global_step=global_step) | |
| with tf.Session(config=self.config) as sess: | |
| tf.global_variables_initializer().run() | |
| # load pre trained weights if load_pretrain==True | |
| if self.load_pretrain: | |
| print('*** loading pre-trained weights from: '+self.pre_train_path+' ***') | |
| if self.load_primary_only: | |
| print('*** loading primary-net only ***') | |
| primary_var = [v for v in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES) if | |
| ('deconv_' not in v.name) and ('_fsn_' not in v.name)] | |
| loader = tf.train.Saver(var_list=primary_var) | |
| else: | |
| loader = tf.train.Saver() | |
| loader.restore(sess, self.pre_train_path) | |
| print("*** Model restore finished, current global step: %d" % global_step.eval()) | |
| # for fine-tuning, choose reset_training_op==True. when resuming training, reset_training_op==False | |
| if self.reset_training_op: | |
| print ("resetting optimizer and global step") | |
| opt_var_list = [optimizer.get_slot(var, name) for name in optimizer.get_slot_names() | |
| for var in tf.global_variables() if optimizer.get_slot(var, name) is not None] | |
| opt_var_list_init = tf.variables_initializer(opt_var_list) | |
| opt_var_list_init.run() | |
| sess.run(global_step.initializer) | |
| # create model saver and file writer | |
| summary_writer = tf.summary.FileWriter(logdir=self.save_log_path, graph=tf.get_default_graph()) | |
| saver = tf.train.Saver() | |
| print('\n*** Start Training ***') | |
| # initialize some variables before training loop | |
| resume_step = global_step.eval() | |
| num_train_images = len(self.img_menpo_list) | |
| batches_in_epoch = int(float(num_train_images) / float(self.batch_size)) | |
| epoch = int(resume_step / batches_in_epoch) | |
| img_inds = self.epoch_inds_shuffle[epoch, :] | |
| log_valid = True | |
| log_valid_images = True | |
| # allocate space for batch images, maps and landmarks | |
| batch_images = np.zeros([self.batch_size, self.image_size, self.image_size, self.c_dim]).astype( | |
| 'float32') | |
| batch_lms = np.zeros([self.batch_size, self.num_landmarks, 2]).astype('float32') | |
| batch_lms_pred = np.zeros([self.batch_size, self.num_landmarks, 2]).astype('float32') | |
| batch_maps_small = np.zeros((self.batch_size, int(self.image_size/4), | |
| int(self.image_size/4), self.num_landmarks)).astype('float32') | |
| batch_maps = np.zeros((self.batch_size, self.image_size, self.image_size, | |
| self.num_landmarks)).astype('float32') | |
| # create gaussians for heatmap generation | |
| gaussian_filt_large = create_gaussian_filter(sigma=self.sigma, win_mult=self.win_mult) | |
| gaussian_filt_small = create_gaussian_filter(sigma=1.*self.sigma/4, win_mult=self.win_mult) | |
| # training loop | |
| for step in range(resume_step, self.train_iter): | |
| j = step % batches_in_epoch # j==0 if we finished an epoch | |
| # if we finished an epoch and this isn't the first step | |
| if step > resume_step and j == 0: | |
| epoch += 1 | |
| img_inds = self.epoch_inds_shuffle[epoch, :] # get next shuffled image inds | |
| log_valid = True | |
| log_valid_images = True | |
| if self.use_epoch_data: # if using pre-augmented data, load epoch directory | |
| epoch_dir = os.path.join(self.epoch_data_dir, str(epoch)) | |
| self.img_menpo_list = load_menpo_image_list( | |
| self.img_path, train_crop_dir=epoch_dir, img_dir_ns=None, mode=self.mode, | |
| bb_dictionary=self.bb_dictionary, image_size=self.image_size, test_data=self.test_data, | |
| augment_basic=False, augment_texture=False, augment_geom=False) | |
| # get batch indices | |
| batch_inds = img_inds[j * self.batch_size:(j + 1) * self.batch_size] | |
| # load batch images, gt maps and landmarks | |
| load_images_landmarks_approx_maps_alloc_once( | |
| self.img_menpo_list, batch_inds, images=batch_images, maps_small=batch_maps_small, | |
| maps=batch_maps, landmarks=batch_lms, image_size=self.image_size, | |
| num_landmarks=self.num_landmarks, scale=self.scale, gauss_filt_large=gaussian_filt_large, | |
| gauss_filt_small=gaussian_filt_small, win_mult=self.win_mult, sigma=self.sigma, | |
| save_landmarks=self.compute_nme) | |
| feed_dict_train = {self.images: batch_images, self.heatmaps: batch_maps, | |
| self.heatmaps_small: batch_maps_small} | |
| # train on batch | |
| sess.run(train_op, feed_dict_train) | |
| # save to log and print status | |
| if step == resume_step or (step + 1) % self.print_every == 0: | |
| # train data log | |
| if self.compute_nme: | |
| batch_maps_pred = sess.run(self.pred_hm_u, {self.images: batch_images}) | |
| batch_heat_maps_to_landmarks_alloc_once( | |
| batch_maps=batch_maps_pred,batch_landmarks=batch_lms_pred, | |
| batch_size=self.batch_size, image_size=self.image_size, | |
| num_landmarks=self.num_landmarks) | |
| train_feed_dict_log = { | |
| self.images: batch_images, self.heatmaps: batch_maps, | |
| self.heatmaps_small: batch_maps_small, self.train_lms: batch_lms, | |
| self.train_pred_lms: batch_lms_pred} | |
| summary, l_p, l_f, l_t, nme = sess.run( | |
| [self.batch_summary_op, self.l2_primary, self.l2_fusion, self.total_loss, | |
| self.nme_loss], | |
| train_feed_dict_log) | |
| print ( | |
| 'epoch: [%d] step: [%d/%d] primary loss: [%.6f] fusion loss: [%.6f]' | |
| ' total loss: [%.6f] NME: [%.6f]' % ( | |
| epoch, step + 1, self.train_iter, l_p, l_f, l_t, nme)) | |
| else: | |
| train_feed_dict_log = {self.images: batch_images, self.heatmaps: batch_maps, | |
| self.heatmaps_small: batch_maps_small} | |
| summary, l_p, l_f, l_t = sess.run( | |
| [self.batch_summary_op, self.l2_primary, self.l2_fusion, self.total_loss], | |
| train_feed_dict_log) | |
| print ( | |
| 'epoch: [%d] step: [%d/%d] primary loss: [%.6f] fusion loss: [%.6f] total loss: [%.6f]' | |
| % (epoch, step + 1, self.train_iter, l_p, l_f, l_t)) | |
| summary_writer.add_summary(summary, step) | |
| # valid data log | |
| if self.valid_size > 0 and (log_valid and epoch % self.log_valid_every == 0) \ | |
| and self.compute_nme: | |
| log_valid = False | |
| self.predict_valid_landmarks_in_batches(self.valid_images_loaded, sess) | |
| valid_feed_dict_log = { | |
| self.valid_lms: self.valid_landmarks_loaded, | |
| self.valid_pred_lms: self.valid_landmarks_pred} | |
| v_summary, v_nme = sess.run([self.valid_summary, self.valid_nme_loss], | |
| valid_feed_dict_log) | |
| summary_writer.add_summary(v_summary, step) | |
| print ( | |
| 'epoch: [%d] step: [%d/%d] valid NME: [%.6f]' % ( | |
| epoch, step + 1, self.train_iter, v_nme)) | |
| # save model | |
| if (step + 1) % self.save_every == 0: | |
| saver.save(sess, os.path.join(self.save_model_path, 'deep_heatmaps'), global_step=step + 1) | |
| print ('model/deep-heatmaps-%d saved' % (step + 1)) | |
| # save images | |
| if step == resume_step or (step + 1) % self.sample_every == 0: | |
| batch_maps_small_pred = sess.run(self.pred_hm_p, {self.images: batch_images}) | |
| if not self.compute_nme: | |
| batch_maps_pred = sess.run(self.pred_hm_u, {self.images: batch_images}) | |
| batch_lms_pred = None | |
| merged_img = merge_images_landmarks_maps_gt( | |
| batch_images.copy(), batch_maps_pred, batch_maps, landmarks=batch_lms_pred, | |
| image_size=self.image_size, num_landmarks=self.num_landmarks, num_samples=self.sample_grid, | |
| scale=self.scale, circle_size=2, fast=self.fast_img_gen) | |
| merged_img_small = merge_images_landmarks_maps_gt( | |
| batch_images.copy(), batch_maps_small_pred, batch_maps_small, | |
| image_size=self.image_size, | |
| num_landmarks=self.num_landmarks, num_samples=self.sample_grid, scale=self.scale, | |
| circle_size=0, fast=self.fast_img_gen) | |
| if self.sample_per_channel: | |
| map_per_channel = map_comapre_channels( | |
| batch_images.copy(), batch_maps_pred, batch_maps, image_size=self.image_size, | |
| num_landmarks=self.num_landmarks, scale=self.scale) | |
| map_per_channel_small = map_comapre_channels( | |
| batch_images.copy(), batch_maps_small_pred, batch_maps_small, image_size=int(self.image_size/4), | |
| num_landmarks=self.num_landmarks, scale=self.scale) | |
| if self.sample_to_log: # save heatmap images to log | |
| if self.sample_per_channel: | |
| summary_img = sess.run( | |
| self.img_summary, {self.log_image_map: np.expand_dims(merged_img, 0), | |
| self.log_map_channels: np.expand_dims(map_per_channel, 0), | |
| self.log_image_map_small: np.expand_dims(merged_img_small, 0), | |
| self.log_map_channels_small: np.expand_dims(map_per_channel_small, 0)}) | |
| else: | |
| summary_img = sess.run( | |
| self.img_summary, {self.log_image_map: np.expand_dims(merged_img, 0), | |
| self.log_image_map_small: np.expand_dims(merged_img_small, 0)}) | |
| summary_writer.add_summary(summary_img, step) | |
| if (self.valid_size >= self.sample_grid) and self.save_valid_images and\ | |
| (log_valid_images and epoch % self.log_valid_every == 0): | |
| log_valid_images = False | |
| batch_maps_small_pred_val,batch_maps_pred_val =\ | |
| sess.run([self.pred_hm_p,self.pred_hm_u], | |
| {self.images: self.valid_images_loaded[:self.sample_grid]}) | |
| merged_img_small = merge_images_landmarks_maps_gt( | |
| self.valid_images_loaded[:self.sample_grid].copy(), batch_maps_small_pred_val, | |
| self.valid_gt_maps_small_loaded, image_size=self.image_size, | |
| num_landmarks=self.num_landmarks, num_samples=self.sample_grid, | |
| scale=self.scale, circle_size=0, fast=self.fast_img_gen) | |
| merged_img = merge_images_landmarks_maps_gt( | |
| self.valid_images_loaded[:self.sample_grid].copy(), batch_maps_pred_val, | |
| self.valid_gt_maps_loaded, image_size=self.image_size, | |
| num_landmarks=self.num_landmarks, num_samples=self.sample_grid, | |
| scale=self.scale, circle_size=2, fast=self.fast_img_gen) | |
| if self.sample_per_channel: | |
| map_per_channel_small = map_comapre_channels( | |
| self.valid_images_loaded[:self.sample_grid].copy(), batch_maps_small_pred_val, | |
| self.valid_gt_maps_small_loaded, image_size=int(self.image_size / 4), | |
| num_landmarks=self.num_landmarks, scale=self.scale) | |
| map_per_channel = map_comapre_channels( | |
| self.valid_images_loaded[:self.sample_grid].copy(), batch_maps_pred, | |
| self.valid_gt_maps_loaded, image_size=self.image_size, | |
| num_landmarks=self.num_landmarks, scale=self.scale) | |
| summary_img = sess.run( | |
| self.img_summary_valid, | |
| {self.log_image_map: np.expand_dims(merged_img, 0), | |
| self.log_map_channels: np.expand_dims(map_per_channel, 0), | |
| self.log_image_map_small: np.expand_dims(merged_img_small, 0), | |
| self.log_map_channels_small: np.expand_dims(map_per_channel_small, 0)}) | |
| else: | |
| summary_img = sess.run( | |
| self.img_summary_valid, | |
| {self.log_image_map: np.expand_dims(merged_img, 0), | |
| self.log_image_map_small: np.expand_dims(merged_img_small, 0)}) | |
| summary_writer.add_summary(summary_img, step) | |
| else: # save heatmap images to directory | |
| sample_path_imgs = os.path.join( | |
| self.save_sample_path, 'epoch-%d-train-iter-%d-1.png' % (epoch, step + 1)) | |
| sample_path_imgs_small = os.path.join( | |
| self.save_sample_path, 'epoch-%d-train-iter-%d-1-s.png' % (epoch, step + 1)) | |
| scipy.misc.imsave(sample_path_imgs, merged_img) | |
| scipy.misc.imsave(sample_path_imgs_small, merged_img_small) | |
| if self.sample_per_channel: | |
| sample_path_ch_maps = os.path.join( | |
| self.save_sample_path, 'epoch-%d-train-iter-%d-3.png' % (epoch, step + 1)) | |
| sample_path_ch_maps_small = os.path.join( | |
| self.save_sample_path, 'epoch-%d-train-iter-%d-3-s.png' % (epoch, step + 1)) | |
| scipy.misc.imsave(sample_path_ch_maps, map_per_channel) | |
| scipy.misc.imsave(sample_path_ch_maps_small, map_per_channel_small) | |
| print('*** Finished Training ***') | |
| def get_image_maps(self, test_image, reuse=None, norm=False): | |
| """ returns heatmaps of input image (menpo image object)""" | |
| self.add_placeholders() | |
| # build model | |
| pred_hm_p, pred_hm_f, pred_hm_u = self.heatmaps_network(self.images, reuse=reuse) | |
| with tf.Session(config=self.config) as sess: | |
| # load trained parameters | |
| saver = tf.train.Saver() | |
| saver.restore(sess, self.test_model_path) | |
| _, model_name = os.path.split(self.test_model_path) | |
| test_image = test_image.pixels_with_channels_at_back().astype('float32') | |
| if norm: | |
| if self.scale is '255': | |
| test_image *= 255 | |
| elif self.scale is '0': | |
| test_image = 2 * test_image - 1 | |
| map_primary, map_fusion, map_upsample = sess.run( | |
| [pred_hm_p, pred_hm_f, pred_hm_u], {self.images: np.expand_dims(test_image, 0)}) | |
| return map_primary, map_fusion, map_upsample | |
| def get_landmark_predictions(self, img_list, pdm_models_dir, clm_model_path, reuse=None, map_to_input_size=False): | |
| """returns dictionary with landmark predictions of each step of the ECpTp algorithm and ECT""" | |
| from thirdparty.face_of_art.pdm_clm_functions import feature_based_pdm_corr, clm_correct | |
| jaw_line_inds = np.arange(0, 17) | |
| left_brow_inds = np.arange(17, 22) | |
| right_brow_inds = np.arange(22, 27) | |
| self.add_placeholders() | |
| # build model | |
| _, _, pred_hm_u = self.heatmaps_network(self.images, reuse=reuse) | |
| with tf.Session(config=self.config) as sess: | |
| # load trained parameters | |
| saver = tf.train.Saver() | |
| saver.restore(sess, self.test_model_path) | |
| _, model_name = os.path.split(self.test_model_path) | |
| e_list = [] | |
| ect_list = [] | |
| ecp_list = [] | |
| ecpt_list = [] | |
| ecptp_jaw_list = [] | |
| ecptp_out_list = [] | |
| for test_image in img_list: | |
| if map_to_input_size: | |
| test_image_transform = test_image[1] | |
| test_image=test_image[0] | |
| # get landmarks for estimation stage | |
| if test_image.n_channels < 3: | |
| test_image_map = sess.run( | |
| pred_hm_u, {self.images: np.expand_dims( | |
| gray2rgb(test_image.pixels_with_channels_at_back()).astype('float32'), 0)}) | |
| else: | |
| test_image_map = sess.run( | |
| pred_hm_u, {self.images: np.expand_dims( | |
| test_image.pixels_with_channels_at_back().astype('float32'), 0)}) | |
| init_lms = heat_maps_to_landmarks(np.squeeze(test_image_map)) | |
| # get landmarks for part-based correction stage | |
| p_pdm_lms = feature_based_pdm_corr(lms_init=init_lms, models_dir=pdm_models_dir, train_type='basic') | |
| # get landmarks for part-based tuning stage | |
| try: # clm may not converge | |
| pdm_clm_lms = clm_correct( | |
| clm_model_path=clm_model_path, image=test_image, map=test_image_map, lms_init=p_pdm_lms) | |
| except: | |
| pdm_clm_lms = p_pdm_lms.copy() | |
| # get landmarks ECT | |
| try: # clm may not converge | |
| ect_lms = clm_correct( | |
| clm_model_path=clm_model_path, image=test_image, map=test_image_map, lms_init=init_lms) | |
| except: | |
| ect_lms = p_pdm_lms.copy() | |
| # get landmarks for ECpTp_out (tune jaw and eyebrows) | |
| ecptp_out = p_pdm_lms.copy() | |
| ecptp_out[left_brow_inds] = pdm_clm_lms[left_brow_inds] | |
| ecptp_out[right_brow_inds] = pdm_clm_lms[right_brow_inds] | |
| ecptp_out[jaw_line_inds] = pdm_clm_lms[jaw_line_inds] | |
| # get landmarks for ECpTp_jaw (tune jaw) | |
| ecptp_jaw = p_pdm_lms.copy() | |
| ecptp_jaw[jaw_line_inds] = pdm_clm_lms[jaw_line_inds] | |
| if map_to_input_size: | |
| ecptp_jaw = test_image_transform.apply(ecptp_jaw) | |
| ecptp_out = test_image_transform.apply(ecptp_out) | |
| ect_lms = test_image_transform.apply(ect_lms) | |
| init_lms = test_image_transform.apply(init_lms) | |
| p_pdm_lms = test_image_transform.apply(p_pdm_lms) | |
| pdm_clm_lms = test_image_transform.apply(pdm_clm_lms) | |
| ecptp_jaw_list.append(ecptp_jaw) # E + p-correction + p-tuning (ECpTp_jaw) | |
| ecptp_out_list.append(ecptp_out) # E + p-correction + p-tuning (ECpTp_out) | |
| ect_list.append(ect_lms) # ECT prediction | |
| e_list.append(init_lms) # init prediction from heatmap network (E) | |
| ecp_list.append(p_pdm_lms) # init prediction + part pdm correction (ECp) | |
| ecpt_list.append(pdm_clm_lms) # init prediction + part pdm correction + global tuning (ECpT) | |
| pred_dict = { | |
| 'E': e_list, | |
| 'ECp': ecp_list, | |
| 'ECpT': ecpt_list, | |
| 'ECT': ect_list, | |
| 'ECpTp_jaw': ecptp_jaw_list, | |
| 'ECpTp_out': ecptp_out_list | |
| } | |
| return pred_dict | |