diff --git a/test.py b/test.py index 58ce26e..821eabc 100644 --- a/test.py +++ b/test.py @@ -34,30 +34,17 @@ def plot3d(depth, mask): # temp -def plot_gt(y_pred, batch): +def plot_gt(depth, gt, mask): import kornia import plotly.graph_objects as go - import pytorch3d from data.render import intrinsic - K = torch.tensor(intrinsic()[None, :, :], device='cuda') - pc = kornia.geometry.depth.depth_to_3d(y_pred - batch['info']['camLoc'][:,0][:,None,None,None].cuda(), K, normalize_points=True) - pc = pc[0,:,batch['depth_mask'][0]].detach().cpu().numpy() - m = batch['mesh_posed'][0, batch['mesh_mask'][0], :] - view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((torch.tensor([torch.pi/2], device='cuda'), - torch.tensor([0.0], device='cuda'), - torch.tensor([0.0], device='cuda')))) - transform = pytorch3d.transforms.Rotate(view) - m = transform.transform_points(m) - view = pytorch3d.transforms.axis_angle_to_matrix(torch.concat((torch.tensor([0.0], device='cuda'), - torch.tensor([torch.pi/2], device='cuda'), - torch.tensor([0.0], device='cuda')))) - transform = pytorch3d.transforms.Rotate(view) - m = transform.transform_points(m) - m = m.cpu().numpy() - + pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True) + pc = pc[0,:,mask[0]].detach().cpu().numpy() + pc_gt = kornia.geometry.depth.depth_to_3d(gt[:,None,:,:], K, normalize_points=True) + pc_gt = pc_gt[0,:,mask[0]].detach().cpu().numpy() go.Figure(data = [go.Scatter3d(x=pc[0,:], y=pc[1,:], z=pc[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'), - go.Scatter3d(x=m[:,0], y=m[:,1], z=m[:,2], marker=dict(size=3, line=dict(width=0)), mode='markers'), + go.Scatter3d(x=pc_gt[0,:], y=pc_gt[1,:], z=pc_gt[2,:], marker=dict(size=3, line=dict(width=0)), mode='markers'), ]).show() pass @@ -113,21 +100,33 @@ def depth_reconst_loss(y_pred, y_true, mask): chamfer, _ = chamfer_pytorch(pred_pc_pc, gt_pc_pc) return chamfer +def normals_reconst_loss(y_pred, normals_gt, mask): + import kornia + from data.render import intrinsic + from utils.data_utils import get_normals_from_depth + K = torch.tensor(intrinsic()[None, :, :], device='cuda') + normal_image = torch.permute(get_normals_from_depth(y_pred, mask, K), [0,3,1,2]) + + mse = loss_fn(normal_image * mask[:,None,:,:], normals_gt * mask[:,None,:,:]) + return mse + def train_step(batch): ### Training model.train() # train mode is on by default after construction # 1. Forward pass - y_pred = model([batch['rgb_image'], batch['view']]) + y_pred, y_pred_normals = model([batch['rgb_image'], batch['view']]) # 2. Calculate loss mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) + mse_normals = loss_fn(y_pred_normals * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:]) # chamfer, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask']) msssim = ms_ssim_module(y_pred * batch['depth_mask'][:,None,:,:] + 5, batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:] + 5) - chamfer_reconst = depth_reconst_loss(y_pred, batch['depth_image'][:,None,:,:], batch['depth_mask']) + chamfer_reconst = depth_reconst_loss(y_pred, batch['depth_image'][:,None,:,:], batch['depth_mask']) + normals_reconst = normals_reconst_loss(y_pred * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:], batch['depth_mask']) # loss = mse #+ 1e-1 * (chamfer + normal) - loss = 1 - msssim + chamfer_reconst + mse + loss = 1 - msssim + mse_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst # plot_gt(y_pred, batch) @@ -148,15 +147,17 @@ def valid_step(batch): 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']]) + pred, pred_normals = 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,:,:]) + mse_normals = loss_fn(pred_normals * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:]) # chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask']) msssim = ms_ssim_module(pred * batch['depth_mask'][:,None,:,:] + 5, batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:] + 5) chamfer_reconst = depth_reconst_loss(pred, batch['depth_image'][:,None,:,:], batch['depth_mask']) + normals_reconst = normals_reconst_loss(pred * batch['depth_mask'][:,None,:,:], batch['normal_image'] * batch['depth_mask'][:,None,:,:], batch['depth_mask']) # loss = mse #+ 1e-1 * (chamfer + normal) - loss = 1 - msssim + chamfer_reconst + mse + loss = 1 - msssim + mse_normals #+ chamfer_reconst + mse + normals_reconst return { 'valid/loss': loss.detach().cpu().numpy() } @@ -187,7 +188,8 @@ if __name__ == '__main__': loss_fn = nn.MSELoss() loss_fn_3d = nn.MSELoss() - optimizer = optim.Adam(model.parameters(), lr=0.01) + optimizer = optim.Adam(model.parameters(), lr=0.001) + lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, 2000, gamma=0.1) # optimizer = optim.SGD(model.parameters(), lr=0.01) iteration = 0 @@ -197,19 +199,22 @@ if __name__ == '__main__': train_loss = train_step(batch) train_accumulator.update(train_loss) - if iteration % 1000 == 0: - print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}") - train_accumulator.reset() - + lr_scheduler.step() iteration += 1 + if iteration % 1000 == 0: + print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}") + train_accumulator.reset() + + # TODO: change when batchsize != 1 if iteration % 100 == 0: for batch in valid_dataloader: valid_loss = valid_step(batch) valid_accumulator.update(valid_loss) print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}") - valid_accumulator.reset() + valid_accumulator.reset() +