import torch
from torch import nn
import math
import scipy.sparse
import numpy as np
from torch.nn.parameter import Parameter
from device import device
from psbody.mesh import Mesh
from graphlib import graph, coarsening, utils, mesh_sampling
import pickle
reference_mesh_file = "./data/template.obj"
ds_factors = [4, 4, 4, 4]
# Generates adjecency matrices A, downsampling matrices D, and upsamling matrices U by sampling
# the mesh 4 times. Each time the mesh is sampled by a factor of 4
reference_mesh = Mesh(filename=reference_mesh_file)
M, A, D, U = mesh_sampling.generate_transform_matrices(reference_mesh,
ds_factors)
pickle.dump([M, A, D, U], open("./data/pai_template.pkl", 'wb'))
#%%
from psbody.mesh import Mesh
import torch
reference_mesh_file = "./data/template.obj"
reference_mesh = Mesh(filename=reference_mesh_file)
mean = torch.load('./data/Processed/sliced/mean.tch')
reference_mesh.v = mean.numpy()
reference_mesh.show()
#%%
def vec2mesh(self, vec):
vec = vec.reshape((self.n_vertex, 3)) * self.std + self.mean
return Mesh(v=vec, f=self.reference_mesh.f)
parser = argparse.ArgumentParser()
parser.add_argument('dir', type=str,)
args = parser.parse_args()
name = ' '.join((args.dir).split('/')).split()[-1]
os.makedirs(os.path.join(args.dir,'unwraps'),exist_ok=True)
step1_make_unwraps.main(os.path.join(args.dir,'frame_data.pkl'),os.path.join(args.dir,'frames'),os.path.join(args.dir,'segmentations'),os.path.join(args.dir,'unwraps'))
step2_segm_vote_gmm.main(os.path.join(args.dir,'unwraps'),os.path.join(args.dir,'segm.png'),os.path.join(args.dir,'gmm.pkl'))
step3_stitch_texture.main(os.path.join(args.dir,'unwraps'), os.path.join(args.dir,'segm.png'), os.path.join(args.dir,'gmm.pkl'), os.path.join(args.dir,name+'_octopus.jpg'),20)
filename = os.path.join(args.dir,name+'.obj')
body = Mesh(filename=filename.replace(name+'.obj',name+'_octopus.obj'))
body_tex = filename.replace('.obj', '_octopus.jpg')
if not os.path.exists(body_tex):
body_tex = 'tex_{}'.format(body_tex)
v, f = body.v, body.f
(mapping, hf) = loop_subdivider(v, f)
hv = mapping.dot(v.ravel()).reshape(-1, 3)
body_hres = Mesh(hv, hf)
vt, ft = np.hstack((body.vt, np.ones((body.vt.shape[0], 1)))), body.ft
(mappingt, hft) = loop_subdivider(vt, ft)
hvt = mappingt.dot(vt.ravel()).reshape(-1, 3)[:, :2]
body_hres.vt, body_hres.ft = hvt, hft
body_hres.set_texture_image(body_tex)
def run_2d_lmk_fitting(tf_model_fname, template_fname, flame_lmk_path,
texture_mapping, target_img_path, target_lmk_path,
out_path):
if 'generic' not in tf_model_fname:
print(
'You are fitting a gender specific model (i.e. female / male). Please make sure you selected the right gender model. Choose the generic model if gender is unknown.'
)
if not os.path.exists(template_fname):
print('Template mesh (in FLAME topology) not found - %s' %
template_fname)
return
if not os.path.exists(flame_lmk_path):
print('FLAME landmark embedding not found - %s ' % flame_lmk_path)
return
if not os.path.exists(target_img_path):
print('Target image not found - s' % target_img_path)
return
if not os.path.exists(target_lmk_path):
print('Landmarks of target image not found - s' % target_lmk_path)
return
if not os.path.exists(out_path):
os.makedirs(out_path)
lmk_face_idx, lmk_b_coords = load_embedding(flame_lmk_path)
target_img = cv2.imread(target_img_path)
lmk_2d = np.load(target_lmk_path)
weights = {}
# Weight of the landmark distance term
weights['lmk'] = 1.0
# Weight of the shape regularizer
weights['shape'] = 1e-3
# Weight of the expression regularizer
weights['expr'] = 1e-3
# Weight of the neck pose (i.e. neck rotationh around the neck) regularizer
weights['neck_pose'] = 100.0
# Weight of the jaw pose (i.e. jaw rotation for opening the mouth) regularizer
weights['jaw_pose'] = 1e-3
# Weight of the eyeball pose (i.e. eyeball rotations) regularizer
weights['eyeballs_pose'] = 10.0
result_mesh, result_scale = fit_lmk2d(target_img, lmk_2d, template_fname,
tf_model_fname, lmk_face_idx,
lmk_b_coords, weights)
if sys.version_info >= (3, 0):
texture_data = np.load(texture_mapping,
allow_pickle=True,
encoding='latin1').item()
else:
texture_data = np.load(texture_mapping, allow_pickle=True).item()
texture_map = compute_texture_map(target_img, result_mesh, result_scale,
texture_data)
out_mesh_fname = os.path.join(
out_path,
os.path.splitext(os.path.basename(target_img_path))[0] + '.obj')
out_img_fname = os.path.join(
out_path,
os.path.splitext(os.path.basename(target_img_path))[0] + '.png')
cv2.imwrite(out_img_fname, texture_map)
result_mesh.set_vertex_colors('white')
result_mesh.vt = texture_data['vt']
result_mesh.ft = texture_data['ft']
result_mesh.set_texture_image(out_img_fname)
result_mesh.write_obj(out_mesh_fname)
np.save(
os.path.join(
out_path,
os.path.splitext(os.path.basename(target_img_path))[0] +
'_scale.npy'), result_scale)
mv = MeshViewers(shape=[1, 2], keepalive=True)
mv[0][0].set_static_meshes([Mesh(result_mesh.v, result_mesh.f)])
mv[0][1].set_static_meshes([result_mesh])
def render_mesh_helper(mesh, t_center, rot=np.zeros(3), v_colors=None, errors=None, error_unit='m', min_dist_in_mm=0.0, max_dist_in_mm=3.0, z_offset=0):
camera_params = {'c': np.array([400, 400]),
'k': np.array([-0.19816071, 0.92822711, 0, 0, 0]),
'f': np.array([4754.97941935 / 2, 4754.97941935 / 2])}
frustum = {'near': 0.01, 'far': 3.0, 'height': 800, 'width': 800}
mesh_copy = Mesh(mesh.v, mesh.f)
# mesh_copy.v[:] = np.matmul(mesh_copy.v[:], rot)
# mesh_copy.v[:] = cv2.Rodrigues(rot)[0].dot((mesh_copy.v-t_center).T).T+t_center
if errors is not None:
intensity = 0.5
unit_factor = get_unit_factor('mm')/get_unit_factor(error_unit)
errors = unit_factor*errors
norm = mpl.colors.Normalize(vmin=min_dist_in_mm, vmax=max_dist_in_mm)
cmap = cm.get_cmap(name='jet')
colormapper = cm.ScalarMappable(norm=norm, cmap=cmap)
rgba_per_v = colormapper.to_rgba(errors)
rgb_per_v = rgba_per_v[:, 0:3]
elif v_colors is not None:
intensity = 0.5
rgb_per_v = v_colors
else:
intensity = 1.5
rgb_per_v = None
tri_mesh = trimesh.Trimesh(vertices=mesh_copy.v, faces=mesh_copy.f, vertex_colors=rgb_per_v)
render_mesh = pyrender.Mesh.from_trimesh(tri_mesh, smooth=True)
scene = pyrender.Scene(ambient_light=[.2, .2, .2], bg_color=[255, 255, 255])
camera = pyrender.IntrinsicsCamera(fx=camera_params['f'][0],
fy=camera_params['f'][1],
cx=camera_params['c'][0],
cy=camera_params['c'][1],
znear=frustum['near'],
zfar=frustum['far'])
scene.add(render_mesh, pose=np.eye(4))
camera_pose = np.eye(4)
camera_pose[:3,3] = np.array([0, 0, 1.0-z_offset])
scene.add(camera, pose=[[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 1, 1],
[0, 0, 0, 1]])
angle = np.pi / 6.0
pos = camera_pose[:3,3]
light_color = np.array([1., 1., 1.])
light = pyrender.PointLight(color=light_color, intensity=intensity)
light_pose = np.eye(4)
light_pose[:3,3] = pos
scene.add(light, pose=light_pose.copy())
light_pose[:3,3] = cv2.Rodrigues(np.array([angle, 0, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3,3] = cv2.Rodrigues(np.array([-angle, 0, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3,3] = cv2.Rodrigues(np.array([0, -angle, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3,3] = cv2.Rodrigues(np.array([0, angle, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
flags = pyrender.RenderFlags.SKIP_CULL_FACES
r = pyrender.OffscreenRenderer(viewport_width=frustum['width'], viewport_height=frustum['height'])
color, _ = r.render(scene, flags=flags)
return color[..., ::-1]
def fit_sources(
dir_tup_list,
tf_model_fname,
template_fname,
weight_reg_shape,
weight_reg_expr,
weight_reg_neck_pos,
weight_reg_jaw_pos,
weight_reg_eye_pos,
showing=False
):
global g_mv
if showing:
g_mv = MeshViewer()
saver = tf.train.import_meta_graph(tf_model_fname + '.meta')
graph = tf.get_default_graph()
tf_model = graph.get_tensor_by_name(u'vertices:0')
with tf.Session() as session:
saver.restore(session, tf_model_fname)
template = Mesh(filename=template_fname)
tf_src = tf.Variable(tf.zeros(template.v.shape, dtype=tf.float64))
# get all params
tf_trans = [x for x in tf.trainable_variables() if 'trans' in x.name][0]
tf_rot = [x for x in tf.trainable_variables() if 'rot' in x.name][0]
tf_pose = [x for x in tf.trainable_variables() if 'pose' in x.name][0]
tf_shape = [x for x in tf.trainable_variables() if 'shape' in x.name][0]
tf_exp = [x for x in tf.trainable_variables() if 'exp' in x.name][0]
def _save_state(*names, **kwargs):
state = dict()
if "trans" in names: state["trans"] = tf_trans.eval()
if "rot" in names: state["rot"] = tf_rot.eval()
if "pose" in names: state["pose"] = tf_pose.eval()
if "shape" in names: state["shape"] = tf_shape.eval()
if "exp" in names: state["exp"] = tf_exp.eval()
if kwargs.get("set_zero", False):
_zero_state(*names)
return state
def _load_state(state):
ops = []
if "trans" in state: ops.append(tf_trans.assign(state["trans"]))
if "rot" in state: ops.append(tf_rot.assign (state["rot"] ))
if "pose" in state: ops.append(tf_pose.assign (state["pose"] ))
if "shape" in state: ops.append(tf_shape.assign(state["shape"]))
if "exp" in state: ops.append(tf_exp.assign (state["exp"] ))
session.run(ops)
def _zero_state(*names):
ops = []
if "trans" in names: ops.append(tf_trans.assign(tf.zeros_like(tf_trans)))
if "rot" in names: ops.append(tf_rot .assign(tf.zeros_like(tf_rot )))
if "pose" in names: ops.append(tf_pose .assign(tf.zeros_like(tf_pose )))
if "shape" in names: ops.append(tf_shape.assign(tf.zeros_like(tf_shape)))
if "exp" in names: ops.append(tf_exp .assign(tf.zeros_like(tf_exp )))
session.run(ops)
mesh_dist = tf.reduce_sum(tf.square(tf.subtract(tf_model, tf_src)))
neck_pose_reg = tf.reduce_sum(tf.square(tf_pose[:3]))
jaw_pose_reg = tf.reduce_sum(tf.square(tf_pose[3:6]))
eye_pose_reg = tf.reduce_sum(tf.square(tf_pose[6:]))
shape_reg = tf.reduce_sum(tf.square(tf_shape))
exp_reg = tf.reduce_sum(tf.square(tf_exp))
reg_term = (
weight_reg_neck_pos * neck_pose_reg +
weight_reg_jaw_pos * jaw_pose_reg +
weight_reg_eye_pos * eye_pose_reg +
weight_reg_shape * shape_reg +
weight_reg_expr * exp_reg
)
# optimizers
optim_shared_rigid = scipy_pt(
loss=mesh_dist,
var_list=[tf_trans, tf_rot],
method='L-BFGS-B',
options={'disp': 0}
)
optim_shared_all = scipy_pt(
loss=mesh_dist+reg_term,
var_list=[tf_trans, tf_rot, tf_pose, tf_shape, tf_exp],
method='L-BFGS-B',
options={'disp': 0}
)
optim_seq = scipy_pt(
loss=mesh_dist+reg_term,
var_list=[tf_shape, tf_exp],
method='L-BFGS-B', options={'disp': 0, 'maxiter': 50}
)
def _fit_sentence(src_dir, dst_dir, prm_dir, last_speaker):
_anchor = os.path.join(dst_dir, "_anchor")
if os.path.exists(_anchor):
print("- Skip " + src_dir)
return
if not os.path.exists(src_dir):
print("- Failed to find " + src_dir)
return
if not os.path.exists(dst_dir): os.makedirs(dst_dir)
if not os.path.exists(prm_dir): os.makedirs(prm_dir)
ply_files = []
for root, _, files in os.walk(src_dir):
for f in files:
if os.path.splitext(f)[1] == ".ply":
ply_files.append(os.path.join(root, f))
ply_files = sorted(ply_files)
# get shared
src_mesh = Mesh(filename=ply_files[0])
session.run(tf.assign(tf_src, src_mesh.v))
speaker = os.path.basename(os.path.dirname(src_dir))
if last_speaker != speaker:
print("- clear speaker information")
_zero_state("trans", "rot", "pose", "shape", "exp")
else:
_zero_state("exp")
stt_dir = os.path.join(os.path.dirname(dst_dir), "state")
if os.path.exists(stt_dir):
state_dict = dict(
trans = np.load(os.path.join(stt_dir, "trans.npy")),
rot = np.load(os.path.join(stt_dir, "rot.npy")),
pose = np.load(os.path.join(stt_dir, "pose.npy")),
shape = np.load(os.path.join(stt_dir, "shape.npy")),
)
_load_state(state_dict)
fitting_mesh = Mesh(session.run(tf_model), src_mesh.f)
fitting_mesh.write_ply(os.path.join(stt_dir, "zero.ply"))
fit_zero_dir = os.path.join(os.path.dirname(os.path.dirname(dst_dir)), "zero_exp")
if not os.path.exists(fit_zero_dir): os.makedirs(fit_zero_dir)
print("- " + speaker + " " + os.path.basename(src_dir))
print(" -> fit shared parameters...")
optim_shared_rigid.minimize(session)
optim_shared_all.minimize(session)
state_dict = _save_state("exp", set_zero=True)
fitting_mesh = Mesh(session.run(tf_model), src_mesh.f)
fitting_mesh.write_ply(os.path.join(fit_zero_dir, "{}.ply".format(speaker)))
_load_state(state_dict)
return
if not os.path.exists(stt_dir): os.makedirs(stt_dir)
np.save(os.path.join(stt_dir, "trans.npy"), tf_trans.eval(), allow_pickle=False)
np.save(os.path.join(stt_dir, "rot.npy"), tf_rot.eval(), allow_pickle=False)
np.save(os.path.join(stt_dir, "pose.npy"), tf_pose.eval(), allow_pickle=False)
np.save(os.path.join(stt_dir, "shape.npy"), tf_shape.eval(), allow_pickle=False)
progress = tqdm(ply_files)
for src_fname in progress:
frame = os.path.basename(src_fname)
progress.set_description(" -> " + frame)
dst_fname = os.path.join(dst_dir, frame)
# param filename
prm_fname = os.path.join(prm_dir, frame)
exp_fname = os.path.splitext(prm_fname)[0] + '_exp.npy'
idn_fname = os.path.splitext(prm_fname)[0] + '_idn.npy'
src_mesh = Mesh(filename=src_fname)
session.run(tf.assign(tf_src, src_mesh.v))
optim_seq.minimize(session)
# save expr
np.save(exp_fname, tf_exp.eval())
np.save(idn_fname, tf_shape.eval())
# state_dict = _save_state("trans", "rot", "pose", "shape", set_zero=True)
# save mesh
fitting_mesh = Mesh(session.run(tf_model), src_mesh.f)
fitting_mesh.write_ply(dst_fname)
# _load_state(state_dict)
# print(tf_shape.eval())
if showing:
g_mv.set_static_meshes([fitting_mesh])
os.system("touch {}".format(_anchor))
return speaker
last_speaker = None
for (src, dst, prm) in dir_tup_list:
last_speaker = _fit_sentence(src, dst, prm, last_speaker)
def cage(length=1, vc=name_to_rgb['black']):
cage_points = np.array([[-1., -1., -1.], [1., 1., 1.], [1., -1., 1.],
[-1., 1., -1.]])
c = Mesh(v=length * cage_points, f=[], vc=vc)
return c
col[(verts[:, 1] < -1.14) & (verts[:, 0] >= 0)] = 1
if display:
ms.set_vertex_colors_from_weights(col)
ms.show()
print('left_foot ', np.where(col)[0].shape)
return col
if __name__ == "__main__":
smplx = get_tpose_smplx()
verts = smplx.vertices.detach().cpu().numpy().squeeze()
faces = smplx.faces
ms = Mesh(v=verts, f=faces)
col = np.zeros((10475, ))
display = False
rfa = cut_right_forearm(display)
col += (rfa * 0.1)
rma = cut_right_midarm(display)
col += (rma * 0.2)
lfa = cut_left_forearm(display)
col += (lfa * 0.3)
lma = cut_left_midarm(display)
col += (lma * 0.4)
avail_items = f.read().splitlines()
avail_items = [k.split('\t') for k in avail_items]
people_names = [k[0] for k in avail_items if k[1] == garment_class]
shape_root = os.path.join(global_var.ROOT, 'neutral_shape_static_pose_new')
smoothing = None
shape_names = ["{:02d}".format(k) for k in range(0, 100)]
for people_name, garment_class in tqdm(avail_items):
shape_static_pose_people = os.path.join(shape_root, people_name)
for shape_name in shape_names:
garment_path = os.path.join(shape_static_pose_people, '{}_{}.obj'.format(shape_name, garment_class))
if not os.path.exists(garment_path):
print("{} doesn't exist".format(garment_path))
try:
m = Mesh(filename=garment_path)
except AttributeError as e:
print(e)
print(garment_path)
exit()
if smoothing is None:
smoothing = DiffusionSmoothing(m.v, m.f, Ltype="cotangent")
steps = [30]
verts_smooth = m.v.copy()
for i, step in enumerate(steps):
smooth_name = '_sm{}'.format(step)
dst_path = os.path.join(shape_static_pose_people, '{}{}_{}.obj'.format(shape_name, smooth_name, garment_class))
if os.path.exists(dst_path):
print("{} exists. Skip".format(dst_path))
data = np.ones(nverts)
data[indices] *= ww
I = csr_matrix((data, (rc, rc)), shape=(nverts, nverts))
A = vstack([L, I])
b = np.vstack((L.dot(mesh.v), tgt_points))
res = spsolve(A.T.dot(A), A.T.dot(b))
mres = Mesh(v=res, f=mesh.f)
return mres
if __name__ == '__main__':
import os
ROOT = "/BS/cpatel/work/data/learn_anim/mixture_exp31/000_0/smooth_TShirtNoCoat/0990/"
body = Mesh(filename=os.path.join(ROOT, "body_160.ply"))
mesh = Mesh(filename=os.path.join(ROOT, "pred_160.ply"))
mesh1 = remove_interpenetration_fast(mesh, body)
mesh1.write_ply("/BS/cpatel/work/proccessed.ply")
mesh.write_ply("/BS/cpatel/work/orig.ply")
body.write_ply("/BS/cpatel/work/body.ply")
# from psbody.mesh import MeshViewers
# mvs = MeshViewers((1, 2))
# mesh1.set_vertex_colors_from_weights(np.linalg.norm(mesh.v - mesh1.v, axis=1))
# mesh.set_vertex_colors_from_weights(np.linalg.norm(mesh.v - mesh1.v, axis=1))
# # mesh1.set_vertex_colors_from_weights(np.zeros(mesh.v.shape[0]))
# # mesh.set_vertex_colors_from_weights(np.zeros(mesh.v.shape[0]))
# mvs[0][0].set_static_meshes([mesh, body])
# mvs[0][1].set_static_meshes([mesh1, body])
def run_tailornet():
gender = 'male'
garment_class = 'pant'
#garment_class = 'shirt'
garment_combine = True
garment_class_pairs = {
#'pant': ['shirt', 't-shirt'],
'pant': ['shirt'],
'short-pant': ['shirt', 't-shirt'],
'skirt': ['shirt', 't-shirt']
}
thetas, betas, gammas = get_single_frame_inputs(garment_class, gender)
# # uncomment the line below to run inference on sequence data
# thetas, betas, gammas = get_sequence_inputs(garment_class, gender)
# load model
tn_runner = get_tn_runner(gender=gender, garment_class=garment_class)
# from trainer.base_trainer import get_best_runner
# tn_runner = get_best_runner("/BS/cpatel/work/data/learn_anim/tn_baseline/{}_{}/".format(garment_class, gender))
smpl = SMPL4Garment(gender=gender)
# make out directory if doesn't exist
if not os.path.isdir(OUT_PATH):
os.mkdir(OUT_PATH)
# run inference
for i, (theta, beta, gamma) in enumerate(zip(thetas, betas, gammas)):
print(i, len(thetas))
# normalize y-rotation to make it front facing
theta_normalized = normalize_y_rotation(theta)
with torch.no_grad():
pred_verts_d = tn_runner.forward(
thetas=torch.from_numpy(theta_normalized[None, :].astype(
np.float32)).cuda(),
betas=torch.from_numpy(beta[None, :].astype(
np.float32)).cuda(),
gammas=torch.from_numpy(gamma[None, :].astype(
np.float32)).cuda(),
)[0].cpu().numpy()
# get garment from predicted displacements
body, pred_gar = smpl.run(beta=beta,
theta=theta,
garment_class=garment_class,
garment_d=pred_verts_d)
# gar_pair_hres = Mesh(
# filename=os.path.join(OUT_PATH,
# "{}_gar_hres_{}_{:04d}.obj".format(gender, garment_class_pairs[garment_class][0], i)))
# gar_pair_hres = remove_interpenetration_fast(gar_pair_hres, body)
# pred_body = remove_interpenetration_fast_custom(body, gar_pair_hres, threshold=0.0010)
# pred_body.write_obj(os.path.join(OUT_PATH, 'deformed_body_shirt_0001.obj'))
#
pred_gar = remove_interpenetration_fast(pred_gar, body)
pred_body, _ = remove_interpenetration_fast_custom(body,
pred_gar,
threshold=0.0005)
pred_body.write_obj(
os.path.join(OUT_PATH, 'deformed_body_pant_shirt.obj'))
hv, hf, mapping = get_hres(pred_gar.v, pred_gar.f)
pred_gar_hres = Mesh(hv, hf)
if garment_combine:
for garment_class_pair in garment_class_pairs[garment_class]:
gar_pair_hres = Mesh(filename=os.path.join(
OUT_PATH, "{}_gar_hres_{}_{:04d}.obj".format(
gender, garment_class_pair, i)))
print(len(body.v), body.f.min(), body.f.max())
print(len(pred_gar_hres.v), pred_gar_hres.f.min(),
pred_gar_hres.f.max())
print(
np.vstack((body.v, pred_gar_hres.v)).shape,
np.min(pred_gar_hres.f + len(body.v)),
np.max(pred_gar_hres.f + len(body.v)))
#gar_pair_hres = remove_interpenetration_fast(gar_pair_hres, pred_gar_hres)
for _ in range(3):
gar_pair_hres = remove_interpenetration_fast(
gar_pair_hres, pred_body)
for _ in range(3):
gar_pair_hres = remove_interpenetration_fast(
gar_pair_hres,
Mesh(
np.vstack((body.v, pred_gar_hres.v)),
np.vstack(
(body.f, pred_gar_hres.f + len(body.v)))))
pred_gar_hres_inverse, _ = remove_interpenetration_fast_custom(
pred_gar_hres,
Mesh(np.vstack((body.v, gar_pair_hres.v)),
np.vstack((body.f, gar_pair_hres.f + len(body.v)))),
threshold=0.000005,
inverse=True)
gar_pair_hres.write_obj(
os.path.join(
OUT_PATH,
"new2_{}_gar_hres_{}_inter_{}_{:04d}.obj".format(
gender, garment_class_pair, garment_class, i)))
# save body and predicted garment
body.write_obj(
os.path.join(OUT_PATH, "{}_body_{:04d}.obj".format(gender, i)))
pred_gar.write_obj(
os.path.join(
OUT_PATH,
"{}_gar_{}_{:04d}.obj".format(gender, garment_class, i)))
pred_gar_hres.write_obj(
os.path.join(
OUT_PATH,
"{}_gar_hres_{}_{:04d}.obj".format(gender, garment_class, i)))
pred_gar_hres_inverse.write_obj(
os.path.join(
OUT_PATH, "{}_gar_hres_{}_{:04d}_inverse.obj".format(
gender, garment_class, i)))
def vis_results(ho,
dorig,
coarse_net,
refine_net,
rh_model,
save=False,
save_dir=None,
rh_model_pkl=None,
vis=True):
# with torch.no_grad():
imw, imh = 1920, 780
cols = len(dorig['bps_object'])
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')
if vis:
mvs = MeshViewers(window_width=imw,
window_height=imh,
shape=[1, cols],
keepalive=True)
# drec_cnet = coarse_net.sample_poses(dorig['bps_object'])
#
# for k in drec_cnet.keys():
# print('drec cnet', k, drec_cnet[k].shape)
# verts_rh_gen_cnet = rh_model(**drec_cnet).vertices
drec_cnet = {}
hand_pose_in = torch.Tensor(ho.hand_pose[3:]).unsqueeze(0)
mano_out_1 = rh_model_pkl(hand_pose=hand_pose_in)
hand_pose_in = mano_out_1.hand_pose
mTc = torch.Tensor(ho.hand_mTc)
approx_global_orient = rotmat2aa(mTc[:3, :3].unsqueeze(0))
if torch.isnan(approx_global_orient).any(): # Using honnotate?
approx_global_orient = torch.Tensor(ho.hand_pose[:3]).unsqueeze(0)
approx_global_orient = approx_global_orient.squeeze(1).squeeze(1)
approx_trans = mTc[:3, 3].unsqueeze(0)
target_verts = torch.Tensor(ho.hand_verts).unsqueeze(0)
pose, trans, rot = util.opt_hand(rh_model, target_verts, hand_pose_in,
approx_trans, approx_global_orient)
# drec_cnet['hand_pose'] = torch.einsum('bi,ij->bj', [hand_pose_in, rh_model_pkl.hand_components])
drec_cnet['transl'] = trans
drec_cnet['global_orient'] = rot
drec_cnet['hand_pose'] = pose
verts_rh_gen_cnet = rh_model(**drec_cnet).vertices
_, h2o, _ = point2point_signed(verts_rh_gen_cnet,
dorig['verts_object'].to(device))
drec_cnet['trans_rhand_f'] = drec_cnet['transl']
drec_cnet['global_orient_rhand_rotmat_f'] = aa2rotmat(
drec_cnet['global_orient']).view(-1, 3, 3)
drec_cnet['fpose_rhand_rotmat_f'] = aa2rotmat(drec_cnet['hand_pose']).view(
-1, 15, 3, 3)
drec_cnet['verts_object'] = dorig['verts_object'].to(device)
drec_cnet['h2o_dist'] = h2o.abs()
print(
'Hand fitting err',
np.linalg.norm(
verts_rh_gen_cnet.squeeze().detach().numpy() - ho.hand_verts, 2,
1).mean())
orig_obj = dorig['mesh_object'][0].v
# print(orig_obj.shape, orig_obj)
# print('Obj fitting err', np.linalg.norm(orig_obj - ho.obj_verts, 2, 1).mean())
drec_rnet = refine_net(**drec_cnet)
mano_out = rh_model(**drec_rnet)
verts_rh_gen_rnet = mano_out.vertices
joints_out = mano_out.joints
if vis:
for cId in range(0, len(dorig['bps_object'])):
try:
from copy import deepcopy
meshes = deepcopy(dorig['mesh_object'])
obj_mesh = meshes[cId]
except:
obj_mesh = points_to_spheres(to_cpu(
dorig['verts_object'][cId]),
radius=0.002,
vc=name_to_rgb['green'])
hand_mesh_gen_cnet = Mesh(v=to_cpu(verts_rh_gen_cnet[cId]),
f=rh_model.faces,
vc=name_to_rgb['pink'])
hand_mesh_gen_rnet = Mesh(v=to_cpu(verts_rh_gen_rnet[cId]),
f=rh_model.faces,
vc=name_to_rgb['gray'])
if 'rotmat' in dorig:
rotmat = dorig['rotmat'][cId].T
obj_mesh = obj_mesh.rotate_vertices(rotmat)
hand_mesh_gen_cnet.rotate_vertices(rotmat)
hand_mesh_gen_rnet.rotate_vertices(rotmat)
# print('rotmat', rotmat)
hand_mesh_gen_cnet.reset_face_normals()
hand_mesh_gen_rnet.reset_face_normals()
# mvs[0][cId].set_static_meshes([hand_mesh_gen_cnet] + obj_mesh, blocking=True)
# mvs[0][cId].set_static_meshes([hand_mesh_gen_rnet,obj_mesh], blocking=True)
mvs[0][cId].set_static_meshes(
[hand_mesh_gen_rnet, hand_mesh_gen_cnet, obj_mesh],
blocking=True)
if save:
save_path = os.path.join(save_dir, str(cId))
makepath(save_path)
hand_mesh_gen_rnet.write_ply(filename=save_path +
'/rh_mesh_gen_%d.ply' % cId)
obj_mesh[0].write_ply(filename=save_path +
'/obj_mesh_%d.ply' % cId)
return verts_rh_gen_rnet, joints_out
def estimate_global_pose(landmarks,
key_vids,
model,
cam,
img,
fix_t=False,
viz=False,
SOLVE_FLATER=True):
'''
Estimates the global rotation and translation.
only diff in estimate_global_pose from single_frame_ferrari is that all animals have the same kp order.
'''
# Redefining part names..
part_names = [
'leftEye', 'rightEye', 'chin', 'frontLeftFoot', 'frontRightFoot',
'backLeftFoot', 'backRightFoot', 'tailStart', 'frontLeftKnee',
'frontRightKnee', 'backLeftKnee', 'backRightKnee', 'leftShoulder',
'rightShoulder', 'frontLeftAnkle', 'frontRightAnkle', 'backLeftAnkle',
'backRightAnkle', 'neck', 'TailTip'
]
# Use shoulder to "knee"(elbow) distance. also tail to "knee" if available.
use_names = [
'neck', 'leftShoulder', 'rightShoulder', 'backLeftKnee',
'backRightKnee', 'tailStart', 'frontLeftKnee', 'frontRightKnee'
]
use_ids = [part_names.index(name) for name in use_names]
# These might not be visible
visible = landmarks[:, 2].astype(bool)
use_ids = [id for id in use_ids if visible[id]]
if len(use_ids) < 3:
print('Frontal?..')
use_names += [
'frontLeftAnkle', 'frontRightAnkle', 'backLeftAnkle',
'backRightAnkle'
]
model.pose[1] = np.pi / 2
init_t = estimate_translation(landmarks, key_vids, cam.f[0].r, model)
use_ids = [part_names.index(name) for name in use_names]
use_ids = [id for id in use_ids if visible[id]]
# Setup projection error:
all_vids = np.hstack([key_vids[id] for id in use_ids])
cam.v = model[all_vids]
keypoints = landmarks[use_ids, :2].astype(float)
# Duplicate keypoints for the # of vertices for that kp.
num_verts_per_kp = [len(key_vids[row_id]) for row_id in use_ids]
j2d = np.vstack([
np.tile(kp, (num_rep, 1))
for kp, num_rep in zip(keypoints, num_verts_per_kp)
])
assert (cam.r.shape == j2d.shape)
# SLOW but correct method!!
# remember which ones belongs together,,
group = np.hstack([
index * np.ones(len(key_vids[row_id]))
for index, row_id in enumerate(use_ids)
])
assignments = np.vstack([group == i for i in np.arange(group[-1] + 1)])
num_points = len(use_ids)
proj_error = (ch.vstack([
cam[choice] if np.sum(choice) == 1 else cam[choice].mean(axis=0)
for choice in assignments
]) - keypoints) / np.sqrt(num_points)
# Fast but not matching average:
# Normalization weight
j2d_norm_weights = np.sqrt(
1. / len(use_ids) *
np.vstack([1. / num * np.ones((num, 1)) for num in num_verts_per_kp]))
proj_error_fast = j2d_norm_weights * (cam - j2d)
if fix_t:
obj = {'cam': proj_error_fast}
else:
obj = {
'cam': proj_error_fast,
'cam_t': 1e1 * (model.trans[2] - init_t[2])
}
# Only estimate body orientation
if fix_t:
free_variables = [model.pose[:3]]
else:
free_variables = [model.trans, model.pose[:3]]
if not SOLVE_FLATER:
obj['feq'] = 1e3 * (cam.f[0] - cam.f[1])
# So it's under control
obj['freg'] = 1e1 * (cam.f[0] - 3000) / 1000.
# here without this cam.f goes negative.. (asking margin of 500)
obj['fpos'] = ch.maximum(0, 500 - cam.f[0])
# cam t also has to be positive!
obj['cam_t_pos'] = ch.maximum(0, 0.01 - model.trans[2])
del obj['cam_t']
free_variables.append(cam.f)
if viz:
import matplotlib.pyplot as plt
plt.ion()
def on_step(_):
plt.figure(1, figsize=(5, 5))
plt.cla()
plt.imshow(img[:, :, ::-1])
img_here = render_mesh(Mesh(model.r, model.f), img.shape[1],
img.shape[0], cam)
plt.imshow(img_here)
plt.scatter(j2d[:, 0], j2d[:, 1], c='w')
plt.scatter(cam.r[:, 0], cam.r[:, 1])
plt.draw()
plt.pause(1e-3)
if 'feq' in obj:
print('flength %.1f %.1f, z %.f' %
(cam.f[0], cam.f[1], model.trans[2]))
else:
on_step = None
from time import time
t0 = time()
init_angles = [[0, 0, 0]] #, [1.5,0,0], [1.5,-1.,0]]
scores = np.zeros(len(init_angles))
for i, angle in enumerate(init_angles):
# Init translation
model.trans[:] = init_t
model.pose[:3] = angle
ch.minimize(obj,
x0=free_variables,
method='dogleg',
callback=on_step,
options={
'maxiter': 100,
'e_3': .0001
})
scores[i] = np.sum(obj['cam'].r**2.)
j = np.argmin(scores)
model.trans[:] = init_t
model.pose[:3] = init_angles[j]
ch.minimize(obj,
x0=free_variables,
method='dogleg',
callback=on_step,
options={
'maxiter': 100,
'e_3': .0001
})
print('Took %g' % (time() - t0))
#import pdb; pdb.set_trace()
if viz:
dist = np.mean(model.r, axis=0)[2]
img_here = render_mesh(Mesh(model.r, model.f), img.shape[1],
img.shape[0], cam)
plt.imshow(img[:, :, ::-1])
plt.imshow(img_here)
return model.pose[:3].r, model.trans.r
np.float32)).unsqueeze(0)
betas_torch = torch.from_numpy(beta[:10].astype(
np.float32)).unsqueeze(0)
smpl_verts = smpl_torch.forward(pose_torch, betas_torch)
transform = smpl_torch.A.detach().numpy()[0]
transform_inv = np.array(
[np.linalg.inv(transform[i]) for i in range(24)])
#rotate root joint using inverse transform
joints = smpl_torch.J_transformed.detach().numpy()[0]
root_joint = np.array([joints[0, 0], joints[0, 1], joints[0, 2], 1])
transformed_root = np.array(
[np.matmul(transform_inv[i], root_joint)[:3] for i in range(24)])
m1 = Mesh(v=smpl_verts.detach().numpy()[0], f=smpl_torch.faces)
mesh = trimesh.Trimesh(m1.v, smpl_torch.faces)
#sample points on mesh surface and displace with sigma = 0.03
boundary_points = []
points = mesh.sample(sample_num)
boundary_points_1 = points + 0.03 * np.random.randn(sample_num, 3)
#sample points in bbox of 110%
bottom_corner, upper_corner = mesh.bounds
bottom_corner = bottom_corner + 0.1 * bottom_corner
upper_corner = upper_corner + 0.1 * upper_corner
x_pt = np.random.uniform(bottom_corner[0], upper_corner[0], sample_num)
y_pt = np.random.uniform(bottom_corner[1], upper_corner[1], sample_num)
z_pt = np.random.uniform(bottom_corner[2], upper_corner[2], sample_num)
boundary_points_2 = np.concatenate(
args = parser.parse_args()
if args.conf is None:
args.conf = os.path.join(os.path.dirname(__file__), 'default.cfg')
print(
'configuration file not specified, trying to load '
'it from current directory', args.conf)
if not os.path.exists(args.conf):
print('Config not found' + args.conf)
config = read_config(args.conf)
print('Initializing parameters')
template_file_path = config['template_fname']
template_mesh = Mesh(filename=template_file_path)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Generating transforms')
M, A, D, U = mesh_operations.generate_transform_matrices(
template_mesh, config['downsampling_factors'])
D_t = [scipy_to_torch_sparse(d).to(device) for d in D]
U_t = [scipy_to_torch_sparse(u).to(device) for u in U]
A_t = [scipy_to_torch_sparse(a).to(device) for a in A]
num_nodes = [len(M[i].v) for i in range(len(M))]
print('Loading Dataset')
if args.data_dir:
data_dir = args.data_dir
def process(self):
train_data, val_data, test_data = [], [], []
train_vertices, val_vertices, test_vertices = [], [], []
for idx, data_file in tqdm(enumerate(self.data_file)):
mesh = Mesh(filename=data_file)
mesh_verts = torch.Tensor(mesh.v)
adjacency = get_vert_connectivity(mesh.v, mesh.f).tocoo()
# edge_index = torch.LongTensor(np.vstack((adjacency.row, adjacency.col)))
edge_index = torch.Tensor(np.vstack((adjacency.row, adjacency.col)))
data = Data(x=mesh_verts, y=mesh_verts, edge_index=edge_index)
if self.split == 'sliced':
if idx % 100 <= 10:
test_data.append(data)
test_vertices.append(mesh.v)
elif idx % 100 <= 20:
val_data.append(data)
val_vertices.append(mesh.v)
else:
train_data.append(data)
train_vertices.append(mesh.v)
elif self.split == 'expression':
if data_file.split('/')[-2] == self.split_term:
test_data.append(data)
test_vertices.append(mesh.v)
else:
train_data.append(data)
train_vertices.append(mesh.v)
elif self.split == 'identity':
if data_file.split('/')[-3] == self.split_term:
test_data.append(data)
test_vertices.append(mesh.v)
else:
train_data.append(data)
train_vertices.append(mesh.v)
elif self.split == 'custom':
test_terms = self.split_term.split('-')
test_flag = False
for term in test_terms:
if term in data_file:
test_flag = True
if test_flag is True:
# print('testing set: {}'.format(data_file))
test_data.append(data)
test_vertices.append(mesh.v)
else:
# print('training set: {}'.format(data_file))
train_data.append(data)
train_vertices.append(mesh.v)
else:
raise Exception('sliced, expression and identity are the only supported split terms')
if self.split != 'sliced':
val_data = test_data[-self.nVal:]
test_data = test_data[:-self.nVal]
val_vertices = test_vertices[-self.nVal:]
test_vertices = test_vertices[:-self.nVal]
mean_train = torch.Tensor(np.mean(train_vertices, axis=0))
std_train = torch.Tensor(np.std(train_vertices, axis=0))
norm_dict = {'mean': mean_train, 'std': std_train}
if self.pre_transform is not None:
print('Transforming data...')
if hasattr(self.pre_transform, 'mean') and hasattr(self.pre_transform, 'std'):
if self.pre_tranform.mean is None:
self.pre_tranform.mean = mean_train
if self.pre_transform.std is None:
self.pre_tranform.std = std_train
train_data = [self.pre_transform(td) for td in train_data]
val_data = [self.pre_transform(td) for td in val_data]
test_data = [self.pre_transform(td) for td in test_data]
print('Saving data...')
torch.save(self.collate(train_data), self.processed_paths[0])
torch.save(self.collate(val_data), self.processed_paths[1])
torch.save(self.collate(test_data), self.processed_paths[2])
torch.save(norm_dict, self.processed_paths[3])
self.save_vertices(np.array(train_vertices), np.array(val_vertices), np.array(test_vertices))
def fit_lmk2d(target_img, target_2d_lmks, model_fname, lmk_face_idx, lmk_b_coords, weights, visualize):
'''
Fit FLAME to 2D landmarks
:param target_2d_lmks target 2D landmarks provided as (num_lmks x 3) matrix
:param model_fname saved FLAME model
:param lmk_face_idx face indices of the landmark embedding in the FLAME topology
:param lmk_b_coords barycentric coordinates of the landmark embedding in the FLAME topology
(i.e. weighting of the three vertices for the trinagle, the landmark is embedded in
:param weights weights of the individual objective functions
:param visualize visualize fitting progress
:return: a mesh with the fitting results
'''
'''
pred_types = {'face': pred_type(slice(0, 17), (0.682, 0.780, 0.909, 0.5)),
'eyebrow1': pred_type(slice(17, 22), (1.0, 0.498, 0.055, 0.4)),
'eyebrow2': pred_type(slice(22, 27), (1.0, 0.498, 0.055, 0.4)),
'nose': pred_type(slice(27, 31), (0.345, 0.239, 0.443, 0.4)),
'nostril': pred_type(slice(31, 36), (0.345, 0.239, 0.443, 0.4)),
'eye1': pred_type(slice(36, 42), (0.596, 0.875, 0.541, 0.3)),
'eye2': pred_type(slice(42, 48), (0.596, 0.875, 0.541, 0.3)),
'lips': pred_type(slice(48, 60), (0.596, 0.875, 0.541, 0.3)),
'teeth': pred_type(slice(60, 68), (0.596, 0.875, 0.541, 0.4))
}
'''
lmks_weights = [[1,1]] * 68
for idx in range(36, 48):
lmks_weights[idx] = [100, 100]
tf_lmks_weights = tf.constant(
lmks_weights,
tf.float64
)
tf_trans = tf.Variable(np.zeros((1,3)), name="trans", dtype=tf.float64, trainable=True)
tf_rot = tf.Variable(np.zeros((1,3)), name="rot", dtype=tf.float64, trainable=True)
tf_pose = tf.Variable(np.zeros((1,12)), name="pose", dtype=tf.float64, trainable=True)
tf_shape = tf.Variable(np.zeros((1,300)), name="shape", dtype=tf.float64, trainable=True)
tf_exp = tf.Variable(np.zeros((1,100)), name="expression", dtype=tf.float64, trainable=True)
smpl = SMPL(model_fname)
tf_model = tf.squeeze(smpl(tf_trans,
tf.concat((tf_shape, tf_exp), axis=-1),
tf.concat((tf_rot, tf_pose), axis=-1)))
with tf.Session() as session:
# session.run(tf.global_variables_initializer())
# Mirror landmark y-coordinates
target_2d_lmks[:,1] = target_img.shape[0]-target_2d_lmks[:,1]
lmks_3d = tf_get_model_lmks(tf_model, smpl.f, lmk_face_idx, lmk_b_coords)
s2d = np.mean(np.linalg.norm(target_2d_lmks-np.mean(target_2d_lmks, axis=0), axis=1))
s3d = tf.reduce_mean(tf.sqrt(tf.reduce_sum(tf.square(lmks_3d-tf.reduce_mean(lmks_3d, axis=0))[:, :2], axis=1)))
tf_scale = tf.Variable(s2d/s3d, dtype=lmks_3d.dtype)
# trans = 0.5*np.array((target_img.shape[0], target_img.shape[1]))/tf_scale
# trans = 0.5 * s3d * np.array((target_img.shape[0], target_img.shape[1])) / s2d
lmks_proj_2d = tf_project_points(lmks_3d, tf_scale, np.zeros(2))
factor = max(max(target_2d_lmks[:,0]) - min(target_2d_lmks[:,0]),max(target_2d_lmks[:,1]) - min(target_2d_lmks[:,1]))
#lmk_dist = weights['lmk']*tf.reduce_sum(tf.square(tf.subtract(lmks_proj_2d, target_2d_lmks))) / (factor ** 2)
lmk_dist = weights['lmk']*tf.reduce_sum(
tf.square(tf.subtract(lmks_proj_2d, target_2d_lmks)) * tf_lmks_weights
) / (factor ** 2)
neck_pose_reg = weights['neck_pose']*tf.reduce_sum(tf.square(tf_pose[:,:3]))
jaw_pose_reg = weights['jaw_pose']*tf.reduce_sum(tf.square(tf_pose[:,3:6]))
eyeballs_pose_reg = weights['eyeballs_pose']*tf.reduce_sum(tf.square(tf_pose[:,6:]))
shape_reg = weights['shape']*tf.reduce_sum(tf.square(tf_shape))
exp_reg = weights['expr']*tf.reduce_sum(tf.square(tf_exp))
session.run(tf.global_variables_initializer())
if visualize:
def on_step(verts, scale, faces, target_img, target_lmks, opt_lmks, lmk_dist=0.0, shape_reg=0.0, exp_reg=0.0, neck_pose_reg=0.0, jaw_pose_reg=0.0, eyeballs_pose_reg=0.0):
import cv2
import sys
import numpy as np
from psbody.mesh import Mesh
from utils.render_mesh import render_mesh
if lmk_dist>0.0 or shape_reg>0.0 or exp_reg>0.0 or neck_pose_reg>0.0 or jaw_pose_reg>0.0 or eyeballs_pose_reg>0.0:
print('lmk_dist: %f, shape_reg: %f, exp_reg: %f, neck_pose_reg: %f, jaw_pose_reg: %f, eyeballs_pose_reg: %f' % (lmk_dist, shape_reg, exp_reg, neck_pose_reg, jaw_pose_reg, eyeballs_pose_reg))
plt_target_lmks = target_lmks.copy()
plt_target_lmks[:, 1] = target_img.shape[0] - plt_target_lmks[:, 1]
for (x, y) in plt_target_lmks:
cv2.circle(target_img, (int(x), int(y)), 4, (0, 0, 255), -1)
plt_opt_lmks = opt_lmks.copy()
plt_opt_lmks[:,1] = target_img.shape[0] - plt_opt_lmks[:,1]
for (x, y) in plt_opt_lmks:
cv2.circle(target_img, (int(x), int(y)), 4, (255, 0, 0), -1)
if sys.version_info >= (3, 0):
rendered_img = render_mesh(Mesh(scale*verts, faces), height=target_img.shape[0], width=target_img.shape[1])
for (x, y) in plt_opt_lmks:
cv2.circle(rendered_img, (int(x), int(y)), 4, (255, 0, 0), -1)
target_img = np.hstack((target_img, rendered_img))
#cv2.imshow('img', target_img)
#cv2.waitKey(10)
else:
def on_step(*_):
pass
print('Optimize rigid transformation')
vars = [tf_scale, tf_trans, tf_rot]
loss = lmk_dist
optimizer = scipy_pt(loss=loss, var_list=vars, method='L-BFGS-B', options={'disp': 1, 'ftol': 5e-6})
optimizer.minimize(session, fetches=[tf_model, tf_scale, tf.constant(smpl.f), tf.constant(target_img), tf.constant(target_2d_lmks), lmks_proj_2d], loss_callback=on_step)
print('Optimize model parameters')
vars = [tf_scale, tf_trans[:2], tf_rot, tf_pose, tf_shape, tf_exp]
loss = lmk_dist + shape_reg + exp_reg + neck_pose_reg + jaw_pose_reg + eyeballs_pose_reg
optimizer = scipy_pt(loss=loss, var_list=vars, method='L-BFGS-B', options={'disp': 0, 'ftol': 1e-7})
optimizer.minimize(session, fetches=[tf_model, tf_scale, tf.constant(smpl.f), tf.constant(target_img), tf.constant(target_2d_lmks), lmks_proj_2d,
lmk_dist, shape_reg, exp_reg, neck_pose_reg, jaw_pose_reg, eyeballs_pose_reg], loss_callback=on_step)
print('Fitting done')
np_verts, np_scale = session.run([tf_model, tf_scale])
return Mesh(np_verts, smpl.f), np_scale
def fit_lmk3d(target_3d_lmks,
template_fname,
tf_model_fname,
lmk_face_idx,
lmk_b_coords,
weights,
show_fitting=True):
'''
Fit FLAME to 3D landmarks
:param target_3d_lmks: target 3D landmarks provided as (num_lmks x 3) matrix
:param template_fname: template mesh in FLAME topology (only the face information are used)
:param tf_model_fname: saved Tensorflow FLAME model
:param lmk_face_idx: face indices of the landmark embedding in the FLAME topology
:param lmk_b_coords: barycentric coordinates of the landmark embedding in the FLAME topology
(i.e. weighting of the three vertices for the trinagle, the landmark is embedded in
:param weights: weights of the individual objective functions
:return: a mesh with the fitting results
'''
template_mesh = Mesh(filename=template_fname)
saver = tf.train.import_meta_graph(tf_model_fname + '.meta')
graph = tf.get_default_graph()
tf_model = graph.get_tensor_by_name(u'vertices:0')
with tf.Session() as session:
saver.restore(session, tf_model_fname)
# Workaround as existing tf.Variable cannot be retrieved back with tf.get_variable
# tf_v_template = [x for x in tf.trainable_variables() if 'v_template' in x.name][0]
tf_trans = [x for x in tf.trainable_variables()
if 'trans' in x.name][0]
tf_rot = [x for x in tf.trainable_variables() if 'rot' in x.name][0]
tf_pose = [x for x in tf.trainable_variables() if 'pose' in x.name][0]
tf_shape = [x for x in tf.trainable_variables()
if 'shape' in x.name][0]
tf_exp = [x for x in tf.trainable_variables() if 'exp' in x.name][0]
lmks = tf_get_model_lmks(tf_model, template_mesh, lmk_face_idx,
lmk_b_coords)
lmk_dist = tf.reduce_sum(
tf.square(1000 * tf.subtract(lmks, target_3d_lmks)))
neck_pose_reg = tf.reduce_sum(tf.square(tf_pose[:3]))
jaw_pose_reg = tf.reduce_sum(tf.square(tf_pose[3:6]))
eyeballs_pose_reg = tf.reduce_sum(tf.square(tf_pose[6:]))
shape_reg = tf.reduce_sum(tf.square(tf_shape))
exp_reg = tf.reduce_sum(tf.square(tf_exp))
# Optimize global transformation first
vars = [tf_trans, tf_rot]
loss = weights['lmk'] * lmk_dist
optimizer = scipy_pt(loss=loss,
var_list=vars,
method='L-BFGS-B',
options={
'disp': 1,
'ftol': 5e-6
})
print('Optimize rigid transformation')
optimizer.minimize(session)
# Optimize for the model parameters
vars = [tf_trans, tf_rot, tf_pose, tf_shape, tf_exp]
loss = weights['lmk'] * lmk_dist + weights['shape'] * shape_reg + weights['expr'] * exp_reg + \
weights['neck_pose'] * neck_pose_reg + weights['jaw_pose'] * jaw_pose_reg + weights['eyeballs_pose'] * eyeballs_pose_reg
optimizer = scipy_pt(loss=loss,
var_list=vars,
method='L-BFGS-B',
options={
'disp': 1,
'ftol': 5e-6
})
print('Optimize model parameters')
optimizer.minimize(session)
print('Fitting done')
if show_fitting:
# Visualize landmark fitting
mv = MeshViewer()
mv.set_static_meshes(
create_lmk_spheres(target_3d_lmks, 0.001, [255.0, 0.0, 0.0]))
mv.set_dynamic_meshes(
[Mesh(session.run(tf_model), template_mesh.f)] +
create_lmk_spheres(session.run(lmks), 0.001,
[0.0, 0.0, 255.0]),
blocking=True)
six.moves.input('Press key to continue')
return Mesh(session.run(tf_model), template_mesh.f)
def _fit_sentence(src_dir, dst_dir, prm_dir, last_speaker):
_anchor = os.path.join(dst_dir, "_anchor")
if os.path.exists(_anchor):
print("- Skip " + src_dir)
return
if not os.path.exists(src_dir):
print("- Failed to find " + src_dir)
return
if not os.path.exists(dst_dir): os.makedirs(dst_dir)
if not os.path.exists(prm_dir): os.makedirs(prm_dir)
ply_files = []
for root, _, files in os.walk(src_dir):
for f in files:
if os.path.splitext(f)[1] == ".ply":
ply_files.append(os.path.join(root, f))
ply_files = sorted(ply_files)
# get shared
src_mesh = Mesh(filename=ply_files[0])
session.run(tf.assign(tf_src, src_mesh.v))
speaker = os.path.basename(os.path.dirname(src_dir))
if last_speaker != speaker:
print("- clear speaker information")
_zero_state("trans", "rot", "pose", "shape", "exp")
else:
_zero_state("exp")
stt_dir = os.path.join(os.path.dirname(dst_dir), "state")
if os.path.exists(stt_dir):
state_dict = dict(
trans = np.load(os.path.join(stt_dir, "trans.npy")),
rot = np.load(os.path.join(stt_dir, "rot.npy")),
pose = np.load(os.path.join(stt_dir, "pose.npy")),
shape = np.load(os.path.join(stt_dir, "shape.npy")),
)
_load_state(state_dict)
fitting_mesh = Mesh(session.run(tf_model), src_mesh.f)
fitting_mesh.write_ply(os.path.join(stt_dir, "zero.ply"))
fit_zero_dir = os.path.join(os.path.dirname(os.path.dirname(dst_dir)), "zero_exp")
if not os.path.exists(fit_zero_dir): os.makedirs(fit_zero_dir)
print("- " + speaker + " " + os.path.basename(src_dir))
print(" -> fit shared parameters...")
optim_shared_rigid.minimize(session)
optim_shared_all.minimize(session)
state_dict = _save_state("exp", set_zero=True)
fitting_mesh = Mesh(session.run(tf_model), src_mesh.f)
fitting_mesh.write_ply(os.path.join(fit_zero_dir, "{}.ply".format(speaker)))
_load_state(state_dict)
return
if not os.path.exists(stt_dir): os.makedirs(stt_dir)
np.save(os.path.join(stt_dir, "trans.npy"), tf_trans.eval(), allow_pickle=False)
np.save(os.path.join(stt_dir, "rot.npy"), tf_rot.eval(), allow_pickle=False)
np.save(os.path.join(stt_dir, "pose.npy"), tf_pose.eval(), allow_pickle=False)
np.save(os.path.join(stt_dir, "shape.npy"), tf_shape.eval(), allow_pickle=False)
progress = tqdm(ply_files)
for src_fname in progress:
frame = os.path.basename(src_fname)
progress.set_description(" -> " + frame)
dst_fname = os.path.join(dst_dir, frame)
# param filename
prm_fname = os.path.join(prm_dir, frame)
exp_fname = os.path.splitext(prm_fname)[0] + '_exp.npy'
idn_fname = os.path.splitext(prm_fname)[0] + '_idn.npy'
src_mesh = Mesh(filename=src_fname)
session.run(tf.assign(tf_src, src_mesh.v))
optim_seq.minimize(session)
# save expr
np.save(exp_fname, tf_exp.eval())
np.save(idn_fname, tf_shape.eval())
# state_dict = _save_state("trans", "rot", "pose", "shape", set_zero=True)
# save mesh
fitting_mesh = Mesh(session.run(tf_model), src_mesh.f)
fitting_mesh.write_ply(dst_fname)
# _load_state(state_dict)
# print(tf_shape.eval())
if showing:
g_mv.set_static_meshes([fitting_mesh])
os.system("touch {}".format(_anchor))
return speaker
def process(self):
train_data, val_data, test_data = [], [], []
train_vertices = []
for key in self.data_file:
for idx, data_file in tqdm(enumerate(self.data_file[key])):
if key == 'lgtd':
mesh = Mesh(filename=data_file[0])
mesh_verts = torch.Tensor(mesh.v)
adjacency = get_vert_connectivity(mesh.v, mesh.f).tocoo()
edge_index = torch.Tensor(
np.vstack((adjacency.row, adjacency.col)))
mesh_m24 = Mesh(filename=data_file[1])
data = Data(x=mesh_verts,
y=torch.Tensor(mesh_m24.v),
edge_index=edge_index)
elif key == 'lgtddx':
mesh = Mesh(filename=data_file[0])
mesh_verts = torch.Tensor(mesh.v)
adjacency = get_vert_connectivity(mesh.v, mesh.f).tocoo()
edge_index = torch.Tensor(
np.vstack((adjacency.row, adjacency.col)))
mesh_m24 = Mesh(filename=data_file[1])
data = Data(x=mesh_verts,
y=torch.Tensor(mesh_m24.v),
edge_index=edge_index,
label=torch.Tensor(data_file[2]))
elif key == 'lgtdvc':
#print(data_file)
mesh = Mesh(filename=data_file[0])
mesh_verts = torch.Tensor(mesh.v)
adjacency = get_vert_connectivity(mesh.v, mesh.f).tocoo()
edge_index = torch.Tensor(
np.vstack((adjacency.row, adjacency.col)))
mesh_fu = Mesh(filename=data_file[1])
data = Data(x=mesh_verts,
y=torch.Tensor(mesh_fu.v),
edge_index=edge_index,
label=torch.Tensor(data_file[2]),
period=torch.Tensor(data_file[3]))
else:
mesh = Mesh(filename=data_file)
mesh_verts = torch.Tensor(mesh.v)
adjacency = get_vert_connectivity(mesh.v, mesh.f).tocoo()
edge_index = torch.Tensor(
np.vstack((adjacency.row, adjacency.col)))
if key == 'ad':
data = Data(x=mesh_verts,
y=mesh_verts,
label=torch.Tensor([1, 0]),
edge_index=edge_index)
elif key == 'cn':
data = Data(x=mesh_verts,
y=mesh_verts,
label=torch.Tensor([0, 1]),
edge_index=edge_index)
elif key == 'all':
data = Data(x=mesh_verts,
y=mesh_verts,
edge_index=edge_index)
if idx % 100 < 10:
test_data.append(data)
#print(data.period)
elif idx % 100 < 20:
val_data.append(data)
else:
train_data.append(data)
train_vertices.append(mesh.v)
print(len(train_data), len(val_data), len(test_data))
mean_train = torch.Tensor(np.mean(train_vertices, axis=0))
std_train = torch.Tensor(np.std(train_vertices, axis=0))
norm_dict = {'mean': mean_train, 'std': std_train}
if self.pre_transform is not None:
if hasattr(self.pre_transform, 'mean') and hasattr(
self.pre_transform, 'std'):
if self.pre_tranform.mean is None:
self.pre_tranform.mean = mean_train
if self.pre_transform.std is None:
self.pre_tranform.std = std_train
train_data = [self.pre_transform(td) for td in train_data]
val_data = [self.pre_transform(td) for td in val_data]
test_data = [self.pre_transform(td) for td in test_data]
torch.save(self.collate(train_data), self.processed_paths[0])
torch.save(self.collate(val_data), self.processed_paths[1])
torch.save(self.collate(test_data), self.processed_paths[2])
torch.save(norm_dict, self.processed_paths[3])
def vis_results(dorig,
coarse_net,
refine_net,
rh_model,
show_gen=True,
show_rec=True,
save=False,
save_dir=None):
with torch.no_grad():
imw, imh = 400, 1000
cols = len(dorig['bps_object'])
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
if show_rec:
mvs = MeshViewers(window_width=imw * cols,
window_height=imh,
shape=[3, cols],
keepalive=True)
drec_cnet = coarse_net(**dorig)
verts_rh_rec_cnet = rh_model(**drec_cnet).vertices
_, h2o, _ = point2point_signed(verts_rh_rec_cnet,
dorig['verts_object'])
drec_cnet['trans_rhand_f'] = drec_cnet['transl']
drec_cnet['global_orient_rhand_rotmat_f'] = aa2rotmat(
drec_cnet['global_orient']).view(-1, 3, 3)
drec_cnet['fpose_rhand_rotmat_f'] = aa2rotmat(
drec_cnet['hand_pose']).view(-1, 15, 3, 3)
drec_cnet['verts_object'] = dorig['verts_object']
drec_cnet['h2o_dist'] = h2o.abs()
drec_rnet = refine_net(**drec_cnet)
verts_rh_rec_rnet = rh_model(**drec_rnet).vertices
for cId in range(0, len(dorig['bps_object'])):
try:
from copy import deepcopy
meshes = deepcopy(dorig['mesh_object'])
obj_mesh = [meshes[cId]]
except:
obj_mesh = points_to_spheres(points=to_cpu(
dorig['verts_object'][cId]),
radius=0.002,
vc=name_to_rgb['green'])
hand_mesh_orig = Mesh(v=to_cpu(dorig['verts_rhand'][cId]),
f=rh_model.faces,
vc=name_to_rgb['blue'])
hand_mesh_rec_cnet = Mesh(v=to_cpu(verts_rh_rec_cnet[cId]),
f=rh_model.faces,
vc=name_to_rgb['green'])
hand_mesh_rec_rnet = Mesh(v=to_cpu(verts_rh_rec_rnet[cId]),
f=rh_model.faces,
vc=name_to_rgb['red'])
if 'rotmat' in dorig:
rotmat = dorig['rotmat'][cId].T
obj_mesh = [obj_mesh[0].rotate_vertices(rotmat)]
hand_mesh_orig.rotate_vertices(rotmat)
hand_mesh_rec_cnet.rotate_vertices(rotmat)
hand_mesh_rec_rnet.rotate_vertices(rotmat)
mvs[0][cId].set_static_meshes([hand_mesh_orig] + obj_mesh,
blocking=True)
mvs[1][cId].set_static_meshes([hand_mesh_rec_cnet] + obj_mesh,
blocking=True)
mvs[2][cId].set_static_meshes([hand_mesh_rec_rnet] + obj_mesh,
blocking=True)
if save:
save_path = os.path.join(save_dir, str(cId))
makepath(save_path)
hand_mesh_rec_rnet.write_ply(filename=save_path +
'/rh_mesh_gen_%d.ply' % cId)
obj_mesh[0].write_ply(filename=save_path +
'/obj_mesh_%d.ply' % cId)
if show_gen:
mvs = MeshViewers(window_width=imw * cols,
window_height=imh,
shape=[2, cols],
keepalive=True)
drec_cnet = coarse_net.sample_poses(dorig['bps_object'])
verts_rh_gen_cnet = rh_model(**drec_cnet).vertices
_, h2o, _ = point2point_signed(verts_rh_gen_cnet,
dorig['verts_object'].to(device))
drec_cnet['trans_rhand_f'] = drec_cnet['transl']
drec_cnet['global_orient_rhand_rotmat_f'] = aa2rotmat(
drec_cnet['global_orient']).view(-1, 3, 3)
drec_cnet['fpose_rhand_rotmat_f'] = aa2rotmat(
drec_cnet['hand_pose']).view(-1, 15, 3, 3)
drec_cnet['verts_object'] = dorig['verts_object'].to(device)
drec_cnet['h2o_dist'] = h2o.abs()
drec_rnet = refine_net(**drec_cnet)
verts_rh_gen_rnet = rh_model(**drec_rnet).vertices
for cId in range(0, len(dorig['bps_object'])):
try:
from copy import deepcopy
meshes = deepcopy(dorig['mesh_object'])
obj_mesh = [meshes[cId]]
except:
obj_mesh = points_to_spheres(to_cpu(
dorig['verts_object'][cId]),
radius=0.002,
vc=name_to_rgb['green'])
hand_mesh_gen_cnet = Mesh(v=to_cpu(verts_rh_gen_cnet[cId]),
f=rh_model.faces,
vc=name_to_rgb['pink'])
hand_mesh_gen_rnet = Mesh(v=to_cpu(verts_rh_gen_rnet[cId]),
f=rh_model.faces,
vc=name_to_rgb['gray'])
if 'rotmat' in dorig:
rotmat = dorig['rotmat'][cId].T
obj_mesh = [obj_mesh[0].rotate_vertices(rotmat)]
hand_mesh_gen_cnet.rotate_vertices(rotmat)
hand_mesh_gen_rnet.rotate_vertices(rotmat)
mvs[0][cId].set_static_meshes([hand_mesh_gen_cnet] + obj_mesh,
blocking=True)
mvs[1][cId].set_static_meshes([hand_mesh_gen_rnet] + obj_mesh,
blocking=True)
if save:
save_path = os.path.join(save_dir, str(cId))
makepath(save_path)
hand_mesh_gen_rnet.write_ply(filename=save_path +
'/rh_mesh_gen_%d.ply' % cId)
obj_mesh[0].write_ply(filename=save_path +
'/obj_mesh_%d.ply' % cId)
# args = parser.parse_args()
def frontalize(vertices, canonical_vertices):
# mesh = Mesh(filename="/home/user/3dfaceRe/center-loss_conv/scripts/template_fwh.obj")
# canonical_vertices = mesh.v
#canonical_vertices = np.load('Data/uv-data/canonical_vertices.npy')
vertices_homo = np.hstack((vertices, np.ones([vertices.shape[0],
1]))) #n x 4
P = np.linalg.lstsq(vertices_homo,
canonical_vertices)[0].T # Affine matrix. 3 x 4
front_vertices = vertices_homo.dot(P.T)
return front_vertices
template = Mesh(filename='./data/template.obj')
nVal = 100 # args.num_valid
root_dir = './' # args.root_dir
dataset = 'data' # args.dataset
name = 'sliced'
data = os.path.join(root_dir, dataset, 'Processed', name)
train = np.load(data + '/train.npy')
# train = torch.tensor(train.astype('float32'))
# mean = torch.mean(train, dim=0)
# std = torch.std(train, dim=0)
# torch.save(mean, './data/Processed/sliced/mean.tch')
# torch.save(std, './data/Processed/sliced/std.tch')
def fit_lmk2d(target_img, target_2d_lmks, template_fname, tf_model_fname,
lmk_face_idx, lmk_b_coords, weights):
'''
Fit FLAME to 2D landmarks
:param target_2d_lmks: target 2D landmarks provided as (num_lmks x 3) matrix
:param template_fname: template mesh in FLAME topology (only the face information are used)
:param tf_model_fname: saved Tensorflow FLAME model
:param lmk_face_idx: face indices of the landmark embedding in the FLAME topology
:param lmk_b_coords: barycentric coordinates of the landmark embedding in the FLAME topology
(i.e. weighting of the three vertices for the trinagle, the landmark is embedded in
:param weights: weights of the individual objective functions
:return: a mesh with the fitting results
'''
template_mesh = Mesh(filename=template_fname)
saver = tf.train.import_meta_graph(tf_model_fname + '.meta')
graph = tf.get_default_graph()
tf_model = graph.get_tensor_by_name(u'vertices:0')
with tf.Session() as session:
saver.restore(session, tf_model_fname)
# Workaround as existing tf.Variable cannot be retrieved back with tf.get_variable
# tf_v_template = [x for x in tf.trainable_variables() if 'v_template' in x.name][0]
tf_trans = [x for x in tf.trainable_variables()
if 'trans' in x.name][0]
tf_rot = [x for x in tf.trainable_variables() if 'rot' in x.name][0]
tf_pose = [x for x in tf.trainable_variables() if 'pose' in x.name][0]
tf_shape = [x for x in tf.trainable_variables()
if 'shape' in x.name][0]
tf_exp = [x for x in tf.trainable_variables() if 'exp' in x.name][0]
# Mirror landmark y-coordinates
target_2d_lmks[:, 1] = target_img.shape[0] - target_2d_lmks[:, 1]
lmks_3d = tf_get_model_lmks(tf_model, template_mesh, lmk_face_idx,
lmk_b_coords)
s2d = np.mean(
np.linalg.norm(target_2d_lmks - np.mean(target_2d_lmks, axis=0),
axis=1))
s3d = tf.reduce_mean(
tf.sqrt(
tf.reduce_sum(
tf.square(lmks_3d -
tf.reduce_mean(lmks_3d, axis=0))[:, :2],
axis=1)))
tf_scale = tf.Variable(s2d / s3d, dtype=lmks_3d.dtype)
# trans = 0.5*np.array((target_img.shape[0], target_img.shape[1]))/tf_scale
# trans = 0.5 * s3d * np.array((target_img.shape[0], target_img.shape[1])) / s2d
lmks_proj_2d = tf_project_points(lmks_3d, tf_scale, np.zeros(2))
factor = max(
max(target_2d_lmks[:, 0]) - min(target_2d_lmks[:, 0]),
max(target_2d_lmks[:, 1]) - min(target_2d_lmks[:, 1]))
lmk_dist = weights['lmk'] * tf.reduce_sum(
tf.square(tf.subtract(lmks_proj_2d, target_2d_lmks))) / (factor**2)
neck_pose_reg = weights['neck_pose'] * tf.reduce_sum(
tf.square(tf_pose[:3]))
jaw_pose_reg = weights['jaw_pose'] * tf.reduce_sum(
tf.square(tf_pose[3:6]))
eyeballs_pose_reg = weights['eyeballs_pose'] * tf.reduce_sum(
tf.square(tf_pose[6:]))
shape_reg = weights['shape'] * tf.reduce_sum(tf.square(tf_shape))
exp_reg = weights['expr'] * tf.reduce_sum(tf.square(tf_exp))
session.run(tf.global_variables_initializer())
def on_step(verts,
scale,
faces,
target_img,
target_lmks,
opt_lmks,
lmk_dist=0.0,
shape_reg=0.0,
exp_reg=0.0,
neck_pose_reg=0.0,
jaw_pose_reg=0.0,
eyeballs_pose_reg=0.0):
import cv2
import sys
import numpy as np
from psbody.mesh import Mesh
from utils.render_mesh import render_mesh
if lmk_dist > 0.0 or shape_reg > 0.0 or exp_reg > 0.0 or neck_pose_reg > 0.0 or jaw_pose_reg > 0.0 or eyeballs_pose_reg > 0.0:
print(
'lmk_dist: %f, shape_reg: %f, exp_reg: %f, neck_pose_reg: %f, jaw_pose_reg: %f, eyeballs_pose_reg: %f'
% (lmk_dist, shape_reg, exp_reg, neck_pose_reg,
jaw_pose_reg, eyeballs_pose_reg))
plt_target_lmks = target_lmks.copy()
plt_target_lmks[:, 1] = target_img.shape[0] - plt_target_lmks[:, 1]
for (x, y) in plt_target_lmks:
cv2.circle(target_img, (int(x), int(y)), 4, (0, 0, 255), -1)
plt_opt_lmks = opt_lmks.copy()
plt_opt_lmks[:, 1] = target_img.shape[0] - plt_opt_lmks[:, 1]
for (x, y) in plt_opt_lmks:
cv2.circle(target_img, (int(x), int(y)), 4, (255, 0, 0), -1)
if sys.version_info >= (3, 0):
rendered_img = render_mesh(Mesh(scale * verts, faces),
height=target_img.shape[0],
width=target_img.shape[1])
for (x, y) in plt_opt_lmks:
cv2.circle(rendered_img, (int(x), int(y)), 4, (255, 0, 0),
-1)
target_img = np.hstack((target_img, rendered_img))
cv2.imshow('img', target_img)
cv2.waitKey(10)
print('Optimize rigid transformation')
vars = [tf_scale, tf_trans, tf_rot]
loss = lmk_dist
optimizer = scipy_pt(loss=loss,
var_list=vars,
method='L-BFGS-B',
options={
'disp': 1,
'ftol': 5e-6
})
optimizer.minimize(session,
fetches=[
tf_model, tf_scale,
tf.constant(template_mesh.f),
tf.constant(target_img),
tf.constant(target_2d_lmks), lmks_proj_2d
],
loss_callback=on_step)
print('Optimize model parameters')
vars = [tf_scale, tf_trans[:2], tf_rot, tf_pose, tf_shape, tf_exp]
loss = lmk_dist + shape_reg + exp_reg + neck_pose_reg + jaw_pose_reg + eyeballs_pose_reg
optimizer = scipy_pt(loss=loss,
var_list=vars,
method='L-BFGS-B',
options={
'disp': 0,
'ftol': 1e-7
})
optimizer.minimize(session,
fetches=[
tf_model, tf_scale,
tf.constant(template_mesh.f),
tf.constant(target_img),
tf.constant(target_2d_lmks), lmks_proj_2d,
lmk_dist, shape_reg, exp_reg, neck_pose_reg,
jaw_pose_reg, eyeballs_pose_reg
],
loss_callback=on_step)
print('Fitting done')
np_verts, np_scale = session.run([tf_model, tf_scale])
return Mesh(np_verts, template_mesh.f), np_scale
(pack(results[i].dot(ch.concatenate((self.J[i, :], [0])))))
for i in range(len(results))
]
result = ch.dstack(results2)
return result, results_global
def compute_r(self):
return self.v.r
def compute_dr_wrt(self, wrt):
if wrt is not self.trans and wrt is not self.betas and wrt is not self.pose and wrt is not self.v_personal and wrt is not self.v_template:
return None
return self.v.dr_wrt(wrt)
if __name__ == '__main__':
from utils.smpl_paths import SmplPaths
dp = SmplPaths(gender='neutral')
smpl = Smpl(dp.get_smpl_file())
from psbody.mesh.meshviewer import MeshViewer
from psbody.mesh import Mesh
mv = MeshViewer()
mv.set_static_meshes([Mesh(smpl.r, smpl.f)])
raw_input("Press Enter to continue...")
## This file contains correspondances between garment vertices and smpl body
fts_file = 'assets/garment_fts.pkl'
vert_indices, fts = pkl.load(open(fts_file, "rb"), encoding="latin1")
fts['naked'] = ft
## Choose any garmet type as source
garment_type = 'TShirtNoCoat'
index = np.random.randint(0, len(gar_dict[garment_type])) ## Randomly pick from the digital wardrobe
path = split(gar_dict[garment_type][index])[0]
garment_org_body_unposed = load_smpl_from_file(join(path, 'registration.pkl'))
garment_org_body_unposed.pose[:] = 0
garment_org_body_unposed.trans[:] = 0
garment_org_body_unposed = Mesh(garment_org_body_unposed.v, garment_org_body_unposed.f)
garment_unposed = Mesh(filename=join(path, garment_type + '.obj'))
garment_tex = join(path, 'multi_tex.jpg')
## Generate random SMPL body (Feel free to set up ur own smpl) as target subject
smpl.pose[:] = np.random.randn(72) *0.05
smpl.betas[:] = np.random.randn(10) *0.01
smpl.trans[:] = 0
tgt_body = Mesh(smpl.r, smpl.f)
vert_inds = vert_indices[garment_type]
garment_unposed.set_texture_image(garment_tex)
new_garment = dress(smpl, garment_org_body_unposed, garment_unposed, vert_inds, garment_tex)
def fit_single_frame(
img,
keypoints,
init_trans,
scan,
scene_name,
body_model,
camera,
joint_weights,
body_pose_prior,
jaw_prior,
left_hand_prior,
right_hand_prior,
shape_prior,
expr_prior,
angle_prior,
result_fn='out.pkl',
mesh_fn='out.obj',
body_scene_rendering_fn='body_scene.png',
out_img_fn='overlay.png',
loss_type='smplify',
use_cuda=True,
init_joints_idxs=(9, 12, 2, 5),
use_face=True,
use_hands=True,
data_weights=None,
body_pose_prior_weights=None,
hand_pose_prior_weights=None,
jaw_pose_prior_weights=None,
shape_weights=None,
expr_weights=None,
hand_joints_weights=None,
face_joints_weights=None,
depth_loss_weight=1e2,
interpenetration=True,
coll_loss_weights=None,
df_cone_height=0.5,
penalize_outside=True,
max_collisions=8,
point2plane=False,
part_segm_fn='',
focal_length_x=5000.,
focal_length_y=5000.,
side_view_thsh=25.,
rho=100,
vposer_latent_dim=32,
vposer_ckpt='',
use_joints_conf=False,
interactive=True,
visualize=False,
save_meshes=True,
degrees=None,
batch_size=1,
dtype=torch.float32,
ign_part_pairs=None,
left_shoulder_idx=2,
right_shoulder_idx=5,
####################
### PROX
render_results=True,
camera_mode='moving',
## Depth
s2m=False,
s2m_weights=None,
m2s=False,
m2s_weights=None,
rho_s2m=1,
rho_m2s=1,
init_mode=None,
trans_opt_stages=None,
viz_mode='mv',
#penetration
sdf_penetration=False,
sdf_penetration_weights=0.0,
sdf_dir=None,
cam2world_dir=None,
#contact
contact=False,
rho_contact=1.0,
contact_loss_weights=None,
contact_angle=15,
contact_body_parts=None,
body_segments_dir=None,
load_scene=False,
scene_dir=None,
**kwargs):
assert batch_size == 1, 'PyTorch L-BFGS only supports batch_size == 1'
body_model.reset_params()
body_model.transl.requires_grad = True
device = torch.device('cuda') if use_cuda else torch.device('cpu')
if visualize:
pil_img.fromarray((img * 255).astype(np.uint8)).show()
if degrees is None:
degrees = [0, 90, 180, 270]
if data_weights is None:
data_weights = [
1,
] * 5
if body_pose_prior_weights is None:
body_pose_prior_weights = [4.04 * 1e2, 4.04 * 1e2, 57.4, 4.78]
msg = ('Number of Body pose prior weights {}'.format(
len(body_pose_prior_weights)) +
' does not match the number of data term weights {}'.format(
len(data_weights)))
assert (len(data_weights) == len(body_pose_prior_weights)), msg
if use_hands:
if hand_pose_prior_weights is None:
hand_pose_prior_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
msg = ('Number of Body pose prior weights does not match the' +
' number of hand pose prior weights')
assert (
len(hand_pose_prior_weights) == len(body_pose_prior_weights)), msg
if hand_joints_weights is None:
hand_joints_weights = [0.0, 0.0, 0.0, 1.0]
msg = ('Number of Body pose prior weights does not match the' +
' number of hand joint distance weights')
assert (
len(hand_joints_weights) == len(body_pose_prior_weights)), msg
if shape_weights is None:
shape_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
msg = ('Number of Body pose prior weights = {} does not match the' +
' number of Shape prior weights = {}')
assert (len(shape_weights) == len(body_pose_prior_weights)), msg.format(
len(shape_weights), len(body_pose_prior_weights))
if use_face:
if jaw_pose_prior_weights is None:
jaw_pose_prior_weights = [[x] * 3 for x in shape_weights]
else:
jaw_pose_prior_weights = map(lambda x: map(float, x.split(',')),
jaw_pose_prior_weights)
jaw_pose_prior_weights = [list(w) for w in jaw_pose_prior_weights]
msg = ('Number of Body pose prior weights does not match the' +
' number of jaw pose prior weights')
assert (
len(jaw_pose_prior_weights) == len(body_pose_prior_weights)), msg
if expr_weights is None:
expr_weights = [1e2, 5 * 1e1, 1e1, .5 * 1e1]
msg = ('Number of Body pose prior weights = {} does not match the' +
' number of Expression prior weights = {}')
assert (len(expr_weights) == len(body_pose_prior_weights)), msg.format(
len(body_pose_prior_weights), len(expr_weights))
if face_joints_weights is None:
face_joints_weights = [0.0, 0.0, 0.0, 1.0]
msg = ('Number of Body pose prior weights does not match the' +
' number of face joint distance weights')
assert (len(face_joints_weights) == len(body_pose_prior_weights)), msg
if coll_loss_weights is None:
coll_loss_weights = [0.0] * len(body_pose_prior_weights)
msg = ('Number of Body pose prior weights does not match the' +
' number of collision loss weights')
assert (len(coll_loss_weights) == len(body_pose_prior_weights)), msg
use_vposer = kwargs.get('use_vposer', True)
vposer, pose_embedding = [
None,
] * 2
if use_vposer:
pose_embedding = torch.zeros([batch_size, 32],
dtype=dtype,
device=device,
requires_grad=True)
vposer_ckpt = osp.expandvars(vposer_ckpt)
vposer, _ = load_vposer(vposer_ckpt, vp_model='snapshot')
vposer = vposer.to(device=device)
vposer.eval()
if use_vposer:
body_mean_pose = torch.zeros([batch_size, vposer_latent_dim],
dtype=dtype)
else:
body_mean_pose = body_pose_prior.get_mean().detach().cpu()
keypoint_data = torch.tensor(keypoints, dtype=dtype)
gt_joints = keypoint_data[:, :, :2]
if use_joints_conf:
joints_conf = keypoint_data[:, :, 2].reshape(1, -1)
# Transfer the data to the correct device
gt_joints = gt_joints.to(device=device, dtype=dtype)
if use_joints_conf:
joints_conf = joints_conf.to(device=device, dtype=dtype)
scan_tensor = None
if scan is not None:
scan_tensor = torch.tensor(scan.get('points'),
device=device,
dtype=dtype).unsqueeze(0)
# load pre-computed signed distance field
sdf = None
sdf_normals = None
grid_min = None
grid_max = None
voxel_size = None
if sdf_penetration:
with open(osp.join(sdf_dir, scene_name + '.json'), 'r') as f:
sdf_data = json.load(f)
grid_min = torch.tensor(np.array(sdf_data['min']),
dtype=dtype,
device=device)
grid_max = torch.tensor(np.array(sdf_data['max']),
dtype=dtype,
device=device)
grid_dim = sdf_data['dim']
voxel_size = (grid_max - grid_min) / grid_dim
sdf = np.load(osp.join(sdf_dir, scene_name + '_sdf.npy')).reshape(
grid_dim, grid_dim, grid_dim)
sdf = torch.tensor(sdf, dtype=dtype, device=device)
if osp.exists(osp.join(sdf_dir, scene_name + '_normals.npy')):
sdf_normals = np.load(
osp.join(sdf_dir, scene_name + '_normals.npy')).reshape(
grid_dim, grid_dim, grid_dim, 3)
sdf_normals = torch.tensor(sdf_normals, dtype=dtype, device=device)
else:
print("Normals not found...")
with open(os.path.join(cam2world_dir, scene_name + '.json'), 'r') as f:
cam2world = np.array(json.load(f))
R = torch.tensor(cam2world[:3, :3].reshape(3, 3),
dtype=dtype,
device=device)
t = torch.tensor(cam2world[:3, 3].reshape(1, 3),
dtype=dtype,
device=device)
# Create the search tree
search_tree = None
pen_distance = None
filter_faces = None
if interpenetration:
from mesh_intersection.bvh_search_tree import BVH
import mesh_intersection.loss as collisions_loss
from mesh_intersection.filter_faces import FilterFaces
assert use_cuda, 'Interpenetration term can only be used with CUDA'
assert torch.cuda.is_available(), \
'No CUDA Device! Interpenetration term can only be used' + \
' with CUDA'
search_tree = BVH(max_collisions=max_collisions)
pen_distance = \
collisions_loss.DistanceFieldPenetrationLoss(
sigma=df_cone_height, point2plane=point2plane,
vectorized=True, penalize_outside=penalize_outside)
if part_segm_fn:
# Read the part segmentation
part_segm_fn = os.path.expandvars(part_segm_fn)
with open(part_segm_fn, 'rb') as faces_parents_file:
face_segm_data = pickle.load(faces_parents_file,
encoding='latin1')
faces_segm = face_segm_data['segm']
faces_parents = face_segm_data['parents']
# Create the module used to filter invalid collision pairs
filter_faces = FilterFaces(
faces_segm=faces_segm,
faces_parents=faces_parents,
ign_part_pairs=ign_part_pairs).to(device=device)
# load vertix ids of contact parts
contact_verts_ids = ftov = None
if contact:
contact_verts_ids = []
for part in contact_body_parts:
with open(os.path.join(body_segments_dir, part + '.json'),
'r') as f:
data = json.load(f)
contact_verts_ids.append(list(set(data["verts_ind"])))
contact_verts_ids = np.concatenate(contact_verts_ids)
vertices = body_model(return_verts=True,
body_pose=torch.zeros((batch_size, 63),
dtype=dtype,
device=device)).vertices
vertices_np = vertices.detach().cpu().numpy().squeeze()
body_faces_np = body_model.faces_tensor.detach().cpu().numpy().reshape(
-1, 3)
m = Mesh(v=vertices_np, f=body_faces_np)
ftov = m.faces_by_vertex(as_sparse_matrix=True)
ftov = sparse.coo_matrix(ftov)
indices = torch.LongTensor(np.vstack((ftov.row, ftov.col))).to(device)
values = torch.FloatTensor(ftov.data).to(device)
shape = ftov.shape
ftov = torch.sparse.FloatTensor(indices, values, torch.Size(shape))
# Read the scene scan if any
scene_v = scene_vn = scene_f = None
if scene_name is not None:
if load_scene:
scene = Mesh(filename=os.path.join(scene_dir, scene_name + '.ply'))
scene.vn = scene.estimate_vertex_normals()
scene_v = torch.tensor(scene.v[np.newaxis, :],
dtype=dtype,
device=device).contiguous()
scene_vn = torch.tensor(scene.vn[np.newaxis, :],
dtype=dtype,
device=device)
scene_f = torch.tensor(scene.f.astype(int)[np.newaxis, :],
dtype=torch.long,
device=device)
# Weights used for the pose prior and the shape prior
opt_weights_dict = {
'data_weight': data_weights,
'body_pose_weight': body_pose_prior_weights,
'shape_weight': shape_weights
}
if use_face:
opt_weights_dict['face_weight'] = face_joints_weights
opt_weights_dict['expr_prior_weight'] = expr_weights
opt_weights_dict['jaw_prior_weight'] = jaw_pose_prior_weights
if use_hands:
opt_weights_dict['hand_weight'] = hand_joints_weights
opt_weights_dict['hand_prior_weight'] = hand_pose_prior_weights
if interpenetration:
opt_weights_dict['coll_loss_weight'] = coll_loss_weights
if s2m:
opt_weights_dict['s2m_weight'] = s2m_weights
if m2s:
opt_weights_dict['m2s_weight'] = m2s_weights
if sdf_penetration:
opt_weights_dict['sdf_penetration_weight'] = sdf_penetration_weights
if contact:
opt_weights_dict['contact_loss_weight'] = contact_loss_weights
keys = opt_weights_dict.keys()
opt_weights = [
dict(zip(keys, vals))
for vals in zip(*(opt_weights_dict[k] for k in keys
if opt_weights_dict[k] is not None))
]
for weight_list in opt_weights:
for key in weight_list:
weight_list[key] = torch.tensor(weight_list[key],
device=device,
dtype=dtype)
# load indices of the head of smpl-x model
with open(osp.join(body_segments_dir, 'body_mask.json'), 'r') as fp:
head_indx = np.array(json.load(fp))
N = body_model.get_num_verts()
body_indx = np.setdiff1d(np.arange(N), head_indx)
head_mask = np.in1d(np.arange(N), head_indx)
body_mask = np.in1d(np.arange(N), body_indx)
# The indices of the joints used for the initialization of the camera
init_joints_idxs = torch.tensor(init_joints_idxs, device=device)
edge_indices = kwargs.get('body_tri_idxs')
# which initialization mode to choose: similar traingles, mean of the scan or the average of both
if init_mode == 'scan':
init_t = init_trans
elif init_mode == 'both':
init_t = (init_trans.to(device) + fitting.guess_init(
body_model,
gt_joints,
edge_indices,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
model_type=kwargs.get('model_type', 'smpl'),
focal_length=focal_length_x,
dtype=dtype)) / 2.0
else:
init_t = fitting.guess_init(body_model,
gt_joints,
edge_indices,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
model_type=kwargs.get(
'model_type', 'smpl'),
focal_length=focal_length_x,
dtype=dtype)
camera_loss = fitting.create_loss('camera_init',
trans_estimation=init_t,
init_joints_idxs=init_joints_idxs,
depth_loss_weight=depth_loss_weight,
camera_mode=camera_mode,
dtype=dtype).to(device=device)
camera_loss.trans_estimation[:] = init_t
loss = fitting.create_loss(loss_type=loss_type,
joint_weights=joint_weights,
rho=rho,
use_joints_conf=use_joints_conf,
use_face=use_face,
use_hands=use_hands,
vposer=vposer,
pose_embedding=pose_embedding,
body_pose_prior=body_pose_prior,
shape_prior=shape_prior,
angle_prior=angle_prior,
expr_prior=expr_prior,
left_hand_prior=left_hand_prior,
right_hand_prior=right_hand_prior,
jaw_prior=jaw_prior,
interpenetration=interpenetration,
pen_distance=pen_distance,
search_tree=search_tree,
tri_filtering_module=filter_faces,
s2m=s2m,
m2s=m2s,
rho_s2m=rho_s2m,
rho_m2s=rho_m2s,
head_mask=head_mask,
body_mask=body_mask,
sdf_penetration=sdf_penetration,
voxel_size=voxel_size,
grid_min=grid_min,
grid_max=grid_max,
sdf=sdf,
sdf_normals=sdf_normals,
R=R,
t=t,
contact=contact,
contact_verts_ids=contact_verts_ids,
rho_contact=rho_contact,
contact_angle=contact_angle,
dtype=dtype,
**kwargs)
loss = loss.to(device=device)
with fitting.FittingMonitor(batch_size=batch_size,
visualize=visualize,
viz_mode=viz_mode,
**kwargs) as monitor:
img = torch.tensor(img, dtype=dtype)
H, W, _ = img.shape
# Reset the parameters to estimate the initial translation of the
# body model
if camera_mode == 'moving':
body_model.reset_params(body_pose=body_mean_pose)
# Update the value of the translation of the camera as well as
# the image center.
with torch.no_grad():
camera.translation[:] = init_t.view_as(camera.translation)
camera.center[:] = torch.tensor([W, H], dtype=dtype) * 0.5
# Re-enable gradient calculation for the camera translation
camera.translation.requires_grad = True
camera_opt_params = [camera.translation, body_model.global_orient]
elif camera_mode == 'fixed':
body_model.reset_params(body_pose=body_mean_pose, transl=init_t)
camera_opt_params = [body_model.transl, body_model.global_orient]
# If the distance between the 2D shoulders is smaller than a
# predefined threshold then try 2 fits, the initial one and a 180
# degree rotation
shoulder_dist = torch.dist(gt_joints[:, left_shoulder_idx],
gt_joints[:, right_shoulder_idx])
try_both_orient = shoulder_dist.item() < side_view_thsh
camera_optimizer, camera_create_graph = optim_factory.create_optimizer(
camera_opt_params, **kwargs)
# The closure passed to the optimizer
fit_camera = monitor.create_fitting_closure(
camera_optimizer,
body_model,
camera,
gt_joints,
camera_loss,
create_graph=camera_create_graph,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
scan_tensor=scan_tensor,
return_full_pose=False,
return_verts=False)
# Step 1: Optimize over the torso joints the camera translation
# Initialize the computational graph by feeding the initial translation
# of the camera and the initial pose of the body model.
camera_init_start = time.time()
cam_init_loss_val = monitor.run_fitting(camera_optimizer,
fit_camera,
camera_opt_params,
body_model,
use_vposer=use_vposer,
pose_embedding=pose_embedding,
vposer=vposer)
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
tqdm.write('Camera initialization done after {:.4f}'.format(
time.time() - camera_init_start))
tqdm.write('Camera initialization final loss {:.4f}'.format(
cam_init_loss_val))
# If the 2D detections/positions of the shoulder joints are too
# close the rotate the body by 180 degrees and also fit to that
# orientation
if try_both_orient:
body_orient = body_model.global_orient.detach().cpu().numpy()
flipped_orient = cv2.Rodrigues(body_orient)[0].dot(
cv2.Rodrigues(np.array([0., np.pi, 0]))[0])
flipped_orient = cv2.Rodrigues(flipped_orient)[0].ravel()
flipped_orient = torch.tensor(flipped_orient,
dtype=dtype,
device=device).unsqueeze(dim=0)
orientations = [body_orient, flipped_orient]
else:
orientations = [body_model.global_orient.detach().cpu().numpy()]
# store here the final error for both orientations,
# and pick the orientation resulting in the lowest error
results = []
body_transl = body_model.transl.clone().detach()
# Step 2: Optimize the full model
final_loss_val = 0
for or_idx, orient in enumerate(tqdm(orientations,
desc='Orientation')):
opt_start = time.time()
new_params = defaultdict(transl=body_transl,
global_orient=orient,
body_pose=body_mean_pose)
body_model.reset_params(**new_params)
if use_vposer:
with torch.no_grad():
pose_embedding.fill_(0)
for opt_idx, curr_weights in enumerate(
tqdm(opt_weights, desc='Stage')):
if opt_idx not in trans_opt_stages:
body_model.transl.requires_grad = False
else:
body_model.transl.requires_grad = True
body_params = list(body_model.parameters())
final_params = list(
filter(lambda x: x.requires_grad, body_params))
if use_vposer:
final_params.append(pose_embedding)
body_optimizer, body_create_graph = optim_factory.create_optimizer(
final_params, **kwargs)
body_optimizer.zero_grad()
curr_weights['bending_prior_weight'] = (
3.17 * curr_weights['body_pose_weight'])
if use_hands:
joint_weights[:, 25:76] = curr_weights['hand_weight']
if use_face:
joint_weights[:, 76:] = curr_weights['face_weight']
loss.reset_loss_weights(curr_weights)
closure = monitor.create_fitting_closure(
body_optimizer,
body_model,
camera=camera,
gt_joints=gt_joints,
joints_conf=joints_conf,
joint_weights=joint_weights,
loss=loss,
create_graph=body_create_graph,
use_vposer=use_vposer,
vposer=vposer,
pose_embedding=pose_embedding,
scan_tensor=scan_tensor,
scene_v=scene_v,
scene_vn=scene_vn,
scene_f=scene_f,
ftov=ftov,
return_verts=True,
return_full_pose=True)
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
stage_start = time.time()
final_loss_val = monitor.run_fitting(
body_optimizer,
closure,
final_params,
body_model,
pose_embedding=pose_embedding,
vposer=vposer,
use_vposer=use_vposer)
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
elapsed = time.time() - stage_start
if interactive:
tqdm.write(
'Stage {:03d} done after {:.4f} seconds'.format(
opt_idx, elapsed))
if interactive:
if use_cuda and torch.cuda.is_available():
torch.cuda.synchronize()
elapsed = time.time() - opt_start
tqdm.write(
'Body fitting Orientation {} done after {:.4f} seconds'.
format(or_idx, elapsed))
tqdm.write(
'Body final loss val = {:.5f}'.format(final_loss_val))
# Get the result of the fitting process
# Store in it the errors list in order to compare multiple
# orientations, if they exist
result = {
'camera_' + str(key): val.detach().cpu().numpy()
for key, val in camera.named_parameters()
}
result.update({
key: val.detach().cpu().numpy()
for key, val in body_model.named_parameters()
})
if use_vposer:
result['pose_embedding'] = pose_embedding.detach().cpu().numpy(
)
body_pose = vposer.decode(pose_embedding,
output_type='aa').view(
1, -1) if use_vposer else None
result['body_pose'] = body_pose.detach().cpu().numpy()
results.append({'loss': final_loss_val, 'result': result})
with open(result_fn, 'wb') as result_file:
if len(results) > 1:
min_idx = (0 if results[0]['loss'] < results[1]['loss'] else 1)
else:
min_idx = 0
pickle.dump(results[min_idx]['result'], result_file, protocol=2)
if save_meshes or visualize:
body_pose = vposer.decode(pose_embedding, output_type='aa').view(
1, -1) if use_vposer else None
model_type = kwargs.get('model_type', 'smpl')
append_wrists = model_type == 'smpl' and use_vposer
if append_wrists:
wrist_pose = torch.zeros([body_pose.shape[0], 6],
dtype=body_pose.dtype,
device=body_pose.device)
body_pose = torch.cat([body_pose, wrist_pose], dim=1)
model_output = body_model(return_verts=True, body_pose=body_pose)
vertices = model_output.vertices.detach().cpu().numpy().squeeze()
import trimesh
out_mesh = trimesh.Trimesh(vertices, body_model.faces, process=False)
out_mesh.export(mesh_fn)
if render_results:
import pyrender
# common
H, W = 1080, 1920
camera_center = np.array([951.30, 536.77])
camera_pose = np.eye(4)
camera_pose = np.array([1.0, -1.0, -1.0, 1.0]).reshape(-1,
1) * camera_pose
camera = pyrender.camera.IntrinsicsCamera(fx=1060.53,
fy=1060.38,
cx=camera_center[0],
cy=camera_center[1])
light = pyrender.DirectionalLight(color=np.ones(3), intensity=2.0)
material = pyrender.MetallicRoughnessMaterial(
metallicFactor=0.0,
alphaMode='OPAQUE',
baseColorFactor=(1.0, 1.0, 0.9, 1.0))
body_mesh = pyrender.Mesh.from_trimesh(out_mesh, material=material)
## rendering body
img = img.detach().cpu().numpy()
H, W, _ = img.shape
scene = pyrender.Scene(bg_color=[0.0, 0.0, 0.0, 0.0],
ambient_light=(0.3, 0.3, 0.3))
scene.add(camera, pose=camera_pose)
scene.add(light, pose=camera_pose)
# for node in light_nodes:
# scene.add_node(node)
scene.add(body_mesh, 'mesh')
r = pyrender.OffscreenRenderer(viewport_width=W,
viewport_height=H,
point_size=1.0)
color, _ = r.render(scene, flags=pyrender.RenderFlags.RGBA)
color = color.astype(np.float32) / 255.0
valid_mask = (color[:, :, -1] > 0)[:, :, np.newaxis]
input_img = img
output_img = (color[:, :, :-1] * valid_mask +
(1 - valid_mask) * input_img)
img = pil_img.fromarray((output_img * 255).astype(np.uint8))
img.save(out_img_fn)
##redering body+scene
body_mesh = pyrender.Mesh.from_trimesh(out_mesh, material=material)
static_scene = trimesh.load(osp.join(scene_dir, scene_name + '.ply'))
trans = np.linalg.inv(cam2world)
static_scene.apply_transform(trans)
static_scene_mesh = pyrender.Mesh.from_trimesh(static_scene)
scene = pyrender.Scene()
scene.add(camera, pose=camera_pose)
scene.add(light, pose=camera_pose)
scene.add(static_scene_mesh, 'mesh')
scene.add(body_mesh, 'mesh')
r = pyrender.OffscreenRenderer(viewport_width=W, viewport_height=H)
color, _ = r.render(scene)
color = color.astype(np.float32) / 255.0
img = pil_img.fromarray((color * 255).astype(np.uint8))
img.save(body_scene_rendering_fn)
def pose_result(self,
verts,
pose_params,
save_obj,
cloth_type=None,
obj_dir=None):
'''
:param verts: [N, 6890, 3]
:param pose_params: [N, 72]
'''
if verts.shape[0] != 1: # minimal shape: pose it to every pose
assert verts.shape[0] == pose_params.shape[
0] # otherwise the number of results should equal the number of pose identities
verts_posed = []
if save_obj:
if not exists(obj_dir):
os.makedirs(obj_dir)
print('saving results as .obj files to {}...'.format(obj_dir))
if verts.shape[0] == 1:
self.smpl_model.v_template[:] = torch.from_numpy(verts[0])
for i in range(len(pose_params)):
# model.pose[:] = pose_params[i]
self.smpl_model.body_pose[:] = torch.from_numpy(
pose_params[i][3:])
self.smpl_model.global_orient[:] = torch.from_numpy(
pose_params[i][:3])
verts_out = self.smpl_model().vertices.detach().cpu().numpy()
verts_posed.append(verts_out)
if save_obj:
if cloth_type is not None:
Mesh(verts_out.squeeze(),
self.smpl_model.faces).write_obj(
join(obj_dir,
'{}_{:0>4d}.obj').format(cloth_type, i))
else:
Mesh(verts_out.squeeze(),
self.smpl_model.faces).write_obj(
join(obj_dir, '{:0>4d}.obj').format(i))
else:
for i in range(len(verts)):
self.smpl_model.v_template[:] = torch.from_numpy(verts[i])
self.smpl_model.body_pose[:] = torch.from_numpy(
pose_params[i][3:])
self.smpl_model.global_orient[:] = torch.from_numpy(
pose_params[i][:3])
verts_out = self.smpl_model().vertices.detach().cpu().numpy()
verts_posed.append(verts_out)
if save_obj:
if cloth_type is not None:
Mesh(verts_out.squeeze(),
self.smpl_model.faces).write_obj(
join(obj_dir,
'{}_{:0>4d}.obj').format(cloth_type, i))
else:
Mesh(verts_out.squeeze(),
self.smpl_model.faces).write_obj(
join(obj_dir, '{:0>4d}.obj').format(i))
return verts_posed
def _render_sequences_helper(self, video_fname, seq_raw_audio,
seq_processed_audio, seq_template, seq_verts,
condition_idx):
def add_image_text(img, text):
font = cv2.FONT_HERSHEY_SIMPLEX
textsize = cv2.getTextSize(text, font, 1, 2)[0]
textX = (img.shape[1] - textsize[0]) // 2
textY = textsize[1] + 10
cv2.putText(img, '%s' % (text), (textX, textY), font, 1,
(0, 0, 255), 2, cv2.LINE_AA)
num_frames = seq_verts.shape[0]
tmp_audio_file = tempfile.NamedTemporaryFile(
'w', suffix='.wav', dir=os.path.dirname(video_fname))
wavfile.write(tmp_audio_file.name, seq_raw_audio['sample_rate'],
seq_raw_audio['audio'])
tmp_video_file = tempfile.NamedTemporaryFile(
'w', suffix='.mp4', dir=os.path.dirname(video_fname))
if int(cv2.__version__[0]) < 3:
print('cv2 < 3')
writer = cv2.VideoWriter(tmp_video_file.name,
cv2.cv.CV_FOURCC(*'mp4v'), 25,
(1600, 800), True)
else:
print('cv2 >= 3')
writer = cv2.VideoWriter(tmp_video_file.name,
cv2.VideoWriter_fourcc(*'mp4v'), 25,
(1600, 800), True)
feed_dict = {
self.speech_features:
np.expand_dims(np.stack(seq_processed_audio), -1),
self.condition_subject_id:
np.repeat(condition_idx, num_frames),
self.is_training:
False,
self.input_template:
np.repeat(seq_template[np.newaxis, :, :, np.newaxis],
num_frames,
axis=0)
}
predicted_vertices, predicted_offset = self.session.run(
[self.output_decoder, self.expression_offset], feed_dict)
predicted_vertices = np.squeeze(predicted_vertices)
center = np.mean(seq_verts[0], axis=0)
for i_frame in range(min(200, num_frames)):
gt_img = render_mesh_helper(
Mesh(seq_verts[i_frame], self.template_mesh.f), center)
# add_image_text(gt_img, 'Captured data')
pred_img = render_mesh_helper(
Mesh(predicted_vertices[i_frame], self.template_mesh.f),
center)
# add_image_text(pred_img, 'VOCA prediction')
img = np.hstack((gt_img, pred_img))
writer.write(img)
writer.release()
cmd = (
'ffmpeg' +
' -i {0} -i {1} -vcodec h264 -ac 2 -strict -2 -channel_layout stereo -pix_fmt yuv420p {2}'
.format(tmp_audio_file.name, tmp_video_file.name,
video_fname)).split()
call(cmd)
def main(args):
if not os.path.exists(args.conf):
print('Config not found' + args.conf)
config = read_config(args.conf)
for k in config.keys():
print(k, config[k])
print('Initializing parameters')
template_file_path = config['template_fname']
template_mesh = Mesh(filename=template_file_path)
if args.checkpoint_dir:
checkpoint_dir = args.checkpoint_dir
else:
checkpoint_dir = config['checkpoint_dir']
if not os.path.exists(checkpoint_dir):
os.makedirs(checkpoint_dir)
visualize = config['visualize']
output_dir = config['visual_output_dir']
if visualize is True and not output_dir:
print(
'No visual output directory is provided. Checkpoint directory will be used to store the visual results'
)
output_dir = checkpoint_dir
if output_dir and not os.path.exists(output_dir):
os.makedirs(output_dir)
eval_flag = config['eval']
lr = config['learning_rate']
lr_decay = config['learning_rate_decay']
weight_decay = config['weight_decay']
total_epochs = config['epoch']
workers_thread = config['workers_thread']
opt = config['optimizer']
batch_size = config['batch_size']
val_losses, accs, durations = [], [], []
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Generating transforms')
M, A, D, U = mesh_operations.generate_transform_matrices(
template_mesh, config['downsampling_factors'])
D_t = [scipy_to_torch_sparse(d).to(device) for d in D]
U_t = [scipy_to_torch_sparse(u).to(device) for u in U]
A_t = [scipy_to_torch_sparse(a).to(device) for a in A]
num_nodes = [len(M[i].v) for i in range(len(M))]
print('Loading Dataset')
if args.data_dir:
data_dir = args.data_dir
else:
data_dir = config['data_dir']
normalize_transform = Normalize()
dataset = ComaDataset(data_dir,
dtype='train',
split=args.split,
split_term=args.split_term,
pre_transform=normalize_transform)
dataset_val = ComaDataset(data_dir,
dtype='val',
split=args.split,
split_term=args.split_term,
pre_transform=normalize_transform)
dataset_test = ComaDataset(data_dir,
dtype='test',
split=args.split,
split_term=args.split_term,
pre_transform=normalize_transform)
train_loader = DataLoader(dataset,
batch_size=batch_size,
shuffle=True,
num_workers=workers_thread)
val_loader = DataLoader(dataset_val,
batch_size=1,
shuffle=True,
num_workers=workers_thread)
test_loader = DataLoader(dataset_test,
batch_size=1,
shuffle=False,
num_workers=workers_thread)
print('Loading model')
start_epoch = 1
coma = Coma(dataset, config, D_t, U_t, A_t, num_nodes)
if opt == 'adam':
optimizer = torch.optim.Adam(coma.parameters(),
lr=lr,
weight_decay=weight_decay)
elif opt == 'sgd':
optimizer = torch.optim.SGD(coma.parameters(),
lr=lr,
weight_decay=weight_decay,
momentum=0.9)
else:
raise Exception('No optimizer provided')
checkpoint_file = config['checkpoint_file']
print(checkpoint_file)
if checkpoint_file:
checkpoint = torch.load(checkpoint_file)
start_epoch = checkpoint['epoch_num']
coma.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
#To find if this is fixed in pytorch
for state in optimizer.state.values():
for k, v in state.items():
if isinstance(v, torch.Tensor):
state[k] = v.to(device)
coma.to(device)
if eval_flag:
val_loss = evaluate(coma, output_dir, test_loader, dataset_test,
template_mesh, device, visualize)
print('val loss', val_loss)
return
best_val_loss = float('inf')
val_loss_history = []
from datetime import datetime
current_time = datetime.now().strftime('%b%d_%H-%M-%S')
log_dir = os.path.join('runs/clae_dxo_de', current_time)
writer = SummaryWriter(log_dir + 'ds2_lr0.04_z8')
for epoch in range(start_epoch, total_epochs + 1):
print("Training for epoch ", epoch)
train_loss = train(coma, train_loader, len(dataset), optimizer, device)
#train_loss = train(coma, train_loader, len(dataset), optimizer, device)
val_loss = evaluate(coma,
output_dir,
val_loader,
dataset_val,
template_mesh,
device,
epoch,
visualize=visualize)
writer.add_scalar('data/train_loss', train_loss, epoch)
writer.add_scalar('data/val_loss', val_loss, epoch)
print('epoch ', epoch, ' Train loss ', train_loss, ' Val loss ',
val_loss)
if val_loss < best_val_loss:
save_model(coma, optimizer, epoch, train_loss, val_loss,
checkpoint_dir)
best_val_loss = val_loss
if epoch == total_epochs or epoch % 100 == 0:
save_model(coma, optimizer, epoch, train_loss, val_loss,
checkpoint_dir)
val_loss_history.append(val_loss)
val_losses.append(best_val_loss)
if opt == 'sgd':
adjust_learning_rate(optimizer, lr_decay)
if torch.cuda.is_available():
torch.cuda.synchronize()
writer.close()
def render_mesh_helper(mesh,
t_center,
rot=np.zeros(3),
tex_img=None,
v_colors=None,
errors=None,
error_unit='m',
min_dist_in_mm=0.0,
max_dist_in_mm=3.0,
z_offset=0):
camera_params = {
'c': np.array([400, 400]),
'k': np.array([-0.19816071, 0.92822711, 0, 0, 0]),
'f': np.array([4754.97941935 / 2, 4754.97941935 / 2])
}
frustum = {'near': 0.01, 'far': 3.0, 'height': 800, 'width': 800}
mesh_copy = Mesh(mesh.v, mesh.f)
mesh_copy.v[:] = cv2.Rodrigues(rot)[0].dot(
(mesh_copy.v - t_center).T).T + t_center
texture_rendering = tex_img is not None and hasattr(
mesh, 'vt') and hasattr(mesh, 'ft')
if texture_rendering:
intensity = 0.5
tex = pyrender.Texture(source=tex_img, source_channels='RGB')
material = pyrender.material.MetallicRoughnessMaterial(
baseColorTexture=tex)
# Workaround as pyrender requires number of vertices and uv coordinates to be the same
temp_filename = '%s.obj' % next(tempfile._get_candidate_names())
mesh.write_obj(temp_filename)
tri_mesh = trimesh.load(temp_filename, process=False)
try:
os.remove(temp_filename)
except:
print('Failed deleting temporary file - %s' % temp_filename)
render_mesh = pyrender.Mesh.from_trimesh(tri_mesh, material=material)
elif errors is not None:
intensity = 0.5
unit_factor = get_unit_factor('mm') / get_unit_factor(error_unit)
errors = unit_factor * errors
norm = mpl.colors.Normalize(vmin=min_dist_in_mm, vmax=max_dist_in_mm)
cmap = cm.get_cmap(name='jet')
colormapper = cm.ScalarMappable(norm=norm, cmap=cmap)
rgba_per_v = colormapper.to_rgba(errors)
rgb_per_v = rgba_per_v[:, 0:3]
elif v_colors is not None:
intensity = 0.5
rgb_per_v = v_colors
else:
intensity = 1.5
rgb_per_v = None
if not texture_rendering:
tri_mesh = trimesh.Trimesh(vertices=mesh_copy.v,
faces=mesh_copy.f,
vertex_colors=rgb_per_v)
render_mesh = pyrender.Mesh.from_trimesh(tri_mesh, smooth=True)
scene = pyrender.Scene(ambient_light=[.2, .2, .2],
bg_color=[255, 255, 255])
camera = pyrender.IntrinsicsCamera(fx=camera_params['f'][0],
fy=camera_params['f'][1],
cx=camera_params['c'][0],
cy=camera_params['c'][1],
znear=frustum['near'],
zfar=frustum['far'])
scene.add(render_mesh, pose=np.eye(4))
camera_pose = np.eye(4)
camera_pose[:3, 3] = np.array([0, 0, 1.0 - z_offset])
scene.add(camera,
pose=[[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, 1], [0, 0, 0, 1]])
angle = np.pi / 6.0
pos = camera_pose[:3, 3]
light_color = np.array([1., 1., 1.])
light = pyrender.PointLight(color=light_color, intensity=intensity)
light_pose = np.eye(4)
light_pose[:3, 3] = pos
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([angle, 0, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([-angle, 0, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([0, -angle, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
light_pose[:3, 3] = cv2.Rodrigues(np.array([0, angle, 0]))[0].dot(pos)
scene.add(light, pose=light_pose.copy())
flags = pyrender.RenderFlags.SKIP_CULL_FACES
try:
r = pyrender.OffscreenRenderer(viewport_width=frustum['width'],
viewport_height=frustum['height'])
color, _ = r.render(scene, flags=flags)
except:
print('pyrender: Failed rendering frame')
color = np.zeros((frustum['height'], frustum['width'], 3),
dtype='uint8')
return color[..., ::-1]
