def get_normalize_mesh(model_file, norm_mesh_sub_dir):
total = 16384
print("trimesh_load:", model_file)
mesh_list = trimesh.load_mesh(model_file, process=False)
if not isinstance(mesh_list, list):
mesh_list = [mesh_list]
area_sum = 0
area_lst = []
for idx, mesh in enumerate(mesh_list):
area = np.sum(mesh.area_faces)
area_lst.append(area)
area_sum += area
area_lst = np.asarray(area_lst)
amount_lst = (area_lst * total / area_sum).astype(np.int32)
points_all = np.zeros((0, 3), dtype=np.float32)
for i in range(amount_lst.shape[0]):
mesh = mesh_list[i]
print("start sample surface of ", mesh.faces.shape[0])
points, index = trimesh.sample.sample_surface(mesh, amount_lst[i])
print("end sample surface")
points_all = np.concatenate([points_all, points], axis=0)
centroid = np.mean(points_all, axis=0)
points_all = points_all - centroid
m = np.max(np.sqrt(np.sum(points_all**2, axis=1)))
obj_file = os.path.join(norm_mesh_sub_dir, "pc_norm.obj")
ori_mesh = pymesh.load_mesh(model_file)
print("centroid, m", centroid, m)
pymesh.save_mesh_raw(obj_file, (ori_mesh.vertices - centroid) / float(m),
ori_mesh.faces)
print("export_mesh", obj_file)
return obj_file, centroid, m
def separate_single_mesh(src_mesh, tar_mesh):
src_mesh = pymesh.load_mesh(src_mesh)
dis_meshes = pymesh.separate_mesh(src_mesh, connectivity_type='auto')
pymesh.save_mesh_raw(tar_mesh+".obj", src_mesh.vertices, src_mesh.faces)
count=0
for dis_mesh in dis_meshes:
print("dis_mesh.vertices.shape",dis_mesh.vertices.shape)
print("dis_mesh.faces.shape",dis_mesh.faces.shape)
pymesh.save_mesh_raw(tar_mesh+"_"+str(count)+".obj", dis_mesh.vertices, dis_mesh.faces)
count+=1
def test_one(opt, cage_shape, new_source, new_source_face, new_target, new_target_face):
states = torch.load(opt.ckpt)
if "states" in states:
states = states["states"]
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-initial.ply"),
states["source_vertices"][0].transpose(
0, 1).detach().cpu(),
states["source_faces"][0].detach().cpu())
# states["template_vertices"] = cage_shape.transpose(1, 2)
# states["source_vertices"] = new_source.transpose(1, 2)
# states["source_faces"] = new_source_face
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-Sa.ply"),
new_source[0].detach().cpu(), new_source_face[0].detach().cpu())
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-Sb.ply"),
new_target[0].detach().cpu(), new_target_face[0].detach().cpu())
net = networks.FixedSourceDeformer(opt, 3, opt.num_point, bottleneck_size=512,
template_vertices=cage_shape.transpose(1, 2), template_faces=states["template_faces"].cuda(),
source_vertices=new_source.transpose(1, 2), source_faces=new_source_face).cuda()
net.eval()
load_network(net, states)
outputs = net(new_target.transpose(1, 2).contiguous())
deformed = outputs["deformed"]
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-Sab.ply"),
deformed[0].detach().cpu(), new_target_face[0].detach().cpu())
def main():
args = parse_args();
mesh = pymesh.load_mesh(args.input_mesh);
offset = np.array(args.offset);
axis = np.array(args.rotation_axis);
angle = math.radians(args.rotation_angle);
rot = pymesh.Quaternion.fromAxisAngle(axis, angle);
rot = rot.to_matrix();
vertices = mesh.vertices;
bbox = mesh.bbox;
centroid = 0.5 * (bbox[0] + bbox[1]);
vertices = np.dot(rot, (vertices - centroid).T).T + centroid + offset;
pymesh.save_mesh_raw(args.output_mesh, vertices, mesh.faces, mesh.voxels);
def main():
args = parse_args()
mesh = pymesh.load_mesh(args.input_mesh)
offset = np.array(args.offset)
axis = np.array(args.rotation_axis)
angle = math.radians(args.rotation_angle)
rot = pymesh.Quaternion.fromAxisAngle(axis, angle)
rot = rot.to_matrix()
vertices = mesh.vertices
bbox = mesh.bbox
centroid = 0.5 * (bbox[0] + bbox[1])
vertices = np.dot(rot, (vertices - centroid).T).T + centroid + offset
pymesh.save_mesh_raw(args.output_mesh, vertices, mesh.faces, mesh.voxels)
def clean_single_mesh(src_mesh, tar_mesh, dist_thresh, num_thresh):
src_mesh_obj = pymesh.load_mesh(src_mesh)
dis_meshes = pymesh.separate_mesh(src_mesh_obj, connectivity_type='auto')
max_mesh_verts = 0
for dis_mesh in dis_meshes:
if dis_mesh.vertices.shape[0] > max_mesh_verts:
max_mesh_verts = dis_mesh.vertices.shape[0]
collection=[]
for dis_mesh in dis_meshes:
if dis_mesh.vertices.shape[0] > max_mesh_verts*num_thresh:
centroid = np.mean(dis_mesh.vertices, axis=0)
if np.sqrt(np.sum(np.square(centroid))) < dist_thresh:
collection.append(dis_mesh)
tar_mesh_obj = pymesh.merge_meshes(collection)
pymesh.save_mesh_raw(tar_mesh, tar_mesh_obj.vertices, tar_mesh_obj.faces)
print("threshes:", str(dist_thresh), str(num_thresh), " clean: ", src_mesh, " create: ",tar_mesh)
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 get_normalize_mesh(model_dir, norm_mesh_sub_dir, target_dir):
norm_file = os.path.join(norm_mesh_sub_dir, "pc_norm.obj")
target_norm_file = os.path.join(target_dir, "pc_norm.obj")
command_str = "cp " + norm_file + " " + target_norm_file
print("command:", command_str)
os.system(command_str)
model_obj = os.path.join(model_dir, "model.obj")
target_model_obj = os.path.join(target_dir, "model.obj")
params = np.loadtxt(os.path.join(norm_mesh_sub_dir, "pc_norm.txt"))
print("trimesh_load:", model_obj)
centroid = params[:3]
m = params[3]
print("centroid, m", centroid, m)
ori_mesh = pymesh.load_mesh(model_obj)
verts = (ori_mesh.vertices - centroid) / float(m)
pymesh.save_mesh_raw(target_model_obj, verts, ori_mesh.faces)
def train():
dataset = build_dataset(opt)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=opt.batch_size,
shuffle=True,
drop_last=True,
num_workers=0,
worker_init_fn=lambda id: np.random.seed(np.random.get_state()[1][0] +
id))
source_shape = dataset.mesh_vertex.unsqueeze(0).to(dtype=torch.float)
source_face = dataset.mesh_face.unsqueeze(0)
cage_shape = dataset.cage_vertex.unsqueeze(0).to(dtype=torch.float)
cage_face = dataset.cage_face.unsqueeze(0)
mesh = Mesh(vertices=cage_shape[0], faces=cage_face[0])
build_gemm(mesh, cage_face[0])
cage_edge_points = torch.from_numpy(get_edge_points(mesh)).cuda()
cage_edges = edge_vertex_indices(cage_face[0])
# network
net = networks.FixedSourceDeformer(
opt,
3,
opt.num_point,
bottleneck_size=opt.bottleneck_size,
template_vertices=cage_shape.transpose(1, 2),
template_faces=cage_face,
source_vertices=source_shape.transpose(1, 2),
source_faces=source_face).cuda()
print(net)
net.apply(weights_init)
if opt.ckpt:
load_network(net, opt.ckpt)
net.train()
all_losses = losses.AllLosses(opt)
# optimizer
optimizer = torch.optim.Adam([{
'params': net.nd_decoder.parameters()
}, {
"params": net.encoder.parameters()
}],
lr=opt.lr)
# train
os.makedirs(opt.log_dir, exist_ok=True)
shutil.copy2(__file__, opt.log_dir)
shutil.copy2(os.path.join(os.path.dirname(__file__), "network2.py"),
opt.log_dir)
shutil.copy2(os.path.join(os.path.dirname(__file__), "common.py"),
opt.log_dir)
shutil.copy2(os.path.join(os.path.dirname(__file__), "losses.py"),
opt.log_dir)
shutil.copy2(os.path.join(os.path.dirname(__file__), "datasets.py"),
opt.log_dir)
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, "t{:06d}_Sa.ply".format(0)),
net.source_vertices[0].transpose(0, 1).detach().cpu().numpy(),
net.source_faces[0].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, "t{:06d}_template.ply".format(0)),
net.template_vertices[0].transpose(0, 1).detach().cpu().numpy(),
net.template_faces[0].detach().cpu())
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
max(int(opt.nepochs * 0.75),
1),
gamma=0.5,
last_epoch=-1)
# train
net.train()
t = 0
start_epoch = 0
warmed_up = False
mvc_weight = opt.mvc_weight
opt.mvc_weight = 0
os.makedirs(opt.log_dir, exist_ok=True)
running_avg_loss = -1
log_file = open(os.path.join(opt.log_dir, "loss_log.txt"), "a")
log_interval = min(max(len(dataloader) // 5, 50), 200)
save_interval = max(opt.nepochs // 10, 1)
with torch.autograd.detect_anomaly():
if opt.epoch:
start_epoch = opt.epoch % opt.nepochs
t += start_epoch * len(dataloader)
for epoch in range(start_epoch, opt.nepochs):
for epoch_t, data in enumerate(dataloader):
progress = epoch_t / len(dataloader) + epoch
warming_up = progress < opt.warmup_epochs
if (opt.deform_template or opt.optimize_template) and (
progress >= opt.warmup_epochs) and (not warmed_up):
if opt.deform_template:
optimizer.add_param_group({
'params':
net.nc_decoder.parameters(),
'lr':
0.1 * opt.lr
})
if opt.optimize_template:
optimizer.add_param_group({
'params': net.template_vertices,
'lr': 0.1 * opt.lr
})
warmed_up = True
# start to compute mvc weight
opt.mvc_weight = mvc_weight
save_network(net,
opt.log_dir,
network_label="net",
epoch_label="warmed_up")
############# get data ###########
data = dataset.uncollate(data)
data["cage_edge_points"] = cage_edge_points
data["cage_edges"] = cage_edges
data["source_shape"] = net.source_vertices.detach()
data["source_face"] = net.source_faces.detach()
############# run network ###########
optimizer.zero_grad()
target_shape_t = data["target_shape"].transpose(1, 2)
sample_idx = None
if "sample_idx" in data:
sample_idx = data["sample_idx"]
if data["source_normals"] is not None:
data["source_normals"] = torch.gather(
data["source_normals"], 1,
sample_idx.unsqueeze(-1).expand(-1, -1, 3))
outputs = net(target_shape_t, sample_idx)
if opt.sfnormal_weight > 0 and ("source_mesh" in data
and "source_mesh" is not None):
if outputs["deformed"].shape[1] == data[
"source_mesh"].shape[1]:
outputs["deformed_hr"] = outputs["deformed"]
else:
outputs["deformed_hr"] = deform_with_MVC(
outputs["cage"].expand(
data["source_mesh"].shape[0], -1, -1).detach(),
outputs["new_cage"], outputs["cage_face"].expand(
data["source_mesh"].shape[0], -1,
-1), data["source_mesh"])
data["source_shape"] = outputs["source_shape"]
############# get losses ###########
current_loss = all_losses(data, outputs, progress)
loss_sum = torch.sum(
torch.stack([v for v in current_loss.values()], dim=0))
if running_avg_loss < 0:
running_avg_loss = loss_sum
else:
running_avg_loss = running_avg_loss + (
loss_sum.item() - running_avg_loss) / (t + 1)
if (t % log_interval
== 0) or (loss_sum > 10 * running_avg_loss):
log_str = "warming up {} e {:03d} t {:05d}: {}".format(
not warmed_up, epoch, t, ", ".join([
"{} {:.3g}".format(k,
v.mean().item())
for k, v in current_loss.items()
]))
print(log_str)
log_file.write(log_str + "\n")
log_outputs(opt, t, outputs, data)
# save_ply(data["target_shape"][0].detach().cpu().numpy(), os.path.join(opt.log_dir,"step-{:06d}-Sb.ply".format(t)))
# save_ply(outputs["deformed"][0].detach().cpu().numpy(), os.path.join(opt.log_dir,"step-{:06d}-Sab.ply".format(t)))
# write_trimesh(os.path.join(opt.log_dir, "step-{:06d}-cage1.ply".format(t)),
# outputs["cage"][0].detach().cpu(), outputs["cage_face"][0].detach().cpu(), binary=True)
# write_trimesh(os.path.join(opt.log_dir, "step-{:06d}-cage2.ply".format(t)),
# outputs["new_cage"][0].detach().cpu(), outputs["cage_face"][0].detach().cpu(), binary=True)
if loss_sum > 100 * running_avg_loss:
logger.info(
"loss ({}) > 10*running_average_loss ({}). Skip without update."
.format(loss_sum, 5 * running_avg_loss))
torch.cuda.empty_cache()
continue
loss_sum.backward()
if opt.alternate_cd:
optimize_C = (progress > opt.warmup_epochs) and (
t % (opt.c_step + opt.d_step)) > opt.d_step
if optimize_C:
net.nd_decoder.zero_grad()
net.encoder.zero_grad()
else:
try:
net.nc_decoder.zero_grad()
except AttributeError:
net.template_vertices.grad.zero_()
# clamp_gradient_norm(net, 1)
optimizer.step()
if (t + 1) % 500 == 0:
save_network(net,
opt.log_dir,
network_label="net",
epoch_label="latest")
t += 1
if (epoch + 1) % save_interval == 0:
save_network(net,
opt.log_dir,
network_label="net",
epoch_label=epoch)
scheduler.step()
log_file.close()
save_network(net, opt.log_dir, network_label="net", epoch_label="final")
test_all(net=net)
def test(net=None, subdir="test"):
opt.phase = "test"
if isinstance(opt.target_model, str):
opt.target_model = [opt.target_model]
if net is None:
states = torch.load(opt.ckpt)
if "states" in states:
states = states["states"]
if opt.template:
cage_shape, cage_face = read_trimesh(opt.template)
cage_shape = torch.from_numpy(
cage_shape[:, :3]).unsqueeze(0).float()
cage_face = torch.from_numpy(cage_face).unsqueeze(0).long()
states["template_vertices"] = cage_shape.transpose(1, 2)
states["template_faces"] = cage_face
if opt.source_model:
source_shape, source_face = read_trimesh(opt.source_model)
source_shape = torch.from_numpy(
source_shape[:, :3]).unsqueeze(0).float()
source_face = torch.from_numpy(source_face).unsqueeze(0).long()
states["source_vertices"] = source_shape.transpose(1, 2)
states["source_faces"] = source_shape
net = networks.FixedSourceDeformer(
opt,
3,
opt.num_point,
bottleneck_size=opt.bottleneck_size,
template_vertices=states["template_vertices"],
template_faces=states["template_faces"],
source_vertices=states["source_vertices"],
source_faces=states["source_faces"]).cuda()
load_network(net, states)
net = net.cuda()
net.eval()
else:
net.eval()
print(net)
test_output_dir = os.path.join(opt.log_dir, subdir)
os.makedirs(test_output_dir, exist_ok=True)
with torch.no_grad():
for target_model in opt.target_model:
assert (os.path.isfile(target_model))
target_face = None
target_shape, target_face = read_trimesh(target_model)
# target_shape = read_ply(target_model)[:,:3]
# target_shape, _, scale = normalize_to_box(target_shape)
# normalize acording to height y axis
# target_shape = target_shape/2*1.7
target_shape = torch.from_numpy(
target_shape[:, :3]).cuda().float().unsqueeze(0)
if target_face is None:
target_face = net.source_faces
else:
target_face = torch.from_numpy(
target_face).cuda().long().unsqueeze(0)
t_filename = os.path.splitext(os.path.basename(target_model))[0]
source_mesh = net.source_vertices.transpose(1, 2).detach()
source_face = net.source_faces.detach()
# furthest sampling
target_shape_sampled = furthest_point_sample(
target_shape, net.source_vertices.shape[2], NCHW=False)[1]
# target_shape_sampled = (target_shape[:, np.random.permutation(target_shape.shape[1]), :]).contiguous()
outputs = net(target_shape_sampled.transpose(1, 2),
None,
cage_only=True)
# deformed = outputs["deformed"]
deformed = deform_with_MVC(
outputs["cage"], outputs["new_cage"],
outputs["cage_face"].expand(outputs["cage"].shape[0], -1,
-1), source_mesh)
b = 0
save_ply(
target_shape_sampled[b].cpu().numpy(),
os.path.join(opt.log_dir, subdir,
"template-{}-Sb.pts".format(t_filename)))
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, subdir,
"template-{}-Sa.ply".format(t_filename)),
source_mesh[0].detach().cpu(), source_face[0].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, subdir,
"template-{}-Sb.ply".format(t_filename)),
target_shape[b].detach().cpu(), target_face[b].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, subdir,
"template-{}-Sab.ply".format(t_filename)),
deformed[b].detach().cpu(), source_face[b].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, subdir,
"template-{}-cage1.ply".format(t_filename)),
outputs["cage"][b].detach().cpu(),
outputs["cage_face"][b].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, subdir,
"template-{}-cage2.ply".format(t_filename)),
outputs["new_cage"][b].detach().cpu(),
outputs["cage_face"][b].detach().cpu())
PairedSurreal.render_result(test_output_dir)
def test_all(net=None, subdir="test"):
opt.phase = "test"
dataset = build_dataset(opt)
if net is None:
source_shape = dataset.mesh_vertex.unsqueeze(0).to(dtype=torch.float)
source_face = dataset.mesh_face.unsqueeze(0)
cage_shape = dataset.cage_vertex.unsqueeze(0).to(dtype=torch.float)
cage_face = dataset.cage_face.unsqueeze(0)
net = networks.FixedSourceDeformer(
opt,
3,
opt.num_point,
bottleneck_size=opt.bottleneck_size,
template_vertices=cage_shape.transpose(1, 2),
template_faces=cage_face,
source_vertices=source_shape.transpose(1, 2),
source_faces=source_face).cuda()
load_network(net, opt.ckpt)
net.eval()
else:
net.eval()
print(net)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=1,
shuffle=False,
drop_last=False,
num_workers=3,
worker_init_fn=lambda id: np.random.seed(np.random.get_state()[1][0] +
id))
chamfer_distance = losses.LabeledChamferDistance(beta=0, gamma=1)
mse_distance = torch.nn.MSELoss()
avg_CD = 0
avg_EMD = 0
test_output_dir = os.path.join(opt.log_dir, 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():
source_mesh = net.source_vertices.transpose(1, 2).detach()
source_face = net.source_faces.detach()
for i, data in enumerate(dataloader):
data = dataset.uncollate(data)
target_shape, target_filename = data["target_shape"], data[
"target_file"]
sample_idx = None
if "sample_idx" in data:
sample_idx = data["sample_idx"]
outputs = net(target_shape.transpose(1, 2), sample_idx)
deformed = outputs["deformed"]
deformed = deform_with_MVC(
outputs["cage"], outputs["new_cage"],
outputs["cage_face"].expand(outputs["cage"].shape[0], -1,
-1), source_mesh)
for b in range(outputs["deformed"].shape[0]):
t_filename = os.path.splitext(target_filename[b])[0]
target_shape_np = target_shape.detach().cpu()[b].numpy()
if data["target_face"] is not None and data[
"target_mesh"] is not None:
pymesh.save_mesh_raw(
os.path.join(
opt.log_dir, subdir,
"template-{}-Sa.ply".format(t_filename)),
source_mesh[0].detach().cpu(),
source_face[0].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(
opt.log_dir, subdir,
"template-{}-Sb.ply".format(t_filename)),
data["target_mesh"][b].detach().cpu(),
data["target_face"][b].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(
opt.log_dir, subdir,
"template-{}-Sab.ply".format(t_filename)),
deformed[b].detach().cpu(),
source_face[b].detach().cpu())
else:
save_ply(
source_mesh[0].detach().cpu(),
os.path.join(
opt.log_dir, subdir,
"template-{}-Sa.ply".format(t_filename)))
save_ply(
target_shape[b].detach().cpu(),
os.path.join(
opt.log_dir, subdir,
"template-{}-Sb.ply".format(t_filename)),
normals=data["target_normals"][b].detach().cpu())
save_ply(
deformed[b].detach().cpu(),
os.path.join(
opt.log_dir, subdir,
"template-{}-Sab.ply".format(t_filename)),
normals=data["target_normals"][b].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(
opt.log_dir, subdir,
"template-{}-cage1.ply".format(t_filename)),
outputs["cage"][b].detach().cpu(),
outputs["cage_face"][b].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(
opt.log_dir, subdir,
"template-{}-cage2.ply".format(t_filename)),
outputs["new_cage"][b].detach().cpu(),
outputs["cage_face"][b].detach().cpu())
log_str = "{}/{} {}".format(i, len(dataloader), t_filename)
print(log_str)
f.write(log_str + "\n")
dataset.render_result(test_output_dir)
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 forward(opt):
"""
Takes an input and a target mesh. Deform input in output and propagate a
manually defined high frequency from the oinput to the output
:return:
"""
my_utils.plant_seeds(randomized_seed=opt.randomize)
os.makedirs(opt.output_dir, exist_ok=True)
trainer = t.Trainer(opt)
trainer.build_dataset_train_for_matching()
trainer.build_dataset_test_for_matching()
trainer.build_network()
trainer.build_losses()
trainer.network.eval()
if opt.eval_list and os.path.isfile(opt.eval_list):
source_target_files = np.loadtxt(opt.eval_list, dtype=str)
source_target_files = source_target_files.tolist()
for i, st in enumerate(source_target_files):
source, target = st
cat1, fname1 = source.split('/')
fname1 = os.path.splitext(fname1)[0]
cat2, fname2 = target.split('/')
fname2 = os.path.splitext(fname2)[0]
if len(opt.shapenetv1_path) > 0:
source_target_files[i] = (os.path.join(opt.shapenetv1_path,
cat1, fname1,
"model.obj"),
os.path.join(opt.shapenetv1_path,
cat2, fname2,
"model.obj"))
elif len(opt.shapenetv2_path) > 0:
source_target_files[i] = (os.path.join(opt.shapenetv2_path,
cat1, fname1, "models",
"model_normalized.obj"),
os.path.join(opt.shapenetv2_path,
cat2, fname2, "models",
"model_normalized.obj"))
elif (opt.eval_source != "" and opt.eval_source[-4:] == ".txt") and (
opt.eval_target != "" and opt.eval_target[-4:] == ".txt"):
source_target_files = [
(figure_2_3.convert_path(opt.shapenetv1_path, opt.eval_source),
figure_2_3.convert_path(opt.shapenetv1_path, opt.eval_target))
]
rot_mat = get_3D_rot_matrix(1, np.pi / 2)
rot_mat_rev = get_3D_rot_matrix(1, -np.pi / 2)
isV2 = len(opt.shapenetv2_path) > 0
for i, source_target in enumerate(source_target_files):
basename = get_model_id(source_target[0], isV2) + "-" + get_model_id(
source_target[1], isV2)
path_deformed = os.path.join(opt.output_dir, basename + "-Sab.ply")
path_source = os.path.join(opt.output_dir, basename + "-Sa.ply")
path_target = os.path.join(opt.output_dir, basename + "-Sb.ply")
mesh_path = source_target[0]
print(mesh_path)
source_mesh_edge = get_shapenet_model.link(mesh_path)
mesh_path = source_target[1]
target_mesh_edge = get_shapenet_model.link(mesh_path)
print("Deforming source in target")
source = source_mesh_edge.vertices
target = target_mesh_edge.vertices
pymesh.save_mesh_raw(path_source,
source,
source_mesh_edge.faces,
ascii=True)
pymesh.save_mesh_raw(path_target,
target,
target_mesh_edge.faces,
ascii=True)
if len(opt.shapenetv2_path) > 0:
source = source.dot(rot_mat)
target = target.dot(rot_mat)
source = torch.from_numpy(source).cuda().float().unsqueeze(0)
target = torch.from_numpy(target).cuda().float().unsqueeze(0)
with torch.no_grad():
source, _, _, _, _ = loss.forward_chamfer(
trainer.network,
source,
target,
local_fix=None,
distChamfer=trainer.distChamfer)
try:
source = source.squeeze().cpu().detach().numpy()
if len(opt.shapenetv2_path) > 0:
source = source.dot(rot_mat_rev)
P2_P1_mesh = pymesh.form_mesh(vertices=source,
faces=source_mesh_edge.faces)
pymesh.save_mesh(path_deformed, P2_P1_mesh, ascii=True)
# print("computing signal tranfer form source to target")
# high_frequencies.high_frequency_propagation(path_source, path_deformed, path_target)
except Exception as e:
print(e)
import pdb
pdb.set_trace()
path_deformed = path_deformed[:-4] + ".pts"
save_pts(path_deformed, source.squeeze().cpu().detach().numpy())
def test_orientation():
m = pymesh.load_mesh(os.path.join(os.path.dirname(__file__), 'source_models/bunny.obj'))
new_mesh = orientation.orient_printed_mesh(m)
pymesh.save_mesh_raw(os.path.join(os.path.dirname(__file__), 'watermarked_models/out_orientation.obj'), new_mesh.vertices, new_mesh.faces)
import pymesh
file_path = "/Users/luchenliu/Desktop/Intern/1_projection/glasses.obj"
mesh = pymesh.load_mesh(file_path)
pymesh.save_mesh_raw(file_path.replace(".obj", ".ply"), mesh.vertices,
mesh.faces, mesh.voxels)
def get_normalize_mesh(model_file, norm_mesh_sub_dir, pnt_dir, ref_sub_dir):
total = 16384 * 50
print("trimesh_load:", model_file)
ref_file = os.path.join(ref_sub_dir, "isosurf.obj")
mesh_list = trimesh.load_mesh(model_file, process=False)
if not isinstance(mesh_list, list):
mesh_list = [mesh_list]
area_sum = 0
area_lst = []
for idx, mesh in enumerate(mesh_list):
area = np.sum(mesh.area_faces)
area_lst.append(area)
area_sum += area
area_lst = np.asarray(area_lst)
amount_lst = (area_lst * total / area_sum).astype(np.int32)
points_all = np.zeros((0, 3), dtype=np.float32)
all_face_normals = np.zeros((0, 3), dtype=np.float32)
all_vert_normals = np.zeros((0, 3, 3), dtype=np.float32)
all_tries = np.zeros((0, 3, 3), dtype=np.float32)
sample_indices = np.zeros((0), dtype=np.int)
for i in range(amount_lst.shape[0]):
mesh = mesh_list[i]
# print("start sample surface of ", mesh.faces.shape[0])
points, index = trimesh.sample.sample_surface(mesh, amount_lst[i])
if not os.path.exists(ref_file):
sample_indices = np.concatenate(
[sample_indices, sample_indices.shape[0] + index], axis=0)
vert_ind = mesh.faces.reshape(-1)
all_tries = np.concatenate(
[all_tries, mesh.vertices[vert_ind].reshape([-1, 3, 3])],
axis=0)
all_face_normals = np.concatenate(
[all_face_normals, mesh.face_normals], axis=0)
all_vert_normals = np.concatenate([
all_vert_normals, mesh.vertex_normals[vert_ind].reshape(
[-1, 3, 3])
],
axis=0)
# print("end sample surface")
points_all = np.concatenate([points_all, points], axis=0)
centroid = np.mean(points_all, axis=0)
points_all = points_all - centroid
m = np.max(np.sqrt(np.sum(points_all**2, axis=1)))
obj_file = os.path.join(norm_mesh_sub_dir, "pc_norm.obj")
param_file = os.path.join(norm_mesh_sub_dir, "pc_norm.txt")
params = np.concatenate([centroid, np.expand_dims(m, axis=0)])
np.savetxt(param_file, params)
print("export_mesh", obj_file)
from_marchingcube = False
if not os.path.exists(ref_file):
ori_mesh = pymesh.load_mesh(model_file)
print("centroid, m", centroid, m)
else:
from_marchingcube = True
mesh_list = trimesh.load_mesh(ref_file, process=False)
print("trimesh_load ref_file:", ref_file)
if not isinstance(mesh_list, list):
mesh_list = [mesh_list]
area_sum = 0
area_lst = []
for idx, mesh in enumerate(mesh_list):
area = np.sum(mesh.area_faces)
area_lst.append(area)
area_sum += area
area_lst = np.asarray(area_lst)
amount_lst = (area_lst * total / area_sum).astype(np.int32)
points_all = np.zeros((0, 3), dtype=np.float32)
for i in range(amount_lst.shape[0]):
mesh = mesh_list[i]
# print("start sample surface of ", mesh.faces.shape[0])
points, index = trimesh.sample.sample_surface(mesh, amount_lst[i])
sample_indices = np.concatenate(
[sample_indices, sample_indices.shape[0] + index], axis=0)
vert_ind = mesh.faces.reshape(-1)
all_tries = np.concatenate(
[all_tries, mesh.vertices[vert_ind].reshape([-1, 3, 3])],
axis=0)
all_face_normals = np.concatenate(
[all_face_normals, mesh.face_normals], axis=0)
all_vert_normals = np.concatenate([
all_vert_normals, mesh.vertex_normals[vert_ind].reshape(
[-1, 3, 3])
],
axis=0)
# print("end sample surface")
points_all = np.concatenate([points_all, points], axis=0)
centroid = np.mean(points_all, axis=0)
points_all = points_all - centroid
m = np.max(np.sqrt(np.sum(points_all**2, axis=1)))
ori_mesh = pymesh.load_mesh(ref_file)
surfpoints = points_all / float(m)
all_tries = (all_tries - centroid[np.newaxis, np.newaxis, :]) / float(m)
print("centroid, m", centroid, m)
verts = (ori_mesh.vertices - centroid) / float(m)
pymesh.save_mesh_raw(obj_file, verts, ori_mesh.faces)
face_norm_surfpnt = save_surface(sample_indices, surfpoints, all_tries,
all_face_normals, all_vert_normals,
pnt_dir)
return all_tries, all_face_normals, all_vert_normals, params, surfpoints, face_norm_surfpnt, from_marchingcube
])
new_face = np.array(new_f)
return new_face
random_sample = np.random.choice(6890, size=6890, replace=False)
random_sample2 = np.random.choice(6890, size=6890, replace=False)
id_mesh = pymesh.load_mesh('./demo_data/13_643.obj')
pose_mesh = pymesh.load_mesh('./demo_data/14_664.obj')
with torch.no_grad():
id_mesh_points = id_mesh.vertices[random_sample]
id_mesh_points = id_mesh_points - (id_mesh.bbox[0] + id_mesh.bbox[1]) / 2
id_mesh_points = torch.from_numpy(id_mesh_points.astype(np.float32)).cuda()
pose_mesh_points = pose_mesh.vertices #[random_sample2]
pose_mesh_points = pose_mesh_points - (pose_mesh.bbox[0] +
pose_mesh.bbox[1]) / 2
pose_mesh_points = torch.from_numpy(pose_mesh_points.astype(
np.float32)).cuda()
pointsReconstructed = net_G(
pose_mesh_points.transpose(0, 1).unsqueeze(0),
id_mesh_points.transpose(0, 1).unsqueeze(0)) # forward pass
new_face = face_reverse(id_mesh.faces)
pymesh.save_mesh_raw('./demo_data/13_664.obj',
pointsReconstructed.cpu().numpy().squeeze(), new_face)
scheduler.step(loss.item())
return deformed_shape
if __name__ == "__main__":
parser = MyOptions()
opt = parser.parse()
opt.log_dir = os.path.dirname(opt.ckpt)
os.makedirs(os.path.join(opt.log_dir, opt.subdir), exist_ok=True)
if opt.use_cage is None:
# optimize initial cage for the new target
cage_v, cage_f = optimize(opt)
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "optimized_template_cage.ply"),
cage_v[0].detach().cpu(), cage_f[0].detach().cpu())
else:
cage_v, cage_f = read_trimesh(opt.use_cage)
cage_v = torch.from_numpy(cage_v[:, :3].astype(np.float32)).cuda()
cage_f = torch.from_numpy(cage_f[:, :3].astype(np.int64)).cuda()
cage_v.unsqueeze_(0)
cage_f.unsqueeze_(0)
# # test using the new source and initial cage
# target_shape_pose, target_face_pose, _ = read_trimesh("/home/mnt/points/data/MPI-FAUST/training/registrations/tr_reg_002.ply")
# target_shape_pose = torch.from_numpy(target_shape_pose[:,:3].astype(np.float32)).cuda()
# target_face_pose = torch.from_numpy(target_face_pose[:,:3].astype(np.int64)).cuda()
# target_shape_pose, _, _ = center_bounding_box(target_shape_pose)
# target_shape_pose.unsqueeze_(0)
# target_face_pose.unsqueeze_(0)
# test_one(opt, cage_v, target_shape, target_face, target_shape_pose, target_face_pose)
def __call__(self, outputMesh):
generate_x = np.array([])
generate_y = np.array([])
generate_z = np.array([])
source_control_x = self.source_vertices_x[0:self.control_num]
source_control_y = self.source_vertices_y[0:self.control_num]
source_control_z = self.source_vertices_z[0:self.control_num]
target_control_x = self.target_vertices_x[0:self.control_num]
target_control_y = self.target_vertices_y[0:self.control_num]
target_control_z = self.target_vertices_z[0:self.control_num]
result_x = source_control_x
result_y = source_control_y
result_z = source_control_z
vertex_num = self.source_vertices_x.size
for i in range(self.control_num, vertex_num):
source_di_x = np.array([])
source_di_y = np.array([])
source_di_z = np.array([])
for j in range(0, self.control_num):
di_x = self.source_vertices_x.item(
i) - self.source_vertices_x.item(j)
di_y = self.source_vertices_y.item(
i) - self.source_vertices_y.item(j)
di_z = self.source_vertices_z.item(
i) - self.source_vertices_z.item(j)
source_di_x = np.append(source_di_x, di_x)
source_di_y = np.append(source_di_y, di_y)
source_di_z = np.append(source_di_z, di_z)
#print(source_di_x.shape)
#print(source_di_y.shape)
#print(source_di_z.shape)
#print(source_control_x.shape)
fitting_x = RBF(source_control_x, source_control_y,
source_control_z, source_di_x)
fitting_y = RBF(source_control_x, source_control_y,
source_control_z, source_di_y)
fitting_z = RBF(source_control_x, source_control_y,
source_control_z, source_di_z)
generate_x = np.add(
fitting_x(target_control_x, target_control_y,
target_control_z), target_control_x)
generate_y = np.add(
fitting_y(target_control_x, target_control_y,
target_control_z), target_control_y)
generate_z = np.add(
fitting_z(target_control_x, target_control_y,
target_control_z), target_control_z)
result_x = np.append(result_x, np.mean(generate_x))
result_y = np.append(result_y, np.mean(generate_y))
result_z = np.append(result_z, np.mean(generate_z))
result_vertices = np.vstack([result_x, result_y, result_z]).T
pm.save_mesh_raw(outputMesh + ".obj",
result_vertices,
self.target_face,
voxels=None)
print("Deformation finished")
check_a = result_vertices.flatten()
check_b = self.source_vertices.flatten()
percentage = 0
for i in range(0, check_a.size):
percentage = percentage + (abs(check_a.item(i) - check_b.item(i)) /
check_b.item(i))
#print(percentage)
percentage = percentage / check_a.size
print("Average error (in percentage):", percentage * 100)
files = glob.glob(os.path.join(phrase_path, '*.off'))
for file in files:
_, filename = os.path.split(file)
new_filename = new_root + '/' + type + '/' + phrase + '/' + filename[:
-4]
# load mesh
mesh = pymesh.load_mesh(file)
### off2ply
vertices = mesh.vertices
faces = mesh.faces
voxels = mesh.voxels
# For surface mesh
pymesh.save_mesh_raw(new_filename + ".ply", vertices, faces)
# For volume mesh
# pymesh.save_mesh_raw(new_filename + ".ply", vertices, faces, voxels)
# In ascii and using float
# pymesh.save_mesh_raw(new_filename + ".ply", vertices, faces, voxels, ascii=True, use_float=True)
### off2npz
# clean up
mesh, _ = pymesh.remove_isolated_vertices(mesh)
mesh, _ = pymesh.remove_duplicated_vertices(mesh)
# get elements
vertices = mesh.vertices.copy()
faces = mesh.faces.copy()
# move to center
def write(self, mesh_file, vertices, faces):
import pymesh
pymesh.save_mesh_raw(mesh_file, vertices, faces)
def test_all(opt, new_cage_shape):
opt.phase = "test"
opt.target_model = None
print(opt.model)
if opt.is_poly:
source_mesh = om.read_polymesh(opt.model)
else:
source_mesh = om.read_trimesh(opt.model)
dataset = build_dataset(opt)
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))
states = torch.load(opt.ckpt)
if "states" in states:
states = states["states"]
# states["template_vertices"] = new_cage_shape.transpose(1, 2)
# states["source_vertices"] = new_source.transpose(1,2)
# states["source_faces"] = new_source_face
# new_source_face = states["source_faces"]
om.write_mesh(os.path.join(opt.log_dir, opt.subdir,
"template-Sa.ply"), source_mesh)
net = networks.FixedSourceDeformer(opt, 3, opt.num_point, bottleneck_size=opt.bottleneck_size,
template_vertices=states["template_vertices"], template_faces=states["template_faces"].cuda(),
source_vertices=states["source_vertices"], source_faces=states["source_faces"]).cuda()
load_network(net, states)
source_points = torch.from_numpy(
source_mesh.points().copy()).float().cuda().unsqueeze(0)
with torch.no_grad():
# source_face = net.source_faces.detach()
for i, data in enumerate(dataloader):
data = dataset.uncollate(data)
target_shape, target_filename = data["target_shape"], data["target_file"]
logger.info("", data["target_file"][0])
sample_idx = None
if "sample_idx" in data:
sample_idx = data["sample_idx"]
outputs = net(target_shape.transpose(1, 2), cage_only=True)
if opt.d_residual:
cage_offset = outputs["new_cage"]-outputs["cage"]
outputs["cage"] = new_cage_shape
outputs["new_cage"] = new_cage_shape+cage_offset
deformed = deform_with_MVC(outputs["cage"], outputs["new_cage"], outputs["cage_face"].expand(
outputs["cage"].shape[0], -1, -1), source_points)
for b in range(deformed.shape[0]):
t_filename = os.path.splitext(target_filename[b])[0]
source_mesh_arr = source_mesh.points()
source_mesh_arr[:] = deformed[0].cpu().detach().numpy()
om.write_mesh(os.path.join(
opt.log_dir, opt.subdir, "template-{}-Sab.obj".format(t_filename)), source_mesh)
# if data["target_face"] is not None and data["target_mesh"] is not None:
# pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-{}-Sa.ply".format(t_filename)),
# source_mesh[0].detach().cpu(), source_face[b].detach().cpu())
pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-{}-Sb.ply".format(t_filename)),
data["target_mesh"][b].detach().cpu(), data["target_face"][b].detach().cpu())
# pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-{}-Sab.ply".format(t_filename)),
# deformed[b].detach().cpu(), source_face[b].detach().cpu())
# else:
# save_ply(source_mesh[0].detach().cpu(), os.path.join(opt.log_dir, opt.subdir,"template-{}-Sa.ply".format(t_filename)))
# save_ply(target_shape[b].detach().cpu(), os.path.join(opt.log_dir, opt.subdir,"template-{}-Sb.ply".format(t_filename)),
# normals=data["target_normals"][b].detach().cpu())
# save_ply(deformed[b].detach().cpu(), os.path.join(opt.log_dir, opt.subdir,"template-{}-Sab.ply".format(t_filename)),
# normals=data["target_normals"][b].detach().cpu())
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, opt.subdir, "template-{}-cage1.ply".format(t_filename)),
outputs["cage"][b].detach().cpu(), outputs["cage_face"][b].detach().cpu(),
)
pymesh.save_mesh_raw(
os.path.join(opt.log_dir, opt.subdir, "template-{}-cage2.ply".format(t_filename)),
outputs["new_cage"][b].detach().cpu(), outputs["cage_face"][b].detach().cpu(),
)
# if opt.opt_lap and deformed.shape[1] == source_mesh.shape[1]:
# deformed = optimize_lap(opt, source_mesh, deformed, source_face)
# for b in range(deformed.shape[0]):
# pymesh.save_mesh_raw(os.path.join(opt.log_dir, opt.subdir, "template-{}-Sab-optlap.ply".format(t_filename)),
# deformed[b].detach().cpu(), source_face[b].detach().cpu())
if i % 20 == 0:
logger.success("[{}/{}] Done".format(i, len(dataloader)))
dataset.render_result(os.path.join(opt.log_dir, opt.subdir))
source_control_z = source_vertices_z[0:control_num]
target_control_x = target_vertices_x[0:control_num]
target_control_y = target_vertices_y[0:control_num]
target_control_z = target_vertices_z[0:control_num]
result_x = target_control_x
result_y = target_control_y
result_z = target_control_z
vertex_num = source_vertices_x.size
for i in range(control_num, vertex_num):
for j in range(0, control_num):
distance_x = source_vertices_x.item(i) - source_vertices_x.item(j)
distance_y = source_vertices_y.item(i) - source_vertices_y.item(j)
distance_z = source_vertices_z.item(i) - source_vertices_z.item(j)
source_distance_x = np.append(source_distance_x, distance_x)
source_distance_y = np.append(source_distance_y, distance_y)
source_distance_z = np.append(source_distance_z, distance_z)
fitting_func_x = Rbf(source_control_x, source_distance_x)
fitting_func_y = Rbf(source_control_y, source_distance_y)
fitting_func_z = Rbf(source_control_z, source_distance_z)
generate_x = np.add(Rbf(source_control_x) + target_control_x)
generate_y = np.add(Rbf(source_control_y) + target_control_y)
generate_z = np.add(Rbf(source_control_z) + target_control_z)
result_x = np.append(result_x, np.mean(generate_x))
result_y = np.append(result_y, np.mean(generate_y))
result_z = np.append(result_z, np.mean(generate_z))
result_vertices = np.vstack([result_x, result_y, result_z]).T
pm.save_mesh_raw("output.obj", result_vertices, target_faces, voxels=None)
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
def save_as_ply(filename, new_filename):
mesh = pymesh.load_mesh(filename)
vertices = mesh.vertices
faces = mesh.faces
pymesh.save_mesh_raw(new_filename, vertices, faces)
from common import read_trimesh
work_dir = sys.argv[1]
output_dir = sys.argv[2]
sources = glob(os.path.join(work_dir, "*Sa.*"))
logger.info("Found {} source files".format(len(sources)))
os.makedirs(output_dir, exist_ok=True)
for source in sources:
target = source.replace("Sa", "Sb")
fn = os.path.basename(target)
fn = fn.replace("Sb", "Sab")
V_t, F_t = read_trimesh(target, clean=True)
V_t = V_t[:,:3]
V_s, F_s = read_trimesh(source, clean=True)
V_s = V_s[:,:3]
bb_max = np.max(V_t, axis=0)
bb_min = np.min(V_t, axis=0)
size_t = (bb_max - bb_min)
bb_max = np.max(V_s, axis=0)
bb_min = np.min(V_s, axis=0)
size_s = (bb_max - bb_min)
V_st = (V_s * size_t/size_s)
pymesh.save_mesh(os.path.join(output_dir, fn.replace("Sab", "Sa")), pymesh.form_mesh(V_s, F_s))
pymesh.save_mesh(os.path.join(output_dir, fn.replace("Sab", "Sb")), pymesh.form_mesh(V_t, F_t))
pymesh.save_mesh_raw(os.path.join(output_dir, fn), V_st, F_s)
