luni/reconstruction
1
0
Fork 0
Code Issues 6 Pull Requests Packages Projects 1 Releases Wiki Activity

ca;cu;ate normal error with chamfer

Browse Source
This commit is contained in:
luni
2023-01-14 18:24:30 +00:00
parent c315f79941
commit 35928ccbff
1 changed files with 12 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test.py
+12 -10
View File
@@ -43,13 +43,15 @@ def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask):
K = torch.tensor(intrinsic()[None, :, :], device='cuda')
pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
pc = [torch.permute(pc[i,:,depth_mask[i]], [1,0]) for i in range(depth_image.shape[0])]
pc_pc = Pointclouds(pc)
mesh_gt = [vertices[i, mesh_mask[i], :] for i in range(vertices.shape[0])]
mesh_gt_pc = Pointclouds(mesh_gt)
pc_norm = get_normals_from_depth(depth_image, depth_mask, K)
mesh_norm = estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50)
pc_norm = [pc_norm[i,depth_mask[i],:] for i in range(pc_norm.shape[0])]
mesh_gt = [vertices[i, mesh_mask[i], :] for i in range(vertices.shape[0])]
pc_pc = Pointclouds(pc, normals=pc_norm)
mesh_gt_pc = Pointclouds(mesh_gt)
mesh_norm = list(torch.unbind(estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50)))
mesh_gt_pc = Pointclouds(mesh_gt, normals=mesh_norm)
chamfer = chamfer_distance(pc_pc, mesh_gt_pc, x_normals=pc_norm, y_normals=mesh_norm)
chamfer = chamfer_distance(pc_pc, mesh_gt_pc)
return chamfer
@@ -73,8 +75,8 @@ def train_step(batch):
# 2. Calculate loss
mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
chamfer, _ = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
loss = mse + 1e-1 * chamfer
chamfer, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
loss = mse + 1e-1 * (chamfer + normal)
# 3. Zero grad optimizer
optimizer.zero_grad()
@@ -90,15 +92,15 @@ def train_step(batch):
def valid_step(batch):
### Testing
model.eval() # put the model in evaluation mode for testing (inference)
model.eval() # put the model in evaluation mode for testing (inference)
# 1. Forward pass
with torch.inference_mode():
pred = model([batch['rgb_image'], batch['view']])
# 2. Calculate the loss
mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
chamfer, _ = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
loss = mse + 1e-1 * chamfer
chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
loss = mse + 1e-1 * (chamfer + normal)
return { 'valid/loss': loss.detach().cpu().numpy() }