def run(in_files, out_dir, ref_file):
ref_v, ref_f = read_trimesh(ref_file)
ref_v = torch.from_numpy(ref_v[:, :3]).float()
ref_f = torch.from_numpy(ref_f).long()
_, ref_area = compute_face_normals_and_areas(ref_v, ref_f)
ref_area = torch.sum(ref_area, dim=-1)
for in_file in in_files:
v, f = read_trimesh(in_file)
v = torch.from_numpy(v[:, :3]).float()
f = torch.from_numpy(f).long()
v, _, _ = center_bounding_box(v)
_, area = compute_face_normals_and_areas(v, f)
area = torch.sum(area, dim=-1)
ratio = torch.sqrt(ref_area/area)
ratio = ratio.unsqueeze(-1).unsqueeze(-1)
v = v * ratio
out_path = os.path.join(out_dir, os.path.basename(in_file))
pymesh.save_mesh_raw(out_path, v.numpy(), f.numpy())
print("saved to {}".format(out_path))
def test(net=None, save_subdir="test"):
opt.phase = "test"
dataset = build_dataset(opt)
if opt.dim == 3:
init_cage_V, init_cage_Fs = loadInitCage([opt.template])
cage_V_t = init_cage_V.transpose(1, 2).detach().cuda()
else:
init_cage_V = generatePolygon(0, 0, 1.5, 0, 0, 0, opt.cage_deg)
init_cage_V = torch.tensor([(x, y) for x, y in init_cage_V],
dtype=torch.float).unsqueeze(0)
cage_V_t = init_cage_V.transpose(1, 2).detach().cuda()
init_cage_Fs = [
torch.arange(opt.cage_deg, dtype=torch.int64).view(1, 1,
-1).cuda()
]
if net is None:
# network
net = networks.NetworkFull(
opt,
dim=opt.dim,
bottleneck_size=opt.bottleneck_size,
template_vertices=cage_V_t,
template_faces=init_cage_Fs[-1],
).cuda()
net.eval()
load_network(net, opt.ckpt)
else:
net.eval()
print(net)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=1,
shuffle=False,
drop_last=False,
collate_fn=tolerating_collate,
num_workers=0,
worker_init_fn=lambda id: np.random.seed(np.random.get_state()[1][0] +
id))
test_output_dir = os.path.join(opt.log_dir, save_subdir)
os.makedirs(test_output_dir, exist_ok=True)
with open(os.path.join(test_output_dir, "eval.txt"), "w") as f:
with torch.no_grad():
for i, data in enumerate(dataloader):
data = dataset.uncollate(data)
############# blending ############
# sample 4 different alpha
if opt.blend_style:
num_alpha = 4
blend_alpha = torch.linspace(
0, 1, steps=num_alpha,
dtype=torch.float32).cuda().reshape(num_alpha, 1)
data["source_shape"] = data["source_shape"].expand(
num_alpha, -1, -1).contiguous()
data["target_shape"] = data["target_shape"].expand(
num_alpha, -1, -1).contiguous()
else:
blend_alpha = 1.0
data["alpha"] = blend_alpha
###################################
source_shape_t = data["source_shape"].transpose(
1, 2).contiguous().detach()
target_shape_t = data["target_shape"].transpose(
1, 2).contiguous().detach()
outputs = net(source_shape_t, target_shape_t, blend_alpha)
deformed = outputs["deformed"]
####################### evaluation ########################
s_filename = os.path.splitext(data["source_file"][0])[0]
t_filename = os.path.splitext(data["target_file"][0])[0]
log_str = "{}/{} {}-{} ".format(i, len(dataloader), s_filename,
t_filename)
print(log_str)
f.write(log_str + "\n")
###################### outputs ############################
for b in range(deformed.shape[0]):
if "source_mesh" in data and data[
"source_mesh"] is not None:
if isinstance(data["source_mesh"][0], str):
source_mesh = om.read_polymesh(
data["source_mesh"][0]).points().copy()
source_mesh = dataset.normalize(
source_mesh, opt.isV2)
source_mesh = torch.from_numpy(
source_mesh.astype(
np.float32)).unsqueeze(0).cuda()
deformed = deform_with_MVC(
outputs["cage"][b:b + 1],
outputs["new_cage"][b:b + 1],
outputs["cage_face"], source_mesh)
else:
deformed = deform_with_MVC(
outputs["cage"][b:b + 1],
outputs["new_cage"][b:b + 1],
outputs["cage_face"], data["source_mesh"])
deformed[b] = center_bounding_box(deformed[b])[0]
if data["source_face"] is not None and data[
"source_mesh"] is not None:
source_mesh = data["source_mesh"][0].detach().cpu()
source_mesh = center_bounding_box(source_mesh)[0]
source_face = data["source_face"][0].detach().cpu()
tosave = pymesh.form_mesh(vertices=source_mesh,
faces=source_face)
pymesh.save_mesh(os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sa.obj".format(s_filename, t_filename)),
tosave,
use_float=True)
tosave = pymesh.form_mesh(
vertices=deformed[0].detach().cpu(),
faces=source_face)
pymesh.save_mesh(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sab-{}.obj".format(
s_filename, t_filename, b)),
tosave,
use_float=True,
)
elif data["source_face"] is None and isinstance(
data["source_mesh"][0], str):
orig_file_path = data["source_mesh"][0]
mesh = om.read_polymesh(orig_file_path)
points_arr = mesh.points()
points_arr[:] = source_mesh[0].detach().cpu().numpy()
om.write_mesh(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sa.obj".format(s_filename, t_filename)),
mesh)
points_arr[:] = deformed[0].detach().cpu().numpy()
om.write_mesh(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sab-{}.obj".format(
s_filename, t_filename, b)), mesh)
else:
# save to "pts" for rendering
save_pts(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sa.pts".format(s_filename, t_filename)),
data["source_shape"][b].detach().cpu())
save_pts(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sab-{}.pts".format(
s_filename, t_filename, b)),
deformed[0].detach().cpu())
if data["target_face"] is not None and data[
"target_mesh"] is not None:
data["target_mesh"][0] = center_bounding_box(
data["target_mesh"][0])[0]
tosave = pymesh.form_mesh(
vertices=data["target_mesh"][0].detach().cpu(),
faces=data["target_face"][0].detach().cpu())
pymesh.save_mesh(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sb.obj".format(s_filename, t_filename)),
tosave,
use_float=True,
)
elif data["target_face"] is None and isinstance(
data["target_mesh"][0], str):
orig_file_path = data["target_mesh"][0]
mesh = om.read_polymesh(orig_file_path)
points_arr = mesh.points()
points_arr[:] = dataset.normalize(
points_arr.copy(), opt.isV2)
om.write_mesh(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sb.obj".format(s_filename, t_filename)),
mesh)
else:
save_pts(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-Sb.pts".format(s_filename, t_filename)),
data["target_shape"][0].detach().cpu())
outputs["cage"][b] = center_bounding_box(
outputs["cage"][b])[0]
outputs["new_cage"][b] = center_bounding_box(
outputs["new_cage"][b])[0]
pymesh.save_mesh_raw(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-cage1-{}.ply".format(
s_filename, t_filename, b)),
outputs["cage"][b].detach().cpu(),
outputs["cage_face"][0].detach().cpu(),
binary=True)
pymesh.save_mesh_raw(
os.path.join(
opt.log_dir, save_subdir,
"{}-{}-cage2-{}.ply".format(
s_filename, t_filename, b)),
outputs["new_cage"][b].detach().cpu(),
outputs["cage_face"][0].detach().cpu(),
binary=True)
dataset.render_result(test_output_dir)
def optimize(opt):
"""
weights are the same with the original source mesh
target=net(old_source)
"""
# load new target
if opt.is_poly:
target_mesh = om.read_polymesh(opt.model)
else:
target_mesh = om.read_trimesh(opt.model)
target_shape_arr = target_mesh.points()
target_shape = target_shape_arr.copy()
target_shape = torch.from_numpy(
target_shape[:, :3].astype(np.float32)).cuda()
target_shape.unsqueeze_(0)
states = torch.load(opt.ckpt)
if "states" in states:
states = states["states"]
cage_v = states["template_vertices"].transpose(1, 2).cuda()
cage_f = states["template_faces"].cuda()
shape_v = states["source_vertices"].transpose(1, 2).cuda()
shape_f = states["source_faces"].cuda()
if os.path.isfile(opt.model.replace(os.path.splitext(opt.model)[1], ".picked")) and os.path.isfile(opt.source_model.replace(os.path.splitext(opt.source_model)[1], ".picked")):
new_label_path = opt.model.replace(os.path.splitext(opt.model)[1], ".picked")
orig_label_path = opt.source_model.replace(os.path.splitext(opt.source_model)[1], ".picked")
logger.info("Loading picked labels {} and {}".format(orig_label_path, new_label_path))
import pandas as pd
new_label = pd.read_csv(new_label_path, delimiter=" ",skiprows=1, header=None)
orig_label = pd.read_csv(orig_label_path, delimiter=" ",skiprows=1, header=None)
orig_label_name = orig_label.iloc[:,5]
new_label_name = new_label.iloc[:,5].tolist()
new_to_orig_idx = []
for i, name in enumerate(new_label_name):
matched_idx = orig_label_name[orig_label_name==name].index
if matched_idx.size == 1:
new_to_orig_idx.append((i, matched_idx[0]))
new_to_orig_idx = np.array(new_to_orig_idx)
if new_label.shape[1] == 10:
new_vidx = new_label.iloc[:,9].to_numpy()[new_to_orig_idx[:,0]]
target_points = target_shape[:, new_vidx, :]
else:
new_label_points = torch.from_numpy(new_label.iloc[:,6:9].to_numpy().astype(np.float32))
target_points = new_label_points.unsqueeze(0).cuda()
target_points, new_vidx, _ = faiss_knn(1, target_points, target_shape, NCHW=False)
target_points = target_points.squeeze(2) # B,N,3
new_label[9] = new_vidx.squeeze(0).squeeze(-1).cpu().numpy()
new_label.to_csv(new_label_path, sep=" ", header=[str(new_label.shape[0])]+[""]*(new_label.shape[1]-1), index=False)
target_points = target_points[:, new_to_orig_idx[:,0], :]
target_points = target_points.cuda()
source_shape, _ = read_trimesh(opt.source_model)
source_shape = torch.from_numpy(source_shape[None, :,:3]).float()
if orig_label.shape[1] == 10:
orig_vidx = orig_label.iloc[:,9].to_numpy()[new_to_orig_idx[:,1]]
source_points = source_shape[:, orig_vidx, :]
else:
orig_label_points = torch.from_numpy(orig_label.iloc[:,6:9].to_numpy().astype(np.float32))
source_points = orig_label_points.unsqueeze(0)
# find the closest point on the original meshes
source_points, new_vidx, _ = faiss_knn(1, source_points, source_shape, NCHW=False)
source_points = source_points.squeeze(2) # B,N,3
orig_label[9] = new_vidx.squeeze(0).squeeze(-1).cpu().numpy()
orig_label.to_csv(orig_label_path, sep=" ", header=[str(orig_label.shape[0])]+[""]*(orig_label.shape[1]-1), index=False)
source_points = source_points[:,new_to_orig_idx[:,1],:]
_, source_center, _ = center_bounding_box(source_shape[0])
source_points -= source_center
source_points = source_points.cuda()
# # shift target so that the belly match
# try:
# orig_bellyUp_idx = orig_label_name[orig_label_name=="bellUp"].index[0]
# orig_bellyUp = orig_label_points[orig_bellyUp_idx, :]
# new_bellyUp_idx = [i for i, i2 in new_to_orig_idx if i2==orig_bellyUp_idx][0]
# new_bellyUp = new_label_points[new_bellyUp_idx,:]
# target_points += (orig_bellyUp - new_bellyUp)
# except Exception as e:
# logger.warn("Couldn\'t match belly to belly")
# traceback.print_exc(file=sys.stdout)
# source_points[0] = center_bounding_box(source_points[0])[0]
elif not os.path.isfile(opt.model.replace(os.path.splitext(opt.model)[1], ".picked")) and os.path.isfile(opt.source_model.replace(os.path.splitext(opt.source_model)[1], ".picked")):
logger.info("Assuming Faust model")
orig_label_path = opt.source_model.replace(os.path.splitext(opt.source_model)[1], ".picked")
logger.info("Loading picked labels {}".format(orig_label_path))
import pandas as pd
orig_label = pd.read_csv(orig_label_path, delimiter=" ",skiprows=1, header=None)
orig_label_name = orig_label.iloc[:,5]
source_shape, _ = read_trimesh(opt.source_model)
source_shape = torch.from_numpy(source_shape[None, :,:3]).cuda().float()
if orig_label.shape[1] == 10:
idx = torch.from_numpy(orig_label.iloc[:,9].to_numpy()).long()
source_points = source_shape[:,idx,:]
target_points = target_shape[:,idx,:]
else:
source_points = torch.from_numpy(orig_label.iloc[:,6:9].to_numpy().astype(np.float32))
source_points = source_points.unsqueeze(0).cuda()
# find the closest point on the original meshes
source_points, idx, _ = faiss_knn(1, source_points, source_shape, NCHW=False)
source_points = source_points.squeeze(2) # B,N,3
idx = idx.squeeze(-1)
target_points = target_shape[:,idx,:]
_, source_center, _ = center_bounding_box(source_shape[0])
source_points -= source_center
elif opt.corres_idx is None and target_shape.shape[1] == shape_v.shape[1]:
logger.info("No correspondence provided, assuming registered Faust models")
# corresp_idx = torch.randint(0, shape_f.shape[1], (100,)).cuda()
corresp_v = torch.unique(torch.randint(0, shape_v.shape[1], (4800,))).cuda()
target_points = torch.index_select(target_shape, 1, corresp_v)
source_points = torch.index_select(shape_v, 1, corresp_v)
target_shape[0], target_center, target_scale = center_bounding_box(target_shape[0])
_, _, source_scale = center_bounding_box(shape_v[0])
target_scale_factor = (source_scale/target_scale)[1]
target_shape *= target_scale_factor
target_points -= target_center
target_points = (target_points*target_scale_factor).detach()
# make sure test use the normalized
target_shape_arr[:] = target_shape[0].cpu().numpy()
om.write_mesh(os.path.join(opt.log_dir, opt.subdir, os.path.splitext(
os.path.basename(opt.model))[0]+"_normalized.obj"), target_mesh)
opt.model = os.path.join(opt.log_dir, opt.subdir, os.path.splitext(
os.path.basename(opt.model))[0]+"_normalized.obj")
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-initial.obj"),
shape_v[0].cpu().numpy(), shape_f[0].cpu().numpy())
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "cage-initial.obj"),
cage_v[0].cpu().numpy(), cage_f[0].cpu().numpy())
save_ply(target_points[0].cpu().numpy(), os.path.join(
opt.log_dir, opt.subdir, "target_points.ply"))
save_ply(source_points[0].cpu().numpy(), os.path.join(
opt.log_dir, opt.subdir, "source_points.ply"))
logger.info("Optimizing for {} corresponding vertices".format(
target_points.shape[1]))
cage_init = cage_v.clone().detach()
lap_loss = MeshLaplacianLoss(torch.nn.MSELoss(reduction="none"), use_cot=True,
use_norm=True, consistent_topology=True, precompute_L=True)
mvc_reg_loss = MVCRegularizer(threshold=50, beta=1.0, alpha=0.0)
cage_v.requires_grad_(True)
optimizer = torch.optim.Adam([cage_v], lr=opt.lr, betas=(0.5, 0.9))
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, int(opt.nepochs*0.4), gamma=0.5, last_epoch=-1)
if opt.dim == 3:
weights_ref = mean_value_coordinates_3D(
source_points, cage_init, cage_f, verbose=False)
else:
raise NotImplementedError
for t in range(opt.nepochs):
optimizer.zero_grad()
weights = mean_value_coordinates_3D(
target_points, cage_v, cage_f, verbose=False)
loss_mvc = torch.mean((weights-weights_ref)**2)
# reg = torch.sum((cage_init-cage_v)**2, dim=-1)*1e-4
reg = 0
if opt.clap_weight > 0:
reg = lap_loss(cage_init, cage_v, face=cage_f)*opt.clap_weight
reg = reg.mean()
if opt.mvc_weight > 0:
reg += mvc_reg_loss(weights)*opt.mvc_weight
# weight regularizer with the shape difference
# dist = torch.sum((source_points - target_points)**2, dim=-1)
# weights = torch.exp(-dist)
# reg = reg*weights*0.1
loss = loss_mvc + reg
if (t+1) % 50 == 0:
print("t {}/{} mvc_loss: {} reg: {}".format(t,
opt.nepochs, loss_mvc.item(), reg.item()))
if loss_mvc.item() < 5e-6:
break
loss.backward()
optimizer.step()
scheduler.step()
return cage_v, cage_f