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

add new losses

Browse Source
This commit is contained in:
luni
2023-01-18 23:19:09 +00:00
parent 9e1c827526
commit 4443989a6e
1 changed files with 35 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test.py
+31 -26
View File
@@ -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'])
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,12 +199,14 @@ if __name__ == '__main__':
train_loss = train_step(batch)
train_accumulator.update(train_loss)
lr_scheduler.step()
iteration += 1
if iteration % 1000 == 0:
print(f"Iteration: {iteration} | Train loss: {train_accumulator.get()['train/loss']}")
train_accumulator.reset()
iteration += 1
# TODO: change when batchsize != 1
if iteration % 100 == 0:
for batch in valid_dataloader:
valid_loss = valid_step(batch)
@@ -213,4 +217,5 @@ if __name__ == '__main__':
pass