import os os.environ["CUDA_VISIBLE_DEVICES"] = '6' 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, mesh, info, root_dist, scale, mask, root_joint_pixel): import kornia import plotly.graph_objects as go from utils.data_utils import get_camera import torchvision as tv cameras, K = get_camera(info['camLoc'][0], [480, 640], depth.get_device()) mask_orig = mask[:,None,:,:] mask, mask_croped = restore(mask[:,None,:,:], scale, root_joint_pixel) depth, _ = restore(depth, scale, root_joint_pixel, (mask_orig.to(torch.bool), mask_croped.to(torch.bool)), root_dist)# * mask pc = depth_to_3d(depth, mask.to(torch.bool), root_joint_pixel, root_dist, K) pc = pc.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() def restore(img, scale, center, masks=None, root_dist=0.0): import torchvision as tv import numpy as np from utils.data_utils import resize_image if masks is None: resizer = tv.transforms.Resize((torch.round(256 * scale).to(torch.int64), torch.round(256 * scale).to(torch.int64)), interpolation=tv.transforms.functional.InterpolationMode.NEAREST) else: resizer = lambda img, mask_croped, mask_resized: resize_image(img, mask_croped[:,0,:,:], mask_resized[:,0,:,:], (torch.round(256 * scale).to(torch.int64), torch.round(256 * scale).to(torch.int64))) cropper = lambda img, off: tv.transforms.functional.crop(img, off[1], off[0], 480, 640) if masks is None: out = resizer(img + root_dist) else: out = resizer(img + root_dist, masks[0], masks[1]) out_croped = out offset = (np.array(np.round([out.shape[2]/2, out.shape[2]/2]))[None,:] - center.cpu().numpy()).astype(np.int64) out = torch.stack([cropper(out[i], offset[i]) for i in range(img.shape[0])]) out = out return out, out_croped def depth_to_3d(depth, depth_mask, root_joint_pixel, root_dist, K): import kornia # pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True) root_joint_image = torch.zeros_like(depth) root_joint_image[:,:,root_joint_pixel[:,1], root_joint_pixel[:,0]] = root_dist root_reprojected = kornia.geometry.depth.depth_to_3d(root_joint_image, K, normalize_points=True)[0,:,root_joint_pixel[:,1], root_joint_pixel[:,0]][:,0] pc = kornia.geometry.depth.depth_to_3d(depth, K, normalize_points=True) pc = pc[0,:,depth_mask[0,0,:,:]] # pc[2,:] -= root_dist # root_reprojected[2] = root_dist pc = pc - root_reprojected[:,None] return pc def plot3d_with_mesh(depth, mesh, info, root_dist, scale, mask, root_joint_pixel, joints, original, original_mask, device='cuda'): import plotly.graph_objects as go import plotly.express as ex from cloth3d.util import zRotMatrix import pytorch3d import numpy as np from utils.data_utils import intrinsic, extrinsic, get_camera import torchvision as tv cameras, K = get_camera(info['camLoc'][0], [480, 640], depth.get_device()) # cameras, K = get_camera(info['camLoc'][0], [640, 480], depth.get_device()) # K = torch.tensor(intrinsic()[None, :, :], device=device) # pc = cameras.unproject_points(depth) # depth = original[None, :, :, :] # mask = original_mask[None, :, :, :] mask_orig = mask[:,None,:,:] mask, mask_croped = restore(mask[:,None,:,:], scale, root_joint_pixel) depth, _ = restore(depth, scale, root_joint_pixel, (mask_orig.to(torch.bool), mask_croped.to(torch.bool)), root_dist)# * mask pc = depth_to_3d(depth, mask.to(torch.bool), root_joint_pixel, root_dist, K) pc = pc.detach().cpu().numpy() mesh = mesh[0] - joints[0,0,:] # zRot = torch.tensor(zRotMatrix(-info['zrot']), device=device) # mesh = torch.matmul(zRot, mesh.T).T # mesh = mesh - info['trans'][0].to(torch.float32).to(device) # zRot = torch.tensor(zRotMatrix(info['zrot']), device=device) # mesh = torch.matmul(zRot, mesh.T).T # view = pytorch3d.transforms.axis_angle_to_matrix(torch.tensor([0.0, 0.0, torch.pi/2], device=device)) # transform = pytorch3d.transforms.Rotate(view) # mesh = transform.transform_points(mesh) mesh = mesh.detach().cpu().numpy() go.Figure(data = [go.Scatter3d(x=pc[2,:], y=pc[0,:], z=-pc[1,:], marker=dict(size=3, line=dict(width=0)), mode='markers'), go.Scatter3d(x=-mesh[:,0], y=mesh[:,1], z=mesh[:,2], marker=dict(size=3, line=dict(width=0)), mode='markers'), ]).show() pass # 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): 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, :, :] 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): # plot3d_with_mesh(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel'], batch['joints'], batch['original'], batch['original_mask']) # plot3d(batch['depth_image'][:,None,:,:], batch['mesh_posed'], batch['info'], batch['root_dist'], batch['scale'], batch['depth_mask'], batch['root_joint_pixel']) ### Training model.train() # train mode is on by default after construction # input = torch.concat((batch['rgb_image'], batch['normal_image']), axis=1) input = batch['rgb_image'] # import plotly.express as ex # ex.imshow(batch['depth_image'][0,:,:].cpu().numpy(), zmin=-1, zmax=5).show() # 1. Forward pass y_pred, y_pred_normals = model([input, 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 = mse #+ 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(): # input = torch.concat((batch['rgb_image'], batch['normal_image']), axis=1) input = batch['rgb_image'] pred, pred_normals = model([input, 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 = mse #+ 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 use_cache = True train_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache) valid_dataset = Dataset('/home/workspace/data/dataset/cloth3d/train', '/home/workspace/data/dataset/cloth3drgb/train', 'frames_00005_single.txt', random_view=random_view, device=device, use_cache=use_cache) 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.ExponentialLR(optimizer, gamma=0.9999, verbose=False) # 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