add chamfer loss

data/dataset.py

@@ -24,10 +24,11 @@ from cloth3d.util import quads2tris
 
 
 class Dataset(Dataset):
-    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True):
+    def __init__(self, folder, folder_rgb, frames_file='frames_all.txt', device='cuda', use_cache=True, random_view=True):
         self.folder = folder
         self.folder_rgb = folder_rgb
         self.device = device
+        self.random_view = random_view
 
         self.transforms = tv.transforms.ToTensor()
         self.frames = read_frames_file(os.path.join(folder, frames_file))
@@ -52,8 +53,9 @@ class Dataset(Dataset):
 
         path = os.path.join(self.folder, frame.id)
         info = self.reader.read_info(frame.id)
+        zrot = info['zrot'] if not self.random_view else (((np.random.random(1) * 2) -1) * np.pi)[0]
         # mesh_posed = readPC2Frame(os.path.join(path, frame.garment + '_registered.pc16'), frame.frame_nr, True)
-        mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot'])
+        mesh_posed = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot)
         # mesh_posed_no_trans = self.reader.read_garment_vertices(frame.id, frame.garment, frame.frame_nr, absolute=False)
 
         mesh_mask = np.load(os.path.join(path, frame.garment + '_mask.npy'))
@@ -61,7 +63,7 @@ class Dataset(Dataset):
         faces = mask_faces(self.body_faces, mesh_mask)
         # faces = quads2tris(self.reader.read_garment_topology(frame.id, frame.garment))
         # body = readPC2Frame(os.path.join(path, 'Body_posed.pc16'), frame.frame_nr, True)
-        body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], info['zrot'])
+        body, body_faces, root_joint, joints, joints_no_trans = self.reader.read_human(frame.id, frame.frame_nr, info['trans'][:,frame.frame_nr], zrot)
         # body_no_trans, _, _, _, _ = self.reader.read_human(frame.id, frame.frame_nr, absolute=False)
         body_faces = quads2tris(body_faces)
 
@@ -85,7 +87,7 @@ class Dataset(Dataset):
 
         start = time.perf_counter()
         # TODO: is it sure I need this K, and not some transformed one?
-        normal_image = get_normals_from_depth(depth_image[None,:,:], depth_mask[None,:,:], K)
+        normal_image = get_normals_from_depth(depth_image[None,None,:,:], depth_mask[None,:,:], K)
         print(f'normals time: {time.perf_counter() - start}')
 
         depth_mask_body = (pix_to_face <= body_faces.shape[0]) * depth_mask
@@ -190,19 +192,19 @@ class Dataset(Dataset):
         # v2i_coverage[v2i[:,0], v2i[:, 1]] = 1
         # ex.imshow(v2i_coverage).show()
 
-        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(self.device).to(torch.float32)
+        rgb_image = tv.io.read_image(os.path.join(self.folder_rgb, frame.id, f'{frame.frame_nr:04d}.png'), mode=tv.io.ImageReadMode.RGB).to(torch.float32)
 
-        view = torch.tensor(np.concatenate((np.array([info['zrot']], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device)
+        view = torch.tensor(np.concatenate((np.array([zrot], dtype=np.float32), info['camLoc'].astype(np.float32), info['trans'][:,frame.frame_nr].astype(np.float32)), axis=0), device=self.device)
 
         # image = self.transforms(image)
-        data = { 'mesh_posed': mesh_posed,
+        data = { 'mesh_posed': torch.tensor(mesh_posed, device=self.device),
                  'depth_image': depth_image, 
                  'mesh_mask': mesh_mask, 
                  'depth_mask': depth_mask, 
                  'normal_image': normal_image,
                  'depth_mask_garment': depth_mask_garment,
                  'depth_mask_body': depth_mask_body,
-                 'rgb_image': rgb_image,
+                 'rgb_image': rgb_image.to(self.device),
                  'view': view,
                 }
 

models/convnext.py

@@ -180,9 +180,12 @@ class ConvNeXtReverseBase2(nn.Module):
         block: Optional[Callable[..., nn.Module]] = None,
         norm_layer: Optional[Callable[..., nn.Module]] = None,
         device: str = 'cuda',
+        use_skip=False,
     ) -> None:
         super().__init__()
 
+        self.use_skip = use_skip
+
         if not block_setting:
             raise ValueError("The block_setting should not be empty")
         elif not (isinstance(block_setting, Sequence) and all([isinstance(s, CNBlockConfig) for s in block_setting])):
@@ -259,10 +262,13 @@ class ConvNeXtReverseBase2(nn.Module):
 
 
     def _forward_impl(self, x) -> Tensor:
-        x, features = x
+        input, features = x
-        x = self.first(x)
+        x = self.first(input)
         for i in range(5 - 2):
-            x = self.stages[i](x + features[2-i])
+            if self.use_skip:
+                x = self.stages[i](torch.concat((x, features[2-i]), dim=1))
+            else:
+                x = self.stages[i](x)
             x = self.upsamples[i](x)
         
         x = self.stages[-2](x)

models/depth_estim.py

@@ -16,18 +16,20 @@ class ViewEncoder(nn.Module):
         self.lin1 = nn.LazyLinear(features).to(device)
         self.act1 = nn.GELU().to(device)
         self.lin2 = nn.Linear(features, features * 15 * 20).to(device)
+        self.act2 = nn.Tanh().to(device)
     
 
     def forward(self, input: torch.Tensor) -> torch.Tensor:
         x: torch.Tensor = self.lin1(input)
         x = self.act1(x)
         x = self.lin2(x)
+        x = self.act2(x)
         x = x.view(x.shape[0],self.features,15,20)
 
         return x
 
 class DepthEstimation(nn.Module):
-    def __init__(self, view_features=128, device='cuda'):
+    def __init__(self, view_features=128, device='cuda', use_skip=False):
         super().__init__()
 
         
@@ -40,16 +42,16 @@ class DepthEstimation(nn.Module):
         
         block_setting_reverse = [
             CNBlockConfig(768 + view_features, 384, 3),
-            CNBlockConfig(384, 192, 9),
+            CNBlockConfig(384 * ( 2 if use_skip else 1), 192, 9),
-            CNBlockConfig(192, 96, 3),
+            CNBlockConfig(192 * ( 2 if use_skip else 1), 96, 3),
-            CNBlockConfig(96, 27, 3),
+            CNBlockConfig(96 * ( 2 if use_skip else 1), 27, 3),
             CNBlockConfig(27, 9, 3),
             CNBlockConfig(9, None, 3),
         ]
 
         self.convnext = ConvNeXt(block_setting, device=device)
         self.flatten = nn.Flatten()
-        self.decoder = ConvNeXt(block_setting_reverse, reversed=True, device=device)
+        self.decoder = ConvNeXt(block_setting_reverse, use_skip=use_skip, reversed=True, device=device)
 
         self.view_encoder = ViewEncoder(view_features)
 

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()
@@ -48,7 +96,9 @@ def valid_step(batch):
         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/data_utils.py

@@ -46,10 +46,11 @@ def read_frames_file(file):
 
 
 def get_normals_from_depth(depth_image, depth_mask, camera_intrinsic):
-    depth_dilated = kornia.morphology.dilation(depth_image[:,None,:,:], torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
+    depth_dilated = kornia.morphology.dilation(depth_image, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
     depth_dilated = kornia.morphology.dilation(depth_dilated, torch.ones((3,3), device=depth_image.get_device()), engine='convolution')
     depth_dilated = torch.permute(depth_dilated, [0,2,3,1])
-    depth_dilated[depth_mask,:] = depth_image[:,:,:,None][depth_mask,:]
+    depth_image = torch.permute(depth_image, [0,2,3,1])
+    depth_dilated[depth_mask,:] = depth_image[depth_mask,:]
     depth_dilated = torch.permute(depth_dilated, [0,3,1,2])
     normals = kornia.geometry.depth.depth_to_normals(depth_dilated, camera_intrinsic, normalize_points=True)
     normals = torch.permute(normals, [0,2,3,1]) * depth_mask[:,:,:,None]

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