def worker(input_file, output_file):
points = load(input_file)
labels = points[:, -1]
save_ply_property(points[:, :3],
labels,
output_file,
normals=points[:, 3:6])
def log_outputs(opt, step, all_outputs, all_inputs):
# Source
color = all_inputs["source_shape"][:,:,1].cpu().numpy()
save_ply_property(os.path.join(opt.log_dir,"step-{:06d}-Sa.ply".format(step)), all_inputs["source_shape"][0].detach().cpu().numpy(), color[0], cmap_name="rainbow")
# Target
save_ply_property(os.path.join(opt.log_dir,"step-{:06d}-Sb.ply".format(step)), all_inputs["target_shape"][0].detach().cpu().numpy(), color[0], cmap_name="rainbow")
for batch in range(0, all_outputs["cage"].shape[0], opt.batch_size):
if batch // opt.batch_size == 0:
tag = "StoT"
elif batch // opt.batch_size == 2:
tag = "StoS"
elif batch // opt.batch_size == 1:
tag = "TtoS"
elif batch // opt.batch_size == 3:
tag = "TtoT"
# deformed and cage
save_ply_property(os.path.join(opt.log_dir,"step-{:06d}-{}-Sab.ply".format(step, tag)),
all_outputs["deformed"][batch].detach().cpu().numpy(), color[batch], cmap_name="rainbow")
write_trimesh(os.path.join(opt.log_dir, "step-{:06d}-{}-cage1.ply".format(step, tag)),
all_outputs["cage"][batch].detach().cpu(), all_outputs["cage_face"][0].detach().cpu(), binary=True)
write_trimesh(os.path.join(opt.log_dir, "step-{:06d}-{}-cage2.ply".format(step, tag)),
all_outputs["new_cage"][batch].detach().cpu(), all_outputs["cage_face"][0].detach().cpu(), binary=True)
# if using network2
if "cage_surface" in all_outputs:
save_ply(os.path.join(opt.log_dir,"step-{:06d}-{}-cage_surface1.ply".format(step, tag)), all_outputs["cage_surface"][batch].detach().cpu().numpy())
save_ply(os.path.join(opt.log_dir,"step-{:06d}-{}-cage_surface2.ply".format(step, tag)), all_outputs["new_cage_surface"][batch].detach().cpu().numpy())
pred = torch.from_numpy(pred).cuda()
gt = torch.from_numpy(gt).cuda()
pred_tensor, centroid, furthest_distance = normalize_point_batch(pred)
gt_tensor, centroid, furthest_distance = normalize_point_batch(gt)
# B, P_predict, 1
cd_forward, cd_backward = nndistance(pred, gt)
# cd_forward, _ = knn_point(1, gt_tensor, pred_tensor)
# cd_backward, _ = knn_point(1, pred_tensor, gt_tensor)
# cd_forward = cd_forward[0, :, 0]
# cd_backward = cd_backward[0, :, 0]
cd_forward = cd_forward.detach().cpu().numpy()[0]
cd_backward = cd_backward.detach().cpu().numpy()[0]
save_ply_property(pred.squeeze(0).detach().cpu().numpy(), cd_forward, pred_path[:-4]+"_cdF.ply", property_max=0.003, cmap_name="jet")
save_ply_property(gt.squeeze(0).detach().cpu().numpy(), cd_backward, pred_path[:-4]+"_cdB.ply", property_max=0.003, cmap_name="jet")
md_value = np.mean(cd_forward)+np.mean(cd_backward)
hd_value = np.max(np.amax(cd_forward, axis=0)+np.amax(cd_backward, axis=0))
cd_backward = np.mean(cd_backward)
cd_forward = np.mean(cd_forward)
# row["CD_forward"] = np.mean(cd_forward)
# row["CD_backwar"] = np.mean(cd_backward)
row["CD"] = cd_forward+cd_backward
row["hausdorff"] = hd_value
avg_md_forward_value += cd_forward
avg_md_backward_value += cd_backward
avg_hd_value += hd_value
if os.path.isfile(pred_path[:-4] + "_point2mesh_distance.xyz"):
def trainShapeOnImage(scene, refScene, opt, baseline=False, benchmark=False):
expr_dir = os.path.join(opt.output, opt.name)
if not os.path.isdir(expr_dir):
os.makedirs(expr_dir)
trainer = Trainer(opt, scene)
trainer.setup(opt, scene.cloud)
logInterval = (1+sum(opt.steps))//20
renderForwardTime = 0.0
lossTime = 0.0
optimizerStep = 0.0
with open(os.path.join(expr_dir, "loss.csv"), 'w') as loss_log:
learnTick = time.time()
for c in range(opt.cycles):
# creat new reference
tb = c*sum(opt.steps)+opt.startingStep
te = (c+1)*sum(opt.steps)+opt.startingStep
t = tb
with torch.no_grad():
# render reference
trainer.create_reference(refScene)
# render prediction
maxDiffs = []
selectedCameras = []
for i in range(len(refScene.cameras)):
trainer.forward(i)
prediction = trainer.predictions[i]
maxDiff, selectedCamera = viewFromError(opt.genCamera, trainer.groundtruths[i][0], prediction.detach()[0],
trainer.model._localPoints.detach(), trainer.model._projPoints.detach(),
trainer.model,
offset=opt.camOffset*(0.997**c))
maxDiffs.append(maxDiff)
selectedCameras.append(selectedCamera)
maxId = torch.stack(maxDiffs, dim=0).argmax()
selectedCamera = selectedCameras[maxId]
# render again
trainer.create_reference(refScene, selectedCamera)
writeScene(refScene, os.path.join(expr_dir, 't%03d_scene_gt.json' % t))
writeCameras(refScene, os.path.join(expr_dir, 't%03d_cameras.ply' % t))
for i, gt in enumerate(trainer.groundtruths):
if gt is not None:
saveAsPng(gt.cpu()[0], os.path.join(expr_dir, 't%03d_cam%d_gt.png' % (t, i)))
trainer.initiate_cycle()
for t in range(tb, te):
if t % logInterval == 0 and not benchmark:
writeScene(scene, os.path.join(expr_dir, 't%03d_scene.json' % t), os.path.join(expr_dir, "t%03d.ply" % t))
trainer.optimize_parameters()
if t % logInterval == 0 and not benchmark:
for i, prediction in enumerate(trainer.predictions):
saveAsPng(prediction.detach().cpu()[0], os.path.join(expr_dir, 't%03d_cam%d' % (t, i) + ".png"))
if not benchmark:
loss_str = ",".join(["%.3f" % v for v in trainer.loss_image])
reg_str = ",".join(["%.3f" % v for v in trainer.loss_reg])
entries = [trainer.modifier] + [loss_str] + [reg_str]
loss_log.write(",".join(entries)+"\n")
print("{:03d} {}: lr {} loss ({}) \n : reg ({})".format(
t, trainer.modifier, trainer.lr, loss_str, reg_str))
trainer.finish_cycle()
# outlier removal
with torch.no_grad():
# re-project
scores = []
for i, gt in enumerate(trainer.groundtruths):
if gt is None:
continue
trainer.model.setCamera(i)
trainer.model.convertToCameraSpace()
projectedPoints = trainer.model.camera.projectPoints(trainer.model.cameraPoints)
score = removeOutlier(gt, projectedPoints, sigma=100)
saveAsPng(gt.cpu()[0], os.path.join(expr_dir, 'clean_gt_cam%d.png' % i))
save_ply_property(trainer.model.localPoints.cpu()[0].detach(), 1-score.cpu(), os.path.join(expr_dir, 'clean_score_cam % d.ply' % i), property_max=1,
normals=trainer.model.localNormals.detach().cpu()[0], cmap_name="gnuplot2")
rendered = trainer.model.render()[0]
saveAsPng(rendered.cpu(), os.path.join(expr_dir, 'clean_pred_cam%d.png' % i))
scores.append(score)
scores = torch.stack(scores, dim=0)
scores = torch.prod(scores, dim=0)
_, indices = torch.topk(scores, int(trainer.model.localPoints.shape[1]*0.99), dim=0)
# _, indices = torch.topk(score, int(trainer.model.localPoints.shape[1]*0.99), dim=0)
newPoints = torch.index_select(trainer.model.localPoints.data, 1, indices)
newNormals = torch.index_select(trainer.model.localNormals.data, 1, indices)
newColors = torch.index_select(trainer.model.pointColors.data, 1, indices)
scene.cloud.localPoints = newPoints[0]
scene.cloud.localNormals = newNormals[0]
scene.cloud.color = newColors[0]
writeScene(scene, os.path.join(expr_dir, 'final_scene.json'), os.path.join(expr_dir, 'final_cloud.ply'))