From c315f799416d4cac4b2cc03f7e7097e9253ec458 Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sat, 14 Jan 2023 17:13:27 +0000 Subject: [PATCH] add chamfer loss --- test.py | 70 +++++++++++++++--- utils/pytorch3d_chamfer.py | 141 +++++++++++++++++++++++++++++++++++++ 2 files changed, 203 insertions(+), 8 deletions(-) create mode 100644 utils/pytorch3d_chamfer.py diff --git a/test.py b/test.py index 76c0b7f..29aa82e 100644 --- a/test.py +++ b/test.py @@ -18,6 +18,52 @@ from models.depth_estim import DepthEstimation from models.convnext import ConvNeXt + + +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() + + +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_pc = Pointclouds(pc) + mesh_gt = [vertices[i, mesh_mask[i], :] for i in range(vertices.shape[0])] + mesh_gt_pc = Pointclouds(mesh_gt) + pc_norm = get_normals_from_depth(depth_image, depth_mask, K) + mesh_norm = estimate_pointcloud_normals(mesh_gt_pc, neighborhood_size=50) + + chamfer = chamfer_distance(pc_pc, mesh_gt_pc, x_normals=pc_norm, y_normals=mesh_norm) + + 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 train_step(batch): ### Training model.train() # train mode is on by default after construction @@ -26,7 +72,9 @@ def train_step(batch): y_pred = model([batch['rgb_image'], batch['view']]) # 2. Calculate loss - loss = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) + mse = loss_fn(y_pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) + chamfer, _ = chamfer_loss(y_pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask']) + loss = mse + 1e-1 * chamfer # 3. Zero grad optimizer optimizer.zero_grad() @@ -46,9 +94,11 @@ def valid_step(batch): # 1. Forward pass with torch.inference_mode(): pred = model([batch['rgb_image'], batch['view']]) - + # 2. Calculate the loss - loss = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) + mse = loss_fn(pred * batch['depth_mask'][:,None,:,:], batch['depth_image'][:,None,:,:] * batch['depth_mask'][:,None,:,:]) + chamfer, _ = chamfer_loss(pred, batch['mesh_posed'], batch['depth_mask'], batch['mesh_mask']) + loss = mse + 1e-1 * chamfer return { 'valid/loss': loss.detach().cpu().numpy() } @@ -63,14 +113,18 @@ def batch_to_device(batch, device): if __name__ == '__main__': device = "cuda" if torch.cuda.is_available() else "cpu" - train_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt') - valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt') - train_dataloader = DataLoader(train_dataset, batch_size=2, drop_last=True, num_workers=0) - valid_dataloader = DataLoader(valid_dataset, batch_size=2, drop_last=True, num_workers=0) + random_view = False + use_skip = False + batch_size = 2 + + 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_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) + model = DepthEstimation(device=device, use_skip=use_skip) loss_fn = nn.MSELoss() optimizer = optim.Adam(model.parameters(), lr=0.001) diff --git a/utils/pytorch3d_chamfer.py b/utils/pytorch3d_chamfer.py new file mode 100644 index 0000000..78f6c77 --- /dev/null +++ b/utils/pytorch3d_chamfer.py @@ -0,0 +1,141 @@ +# taken from pytorch3d.loss.chamfer, but only one directional + +from typing import Union + +import torch +import torch.nn.functional as F +from pytorch3d.ops.knn import knn_gather, knn_points +from pytorch3d.loss.chamfer import _handle_pointcloud_input, _validate_chamfer_reduction_inputs + + +def chamfer_distance( + x, + y, + x_lengths=None, + y_lengths=None, + x_normals=None, + y_normals=None, + weights=None, + batch_reduction: Union[str, None] = "mean", + point_reduction: str = "mean", + norm: int = 2, +): + """ + Chamfer distance between two pointclouds x and y. + + Args: + x: FloatTensor of shape (N, P1, D) or a Pointclouds object representing + a batch of point clouds with at most P1 points in each batch element, + batch size N and feature dimension D. + y: FloatTensor of shape (N, P2, D) or a Pointclouds object representing + a batch of point clouds with at most P2 points in each batch element, + batch size N and feature dimension D. + x_lengths: Optional LongTensor of shape (N,) giving the number of points in each + cloud in x. + y_lengths: Optional LongTensor of shape (N,) giving the number of points in each + cloud in y. + x_normals: Optional FloatTensor of shape (N, P1, D). + y_normals: Optional FloatTensor of shape (N, P2, D). + weights: Optional FloatTensor of shape (N,) giving weights for + batch elements for reduction operation. + batch_reduction: Reduction operation to apply for the loss across the + batch, can be one of ["mean", "sum"] or None. + point_reduction: Reduction operation to apply for the loss across the + points, can be one of ["mean", "sum"]. + norm: int indicates the norm used for the distance. Supports 1 for L1 and 2 for L2. + + Returns: + 2-element tuple containing + + - **loss**: Tensor giving the reduced distance between the pointclouds + in x and the pointclouds in y. + - **loss_normals**: Tensor giving the reduced cosine distance of normals + between pointclouds in x and pointclouds in y. Returns None if + x_normals and y_normals are None. + """ + _validate_chamfer_reduction_inputs(batch_reduction, point_reduction) + + if not ((norm == 1) or (norm == 2)): + raise ValueError("Support for 1 or 2 norm.") + + x, x_lengths, x_normals = _handle_pointcloud_input(x, x_lengths, x_normals) + y, y_lengths, y_normals = _handle_pointcloud_input(y, y_lengths, y_normals) + + return_normals = x_normals is not None and y_normals is not None + + N, P1, D = x.shape + P2 = y.shape[1] + + # Check if inputs are heterogeneous and create a lengths mask. + is_x_heterogeneous = (x_lengths != P1).any() + x_mask = ( + torch.arange(P1, device=x.device)[None] >= x_lengths[:, None] + ) # shape [N, P1] + + if y.shape[0] != N or y.shape[2] != D: + raise ValueError("y does not have the correct shape.") + if weights is not None: + if weights.size(0) != N: + raise ValueError("weights must be of shape (N,).") + if not (weights >= 0).all(): + raise ValueError("weights cannot be negative.") + if weights.sum() == 0.0: + weights = weights.view(N, 1) + if batch_reduction in ["mean", "sum"]: + return ( + (x.sum((1, 2)) * weights).sum() * 0.0, + (x.sum((1, 2)) * weights).sum() * 0.0, + ) + return ((x.sum((1, 2)) * weights) * 0.0, (x.sum((1, 2)) * weights) * 0.0) + + cham_norm_x = x.new_zeros(()) + + x_nn = knn_points(x, y, lengths1=x_lengths, lengths2=y_lengths, norm=norm, K=1) + + cham_x = x_nn.dists[..., 0] # (N, P1) + + if is_x_heterogeneous: + cham_x[x_mask] = 0.0 + + if weights is not None: + cham_x *= weights.view(N, 1) + + if return_normals: + # Gather the normals using the indices and keep only value for k=0 + x_normals_near = knn_gather(y_normals, x_nn.idx, y_lengths)[..., 0, :] + + cham_norm_x = 1 - torch.abs( + F.cosine_similarity(x_normals, x_normals_near, dim=2, eps=1e-6) + ) + + if is_x_heterogeneous: + cham_norm_x[x_mask] = 0.0 + + if weights is not None: + cham_norm_x *= weights.view(N, 1) + + # Apply point reduction + cham_x = cham_x.sum(1) # (N,) + if return_normals: + cham_norm_x = cham_norm_x.sum(1) # (N,) + if point_reduction == "mean": + x_lengths_clamped = x_lengths.clamp(min=1) + cham_x /= x_lengths_clamped + if return_normals: + cham_norm_x /= x_lengths_clamped + + if batch_reduction is not None: + # batch_reduction == "sum" + cham_x = cham_x.sum() + if return_normals: + cham_norm_x = cham_norm_x.sum() + if batch_reduction == "mean": + div = weights.sum() if weights is not None else max(N, 1) + cham_x /= div + if return_normals: + cham_norm_x /= div + + cham_dist = cham_x + cham_normals = cham_norm_x if return_normals else None + + return cham_dist, cham_normals