add boundary optimization

data/registration_weights.py

@@ -10,7 +10,7 @@ from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal
 from pymeshfix import MeshFix
 import scipy.spatial as spatial
 
-os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from data.cloth3d_io import (
@@ -30,7 +30,7 @@ from data.cloth3d_io import (
 from smpl.smpl_torch import SMPL
 # from custom_utils.custom_utils import _maskFaces
 
-device = 'cpu'
+device = 'cuda'
 smpl = SMPL(device, 'smpl')
 
 
@@ -232,9 +232,9 @@ def register(folder, frame):
     _, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
     faces = quads2tris(faces)
     template = np.load('temp/mapping_cloth.npy')
-    mapping_mask = np.load('temp/mapping_mask.npy')
     mapping_faces = np.load('temp/mapping_faces.npy')
     mapping_weights = np.load('temp/mapping_weights.npy')
+    mapping_mask = np.load('temp/mapping_mask.npy')
 
     mapping_faces, mapping_weights = mapping_faces[mapping_mask], mapping_weights[mapping_mask]
 
@@ -501,9 +501,139 @@ def get_boundaries(folder, frame):
 
     pass
 
+def optimize_boundaries(folder, frame):
+    id, frame_nr, garment = frame.split()
+    frame_nr = int(frame_nr)
+    registered_cloth = np.load('temp/registered_cloth.npy')
+    registered_faces = np.load('temp/registered_faces.npy')
+    registered_weights = np.load('temp/registered_weights.npy')
+    _, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
+    faces = quads2tris(faces)
+    template = np.load('temp/mapping_cloth.npy')
+    mapping_mask = np.load('temp/mapping_mask.npy')
+    mapping_faces = np.load('temp/mapping_faces.npy')
+    mapping_weights = np.load('temp/mapping_weights.npy')
+
+    mapping_faces, mapping_weights = mapping_faces[mapping_mask], mapping_weights[mapping_mask]
+
+    template_border = np.load('temp/template_border.npy')
+    registered_border = np.load('temp/registered_border.npy')
+
+
+    template = torch.tensor(template).to(device).to(torch.float32)
+    faces = torch.tensor(faces).to(device)
+    registered_faces = torch.tensor(registered_faces).to(device)
+    mapping_faces = torch.tensor(mapping_faces).to(device)
+    registered_weights = torch.tensor(registered_weights).to(device).to(torch.float32)
+    registered_cloth = torch.tensor(registered_cloth).to(device).to(torch.float32)
+
+
+    mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
+    avg_edge_length = mesh_edge_lengths(mesh2).mean()
+
+    optim_mask = torch.zeros(registered_cloth.shape[0]).to(device).to(torch.bool)
+    optim_mask[registered_border] = True
+
+    mapping_weights_masked = registered_weights[optim_mask]
+    mapping_weights_nonmasked = registered_weights[~optim_mask]
+    mapping_weights_masked.requires_grad = True
+
+    # Define the optimizer and learning rate
+    optimizer = torch.optim.Adam([mapping_weights_masked], lr=0.001)
+    # optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001)
+
+    # Define the loss functions
+    def laplacian_smoothing_loss(mesh):
+        return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
+
+    def edge_loss(mesh):
+        return mesh_edge_loss(mesh, target_length=avg_edge_length)
+
+    def normal_consistency_loss(mesh):
+        return mesh_normal_consistency(mesh)
+
+    # Optimization loop
+    for i in range(4000):
+        optimizer.zero_grad()
+
+        # find boundaries with the find holes function?
+        registered_masked = (template[mapping_faces[optim_mask],:] * mapping_weights_masked[:,:,None]).sum(1)
+        registered_nonmasked = (template[mapping_faces[~optim_mask],:] * mapping_weights_nonmasked[:,:,None]).sum(1)
+        # registered = torch.cat((registered_masked, registered_nonmasked), 0)
+        registered = torch.zeros_like(registered_weights)
+        registered[optim_mask] = registered_masked
+        registered[~optim_mask] = registered_nonmasked
+
+        # Forward pass
+        mesh1_out = Meshes(verts=registered[None,:,:], faces=registered_faces[None,:,:])
+        # ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()),
+        #                                 )
+        ch_bound_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,registered_border,:], torch.stack(mesh2.verts_list())[:, template_border,:],
+                                        )
+        # loss: torch.Tensor = ch_loss * 1 + ch_bound_loss * 20
+        loss: torch.Tensor = ch_bound_loss * 50
+        lap_loss = laplacian_smoothing_loss(mesh1_out)
+        e_loss = edge_loss(mesh1_out)
+        # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
+        norm_loss = normal_consistency_loss(mesh1_out)
+        loss += lap_loss * 0.05 + e_loss * 2 + norm_loss * 0.1
+
+        # Backward pass
+        loss.backward()
+        optimizer.step()
+
+        # Print the loss
+        if i % 10 == 0:
+            print(f"it: {i} : loss: {loss:.5f} bound_chamfer: {ch_bound_loss:.6f}")
+            # print(f"it: {i} : loss: {loss:.5f} bound_chamfer: {ch_bound_loss:.6f} chamfer: {ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}")
+    
+    registered = mesh1_out.verts_packed().detach().cpu().numpy()
+    template = mesh2.verts_packed().detach().cpu().numpy()
+    faces = faces.detach().cpu().numpy()
+    optim_mask = optim_mask.detach().cpu().numpy()
+    registered_faces = registered_faces.detach().cpu().numpy()
+    # faces, F = F, faces
+    # verts, star_body = star_body, verts
+
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            x=template[:,0], y=template[:,1], z=template[:,2],
+            i=faces[:,0], j=faces[:,1], k=faces[:,2],
+            color='lightblue', opacity=0.80),
+        go.Mesh3d(
+            x=registered[:,0], y=registered[:,1], z=registered[:,2],
+            i=registered_faces[:,0], j=registered_faces[:,1], k=registered_faces[:,2],
+            color='lightpink', opacity=0.50),   
+
+    ])
+    fig.update_layout(scene = dict(
+        xaxis = dict(range=[-1,1],),
+        yaxis = dict(range=[-1,1],),
+        zaxis = dict(range=[-1,1],),),)
+    fig.update_layout(scene_aspectmode='cube')
+    fig.show()
+
+
+    fig = go.Figure(data=[
+        go.Scatter3d(
+            x=registered[:,0], y=registered[:,1], z=registered[:,2],
+            mode='markers', marker=dict(size=4, color='red')),
+        go.Scatter3d(
+            x=template[:,0], y=template[:,1], z=template[:,2],
+            mode='markers', marker=dict(size=2, color='blue')),
+    ])
+    fig.update_layout(scene = dict(
+        xaxis = dict(range=[-1,1],),
+        yaxis = dict(range=[-1,1],),
+        zaxis = dict(range=[-1,1],),),)
+    fig.update_layout(scene_aspectmode='cube')
+    fig.show()
+
+    pass
+
 if __name__ == "__main__":
-    # folder = '/data-net/code/cloth3d/train'
-    folder = r'D:\Work\code\cloth3d\train'
+    folder = '/data-net/code/cloth3d/train'
+    # folder = r'D:\Work\code\cloth3d\train'
     # frames = [f'00761_{i}' for i in range(300)]
     # folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
     frames = get_frames(folder, 'frames_all.txt')
@@ -529,6 +659,7 @@ if __name__ == "__main__":
         
 
         # register(folder, frame)
-        get_boundaries(folder, frame)
+        # get_boundaries(folder, frame)
+        optimize_boundaries(folder, frame)
         break
         pass