221 lines
9.0 KiB
Python
221 lines
9.0 KiB
Python
import os
|
|
os.environ["CUDA_VISIBLE_DEVICES"] = '0'
|
|
|
|
import sys
|
|
|
|
import torch
|
|
|
|
import torch.optim as optim
|
|
import torch.nn as nn
|
|
|
|
from torch.utils.data import DataLoader
|
|
|
|
sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
|
|
|
|
from data.dataset import Dataset
|
|
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):
|
|
import kornia
|
|
import plotly.graph_objects as go
|
|
from data.render import intrinsic
|
|
K = torch.tensor(intrinsic()[None, :, :], device='cuda')
|
|
pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True)
|
|
pc = pc[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'),
|
|
]).show()
|
|
|
|
|
|
# temp
|
|
def plot_gt(depth, gt, mask):
|
|
import kornia
|
|
import plotly.graph_objects as go
|
|
from data.render import intrinsic
|
|
K = torch.tensor(intrinsic()[None, :, :], device='cuda')
|
|
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=pc_gt[0,:], y=pc_gt[1,:], z=pc_gt[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
|
|
from data.render import intrinsic
|
|
from pytorch3d.structures import Pointclouds
|
|
from utils.pytorch3d_chamfer import chamfer_distance
|
|
from pytorch3d.ops.points_normals import estimate_pointcloud_normals
|
|
from utils.data_utils import get_normals_from_depth
|
|
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_norm = get_normals_from_depth(depth_image, depth_mask, K)
|
|
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)
|
|
|
|
return chamfer
|
|
|
|
|
|
# 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
|
|
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')
|
|
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 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, 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 + mse_normals #+ chamfer_reconst + mse + mse_normals + normals_reconst
|
|
|
|
# plot_gt(y_pred, batch)
|
|
|
|
# 3. Zero grad optimizer
|
|
optimizer.zero_grad()
|
|
|
|
# 4. Loss backward
|
|
loss.backward()
|
|
|
|
# 5. Step the optimizer
|
|
optimizer.step()
|
|
|
|
return { 'train/loss': loss.detach().cpu().numpy() }
|
|
|
|
|
|
def valid_step(batch):
|
|
### Testing
|
|
model.eval() # put the model in evaluation mode for testing (inference)
|
|
# 1. Forward pass
|
|
with torch.inference_mode():
|
|
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 + mse_normals #+ chamfer_reconst + mse + normals_reconst
|
|
|
|
|
|
return { 'valid/loss': loss.detach().cpu().numpy() }
|
|
|
|
|
|
def batch_to_device(batch, device):
|
|
filter = {'mesh_posed'}
|
|
for key in batch:
|
|
if key in filter:
|
|
batch[key] = batch[key].to(device)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
|
|
random_view = False
|
|
use_skip = False
|
|
batch_size = 1
|
|
|
|
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()
|
|
valid_accumulator = Accumulator()
|
|
|
|
model = DepthEstimation(device=device, use_skip=use_skip)
|
|
|
|
loss_fn = nn.MSELoss()
|
|
loss_fn_3d = nn.MSELoss()
|
|
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
|
|
while True:
|
|
|
|
for batch in train_dataloader:
|
|
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()
|
|
|
|
# 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()
|
|
|
|
|
|
|
|
|
|
pass |