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 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"])
rand_para_shape = np.random.normal(
size=[1, basis3dmm["basis_shape"].shape[0]])
rand_para_uv = np.random.normal(
size=[1, basis3dmm["uv"]["basis"].shape[0]])
ver_shape = np_get_geometry(basis3dmm, rand_para_shape)[0]
uv_texture, _ = construct(basis3dmm["uv"], rand_para_uv, 512)
print(basis3dmm["vt_list"].shape, basis3dmm["tri"].shape,
basis3dmm["tri_vt"].shape)
# save all files
write_obj(
"rand_shape.obj",
ver_shape,
basis3dmm["vt_list"],
basis3dmm["tri"],
basis3dmm["tri_vt"],
"face.mtl",
)
Image.fromarray(uv_texture.astype(np.uint8)).save("rand_uv.png")
# load regional pyramid bases (it takes a long time because files are huge)
basis3dmm = load_3dmm_basis(
"../files/AI-NExT-Shape-NoAug.mat",
"../files/AI-NExT-AlbedoNormal-RPB",
is_whole_uv=False,
limit_dim=-1,
)
rand_para_uv_dict = {}
for region_name in basis3dmm["uv"]: