add chamfer loss

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')
+    random_view = False
-    valid_dataset = Dataset('dataset/CLOTH3D/train', 'dataset/CLOTH3D2RGB/train', 'frames_small.txt')
+    use_skip = False
-    train_dataloader = DataLoader(train_dataset, batch_size=2, drop_last=True, num_workers=0)
+    batch_size = 2
-    valid_dataloader = DataLoader(valid_dataset, batch_size=2, drop_last=True, num_workers=0)
+
+    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)

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