def RGB_opt(_):
os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.GPU_NO
# load 3DMM
if FLAGS.is_bfm is False:
basis3dmm = load_3dmm_basis(
FLAGS.basis3dmm_path,
FLAGS.uv_path,
)
para_shape_shape = basis3dmm["basis_shape"].shape[0]
para_tex_shape = basis3dmm["uv"]["basis"].shape[0]
else:
basis3dmm = load_3dmm_basis_bfm(FLAGS.basis3dmm_path)
para_shape_shape = basis3dmm["basis_shape"].shape[0]
para_tex_shape = basis3dmm["basis_tex"].shape[0]
# load RGB data
info = RGB_load.load_rgb_data(FLAGS.base_dir, FLAGS.project_type,
FLAGS.num_of_img)
imageH = info["img_list"].shape[1]
imageW = info["img_list"].shape[2]
# build graph
var_list = define_variable(FLAGS.num_of_img, imageH, imageW,
para_shape_shape, para_tex_shape, info)
out_list = build_RGB_opt_graph(var_list, basis3dmm, imageH, imageW)
# summary_op
summary_op = tf.compat.v1.summary.merge_all()
summary_writer = tf.compat.v1.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.compat.v1.ConfigProto()
# config.gpu_options.per_process_gpu_memory_fraction=0.5
config.gpu_options.allow_growth = True
with tf.compat.v1.Session(config=config) as sess:
sess.run(tf.compat.v1.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,
)
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["img_ori_list"], # ? 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"])
def run(prefit_dir, modle_path):
print("---- step3 start -----")
print("running base:", prefit_dir)
# load models
if FLAGS.is_bfm:
mu_shape, pcev_shape, sigma_shape, tri_h_used, kp_86_h = load_pca_models_BFM(
FLAGS.basis3dmm_path
)
else:
mu_shape, pcev_shape, sigma_shape, tri_h_used, kp_86_h = load_pca_models(
modle_path, FLAGS.basis3dmm_path
)
# load datas
(
K,
img_select,
lmk3d_select,
pt3d_select,
pose_all,
pt3d_remove_outliers,
trans_base_2_camera,
depth_ori_select,
) = load_from_npz(prefit_dir)
# sparse fit
if FLAGS.is_bfm:
print("BFM")
lambda_project = 5000
lambda_reg_paras = 2000000000
lambda_3d_key = 1000
lambda_pt2d = 100
else:
lambda_project = 5000
lambda_reg_paras = 1000
lambda_3d_key = 1000
lambda_pt2d = 100
shape = fit_para_shape_with_kp3d(
trans_base_2_camera,
img_select,
depth_ori_select,
lmk3d_select,
K,
K,
pt3d_remove_outliers,
mu_shape,
pcev_shape,
sigma_shape,
kp_86_h,
lambda_reg_paras,
lambda_3d_key,
lambda_pt2d,
lambda_project,
tri_h_used,
)
# save
prefit_head = mu_shape + pcev_shape.dot(shape)
prefit_head = np.reshape(prefit_head, (3, -1), order="F")
np.save(os.path.join(prefit_dir, "para_shape_init.npy"), shape)
write_ply(
os.path.join(prefit_dir, "prefit_head.ply"),
prefit_head.transpose(),
tri_h_used.transpose(),
None,
True,
)
print("---- step3 succeed -----")
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"])
def run(prefit_dir):
print("---- step2 start -----")
print("running base:", prefit_dir)
(
img_select,
depth_select,
lmk3d_select,
pt3d_select,
first_trans_select,
K,
) = load_from_npz(prefit_dir)
# 1. get mask depth
_, _, depth_ori_select = crop_depth_for_fusion(
img_select, depth_select, lmk3d_select
)
# 2. fusion 3d points
pt3d_remove_outliers_camera = find_3d_keypoints_from_landmark_and_depth_86(
first_trans_select, pt3d_select, lmk3d_select, depth_ori_select, img_select, K
)
# get trans
trans_base_2_camera = get_trans_base_to_camera(
pt3d_remove_outliers_camera.transpose(), FLAGS.is_bfm
)
if FLAGS.is_bfm:
thr1 = 20.0
thr2 = 20.0
else:
thr1 = 2.0
thr2 = 2.0
pt3d_remove_outliers = PoseTools.backtrans(
pt3d_remove_outliers_camera, trans_base_2_camera
)
bad_idx = pt3d_remove_outliers_camera[2, :] < 1
pt3d_remove_outliers[:, bad_idx] = -10000
# # ues y and z coord
s1 = pt3d_remove_outliers[1:, 0:8]
s2 = pt3d_remove_outliers[1:, np.array(range(16, 8, -1))]
loss = np.sum(np.array(abs(abs(s1) - abs(s2))), axis=0)
mid = np.median(loss[loss < thr1])
error = np.array(loss > min(mid + 0.5, thr2)).astype(np.int32)
error_countour = np.concatenate(
(np.concatenate((error, [0])), error[np.array(range(7, -1, -1))])
)
bad_idx = np.where(error_countour[:] > 0)
pt3d_remove_outliers[:, bad_idx] = -10000
# save
np.savez(
os.path.join(prefit_dir, "step2_fusion.npz"),
pt3d_remove_outliers=pt3d_remove_outliers,
trans_base_2_camera=trans_base_2_camera,
first_trans_select=first_trans_select,
depth_ori_select=depth_ori_select,
)
temp = pt3d_remove_outliers[:, pt3d_remove_outliers[2, :] > -100] # 3 * n
write_ply(
os.path.join(prefit_dir, "vis_pt3d_remove_outliers.ply"),
temp.transpose(),
None,
None,
True,
)
print("---- step2 succeed -----")