def run(self, image_path, landmark_path, mask_path):
# The implementation is potentially able to optimize with images(batch_size>1),
# here we show the example with a single image fitting
# Photometric optimization is sensitive to the hair or glass occlusions,
# therefore we use a face segmentation network to mask the skin region out.
images = []
landmarks = []
image_masks = []
# check and make the save dir, skip if exists
image_name = os.path.basename(image_path)[:-4]
savefolder = os.path.sep.join([self.config.save_dir, image_name])
util.check_mkdir(savefolder)
if os.path.exists(os.path.join(savefolder, image_name + '.npy')):
print('Skipping for {}...'.format(image_name))
return
# read and process the image
image = cv2.resize(cv2.imread(image_path),
(config.cropped_size, config.cropped_size)).astype(
np.float32) / 255.
image = image[:, :, [2, 1, 0]].transpose(2, 0, 1)
images.append(torch.from_numpy(image[None, :, :, :]).to(self.device))
# read and process the mask
image_mask = np.load(mask_path, allow_pickle=True)
image_mask = image_mask[..., None].astype(np.float32)
image_mask = image_mask.transpose(2, 0, 1)
image_mask_bn = np.zeros_like(image_mask)
image_mask_bn[np.where(image_mask != 0)] = 1.
image_masks.append(
torch.from_numpy(image_mask_bn[None, :, :, :]).to(self.device))
# load and normalize the landmark
landmark = np.load(landmark_path).astype(np.float32)
landmark[:, 0] = landmark[:, 0] / float(image.shape[2]) * 2 - 1
landmark[:, 1] = landmark[:, 1] / float(image.shape[1]) * 2 - 1
landmarks.append(
torch.from_numpy(landmark)[None, :, :].float().to(self.device))
# resize and create batch
images = torch.cat(images, dim=0)
images = F.interpolate(
images, [self.image_size, self.image_size]) # <224, 224>
image_masks = torch.cat(image_masks, dim=0)
image_masks = F.interpolate(image_masks,
[self.image_size, self.image_size])
landmarks = torch.cat(landmarks, dim=0)
# optimize
single_params = self.optimize(images, landmarks, image_masks,
savefolder)
if self.config.save_all_output:
self.render.save_obj(
filename=os.path.join(savefolder, image_name + '.obj'),
vertices=torch.from_numpy(single_params['verts'][0]).to(
self.device),
textures=torch.from_numpy(single_params['albedos'][0]).to(
self.device))
np.save(os.path.join(savefolder, image_name + '.npy'), single_params)
def run(self, imagepath, landmarkpath):
# The implementation is potentially able to optimize with images(batch_size>1),
# The implementation is potentially able to optimize with images(batch_size>1),
# here we show the example with a single image fitting
images = []
landmarks = []
image_masks = []
image_name = os.path.basename(imagepath)[:-4]
savefile = os.path.sep.join(
[self.config.savefolder, image_name + '.npy'])
# photometric optimization is sensitive to the hair or glass occlusions,
# therefore we use a face segmentation network to mask the skin region out.
image_mask_folder = './FFHQ_seg/'
image_mask_path = os.path.sep.join(
[image_mask_folder, image_name + '.npy'])
image = cv2.resize(
cv2.imread(imagepath),
(self.config.cropped_size, self.config.cropped_size)).astype(
np.float32) / 255.
image = image[:, :, [2, 1, 0]].transpose(2, 0, 1)
images.append(torch.from_numpy(image[None, :, :, :]).to(self.device))
image_mask = np.load(image_mask_path, allow_pickle=True)
image_mask = image_mask[..., None].astype('float32')
image_mask = image_mask.transpose(2, 0, 1)
image_mask_bn = np.zeros_like(image_mask)
image_mask_bn[np.where(image_mask != 0)] = 1.
image_masks.append(
torch.from_numpy(image_mask_bn[None, :, :, :]).to(self.device))
landmark = np.load(landmarkpath).astype(np.float32)
landmark[:, 0] = landmark[:, 0] / float(image.shape[2]) * 2 - 1
landmark[:, 1] = landmark[:, 1] / float(image.shape[1]) * 2 - 1
landmarks.append(
torch.from_numpy(landmark)[None, :, :].float().to(self.device))
images = torch.cat(images, dim=0)
images = F.interpolate(images, [self.image_size, self.image_size])
image_masks = torch.cat(image_masks, dim=0)
image_masks = F.interpolate(image_masks,
[self.image_size, self.image_size])
landmarks = torch.cat(landmarks, dim=0)
savefolder = os.path.sep.join([self.config.savefolder, image_name])
util.check_mkdir(savefolder)
# optimize
single_params = self.optimize(images, landmarks, image_masks,
savefolder)
self.render.save_obj(filename=savefile[:-4] + '.obj',
vertices=torch.from_numpy(
single_params['verts'][0]).to(self.device),
textures=torch.from_numpy(
single_params['albedos'][0]).to(self.device))
np.save(savefile, single_params)
def main(config):
# init the env
env_config = set_env(config.gpu)
# get the char list
char_list, char_id_map, id_char_map, char_num = get_char_list()
# check the path
check_mkdir(config.image_dir)
check_mkdir(config.tf_data_dir)
check_mkdir(config.check_point_dir)
# generate the code images
for i in range(config.image_num):
gen_code_image(char_list, config.image_dir, 4)
# generate the tfrecord file
tfrecord_generator = GenTfrecodFile(config.image_dir,
config.tf_data_dir,
char_id_map,
num_test=500,
num_validate=500)
tfrecord_generator.gen_tfrecord()
def main(config):
# init the env
env_config = set_env(config.gpu)
# check the path
check_mkdir(config.image_dir)
check_mkdir(config.tf_data_dir)
check_mkdir(config.check_point_dir)
# get the char list
char_list, char_id_map, id_char_map, char_num = get_char_list()
# train tfrecord
train_tfrecord = '%s/%s' % (config.tf_data_dir, 'train.tfrecords')
# input and output placeholder
x = tf.placeholder(tf.float32, [None, 224, 224])
y0 = tf.placeholder(tf.float32, [None])
y1 = tf.placeholder(tf.float32, [None])
y2 = tf.placeholder(tf.float32, [None])
y3 = tf.placeholder(tf.float32, [None])
# get the input and putput data
tfrecord_generator = GenTfrecodFile(config.image_dir, config.tf_data_dir, char_id_map, num_test=500, num_validate=500)
image, image_raw, label0, label1, label2, label3 = tfrecord_generator.read_and_decode(train_tfrecord)
image_batch, image_raw_batch,label_batch0, label_batch1, label_batch2, label_batch3 = tf.train.shuffle_batch(
[image, image_raw,label0, label1, label2, label3], batch_size = config.batch_size,
capacity = 200000, min_after_dequeue=10000, num_threads=1)
# the net
train_network_fn = nets_factory.get_network_fn(
'alexnet_v2',
num_classes=char_num,
weight_decay=0.0005,
is_training=True)
with tf.Session(config = env_config) as sess:
# inputs: a tensor of size [batch_size, height, width, channels]
X = tf.reshape(x, [config.batch_size, 224, 224, 1])
# 数据输入网络得到输出值
logits0,logits1,logits2,logits3,end_points = train_network_fn(X)
# 把标签转成one_hot的形式
one_hot_labels0 = tf.one_hot(indices=tf.cast(y0, tf.int32), depth=char_num)
one_hot_labels1 = tf.one_hot(indices=tf.cast(y1, tf.int32), depth=char_num)
one_hot_labels2 = tf.one_hot(indices=tf.cast(y2, tf.int32), depth=char_num)
one_hot_labels3 = tf.one_hot(indices=tf.cast(y3, tf.int32), depth=char_num)
# 计算loss
loss0 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits0,labels=one_hot_labels0))
loss1 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits1,labels=one_hot_labels1))
loss2 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits2,labels=one_hot_labels2))
loss3 = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits3,labels=one_hot_labels3))
# 计算总的loss
total_loss = (loss0+loss1+loss2+loss3)/4.0
# 优化total_loss
optimizer = tf.train.AdamOptimizer(learning_rate=config.learning_rate).minimize(total_loss)
# 计算准确率
correct_prediction0 = tf.equal(tf.argmax(one_hot_labels0,1),tf.argmax(logits0,1))
accuracy0 = tf.reduce_mean(tf.cast(correct_prediction0,tf.float32))
correct_prediction1 = tf.equal(tf.argmax(one_hot_labels1,1),tf.argmax(logits1,1))
accuracy1 = tf.reduce_mean(tf.cast(correct_prediction1,tf.float32))
correct_prediction2 = tf.equal(tf.argmax(one_hot_labels2,1),tf.argmax(logits2,1))
accuracy2 = tf.reduce_mean(tf.cast(correct_prediction2,tf.float32))
correct_prediction3 = tf.equal(tf.argmax(one_hot_labels3,1),tf.argmax(logits3,1))
accuracy3 = tf.reduce_mean(tf.cast(correct_prediction3,tf.float32))
# 用于保存模型
saver = tf.train.Saver(max_to_keep=4)
# 初始化
sess.run(tf.global_variables_initializer())
# 创建一个协调器，管理线程
coord = tf.train.Coordinator()
# 启动QueueRunner, 此时文件名队列已经进队
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
for i in range(config.epoch_size):
b_image, b_image_raw, b_label0, b_label1 ,b_label2 ,b_label3 = sess.run([image_batch,
image_raw_batch,
label_batch0,
label_batch1,
label_batch2,
label_batch3])
# 优化模型
sess.run(optimizer, feed_dict={x: b_image, y0:b_label0, y1: b_label1, y2: b_label2, y3: b_label3})
# 每迭代20次计算一次loss和准确率
if i % 20 == 0:
acc0,acc1,acc2,acc3,loss_ = sess.run([accuracy0,accuracy1,accuracy2,accuracy3,total_loss],feed_dict={x: b_image,
y0: b_label0,
y1: b_label1,
y2: b_label2,
y3: b_label3})
print ("Iter:%d Loss:%.3f Accuracy:%.2f,%.2f,%.2f,%.2f Learning_rate:%.6f" % (i,loss_,acc0,acc1,acc2,acc3,config.learning_rate))
# 保存模型
if i % 300 == 0:
saver.save(sess, config.check_point_dir, global_step=i)
if acc0 > 0.90 and acc1 > 0.90 and acc2 > 0.90 and acc3 > 0.90:
break
# 通知其他线程关闭
coord.request_stop()
# 其他所有线程关闭之后，这一函数才能返回
coord.join(threads)
def optimize(self, images, landmarks, image_masks, savefolder=None):
bz = images.shape[0]
pose = nn.Parameter(
torch.zeros(bz, self.config.pose_params).float().to(self.device))
exp = nn.Parameter(
torch.zeros(bz,
self.config.expression_params).float().to(self.device))
shape = nn.Parameter(
torch.zeros(bz, self.config.shape_params).float().to(self.device))
tex = nn.Parameter(
torch.zeros(bz, self.config.tex_params).float().to(self.device))
cam = torch.zeros(bz, self.config.camera_params)
cam[:, 0] = 5.
cam = nn.Parameter(cam.float().to(self.device))
lights = nn.Parameter(torch.zeros(bz, 9, 3).float().to(self.device))
e_opt = torch.optim.Adam([shape, exp, pose, cam, tex, lights],
lr=self.config.e_lr,
weight_decay=self.config.e_wd)
e_opt_rigid = torch.optim.Adam([pose, cam],
lr=self.config.e_lr,
weight_decay=self.config.e_wd)
gt_landmark = landmarks
# rigid fitting of pose and camera with 51 static face landmarks,
# this is due to the non-differentiable attribute of contour landmarks trajectory
for k in range(200):
losses = {}
vertices, landmarks2d, landmarks3d = self.flame(
shape_params=shape, expression_params=exp, pose_params=pose)
trans_vertices = util.batch_orth_proj(vertices, cam)
trans_vertices[..., 1:] = -trans_vertices[..., 1:]
landmarks2d = util.batch_orth_proj(landmarks2d, cam)
landmarks2d[..., 1:] = -landmarks2d[..., 1:]
landmarks3d = util.batch_orth_proj(landmarks3d, cam)
landmarks3d[..., 1:] = -landmarks3d[..., 1:]
losses['landmark'] = util.l2_distance(
landmarks2d[:, 17:, :2], gt_landmark[:,
17:, :2]) * config.w_lmks
all_loss = 0.
for key in losses.keys():
all_loss = all_loss + losses[key]
losses['all_loss'] = all_loss
e_opt_rigid.zero_grad()
all_loss.backward()
e_opt_rigid.step()
if self.config.verbose:
loss_info = '----iter: {}, time: {}\n'.format(
k,
datetime.datetime.now().strftime('%Y-%m-%d-%H:%M:%S'))
for key in losses.keys():
loss_info = loss_info + '{}: {}, '.format(
key, float(losses[key]))
if k % self.config.print_freq == 0:
print(loss_info)
if self.config.save_all_output and k % self.config.print_freq == 0:
grids = {}
visind = range(bz) # [0]
grids['images'] = torchvision.utils.make_grid(
images[visind]).detach().cpu()
grids['landmarks_gt'] = torchvision.utils.make_grid(
util.tensor_vis_landmarks(images[visind],
landmarks[visind]))
grids['landmarks2d'] = torchvision.utils.make_grid(
util.tensor_vis_landmarks(images[visind],
landmarks2d[visind]))
grids['landmarks3d'] = torchvision.utils.make_grid(
util.tensor_vis_landmarks(images[visind],
landmarks3d[visind]))
grid = torch.cat(list(grids.values()), 1)
grid_image = (grid.numpy().transpose(1, 2, 0).copy() *
255)[:, :, [2, 1, 0]]
grid_image = np.minimum(np.maximum(grid_image, 0),
255).astype(np.uint8)
savefolder_iter = os.path.join(
savefolder, 'inter_steps'
) # use a separate folder to save intermediate results
util.check_mkdir(savefolder_iter)
cv2.imwrite('{}/{}.jpg'.format(savefolder_iter, k), grid_image)
# non-rigid fitting of all the parameters with 68 face landmarks, photometric loss and regularization terms.
for k in range(200, 1500):
losses = {}
vertices, landmarks2d, landmarks3d = self.flame(
shape_params=shape, expression_params=exp, pose_params=pose)
trans_vertices = util.batch_orth_proj(vertices, cam)
trans_vertices[..., 1:] = -trans_vertices[..., 1:]
landmarks2d = util.batch_orth_proj(landmarks2d, cam)
landmarks2d[..., 1:] = -landmarks2d[..., 1:]
landmarks3d = util.batch_orth_proj(landmarks3d, cam)
landmarks3d[..., 1:] = -landmarks3d[..., 1:]
losses['landmark'] = util.l2_distance(
landmarks2d[:, :, :2], gt_landmark[:, :, :2]) * config.w_lmks
losses['shape_reg'] = (torch.sum(shape**2) /
2) * config.w_shape_reg # *1e-4
losses['expression_reg'] = (torch.sum(exp**2) /
2) * config.w_expr_reg # *1e-4
losses['pose_reg'] = (torch.sum(pose**2) / 2) * config.w_pose_reg
## render
albedos = self.flametex(tex) / 255.
ops = self.render(vertices, trans_vertices, albedos, lights)
predicted_images = ops['images']
losses['photometric_texture'] = (
image_masks *
(ops['images'] - images).abs()).mean() * config.w_pho
all_loss = 0.
for key in losses.keys():
all_loss = all_loss + losses[key]
losses['all_loss'] = all_loss
e_opt.zero_grad()
all_loss.backward()
e_opt.step()
if self.config.verbose:
loss_info = '----iter: {}, time: {}\n'.format(
k,
datetime.datetime.now().strftime('%Y-%m-%d-%H:%M:%S'))
for key in losses.keys():
loss_info = loss_info + '{}: {}, '.format(
key, float(losses[key]))
if k % self.config.print_freq == 0:
print(loss_info)
# visualize
if self.config.save_all_output and k % self.config.print_freq == 0:
grids = {}
visind = range(bz) # [0]
grids['images'] = torchvision.utils.make_grid(
images[visind]).detach().cpu()
grids['landmarks_gt'] = torchvision.utils.make_grid(
util.tensor_vis_landmarks(images[visind],
landmarks[visind]))
grids['landmarks2d'] = torchvision.utils.make_grid(
util.tensor_vis_landmarks(images[visind],
landmarks2d[visind]))
grids['landmarks3d'] = torchvision.utils.make_grid(
util.tensor_vis_landmarks(images[visind],
landmarks3d[visind]))
grids['albedoimage'] = torchvision.utils.make_grid(
(ops['albedo_images'])[visind].detach().cpu())
grids['render'] = torchvision.utils.make_grid(
predicted_images[visind].detach().float().cpu())
shape_images = self.render.render_shape(
vertices, trans_vertices, images)
grids['shape'] = torchvision.utils.make_grid(
F.interpolate(shape_images[visind],
[224, 224])).detach().float().cpu()
# grids['tex'] = torchvision.utils.make_grid(F.interpolate(albedos[visind], [224, 224])).detach().cpu()
grid = torch.cat(list(grids.values()), 1)
grid_image = (grid.numpy().transpose(1, 2, 0).copy() *
255)[:, :, [2, 1, 0]]
grid_image = np.minimum(np.maximum(grid_image, 0),
255).astype(np.uint8)
savefolder_iter = os.path.join(
savefolder, 'inter_steps'
) # use a separate folder to save intermediate results
util.check_mkdir(savefolder_iter)
cv2.imwrite('{}/{}.jpg'.format(savefolder_iter, k), grid_image)
single_params = {
'shape': shape.detach().cpu().numpy(),
'exp': exp.detach().cpu().numpy(),
'pose': pose.detach().cpu().numpy(),
'cam': cam.detach().cpu().numpy(),
'verts': trans_vertices.detach().cpu().numpy(),
'albedos': albedos.detach().cpu().numpy(),
'tex': tex.detach().cpu().numpy(),
'lit': lights.detach().cpu().numpy()
}
return single_params
config[
'img_dir'] = '/data/yunfan.liu/Data_Preparation_Face_Swapping_Reenactment/sample/synthetic_dataset/{}/origin_image'.format(
config['cropped_size'])
config[
'ldmk_dir'] = '/data/yunfan.liu/Data_Preparation_Face_Swapping_Reenactment/sample/synthetic_dataset/{}/ldmk'.format(
config['cropped_size'])
config[
'mask_dir'] = '/data/yunfan.liu/Data_Preparation_Face_Swapping_Reenactment/sample/synthetic_dataset/{}/mask'.format(
config['cropped_size'])
config[
'save_dir'] = '/data/yunfan.liu/Data_Preparation_Face_Swapping_Reenactment/sample/synthetic_dataset/{}/3DMM_params'.format(
config['cropped_size'])
config = util.dict2obj(config)
util.check_mkdir(config.save_dir)
img_list = os.listdir(
config.img_dir)[args.batch_ind * args.batch_size:(args.batch_ind + 1) *
args.batch_size]
for img_ind in tqdm(range(len(img_list))):
img_name = img_list[img_ind]
img_path = os.path.join(config.img_dir, img_name)
# check for landmark and mask ()
img_prefix = img_name[:-4]
ldmk_save_path = os.path.join(config.ldmk_dir,
img_prefix + '_ldmk.npy')
ldmk = read_ldmk_file(ldmk_save_path, img_path)
mask_save_path = os.path.join(config.mask_dir,
img_prefix + '_mask.npy')
'shape_params': 100,
'expression_params': 50,
'pose_params': 6,
'tex_params': 50,
'use_face_contour': True,
'cropped_size': 256,
'batch_size': 1,
'image_size': 224,
'e_lr': 0.005,
'e_wd': 0.0001,
'savefolder': './test_results/',
# weights of losses and reg terms
'w_pho': 8,
'w_lmks': 1,
'w_shape_reg': 1e-4,
'w_expr_reg': 1e-4,
'w_pose_reg': 0,
}
config = util.dict2obj(config)
util.check_mkdir(config.savefolder)
config.batch_size = 1
fitting = PhotometricFitting(config, device=device_name)
input_folder = './FFHQ'
imagepath = os.path.sep.join([input_folder, image_name + '.png'])
landmarkpath = os.path.sep.join([input_folder, image_name + '.npy'])
fitting.run(imagepath, landmarkpath)
def main(config):
# init the env
env_config = set_env(config.gpu)
# check the path
check_mkdir(config.image_dir)
check_mkdir(config.tf_data_dir)
check_mkdir(config.check_point_dir)
# get the char list
char_list, char_id_map, id_char_map, char_num = get_char_list()
# train tfrecord
train_tfrecord = '%s/%s' % (config.tf_data_dir, 'test.tfrecords')
# input and output placeholder
x = tf.placeholder(tf.float32, [None, 224, 224])
y0 = tf.placeholder(tf.float32, [None])
y1 = tf.placeholder(tf.float32, [None])
y2 = tf.placeholder(tf.float32, [None])
y3 = tf.placeholder(tf.float32, [None])
# get the input and putput data
tfrecord_generator = GenTfrecodFile(config.image_dir,
config.tf_data_dir,
char_id_map,
num_test=500,
num_validate=500)
image, image_raw, label0, label1, label2, label3 = tfrecord_generator.read_and_decode(
train_tfrecord)
image_batch, image_raw_batch, label_batch0, label_batch1, label_batch2, label_batch3 = tf.train.shuffle_batch(
[image, image_raw, label0, label1, label2, label3],
batch_size=config.batch_size,
capacity=200000,
min_after_dequeue=10000,
num_threads=1)
print(image_batch[0].shape)
# the net
train_network_fn = nets_factory.get_network_fn('alexnet_v2',
num_classes=char_num,
weight_decay=0.0005,
is_training=True)
with tf.Session(config=env_config) as sess:
# inputs: a tensor of size [batch_size, height, width, channels]
X = tf.reshape(x, [config.batch_size, 224, 224, 1])
# 数据输入网络得到输出值
logits0, logits1, logits2, logits3, end_points = train_network_fn(X)
predict0 = tf.reshape(logits0, [-1, char_num])
predict0 = tf.argmax(predict0, 1)
predict1 = tf.reshape(logits1, [-1, char_num])
predict1 = tf.argmax(predict1, 1)
predict2 = tf.reshape(logits2, [-1, char_num])
predict2 = tf.argmax(predict2, 1)
predict3 = tf.reshape(logits3, [-1, char_num])
predict3 = tf.argmax(predict3, 1)
# 初始化
sess.run(tf.global_variables_initializer())
# 载入训练好的模型
saver = tf.train.Saver()
#checkpoint = tf.train.latest_checkpoint(check_point_dir)
#saver.restore(sess,checkpoint)
load_check_point(config.check_point_dir, sess, saver)
# 创建一个协调器，管理线程
coord = tf.train.Coordinator()
# 启动QueueRunner, 此时文件名队列已经进队
threads = tf.train.start_queue_runners(sess=sess, coord=coord)
for i in range(10):
# 获取一个批次的数据和标签
b_image, b_image_raw, b_label0, b_label1, b_label2, b_label3 = sess.run(
[
image_batch, image_raw_batch, label_batch0, label_batch1,
label_batch2, label_batch3
])
# 打印标签
print('label:', id_char_map[b_label0[0]], id_char_map[b_label1[0]],
id_char_map[b_label2[0]], id_char_map[b_label3[0]])
# 预测
label0, label1, label2, label3 = sess.run(
[predict0, predict1, predict2, predict3],
feed_dict={x: b_image})
# 打印预测值
print('predict:', id_char_map[label0[0]], id_char_map[label1[0]],
id_char_map[label2[0]], id_char_map[label3[0]])
# 通知其他线程关闭
coord.request_stop()
def main(config):
# init the env
env_config = set_env(config.gpu)
# check the path
check_mkdir(config.image_dir)
check_mkdir(config.tf_data_dir)
check_mkdir(config.check_point_dir)
# get the char list
char_list, char_id_map, id_char_map, char_num = get_char_list()
# train tfrecord
train_tfrecord = '%s/%s' % (config.tf_data_dir, 'test.tfrecords')
# input and output placeholder
x = tf.placeholder(tf.float32, [None, 224, 224])
y0 = tf.placeholder(tf.float32, [None])
y1 = tf.placeholder(tf.float32, [None])
y2 = tf.placeholder(tf.float32, [None])
y3 = tf.placeholder(tf.float32, [None])
# get the input and putput data
image_data = Image.open(config.img_file)
#根据模型的结构resize 为了配合alexnet
image_data = image_data.resize((224, 224))
#灰度化
image_data = np.array(image_data.convert('L'))
image_data = image_data / 255.0
image_data = image_data - 0.5
image_data = image_data * 2.0
image_data = image_data.reshape((1, 224, 224))
print(image_data.shape)
print(type(image_data))
# the net
train_network_fn = nets_factory.get_network_fn('alexnet_v2',
num_classes=char_num,
weight_decay=0.0005,
is_training=True)
with tf.Session(config=env_config) as sess:
# inputs: a tensor of size [batch_size, height, width, channels]
X = tf.reshape(x, [config.batch_size, 224, 224, 1])
# 数据输入网络得到输出值
logits0, logits1, logits2, logits3, end_points = train_network_fn(X)
predict0 = tf.reshape(logits0, [-1, char_num])
predict0 = tf.argmax(predict0, 1)
predict1 = tf.reshape(logits1, [-1, char_num])
predict1 = tf.argmax(predict1, 1)
predict2 = tf.reshape(logits2, [-1, char_num])
predict2 = tf.argmax(predict2, 1)
predict3 = tf.reshape(logits3, [-1, char_num])
predict3 = tf.argmax(predict3, 1)
# 初始化
sess.run(tf.global_variables_initializer())
# 载入训练好的模型
saver = tf.train.Saver()
#checkpoint = tf.train.latest_checkpoint(check_point_dir)
#saver.restore(sess,checkpoint)
load_check_point(config.check_point_dir, sess, saver)
label0, label1, label2, label3 = sess.run(
[predict0, predict1, predict2, predict3],
feed_dict={x: image_data})
# 打印预测值
print('predict:', id_char_map[label0[0]], id_char_map[label1[0]],
id_char_map[label2[0]], id_char_map[label3[0]])
def run(self, img, net, rect_detect, landmark_detect, rect_thresh,
save_name, savefolder):
# The implementation is potentially able to optimize with images(batch_size>1),
# here we show the example with a single image fitting
images = []
landmarks = []
image_masks = []
bbox = rect_detect.extract(img, rect_thresh)
if len(bbox) > 0:
crop_image, new_bbox = self.crop_img(img, bbox[0])
#plt.imshow(crop_image)
#plt.show()
#pdb.set_trace()
resize_img, landmark = landmark_detect.extract(
[crop_image, [new_bbox]])
landmark = landmark[0]
landmark[:,
0] = landmark[:, 0] / float(resize_img.shape[1]) * 2 - 1
landmark[:,
1] = landmark[:, 1] / float(resize_img.shape[0]) * 2 - 1
landmarks.append(
torch.from_numpy(landmark)[None, :, :].double().to(
self.device))
image = cv2.resize(
crop_image,
(self.config.cropped_size, self.config.cropped_size)).astype(
np.float32) / 255.
image = image[:, :, [2, 1, 0]].transpose(2, 0, 1)
images.append(
torch.from_numpy(image[None, :, :, :]).double().to(
self.device))
image_mask = face_seg(crop_image, net)
image_mask = cv2.resize(
image_mask,
(self.config.cropped_size, self.config.cropped_size))
image_mask = image_mask[..., None].astype('float32')
image_mask = image_mask.transpose(2, 0, 1)
image_mask_bn = np.zeros_like(image_mask)
image_mask_bn[np.where(image_mask != 0)] = 1.
image_masks.append(
torch.from_numpy(image_mask_bn[None, :, :, :]).double().to(
self.device))
images = torch.cat(images, dim=0)
images = F.interpolate(images, [self.image_size, self.image_size])
image_masks = torch.cat(image_masks, dim=0)
image_masks = F.interpolate(image_masks,
[self.image_size, self.image_size])
landmarks = torch.cat(landmarks, dim=0)
util.check_mkdir(savefolder)
save_name = os.path.join(savefolder, save_name)
images = images.float()
landmarks = landmarks.float()
# optimize
single_params = self.optimize(images, landmarks, image_masks,
savefolder)
self.render.save_obj(
filename=save_name,
vertices=torch.from_numpy(single_params['verts'][0]).to(
self.device),
textures=torch.from_numpy(single_params['albedos'][0]).to(
self.device))
np.save(save_name, single_params)
