From 35928ccbff31d25daa4679f523ca4b1e67fcbd3c Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 14 Jan 2023 18:24:30 +0000 Subject: [PATCH] ca;cu;ate normal error with chamfer --- test.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/test.py b/test.py index 29aa82e..eaf973c 100644 --- a/test.py +++ b/test.py @@ -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() }