def RGBD_opt():
os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.GPU_NO
# load data
basis3dmm = load_3dmm_basis_bfm(FLAGS.basis3dmm_path)
para_tex_init = np.zeros([1, basis3dmm["basis_tex"].shape[0]]) # 1 * 80
np.save((os.path.join(FLAGS.prefit_dir, "para_tex_init.npy")), para_tex_init)
# load_RGBD_data, sequence index is: mid -- left -- right -- up
info = RGBD_load.load_and_preprocess_RGBD_data(
FLAGS.prefit_dir, FLAGS.prepare_dir, basis3dmm
)
imageH = info["height"]
imageW = info["width"]
para_shape_shape = info["para_shape"].shape[1]
para_tex_shape = info["para_tex"].shape[1]
# build graph
var_list = define_variable(
FLAGS.num_of_img, imageH, imageW, para_shape_shape, para_tex_shape, info
)
out_list = build_RGBD_opt_graph(var_list, basis3dmm, imageH, imageW)
# summary_op
summary_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(FLAGS.summary_dir)
if os.path.exists(FLAGS.summary_dir) is False:
os.makedirs(FLAGS.summary_dir)
# start opt
config = tf.ConfigProto()
# config.gpu_options.per_process_gpu_memory_fraction=0.5
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
import time
starttime = time.time()
for step in range(FLAGS.train_step):
if (step % FLAGS.log_step == 0) | (step == FLAGS.train_step - 1):
out_summary = sess.run(summary_op)
summary_writer.add_summary(out_summary, step)
print(step)
endtime = time.time()
print("time:" + str(endtime - starttime))
starttime = time.time()
sess.run(out_list["train_op"])
para_tex_out = sess.run(out_list["para_tex"])
para_tex_out = np.reshape(para_tex_out[0], [1, -1])
np.save((os.path.join(FLAGS.prefit_dir, "para_tex_init.npy")), para_tex_out)
def RGBD_opt(_):
os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.GPU_NO
# load data
if FLAGS.is_bfm is False:
basis3dmm = load_3dmm_basis(
FLAGS.basis3dmm_path,
FLAGS.uv_path,
)
else:
basis3dmm = load_3dmm_basis_bfm(FLAGS.basis3dmm_path)
# load_RGBD_data, sequence index is: mid -- left -- right -- up
info = RGBD_load.load_and_preprocess_RGBD_data(
FLAGS.prefit_dir, FLAGS.prepare_dir, basis3dmm
)
imageH = info["height"]
imageW = info["width"]
para_shape_shape = info["para_shape"].shape[1]
para_tex_shape = info["para_tex"].shape[1]
# build graph
var_list = define_variable(
FLAGS.num_of_img, imageH, imageW, para_shape_shape, para_tex_shape, info
)
out_list = build_RGBD_opt_graph(var_list, basis3dmm, imageH, imageW)
# summary_op
summary_op = tf.summary.merge_all()
summary_writer = tf.summary.FileWriter(FLAGS.summary_dir)
if os.path.exists(FLAGS.summary_dir) is False:
os.makedirs(FLAGS.summary_dir)
if os.path.exists(FLAGS.out_dir) is False:
os.makedirs(FLAGS.out_dir)
# start opt
config = tf.ConfigProto()
# config.gpu_options.per_process_gpu_memory_fraction=0.5
config.gpu_options.allow_growth = True
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
import time
starttime = time.time()
for step in range(FLAGS.train_step):
if (step % FLAGS.log_step == 0) | (step == FLAGS.train_step - 1):
out_summary = sess.run(summary_op)
summary_writer.add_summary(out_summary, step)
print("step: " + str(step))
endtime = time.time()
print("time:" + str(endtime - starttime))
starttime = time.time()
if step == FLAGS.train_step - 1 and FLAGS.save_ply:
print("output_final_result...")
out_para_shape, out_ver_xyz, out_tex = sess.run(
[out_list["para_shape"], out_list["ver_xyz"], out_list["tex"]]
)
# output ply
v_xyz = out_ver_xyz[0]
if FLAGS.is_bfm is False:
uv_map = out_tex[0] * 255.0
uv_size = uv_map.shape[0]
v_rgb = np.zeros_like(v_xyz) + 200 # N x 3
for (v1, v2, v3), (t1, t2, t3) in zip(
basis3dmm["tri"], basis3dmm["tri_vt"]
):
v_rgb[v1] = uv_map[
int((1.0 - basis3dmm["vt_list"][t1][1]) * uv_size),
int(basis3dmm["vt_list"][t1][0] * uv_size),
]
v_rgb[v2] = uv_map[
int((1.0 - basis3dmm["vt_list"][t2][1]) * uv_size),
int(basis3dmm["vt_list"][t2][0] * uv_size),
]
v_rgb[v3] = uv_map[
int((1.0 - basis3dmm["vt_list"][t3][1]) * uv_size),
int(basis3dmm["vt_list"][t3][0] * uv_size),
]
write_obj(
os.path.join(FLAGS.out_dir, "face.obj"),
v_xyz,
basis3dmm["vt_list"],
basis3dmm["tri"].astype(np.int32),
basis3dmm["tri_vt"].astype(np.int32),
)
else:
v_rgb = out_tex[0] * 255.0
write_ply(
os.path.join(FLAGS.out_dir, "face.ply"),
v_xyz,
basis3dmm["tri"],
v_rgb.astype(np.uint8),
True,
)
## add head
if FLAGS.is_bfm is False:
print("-------------------start add head-------------------")
HeadModel = np.load(
FLAGS.info_for_add_head, allow_pickle=True
).item()
vertex = read_obj(os.path.join(FLAGS.out_dir, "face.obj"))
vertex = vertex.transpose()
vertex_fit_h = vertex[:, HeadModel["head_h_idx"]]
pca_info_h = AddHeadTool.transfer_PCA_format_for_add_head(
basis3dmm, HeadModel
)
vertex_output_coord = AddHeadTool.fix_back_head(
vertex_fit_h,
HeadModel,
pca_info_h,
FLAGS.is_add_head_mirrow,
FLAGS.is_add_head_male,
)
write_obj(
os.path.join(FLAGS.out_dir, "head.obj"),
vertex_output_coord.transpose(),
HeadModel["head_vt_list"],
HeadModel["head_tri"],
HeadModel["head_tri_vt"],
)
print("-------------------add head successfully-------------------")
out_diffuse, out_proj_xyz, out_ver_norm = sess.run(
[out_list["diffuse"], out_list["proj_xyz"], out_list["ver_norm"]]
)
out_diffuse = out_diffuse * 255.0 # RGB 0-255
scio.savemat(
os.path.join(FLAGS.out_dir, "out_for_texture.mat"),
{
"ori_img": info["ori_img"], # ? x ?
"diffuse": out_diffuse, # 300 x 300
"seg": info["seg_list"], # 300 x 300
"proj_xyz": out_proj_xyz, # in 300 x 300 img
"ver_norm": out_ver_norm,
},
)
sess.run(out_list["train_op"])