backup
This commit is contained in:
@@ -17,7 +17,8 @@ from utils.utils import Accumulator
|
||||
from models.depth_estim import DepthEstimation
|
||||
from models.convnext import ConvNeXt
|
||||
|
||||
|
||||
from pytorch_msssim import ssim, ms_ssim, SSIM, MS_SSIM
|
||||
ms_ssim_module = MS_SSIM(data_range=10, size_average=True, channel=1)
|
||||
|
||||
|
||||
def plot3d(depth, mask):
|
||||
@@ -32,6 +33,36 @@ def plot3d(depth, mask):
|
||||
]).show()
|
||||
|
||||
|
||||
# temp
|
||||
def plot_gt(y_pred, batch):
|
||||
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()
|
||||
|
||||
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'),
|
||||
]).show()
|
||||
|
||||
pass
|
||||
|
||||
|
||||
def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask):
|
||||
import kornia
|
||||
import pytorch3d
|
||||
@@ -56,15 +87,32 @@ def chamfer_loss(depth_image, vertices, depth_mask, mesh_mask):
|
||||
return chamfer
|
||||
|
||||
|
||||
def normal_loss(depth_image, depth_mask, normal_image_gt):
|
||||
# def normal_loss(depth_image, depth_mask, normal_image_gt):
|
||||
# import kornia
|
||||
# import pytorch3d
|
||||
# from data.render import intrinsic
|
||||
# K = torch.tensor(intrinsic()[None, :, :], device='cuda')
|
||||
# pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
|
||||
# pc = torch.stack([pc[i,:,depth_mask[i]] for i in range(depth_image.shape[0])])
|
||||
# normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
|
||||
# pass
|
||||
|
||||
def depth_reconst_loss(y_pred, y_true, mask):
|
||||
import kornia
|
||||
import pytorch3d
|
||||
from data.render import intrinsic
|
||||
from pytorch3d.structures import Pointclouds
|
||||
from pytorch3d.loss.chamfer import chamfer_distance as chamfer_pytorch
|
||||
# from utils.pytorch3d_chamfer import chamfer_distance
|
||||
K = torch.tensor(intrinsic()[None, :, :], device='cuda')
|
||||
pc = kornia.geometry.depth.depth_to_3d(depth_image, K, normalize_points=True)
|
||||
pc = torch.stack([pc[i,:,depth_mask[i]] for i in range(depth_image.shape[0])])
|
||||
normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
|
||||
pass
|
||||
pred_pc = kornia.geometry.depth.depth_to_3d(y_pred, K, normalize_points=True)
|
||||
pred_pc = [torch.permute(pred_pc[i, :, mask[i]], [1,0]) for i in range(y_pred.shape[0])]
|
||||
pred_pc_pc = Pointclouds(pred_pc)
|
||||
gt_pc = kornia.geometry.depth.depth_to_3d(y_true, K, normalize_points=True)
|
||||
gt_pc = [torch.permute(gt_pc[i, :, mask[i]], [1,0]) for i in range(y_true.shape[0])]
|
||||
gt_pc_pc = Pointclouds(gt_pc)
|
||||
chamfer, _ = chamfer_pytorch(pred_pc_pc, gt_pc_pc)
|
||||
return chamfer
|
||||
|
||||
|
||||
def train_step(batch):
|
||||
### Training
|
||||
@@ -75,8 +123,13 @@ 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, normal = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
|
||||
loss = mse + 1e-1 * (chamfer + normal)
|
||||
# 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'])
|
||||
# loss = mse #+ 1e-1 * (chamfer + normal)
|
||||
loss = 1 - msssim + chamfer_reconst + mse
|
||||
|
||||
# plot_gt(y_pred, batch)
|
||||
|
||||
# 3. Zero grad optimizer
|
||||
optimizer.zero_grad()
|
||||
@@ -99,8 +152,12 @@ def valid_step(batch):
|
||||
|
||||
# 2. Calculate the loss
|
||||
mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:])
|
||||
chamfer, normal = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask'])
|
||||
loss = mse + 1e-1 * (chamfer + normal)
|
||||
# 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'])
|
||||
# loss = mse #+ 1e-1 * (chamfer + normal)
|
||||
loss = 1 - msssim + chamfer_reconst + mse
|
||||
|
||||
|
||||
return { 'valid/loss': loss.detach().cpu().numpy() }
|
||||
|
||||
@@ -117,10 +174,10 @@ if __name__ == '__main__':
|
||||
|
||||
random_view = False
|
||||
use_skip = False
|
||||
batch_size = 2
|
||||
batch_size = 1
|
||||
|
||||
train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt', random_view=random_view, device=device)
|
||||
valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt', random_view=random_view, device=device)
|
||||
train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_single.txt', random_view=random_view, device=device)
|
||||
valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_single.txt', random_view=random_view, device=device)
|
||||
train_dataloader = DataLoader(train_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
|
||||
valid_dataloader = DataLoader(valid_dataset, batch_size=batch_size, drop_last=True, num_workers=0)
|
||||
train_accumulator = Accumulator()
|
||||
@@ -129,7 +186,9 @@ if __name__ == '__main__':
|
||||
model = DepthEstimation(device=device, use_skip=use_skip)
|
||||
|
||||
loss_fn = nn.MSELoss()
|
||||
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
||||
loss_fn_3d = nn.MSELoss()
|
||||
optimizer = optim.Adam(model.parameters(), lr=0.01)
|
||||
# optimizer = optim.SGD(model.parameters(), lr=0.01)
|
||||
|
||||
iteration = 0
|
||||
while True:
|
||||
@@ -144,11 +203,12 @@ if __name__ == '__main__':
|
||||
|
||||
iteration += 1
|
||||
|
||||
for batch in valid_dataloader:
|
||||
valid_loss = valid_step(batch)
|
||||
valid_accumulator.update(valid_loss)
|
||||
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']}")
|
||||
print(f"Iteration: {iteration} | Valid loss: {valid_accumulator.get()['valid/loss']}")
|
||||
valid_accumulator.reset()
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user