# Print output
if print_every > 0 and (it % print_every) == 0:
print('[Epoch %02d] it=%03d, loss=%.4f' % (epoch_it, it, loss))
with train_summary_writer.as_default():
tf.summary.scalar('loss', loss, step=it)
# Visualize output
if visualize_every > 0 and (it % visualize_every) == 0:
print('Visualizing')
trainer.visualize(data_vis)
# Save checkpoint
if (checkpoint_every > 0 and (it % checkpoint_every) == 0):
print('Saving checkpoint')
checkpoint_io.save('model/model.ckpt',
epoch_it=epoch_it,
it=it,
loss_val_best=metric_val_best)
# Backup if necessary
if (backup_every > 0 and (it % backup_every) == 0):
print('Backup checkpoint')
checkpoint_io.save('backup/model_%d.ckpt' % it,
epoch_it=epoch_it,
it=it,
loss_val_best=metric_val_best)
# Run validation
if validate_every > 0 and (it % validate_every) == 0:
print("evaluate")
eval_dict = trainer.evaluate(val_loader)
print("eval_dict")
metric_val = eval_dict[model_selection_metric]
print('validation metric (%s): %.4f' %
logger.add_scalar('train/loss', loss, it)
# Print output
if print_every > 0 and (it % print_every) == 0:
print('[Epoch %02d] it=%03d, loss=%.4f' % (epoch_it, it, loss))
# Visualize output
if visualize_every > 0 and (it % visualize_every) == 0:
print('Visualizing')
trainer.visualize(data_vis)
# Save checkpoint
if (checkpoint_every > 0 and (it % checkpoint_every) == 0):
print('Saving checkpoint')
checkpoint_io.save('model.pt',
epoch_it=epoch_it,
it=it,
loss_val_best=metric_val_best)
# Backup if necessary
if (backup_every > 0 and (it % backup_every) == 0):
print('Backup checkpoint')
checkpoint_io.save('model_%d.pt' % it,
epoch_it=epoch_it,
it=it,
loss_val_best=metric_val_best)
# Run validation
if validate_every > 0 and (it % validate_every) == 0:
eval_dict = trainer.evaluate(val_loader)
metric_val = eval_dict[model_selection_metric]
print('Validation metric (%s): %.4f' %
(model_selection_metric, metric_val))
def run(pointcloud_path, out_dir, decoder_type='siren', resume=True, **kwargs):
"""
test_implicit_siren_noisy_wNormals
"""
device = torch.device('cuda:0')
if not os.path.exists(out_dir):
os.makedirs(out_dir)
# data
points, normals = np.split(read_ply(pointcloud_path).astype('float32'),
(3, ),
axis=1)
pmax, pmin = points.max(axis=0), points.min(axis=0)
scale = (pmax - pmin).max()
pcenter = (pmax + pmin) / 2
points = (points - pcenter) / scale * 1.5
scale_mat = scale_mat_inv = np.identity(4)
scale_mat[[0, 1, 2], [0, 1, 2]] = 1 / scale * 1.5
scale_mat[[0, 1, 2], [3, 3, 3]] = -pcenter / scale * 1.5
scale_mat_inv = np.linalg.inv(scale_mat)
normals = normals @ np.linalg.inv(scale_mat[:3, :3].T)
object_bounding_sphere = np.linalg.norm(points, axis=1).max()
pcl = trimesh.Trimesh(vertices=points,
vertex_normals=normals,
process=False)
pcl.export(os.path.join(out_dir, "input_pcl.ply"), vertex_normal=True)
assert (np.abs(points).max() < 1)
dataset = torch.utils.data.TensorDataset(torch.from_numpy(points),
torch.from_numpy(normals))
batch_size = 5000
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
num_workers=1,
shuffle=True,
collate_fn=tolerating_collate,
)
gt_surface_pts_all = torch.from_numpy(points).unsqueeze(0).float()
gt_surface_normals_all = torch.from_numpy(normals).unsqueeze(0).float()
gt_surface_normals_all = F.normalize(gt_surface_normals_all, dim=-1)
if kwargs['use_off_normal_loss']:
# subsample from pointset
sub_idx = torch.randperm(gt_surface_normals_all.shape[1])[:20000]
gt_surface_pts_sub = torch.index_select(gt_surface_pts_all, 1,
sub_idx).to(device=device)
gt_surface_normals_sub = torch.index_select(gt_surface_normals_all, 1,
sub_idx).to(device=device)
gt_surface_normals_sub = denoise_normals(gt_surface_pts_sub,
gt_surface_normals_sub,
neighborhood_size=30)
if decoder_type == 'siren':
decoder_params = {
'dim': 3,
"out_dims": {
'sdf': 1
},
"c_dim": 0,
"hidden_size": 256,
'n_layers': 3,
"first_omega_0": 30,
"hidden_omega_0": 30,
"outermost_linear": True,
}
decoder = Siren(**decoder_params)
# pretrained_model_file = os.path.join('data', 'trained_model', 'siren_l{}_c{}_o{}.pt'.format(
# decoder_params['n_layers'], decoder_params['hidden_size'], decoder_params['first_omega_0']))
# loaded_state_dict = torch.load(pretrained_model_file)
# decoder.load_state_dict(loaded_state_dict)
elif decoder_type == 'sdf':
decoder_params = {
'dim': 3,
"out_dims": {
'sdf': 1
},
"c_dim": 0,
"hidden_size": 512,
'n_layers': 8,
'bias': 1.0,
}
decoder = SDF(**decoder_params)
else:
raise ValueError
print(decoder)
decoder = decoder.to(device)
# training
total_iter = 30000
optimizer = torch.optim.Adam(decoder.parameters(), lr=1e-4)
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, [10000, 20000],
gamma=0.5)
shape = Shape(gt_surface_pts_all.cuda(),
n_points=gt_surface_pts_all.shape[1] // 16,
normals=gt_surface_normals_all.cuda())
# initialize siren with sphere_initialization
checkpoint_io = CheckpointIO(out_dir, model=decoder, optimizer=optimizer)
load_dict = dict()
if resume:
models_avail = [f for f in os.listdir(out_dir) if f[-3:] == '.pt']
if len(models_avail) > 0:
models_avail.sort()
load_dict = checkpoint_io.load(models_avail[-1])
it = load_dict.get('it', 0)
if it > 0:
try:
iso_point_files = [
f for f in os.listdir(out_dir) if f[-7:] == 'iso.ply'
]
iso_point_iters = [
int(os.path.basename(f[:-len('_iso.ply')]))
for f in iso_point_files
]
iso_point_iters = np.array(iso_point_iters)
idx = np.argmax(iso_point_iters[(iso_point_iters - it) <= 0])
iso_point_file = np.array(iso_point_files)[(iso_point_iters -
it) <= 0][idx]
iso_points = torch.from_numpy(
read_ply(os.path.join(out_dir, iso_point_file))[..., :3])
shape.points = iso_points.to(device=shape.points.device).view(
1, -1, 3)
print('Loaded iso-points from %s' % iso_point_file)
except Exception as e:
pass
# loss
eikonal_loss = NormalLengthLoss(reduction='mean')
# start training
# save_ply(os.path.join(out_dir, 'in_iso_points.ply'), (to_homogen(shape.points).cpu().detach().numpy() @ scale_mat_inv.T)[...,:3].reshape(-1,3))
save_ply(os.path.join(out_dir, 'in_iso_points.ply'),
shape.points.cpu().view(-1, 3))
# autograd.set_detect_anomaly(True)
iso_points = shape.points
iso_points_normal = None
while True:
if (it > total_iter):
checkpoint_io.save('model_{:04d}.pt'.format(it), it=it)
mesh = get_surface_high_res_mesh(
lambda x: decoder(x).sdf.squeeze(), resolution=512)
mesh.apply_transform(scale_mat_inv)
mesh.export(os.path.join(out_dir, "final.ply"))
break
for batch in data_loader:
gt_surface_pts, gt_surface_normals = batch
gt_surface_pts.unsqueeze_(0)
gt_surface_normals.unsqueeze_(0)
gt_surface_pts = gt_surface_pts.to(device=device).detach()
gt_surface_normals = gt_surface_normals.to(device=device).detach()
optimizer.zero_grad()
decoder.train()
loss = defaultdict(float)
lambda_surface_sdf = 1e3
lambda_surface_normal = 1e2
if kwargs['warm_up'] >= 0 and it >= kwargs['warm_up']:
lambda_surface_sdf = kwargs['lambda_surface_sdf']
lambda_surface_normal = kwargs['lambda_surface_normal']
# debug
if (it - kwargs['warm_up']) % 1000 == 0:
# generate iso surface
with torch.autograd.no_grad():
box_size = (object_bounding_sphere * 2 + 0.2, ) * 3
imgs = plot_cuts(
lambda x: decoder(x).sdf.squeeze().detach(),
box_size=box_size,
max_n_eval_pts=10000,
thres=0.0,
imgs_per_cut=1,
save_path=os.path.join(out_dir, '%010d_iso.html' % it))
mesh = get_surface_high_res_mesh(
lambda x: decoder(x).sdf.squeeze(), resolution=200)
mesh.apply_transform(scale_mat_inv)
mesh.export(os.path.join(out_dir, '%010d_mesh.ply' % it))
if it % 2000 == 0:
checkpoint_io.save('model.pt', it=it)
pred_surface_grad = gradient(gt_surface_pts.clone(),
lambda x: decoder(x).sdf)
# every once in a while update shape and points
# sample points in space and on the shape
# use iso points to weigh data points loss
weights = 1.0
if kwargs['warm_up'] >= 0 and it >= kwargs['warm_up']:
if it == kwargs['warm_up'] or kwargs['resample_every'] > 0 and (
it -
kwargs['warm_up']) % kwargs['resample_every'] == 0:
# if shape.points.shape[1]/iso_points.shape[1] < 1.0:
# idx = fps(iso_points.view(-1,3), torch.zeros(iso_points.shape[1], dtype=torch.long, device=iso_points.device), shape.points.shape[1]/iso_points.shape[1])
# iso_points = iso_points.view(-1,3)[idx].view(1,-1,3)
iso_points = shape.get_iso_points(
iso_points + 0.1 * (torch.rand_like(iso_points) - 0.5),
decoder,
ear=kwargs['ear'],
outlier_tolerance=kwargs['outlier_tolerance'])
# iso_points = shape.get_iso_points(shape.points, decoder, ear=kwargs['ear'], outlier_tolerance=kwargs['outlier_tolerance'])
iso_points_normal = estimate_pointcloud_normals(
iso_points.view(1, -1, 3), 8, False)
if kwargs['denoise_normal']:
iso_points_normal = denoise_normals(iso_points,
iso_points_normal,
num_points=None)
iso_points_normal = iso_points_normal.view_as(
iso_points)
elif iso_points_normal is None:
iso_points_normal = estimate_pointcloud_normals(
iso_points.view(1, -1, 3), 8, False)
# iso_points = resample_uniformly(iso_points.view(1,-1,3))
# TODO: use gradient from network or neighborhood?
iso_points_g = gradient(iso_points.clone(),
lambda x: decoder(x).sdf)
if it == kwargs['warm_up'] or kwargs['resample_every'] > 0 and (
it -
kwargs['warm_up']) % kwargs['resample_every'] == 0:
# save_ply(os.path.join(out_dir, '%010d_iso.ply' % it), (to_homogen(iso_points).cpu().detach().numpy() @ scale_mat_inv.T)[...,:3].reshape(-1,3), normals=iso_points_g.view(-1,3).detach().cpu())
save_ply(os.path.join(out_dir, '%010d_iso.ply' % it),
iso_points.cpu().detach().view(-1, 3),
normals=iso_points_g.view(-1, 3).detach().cpu())
if kwargs['weight_mode'] == 1:
weights = get_iso_bilateral_weights(
gt_surface_pts, gt_surface_normals, iso_points,
iso_points_g).detach()
elif kwargs['weight_mode'] == 2:
weights = get_laplacian_weights(gt_surface_pts,
gt_surface_normals,
iso_points,
iso_points_g).detach()
elif kwargs['weight_mode'] == 3:
weights = get_heat_kernel_weights(gt_surface_pts,
gt_surface_normals,
iso_points,
iso_points_g).detach()
if (it - kwargs['warm_up']
) % 1000 == 0 and kwargs['weight_mode'] != -1:
print("min {:.4g}, max {:.4g}, std {:.4g}, mean {:.4g}".
format(weights.min(), weights.max(), weights.std(),
weights.mean()))
colors = scaler_to_color(1 -
weights.view(-1).cpu().numpy(),
cmap='Reds')
save_ply(
os.path.join(out_dir, '%010d_batch_weight.ply' % it),
(to_homogen(gt_surface_pts).cpu().detach().numpy()
@ scale_mat_inv.T)[..., :3].reshape(-1, 3),
colors=colors)
sample_idx = torch.randperm(
iso_points.shape[1])[:min(batch_size, iso_points.shape[1])]
iso_points_sampled = iso_points.detach()[:, sample_idx, :]
# iso_points_sampled = iso_points.detach()
iso_points_sdf = decoder(iso_points_sampled.detach()).sdf
loss_iso_points_sdf = iso_points_sdf.abs().mean(
) * kwargs['lambda_iso_sdf'] * iso_points_sdf.nelement() / (
iso_points_sdf.nelement() + 8000)
loss['loss_sdf_iso'] = loss_iso_points_sdf.detach()
loss['loss'] += loss_iso_points_sdf
# TODO: predict iso_normals from local_frame
iso_normals_sampled = iso_points_normal.detach()[:,
sample_idx, :]
iso_g_sampled = iso_points_g[:, sample_idx, :]
loss_normals = torch.mean(
(1 - F.cosine_similarity(
iso_normals_sampled, iso_g_sampled, dim=-1).abs())
) * kwargs['lambda_iso_normal'] * iso_points_sdf.nelement() / (
iso_points_sdf.nelement() + 8000)
# loss_normals = torch.mean((1 - F.cosine_similarity(iso_points_normal, iso_points_g, dim=-1).abs())) * kwargs['lambda_iso_normal']
loss['loss_normal_iso'] = loss_normals.detach()
loss['loss'] += loss_normals
idx = torch.randperm(gt_surface_pts.shape[1]).to(
device=gt_surface_pts.device)[:(gt_surface_pts.shape[1] // 2)]
tmp = torch.index_select(gt_surface_pts, 1, idx)
space_pts = torch.cat([
torch.rand_like(tmp) * 2 - 1,
torch.randn_like(tmp, device=tmp.device, dtype=tmp.dtype) * 0.1
+ tmp
],
dim=1)
space_pts.requires_grad_(True)
pred_space_sdf = decoder(space_pts).sdf
pred_space_grad = torch.autograd.grad(
pred_space_sdf, [space_pts], [torch.ones_like(pred_space_sdf)],
create_graph=True)[0]
# 1. eikonal term
loss_eikonal = (
eikonal_loss(pred_surface_grad) +
eikonal_loss(pred_space_grad)) * kwargs['lambda_eikonal']
loss['loss_eikonal'] = loss_eikonal.detach()
loss['loss'] += loss_eikonal
# 2. SDF loss
# loss on iso points
pred_surface_sdf = decoder(gt_surface_pts).sdf
loss_sdf = torch.mean(
weights * pred_surface_sdf.abs()) * lambda_surface_sdf
if kwargs['warm_up'] >= 0 and it >= kwargs['warm_up'] and kwargs[
'lambda_iso_sdf'] != 0:
# loss_sdf = 0.5 * loss_sdf
loss_sdf = loss_sdf * pred_surface_sdf.nelement() / (
pred_surface_sdf.nelement() + iso_points_sdf.nelement())
if kwargs['use_sal_loss'] and iso_points is not None:
dists, idxs, _ = knn_points(space_pts.view(1, -1, 3),
iso_points.view(1, -1, 3).detach(),
K=1)
dists = dists.view_as(pred_space_sdf)
idxs = idxs.view_as(pred_space_sdf)
loss_inter = ((eps_sqrt(dists).sqrt() - pred_space_sdf.abs())**
2).mean() * kwargs['lambda_inter_sal']
else:
alpha = (it / total_iter + 1) * 100
loss_inter = torch.exp(
-alpha *
pred_space_sdf.abs()).mean() * kwargs['lambda_inter_sdf']
loss_sald = torch.tensor(0.0).cuda()
# prevent wrong closing for open mesh
if kwargs['use_off_normal_loss'] and it < 1000:
dists, idxs, _ = knn_points(space_pts.view(1, -1, 3),
gt_surface_pts_sub.view(1, -1,
3).cuda(),
K=1)
knn_normal = knn_gather(
gt_surface_normals_sub.cuda().view(1, -1, 3),
idxs).view(1, -1, 3)
direction_correctness = -F.cosine_similarity(
knn_normal, pred_space_grad, dim=-1)
direction_correctness[direction_correctness < 0] = 0
loss_sald = torch.mean(
direction_correctness * torch.exp(-2 * dists)) * 2
# 3. normal direction
loss_normals = torch.mean(weights * (1 - F.cosine_similarity(
gt_surface_normals, pred_surface_grad, dim=-1))
) * lambda_surface_normal
if kwargs['warm_up'] >= 0 and it >= kwargs['warm_up'] and kwargs[
'lambda_iso_normal'] != 0:
# loss_normals = 0.5 * loss_normals
loss_normals = loss_normals * gt_surface_normals.nelement() / (
gt_surface_normals.nelement() +
iso_normals_sampled.nelement())
loss['loss_sdf'] = loss_sdf.detach()
loss['loss_inter'] = loss_inter.detach()
loss['loss_normals'] = loss_normals.detach()
loss['loss_sald'] = loss_sald
loss['loss'] += loss_sdf
loss['loss'] += loss_inter
loss['loss'] += loss_sald
loss['loss'] += loss_normals
loss['loss'].backward()
torch.nn.utils.clip_grad_norm_(decoder.parameters(), max_norm=1.)
optimizer.step()
scheduler.step()
if it % 20 == 0:
print("iter {:05d} {}".format(
it, ', '.join([
'{}: {}'.format(k, v.item()) for k, v in loss.items()
])))
it += 1
