registration progress save

data/registration.py

@@ -7,6 +7,9 @@ import torch
 from torch.nn.functional import mse_loss
 from pytorch3d.structures import Meshes
 from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
+from pymeshfix import MeshFix
+
+# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
 from data.cloth3d_io import (
@@ -26,7 +29,40 @@ from data.cloth3d_io import (
 from smpl.smpl_torch import SMPL
 # from custom_utils.custom_utils import _maskFaces
 
-smpl = SMPL('cuda', 'smpl')
+device = 'cuda'
+smpl = SMPL(device, 'smpl')
+
+
+import torch.nn.functional as F
+from pytorch3d.structures import Meshes
+
+def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
+    """
+    Given two PyTorch3D meshes and a set of vertex indices for mesh1,
+    find the closest vertices from mesh2 to these vertices.
+
+    Args:
+        mesh1: A PyTorch3D Meshes object representing the first mesh.
+        mesh2: A PyTorch3D Meshes object representing the second mesh.
+        vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
+
+    Returns:
+        A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
+    """
+    # Extract the vertices of mesh1 and mesh2.
+    vertices1 = mesh1.verts_packed()
+    vertices2 = mesh2.verts_packed()
+
+    # Extract the vertices of mesh1 that correspond to the given indices.
+    vertices1_selected = vertices1[vertex_indices]
+
+    # Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
+    distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
+
+    # Find the indices of the closest vertices in mesh2.
+    closest_indices = torch.argmin(distances, dim=1)
+
+    return closest_indices
 
 
 def _maskFaces(V, F, mask, return_unfixed=False):
@@ -90,6 +126,8 @@ def build_lookup_table(repeated):
 from data.cloth3d_io import readPC2Frame
 import scipy.io as sio
 import scipy
+import scipy.spatial as spatial
+
 def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
     pc16_path = os.path.join(folder, id, garment + '.pc16')
     verts = readPC2Frame(pc16_path, frame, True)
@@ -97,8 +135,8 @@ def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorie
     map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
 
     if deorient:
-        R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
-        # R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
         verts = R.apply(verts)
 
     return verts, map
@@ -108,77 +146,293 @@ def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
     verts = readPC2Frame(pc16_path, frame, True)
 
     if deorient:
-        R = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
-        # R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
         verts = R.apply(verts)
 
     return verts
 
+def edge_loss_custom(x, y, edges):
+    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
+    x_e = torch.sum(x_e ** 2, -1)
+    x_e = torch.sqrt(x_e + 1e-9)
+    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
+    y_e = torch.sum(y_e ** 2, -1)
+    y_e = torch.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    loss = torch.mean(d_e ** 2, dim=-1)[0]
+    return loss
+
+
+def crotch(mesh, body, crotch_idx, nr=75):
+    crotch_middle_on_body = crotch_idx
+    # create kdtree for body and mesh
+    body_tree = spatial.cKDTree(body)
+    mesh_tree = spatial.cKDTree(mesh)
+    # find closest point on mesh from crotch middle on body along the z axis
+
+    crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
+
+    crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
+    crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
+
+    body_verts = body[crotch_body]
+    mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
+    mesh_verts = spatial.cKDTree(mesh_verts)
+    closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
+    crotch_mesh = crotch_mesh[closest_mesh_verts]
+
+    return crotch_body, crotch_mesh
+
+
+def find_closest_point_along_z(point, verts, eps=0.01):
+    min_dist = np.inf
+    closest_point = None
+    for i, v in enumerate(verts):
+        dist = np.linalg.norm(point - v)
+        if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
+            min_dist = dist
+            closest_point = i
+    return closest_point
+
+
+def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
+    """
+    Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
+
+    Args:
+        mesh: A PyTorch3D Meshes object representing the mesh.
+
+    Returns:
+        A tensor of shape (E,) representing the lengths of all edges in the mesh,
+        where E is the total number of edges in the mesh.
+    """
+    # Extract the vertices and faces of the mesh.
+    vertices = mesh.verts_packed()
+    faces = mesh.faces_packed()
+
+    # Compute the pairwise distances between all pairs of vertices that form an edge.
+    edge_indices = faces.view(-1, 3)
+    edge_vertices = vertices[edge_indices]
+    edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
+
+    return edge_lengths
+
+
 def register(folder, frame):
     id, frame_nr, garment = frame.split()
     frame_nr = int(frame_nr)
     info = loadInfo(os.path.join(folder, id, 'info.mat'))
     raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
+    # convert it to indices
     mask = raw_mask > 0
     mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
     body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
     template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
     faces = quads2tris(faces)
-    pose = torch.zeros(72)
+    
-    shape = torch.zeros(11)
+    # boundaries = raw_mask == 2
-    shape[-1] = 1
+    # # invert map and use it to invert boudaries
+    # map_inv = np.zeros(map.max()+1, dtype=np.int32)
+    # map_inv[map] = np.arange(len(map))
+    # boundaries = boundaries[map_inv]
+    # # get the boundary indices
+    
+    # mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
+
+
+    
+    # pose = torch.zeros(1, 72).to('cuda')
+    shape = torch.zeros(1, 11).to(device)
+    shape[:, :10] = torch.tensor(info['shape']).to(device)
+    shape[:, -1] = torch.tensor(info['gender']).to(device)
     # for star pose
-    pose[:,5] = 3.1415/4
+    # pose[:, 5] = 3.1415/4
-    pose[:,8] = -3.1415/4
+    # pose[:,8] = -3.1415/4
-    star_body = smpl(pose.unsqueeze(0), shape.unsqueeze(0))[0][0].cpu().numpy()
+    STAR_POSE = np.zeros((24, 3), np.float32)
+    STAR_POSE[1,2] = .15
+    STAR_POSE[2,2] = -.15
+    pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
+    star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
+    # rotate star body with 90 degrees around x axis
+    R = spatial.transform.Rotation.from_euler('x', np.pi/2)
+    star_body = R.apply(star_body).astype(np.float32)
 
     is_skirt = (garment == 'Dress' or garment == 'Skirt')
     skirt_mod = '_skirt' if is_skirt else ''
     F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
 
-    plot(template, faces)
+    ##################### mask body #####################
+    new_crotch = 4729
+
+    # star_body = star_body[mask]
+    # new_crotch = np.zeros((14475), bool)
+    # new_crotch[4729] = True
+    # new_crotch = new_crotch[mask]
+    # new_crotch = np.where(new_crotch)[0][0]
+    # _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
+
+    # plot(template, faces)
     # v = mesh_posed[mask,:]
     
+    ##################### boundaries #####################
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    # mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    ##################### crotch #####################
+    crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
+
+
+
+
+    # export star_body as obj
+    # writeOBJ('temp/star_body.obj', star_body, F)
+
+
+    template = torch.tensor(template).to(device)
+    template_orig = template.clone()
+    faces = torch.tensor(faces).to(device)
+    star_body = torch.tensor(star_body).to(device)
+    F = torch.tensor(F).to(device)
+
+    # verts = template.cpu().numpy()
+    # star_body = star_body.cpu().numpy()
+    # faces = faces.detach().cpu().numpy()
+    # F = F.detach().cpu().numpy()
+    # fig = go.Figure(data=[
+    #     go.Mesh3d(
+    #         x=verts[:,0], y=verts[:,1], z=verts[:,2],
+    #         i=faces[:,0], j=faces[:,1], k=faces[:,2],
+    #         color='lightpink', opacity=0.50),
+    #     go.Mesh3d(
+    #         x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
+    #         i=F[:,0], j=F[:,1], k=F[:,2],
+    #         color='lightblue', opacity=0.50),
+    #     go.Scatter3d(
+    #         x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+    #         mode='markers', marker=dict(size=2, color='red')),
+    #     go.Scatter3d(
+    #         x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+    #         mode='markers', marker=dict(size=2, color='blue')),
+    # ])
+    # # make it cuve and limit to -1 and -1
+    # fig.update_layout(scene = dict(
+    #     xaxis = dict(range=[-1,1],),
+    #     yaxis = dict(range=[-1,1],),
+    #     zaxis = dict(range=[-1,1],),),)
+    # fig.show()
+
     # Define the meshes to optimize
     mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
-    mesh2 = Meshes(verts=star_body, faces=F[None,:,:])
+    mesh2 = Meshes(verts=star_body[None,:,:], faces=F[None,:,:])
+
+    mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
+
+    avg_edge_length = mesh_edge_lengths(mesh1).mean()
 
     # Define the optimizer and learning rate
-    optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.01)
+    mesh1.verts_packed().requires_grad = True
+    optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
+    # 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")
+        return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot
 
-    def edge_loss(mesh1, mesh2):
+    def edge_loss(mesh):
-        return mesh_edge_loss(mesh1, mesh2)
+        return mesh_edge_loss(mesh, target_length=avg_edge_length)
 
-    def normal_consistency_loss(mesh1, mesh2):
+    def normal_consistency_loss(mesh):
-        return mesh_normal_consistency(mesh1, mesh2)
+        return mesh_normal_consistency(mesh)
 
     # Optimization loop
-    for i in range(100):
+    for i in range(10000):
         optimizer.zero_grad()
 
+        # find boundaries with the find holes function?
+
         # Forward pass
         mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
-        loss = chamfer_distance(mesh1_out, mesh2)
+        ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
-        loss += laplacian_smoothing_loss(mesh1_out)
+        bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
-        loss += edge_loss(mesh1_out)
+        # crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
-        loss += normal_consistency_loss(mesh1_out)
+        crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
+        loss: torch.Tensor = ch_loss * 10 +  (bound_ch_loss + crotch_ch_loss) * 5
+        # loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
+        loss += laplacian_smoothing_loss(mesh1_out) * 0.05
+        loss += edge_loss(mesh1_out) * 100
+        # loss += edge_loss_custom(torch.stack(mesh1_out.verts_list()), template_orig[None,:,:], mesh1.edges_packed()) * 100
+        loss += normal_consistency_loss(mesh1_out) * 0.1
 
         # Backward pass
         loss.backward()
         optimizer.step()
 
         # Print the loss
-        print(f"Loss: {loss.item()}")
+        if i % 10 == 0:
+            print(f"Loss: {loss.item()} chamfer: {ch_loss.item()} ")
+
+    # create a plotly plot with the body and the mesh1_out
+    verts = mesh1_out.verts_packed().detach().cpu().numpy()
+    star_body = star_body.detach().cpu().numpy()
+    faces = faces.detach().cpu().numpy()
+    F = F.detach().cpu().numpy()
+    mesh1_boundaries = mesh1_boundaries.detach().cpu().numpy()
+    mesh2_boundaries = mesh2_boundaries.detach().cpu().numpy()
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            x=verts[:,0], y=verts[:,1], z=verts[:,2],
+            i=faces[:,0], j=faces[:,1], k=faces[:,2],
+            color='lightpink', opacity=0.80),
+        go.Mesh3d(
+            x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
+            i=F[:,0], j=F[:,1], k=F[:,2],
+            color='lightblue', opacity=0.50),
+        go.Scatter3d(
+            x=verts[mesh1_boundaries,0], y=verts[mesh1_boundaries,1], z=verts[mesh1_boundaries,2],
+            mode='markers', marker=dict(size=2, color='red')),
+        go.Scatter3d(
+            x=star_body[mesh2_boundaries,0], y=star_body[mesh2_boundaries,1], z=star_body[mesh2_boundaries,2],
+            mode='markers', marker=dict(size=2, color='blue')),   
+        go.Scatter3d(
+            x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+            mode='markers', marker=dict(size=2, color='red')),
+        go.Scatter3d(
+            x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+            mode='markers', marker=dict(size=2, color='blue')),
+    ])
+    # make it cuve and limit to -1 and -1
+    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/dataset/cloth3d/train'
+    folder = '/data-net/code/cloth3d/train'
-    # folder = r'D:\Work\Dataset\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')
@@ -203,5 +457,5 @@ if __name__ == "__main__":
         
 
         register(folder, frame)
-
+        break
         pass

data/registration_post.py

@@ -0,0 +1,482 @@
+import os
+import sys
+import random
+import numpy as np
+import plotly.graph_objects as go
+import torch
+from torch.nn.functional import mse_loss
+from pytorch3d.structures import Meshes
+from pytorch3d.loss import mesh_laplacian_smoothing, mesh_edge_loss, mesh_normal_consistency, chamfer_distance
+from pymeshfix import MeshFix
+
+# os.environ["CUDA_LAUNCH_BLOCKING"] = '1'
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+from data.cloth3d_io import (
+    read_garment_mask,
+    read_garment_template,
+    read_posed_garment_frame,
+    loadInfo,
+    read_posed_body_frame,
+    read_posed_body,
+    read_posed_garment,
+    read_curvatures,
+    read_curvature_frame,
+    writeOBJ,
+    readOBJ,
+    quads2tris
+)
+from smpl.smpl_torch import SMPL
+# from custom_utils.custom_utils import _maskFaces
+
+device = 'cuda'
+smpl = SMPL(device, 'smpl')
+
+
+import torch.nn.functional as F
+from pytorch3d.structures import Meshes
+
+def find_closest_vertices(mesh1: Meshes, mesh2: Meshes, vertex_indices: torch.Tensor) -> torch.Tensor:
+    """
+    Given two PyTorch3D meshes and a set of vertex indices for mesh1,
+    find the closest vertices from mesh2 to these vertices.
+
+    Args:
+        mesh1: A PyTorch3D Meshes object representing the first mesh.
+        mesh2: A PyTorch3D Meshes object representing the second mesh.
+        vertex_indices: A tensor of shape (N,) representing the indices of the vertices in mesh1.
+
+    Returns:
+        A tensor of shape (N,) representing the indices of the closest vertices in mesh2.
+    """
+    # Extract the vertices of mesh1 and mesh2.
+    vertices1 = mesh1.verts_packed()
+    vertices2 = mesh2.verts_packed()
+
+    # Extract the vertices of mesh1 that correspond to the given indices.
+    vertices1_selected = vertices1[vertex_indices]
+
+    # Compute the pairwise distances between the selected vertices of mesh1 and all vertices of mesh2.
+    distances = F.pairwise_distance(vertices1_selected.unsqueeze(1), vertices2.unsqueeze(0), p=2)
+
+    # Find the indices of the closest vertices in mesh2.
+    closest_indices = torch.argmin(distances, dim=1)
+
+    return closest_indices
+
+
+def _maskFaces(V, F, mask, return_unfixed=False):
+    V_mask = mask
+    F_mask = np.take(V_mask, F.reshape(-1), axis=0)
+    F_mask = F_mask.reshape((-1, 3)).sum(axis=1) == 3
+
+    V_ind = np.zeros(V_mask.shape)
+    V_ind[V_mask] = np.arange(V_mask.sum())
+    F_new = np.take(V_ind, F[F_mask, :].reshape(-1), axis=0).reshape((-1, 3))
+    V_new = V[V_mask, :]
+
+    if return_unfixed:
+        return V_new, F_new, F[F_mask, :]
+    else:
+        return V_new, F_new
+
+def get_frames(folder, frames_file):
+    frames = []
+    with open(os.path.join(folder, frames_file), 'r') as f:
+        for line in f:
+            frames.append(line.strip())
+    return frames
+
+
+def plot(mesh, faces):
+    fig = go.Figure()
+
+    fig.add_trace(go.Mesh3d(
+                x=mesh[:,0],
+                y=mesh[:,1],
+                z=mesh[:,2],
+                i = faces[:,0],
+                j = faces[:,1],
+                k = faces[:,2]
+            ),)
+    
+    fig.update_layout(scene_aspectmode='manual', scene_aspectratio=dict(x=1, y=1, z=1),
+                    )
+    helper = dict(xaxis = dict(nticks=4, range=[-1,1],),
+                yaxis = dict(nticks=4, range=[-1,1],),
+                zaxis = dict(nticks=4, range=[-1,1],),)
+    fig.update_layout(scene = helper,
+                        )
+
+    fig.show()
+
+
+def build_lookup_table(repeated):
+    # iterate through the list and create a dict from each sub array in the list where each dict has one element of the sub array as key and the other elements as value as a list and repeat for each element
+
+    lookup = {}
+    for i in range(len(repeated)):
+        for j in range(len(repeated[i])):
+            if len(repeated[i]) == 1:
+                continue
+            lookup[repeated[i][j]] = [repeated[i][k] for k in range(len(repeated[i]))] # if k != j
+    return lookup
+
+
+from data.cloth3d_io import readPC2Frame
+import scipy.io as sio
+import scipy
+import scipy.spatial as spatial
+
+def local_read_posed_garment_frame(folder, id, garment, frame, pose=None, deorient=False):
+    pc16_path = os.path.join(folder, id, garment + '.pc16')
+    verts = readPC2Frame(pc16_path, frame, True)
+
+    map = sio.loadmat(os.path.join(folder, id, garment + '_map.mat'))['map'].reshape((-1))
+
+    if deorient:
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
+    return verts, map
+
+def local_read_posed_body_frame(folder, id, frame, pose=None, deorient=False):
+    pc16_path = os.path.join(folder, id, 'Body_posed.pc16')
+    verts = readPC2Frame(pc16_path, frame, True)
+
+    if deorient:
+        R = spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
+    return verts
+
+def edge_loss_custom(x, y, edges):
+    x_e = x[:,edges[:,0],:] - x[:,edges[:,1],:]
+    x_e = torch.sum(x_e ** 2, -1)
+    x_e = torch.sqrt(x_e + 1e-9)
+    y_e = y[:,edges[:,0],:] - y[:,edges[:,1],:]
+    y_e = torch.sum(y_e ** 2, -1)
+    y_e = torch.sqrt(y_e + 1e-9)
+    d_e = y_e - x_e
+    loss = torch.mean(d_e ** 2, dim=-1)[0]
+    return loss
+
+
+def crotch(mesh, body, crotch_idx, nr=75):
+    crotch_middle_on_body = crotch_idx
+    # create kdtree for body and mesh
+    body_tree = spatial.cKDTree(body)
+    mesh_tree = spatial.cKDTree(mesh)
+    # find closest point on mesh from crotch middle on body along the z axis
+
+    crotch_middle_on_mesh = find_closest_point_along_z(body[crotch_middle_on_body], mesh)
+
+    crotch_body = body_tree.query(body[crotch_middle_on_body], k=nr)[1]
+    crotch_mesh = mesh_tree.query(mesh[crotch_middle_on_mesh], k=nr)[1]
+
+    body_verts = body[crotch_body]
+    mesh_verts = mesh[crotch_mesh] - (mesh[crotch_middle_on_mesh] - body[crotch_middle_on_body])
+    mesh_verts = spatial.cKDTree(mesh_verts)
+    closest_mesh_verts = mesh_verts.query(body_verts, k=1)[1]
+    crotch_mesh = crotch_mesh[closest_mesh_verts]
+
+    return crotch_body, crotch_mesh
+
+
+def find_closest_point_along_z(point, verts, eps=0.01):
+    min_dist = np.inf
+    closest_point = None
+    for i, v in enumerate(verts):
+        dist = np.linalg.norm(point - v)
+        if dist < min_dist and np.abs(point[0] - v[0]) < eps and np.abs(point[1] - v[1]) < eps:
+            min_dist = dist
+            closest_point = i
+    return closest_point
+
+
+def mesh_edge_lengths(mesh: Meshes) -> torch.Tensor:
+    """
+    Given a PyTorch3D Meshes object, compute the lengths of all edges in the mesh.
+
+    Args:
+        mesh: A PyTorch3D Meshes object representing the mesh.
+
+    Returns:
+        A tensor of shape (E,) representing the lengths of all edges in the mesh,
+        where E is the total number of edges in the mesh.
+    """
+    # Extract the vertices and faces of the mesh.
+    vertices = mesh.verts_packed()
+    faces = mesh.faces_packed()
+
+    # Compute the pairwise distances between all pairs of vertices that form an edge.
+    edge_indices = faces.view(-1, 3)
+    edge_vertices = vertices[edge_indices]
+    edge_lengths = torch.norm(edge_vertices[:, 1:] - edge_vertices[:, :-1], dim=2)
+
+    return edge_lengths
+
+# try to optimize the already registered garment ??????????
+
+def register(folder, frame):
+    id, frame_nr, garment = frame.split()
+    frame_nr = int(frame_nr)
+    info = loadInfo(os.path.join(folder, id, 'info.mat'))
+    raw_mask = read_garment_mask(folder, id, garment, return_raw=True)
+    # convert it to indices
+    mask = raw_mask > 0
+    mesh_posed, map = local_read_posed_garment_frame(folder, id, garment, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
+    body_posed = local_read_posed_body_frame(folder, id, frame_nr, pose=info['poses'][:,frame_nr], deorient=True)
+    template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
+    faces = quads2tris(faces)
+    
+    boundaries = raw_mask == 2
+    registered_boundaries = torch.tensor(np.where(boundaries[mask])[0]).to(device)
+    # invert map and use it to invert boudaries
+    # map_inv = np.zeros(map.max()+1, dtype=np.int32)
+    # map_inv[map] = np.arange(len(map))
+    # boundaries = boundaries[map_inv]
+    # # get the boundary indices
+    # mesh1_boundaries = torch.tensor(np.where(boundaries)[0]).to(device)
+
+
+    
+    # # pose = torch.zeros(1, 72).to('cuda')
+    # shape = torch.zeros(1, 11).to(device)
+    # shape[:, :10] = torch.tensor(info['shape']).to(device)
+    # shape[:, -1] = torch.tensor(info['gender']).to(device)
+    # # for star pose
+    # # pose[:, 5] = 3.1415/4
+    # # pose[:,8] = -3.1415/4
+    # STAR_POSE = np.zeros((24, 3), np.float32)
+    # STAR_POSE[1,2] = .15
+    # STAR_POSE[2,2] = -.15
+    # pose = torch.tensor(STAR_POSE.flatten()[None,:]).to(device)
+    # star_body = smpl(pose, shape, remove=True)[0][0].cpu().numpy() - smpl(pose, shape, remove=True)[4][0].cpu().numpy()
+    # # rotate star body with 90 degrees around x axis
+    # R = spatial.transform.Rotation.from_euler('x', np.pi/2)
+    # star_body = R.apply(star_body).astype(np.float32)
+
+    is_skirt = (garment == 'Dress' or garment == 'Skirt')
+    skirt_mod = '_skirt' if is_skirt else ''
+    F = np.load(f'smpl/smpl_faces{skirt_mod}.npy').astype(np.int32)
+
+    ##################### mask body #####################
+    new_crotch = 4729
+
+    # star_body = star_body[mask]
+    new_crotch = np.zeros((14475), bool)
+    new_crotch[4729] = True
+    new_crotch = new_crotch[mask]
+    new_crotch = np.where(new_crotch)[0][0]
+    _, F = _maskFaces(np.zeros((14475, 3)), F, mask)
+    registered = template[map][mask]
+
+    # plot(template, faces)
+    # v = mesh_posed[mask,:]
+    
+    ##################### boundaries #####################
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    template_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    # meshfix = MeshFix(template, faces)
+    # holes = meshfix.extract_holes()
+    # new_border = []
+    # for border in np.array(holes.points):
+    #     # find the index of the vertex in the original mesh
+    #     ind = np.argwhere(np.all(template == border, axis=1))
+    #     new_border.extend(ind.ravel())
+    # mesh1_boundaries = torch.tensor(new_border).to(device)
+    # np.save('temp/body_boundaries.npy', new_border)
+    # mesh1_boundaries = torch.tensor(np.load('temp/mesh1_boundaries.npy')).to(device)
+
+    ##################### crotch #####################
+    # crotch_body, crotch_mesh = crotch(template, star_body, crotch_idx=new_crotch, nr=50)
+
+
+
+
+    # export star_body as obj
+    # writeOBJ('temp/star_body.obj', star_body, F)
+
+
+    template = torch.tensor(template).to(device)
+    template_orig = template.clone()
+    faces = torch.tensor(faces).to(device)
+    registered = torch.tensor(registered).to(device)
+    F = torch.tensor(F).to(device)
+
+    # verts = template.cpu().numpy()
+    # star_body = star_body.cpu().numpy()
+    # faces = faces.detach().cpu().numpy()
+    # F = F.detach().cpu().numpy()
+    # fig = go.Figure(data=[
+    #     go.Mesh3d(
+    #         x=verts[:,0], y=verts[:,1], z=verts[:,2],
+    #         i=faces[:,0], j=faces[:,1], k=faces[:,2],
+    #         color='lightpink', opacity=0.50),
+    #     go.Mesh3d(
+    #         x=star_body[:,0], y=star_body[:,1], z=star_body[:,2],
+    #         i=F[:,0], j=F[:,1], k=F[:,2],
+    #         color='lightblue', opacity=0.50),
+    #     go.Scatter3d(
+    #         x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+    #         mode='markers', marker=dict(size=2, color='red')),
+    #     go.Scatter3d(
+    #         x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+    #         mode='markers', marker=dict(size=2, color='blue')),
+    # ])
+    # # make it cuve and limit to -1 and -1
+    # fig.update_layout(scene = dict(
+    #     xaxis = dict(range=[-1,1],),
+    #     yaxis = dict(range=[-1,1],),
+    #     zaxis = dict(range=[-1,1],),),)
+    # fig.show()
+
+    # Define the meshes to optimize
+    mesh1 = Meshes(verts=registered[None,:,:], faces=F[None,:,:])
+    mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:])
+    mesh1_boundaries = registered_boundaries
+    mesh2_boundaries = template_boundaries
+
+    # mesh2_boundaries = find_closest_vertices(mesh1, mesh2, mesh1_boundaries)
+
+    avg_edge_length = mesh_edge_lengths(mesh1).mean()
+
+    # Define the optimizer and learning rate
+    mesh1.verts_packed().requires_grad = True
+    optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.0001)
+    # 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(10000):
+        optimizer.zero_grad()
+
+        # find boundaries with the find holes function?
+
+        # Forward pass
+        mesh1_out = Meshes(verts=[mesh1.verts_packed()], faces=[mesh1.faces_packed()])
+        ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list()), torch.stack(mesh2.verts_list()))
+        bound_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,mesh1_boundaries,:], torch.stack(mesh2.verts_list())[:,mesh2_boundaries,:])
+        # crotch_ch_loss, _ = chamfer_distance(torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:], torch.stack(mesh2.verts_list())[:,crotch_body,:],)
+        # crotch_ch_loss = torch.mean((torch.stack(mesh1_out.verts_list())[:,crotch_mesh,:] - torch.stack(mesh2.verts_list())[:,crotch_body,:])**2, dim=2).mean()
+        loss: torch.Tensor = ch_loss * 3 +  (bound_ch_loss) * 20
+        # loss: torch.Tensor = ch_loss + ch_norm_loss + (bound_ch_loss + bound_ch_norm_loss) * 1
+        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 * 100 + norm_loss * 0.1
+
+        # Backward pass
+        loss.backward()
+        optimizer.step()
+
+        # Print the loss
+        if i % 10 == 0:
+            print(f"it: {i} : loss: {loss:.5f} chamfer: {ch_loss:.6f} bound_chamfer: {bound_ch_loss:.6f} laplacian: {lap_loss:.6f} edge: {e_loss:.6f} normal: {norm_loss:.5f}")
+    
+    # create a plotly plot with the body and the mesh1_out
+    registered = mesh1_out.verts_packed().detach().cpu().numpy()
+    template = mesh2.verts_packed().detach().cpu().numpy()
+    faces = faces.detach().cpu().numpy()
+    F = F.detach().cpu().numpy()
+    # faces, F = F, faces
+    # verts, star_body = star_body, verts
+    mesh1_boundaries, mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy(), mesh2_boundaries.detach().cpu().numpy()
+    # mesh2_boundaries = mesh1_boundaries.detach().cpu().numpy()
+    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=F[:,0], j=F[:,1], k=F[:,2],
+            color='lightpink', opacity=0.50),
+        go.Scatter3d(
+            x=registered[mesh1_boundaries,0], y=registered[mesh1_boundaries,1], z=registered[mesh1_boundaries,2],
+            mode='markers', marker=dict(size=2, color='red')),
+        go.Scatter3d(
+            x=template[mesh2_boundaries,0], y=template[mesh2_boundaries,1], z=template[mesh2_boundaries,2],
+            mode='markers', marker=dict(size=2, color='blue')),   
+        # go.Scatter3d(
+        #     x=verts[crotch_mesh,0], y=verts[crotch_mesh,1], z=verts[crotch_mesh,2],
+        #     mode='markers', marker=dict(size=2, color='red')),
+        # go.Scatter3d(
+        #     x=star_body[crotch_body,0], y=star_body[crotch_body,1], z=star_body[crotch_body,2],
+        #     mode='markers', marker=dict(size=2, color='blue')),
+    ])
+    # make it cuve and limit to -1 and -1
+    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()
+
+    # create a plotly wireframe plot of registered
+    fig = go.Figure()
+    for edge in mesh1_out.edges_packed().detach().cpu().numpy():
+        fig.add_trace(go.Scatter3d(
+            x=registered[edge,0], y=registered[edge,1], z=registered[edge,2],
+            mode='lines', line=dict(color='red', width=1)))
+    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'
+    # 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')
+    root = '/'.join(folder.split('/')[:-1])
+    type = folder.split('/')[-1]
+
+    use_no_orient = False
+    no_shape = False
+    is_skirt = False
+
+    skirt_modifier = '_skirt' if is_skirt else ''
+    orient_mod = '_no_orient' if use_no_orient else ''
+    shape_mod = '_no_shape' if no_shape else ''
+
+    # mask = np.load('/data-net/dataset/00761/mesh_mask/00761.npy')
+    # faces = np.load('/data-net/dataset/00761/faces/00761.npy')
+
+    for frame in frames:
+        frame = random.choice(frames)
+        # frame = '00761 0 Tshirt'
+        frame = '01893 256 Jumpsuit'
+        
+
+        register(folder, frame)
+        break
+        pass

data/registration_weights.py

@@ -30,7 +30,7 @@ from data.cloth3d_io import (
 from smpl.smpl_torch import SMPL
 # from custom_utils.custom_utils import _maskFaces
 
-device = 'cuda'
+device = 'cpu'
 smpl = SMPL(device, 'smpl')
 
 
@@ -315,7 +315,7 @@ def register(folder, frame):
         return mesh_normal_consistency(mesh)
 
     # Optimization loop
-    for i in range(4000):
+    for i in range(1):
         optimizer.zero_grad()
 
         # find boundaries with the find holes function?
@@ -430,20 +430,139 @@ def register(folder, frame):
     pass
 
 
+def order(verts, edges):
+    # Choose a starting vertex
+    start_vertex = 0
+
+    # Initialize the visited and ordered arrays
+    visited = np.zeros(verts.shape[0], dtype=bool)
+    ordered_vertices = []
+    ordered_edges = []
+
+    # Define the depth-first search function
+    def dfs(vertex):
+        visited[vertex] = True
+        ordered_vertices.append(vertex)
+        for i in range(edges.shape[0]):
+            if edges[i, 0] == vertex and not visited[edges[i, 1]]:
+                ordered_edges.append(i)
+                dfs(edges[i, 1])
+            elif edges[i, 1] == vertex and not visited[edges[i, 0]]:
+                ordered_edges.append(i)
+                dfs(edges[i, 0])
+
+    # Perform depth-first search starting from the chosen vertex
+    dfs(start_vertex)
+
+    return ordered_vertices, ordered_edges
+
+def reindex_edges(edges, mask):
+    lines_new = np.zeros_like(edges)
+    for i in range(edges.shape[0]):
+        line_1 = edges[i,0]
+        line_2 = edges[i,1]
+
+        line_1_new = line_1 - np.count_nonzero(~mask[:line_1])
+        line_2_new = line_2 - np.count_nonzero(~mask[:line_2])
+
+        lines_new[i] = [line_1_new, line_2_new]
+
+    lines_new = np.asarray(lines_new)
+    return lines_new
+
+def reindex_vertices(verts, mask):
+    verts_new = np.zeros_like(verts)
+    for i in range(verts.shape[0]):
+        vert = verts[i]
+        vert_new = vert - np.count_nonzero(~mask[:vert])
+
+        verts_new[i] = vert_new
+
+    verts_new = np.asarray(verts_new)
+    return verts_new
+
+
+def generate_edges(registered_cloth, registered_holes, registered_border):
+    lines = registered_holes.lines
+    lines_filtered = []
+    for i in range(0, lines.shape[0]-2, 3):
+        ind = np.argwhere(np.all(registered_cloth == registered_holes.points[lines[i]], axis=1))
+        ind2 = np.argwhere(np.all(registered_cloth == registered_holes.points[lines[i+1]], axis=1))
+        if ind.ravel() in registered_border and ind2.ravel() in registered_border:
+            lines_filtered.append(lines[i+1:i+3])
+    lines = np.array(lines_filtered)
+
+    return lines
+
+
+def reindex_lines(lines, registered_border_mask):
+    lines_new = np.zeros_like(lines)
+    for i in range(lines.shape[0]):
+        line_1 = lines[i,0]
+        line_2 = lines[i,1]
+
+        line_1_new = line_1 - np.count_nonzero(~registered_border_mask[:line_1])
+        line_2_new = line_2 - np.count_nonzero(~registered_border_mask[:line_2])
+
+        lines_new[i] = [line_1_new, line_2_new]
+
+    lines_new = np.asarray(lines_new)
+    return lines_new
+
+
+def get_boundary_components(registered_border, lines_new):
+    from scipy.sparse import csr_matrix
+    from scipy.sparse.csgraph import connected_components
+    n = len(registered_border)
+    adj_matrix = csr_matrix((np.ones(lines_new.shape[0]), (lines_new[:, 0], lines_new[:, 1])), shape=(n, n))
+
+    # Find the connected components of the graph
+    n_components, labels = connected_components(adj_matrix)
+
+    registered_border = np.array(registered_border)
+    bounds = []
+    bounds_edges_global = []
+    bounds_edges_local = []
+    for i in range(n_components):
+        ind = np.argwhere(labels == i)[:,0]
+        if len(ind) < 10:
+            continue
+
+        bounds.append(registered_border[ind])
+
+        temp_global = []
+        temp_local = []
+        for j in range(lines_new.shape[0]):
+            if lines_new[j, 0] in ind and lines_new[j, 1] in ind:
+                # if lines_new[j, 0] < lines_new[j, 1]:
+                temp_global.append(registered_border[lines_new[j]])
+                temp_local.append(lines_new[j])
+                # else:
+                #     temp.append(registered_border[np.asarray([lines_new[j, 1], lines_new[j, 0]])])
+        bounds_edges_global.append(np.asarray(temp_global))
+        bounds_edges_local.append(np.asarray(temp_local))
+
+    return bounds, bounds_edges_global, bounds_edges_local, labels
+
+
 def get_boundaries(folder, frame):
     from pymeshfix import MeshFix
     from scipy.spatial import KDTree
 
+
     id, frame_nr, garment = frame.split()
     frame_nr = int(frame_nr)
-    registered_cloth = np.load('temp/registered_cloth.npy')
+    registered_cloth = np.load('temp/matlab_match_registered.npy')
-    registered_faces = np.load('temp/registered_faces.npy')
+    registered_faces = np.load('temp/matlab_match_faces.npy')
-    registred_weights = np.load('temp/registered_weights.npy')
+    registered_weight_faces = np.load('temp/matlab_match_weights_faces.npy')
+    registred_weights = np.load('temp/matlab_match_weights_weights.npy')
     _, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
     faces = quads2tris(faces)
     template = np.load('temp/mapping_cloth.npy')
 
-    registered_meshfix = MeshFix(registered_cloth, registered_faces)
+    registered_cloth = (template[registered_weight_faces] * registred_weights[:,:,None]).sum(1)
+
+    registered_meshfix = MeshFix(registered_cloth, registered_faces.astype(np.int32))
     registered_holes = registered_meshfix.extract_holes()
 
     template_meshfix = MeshFix(template, faces)
@@ -457,17 +576,57 @@ def get_boundaries(folder, frame):
     template_border_tree = KDTree(template[template_border])
 
     registered_border = []
-    for border in np.array(registered_holes.points):
+    registered_border_mask = np.zeros(registered_holes.points.shape[0], dtype=bool)
+    for i, border in enumerate(np.array(registered_holes.points)):
         q = template_border_tree.query(border)
         if q[0] > 0.02:
             continue
         # find the index of the vertex in the original mesh
         ind = np.argwhere(np.all(registered_cloth == border, axis=1))
         registered_border.extend(ind.ravel())
+        registered_border_mask[i] = True
 
     np.save('temp/registered_border.npy', registered_border)
     np.save('temp/template_border.npy', template_border)
     
+    lines = generate_edges(registered_cloth, registered_holes, registered_border)
+    lines_new = reindex_lines(lines, registered_border_mask)
+    bounds, bounds_edges_global, bounds_edges_local, labels = get_boundary_components(registered_border, lines_new)
+
+    for k in range(labels.max()+1):
+        bounds_edges_reindexed = reindex_edges(bounds_edges_local[k].copy(), labels == k)
+        verts_order, edges_order = order(bounds[k], bounds_edges_reindexed)
+        traces = []
+        for i in range(bounds_edges_global[k].shape[0]):
+            traces.append(go.Scatter3d(
+                x=registered_cloth[bounds_edges_global[k][i],0], y=registered_cloth[bounds_edges_global[k][i],1], z=registered_cloth[bounds_edges_global[k][i],2],
+                mode='lines', line=dict(color='red', width=1)))
+        for i in range(15):
+            traces.append(go.Scatter3d(
+                x=registered_cloth[bounds_edges_global[k][edges_order][i],0], y=registered_cloth[bounds_edges_global[k][edges_order][i],1], z=registered_cloth[bounds_edges_global[k][edges_order][i],2],
+                mode='lines', line=dict(color='blue', width=1)))
+        fig = go.Figure(data=traces)
+        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()
+
+
+
+    k = 0
+    traces = [go.Scatter3d(
+            x=registered_cloth[bounds[k],0], y=registered_cloth[bounds[k],1], z=registered_cloth[bounds[k],2],
+            mode='markers', line=dict(color='red', width=1))]
+    fig = go.Figure(data=traces)
+    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()
+
     # plot registered cloth and tempalte as scatterplots and add another two scatter plots for the two boundary sets
     fig = go.Figure(data=[
         go.Scatter3d(
@@ -501,6 +660,7 @@ def get_boundaries(folder, frame):
 
     pass
 
+
 def optimize_boundaries(folder, frame):
     id, frame_nr, garment = frame.split()
     frame_nr = int(frame_nr)
@@ -519,7 +679,6 @@ def optimize_boundaries(folder, frame):
     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)
@@ -527,7 +686,6 @@ def optimize_boundaries(folder, frame):
     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()
 
@@ -553,7 +711,7 @@ def optimize_boundaries(folder, frame):
         return mesh_normal_consistency(mesh)
 
     # Optimization loop
-    for i in range(4000):
+    for i in range(1):
         optimizer.zero_grad()
 
         # find boundaries with the find holes function?
@@ -576,7 +734,7 @@ def optimize_boundaries(folder, frame):
         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
+        loss += lap_loss * 0.5 + e_loss * 2 + norm_loss * 0.1
 
         # Backward pass
         loss.backward()
@@ -584,8 +742,8 @@ def optimize_boundaries(folder, frame):
 
         # 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}")
-            # 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}")
+            print(f"it: {i} : loss: {loss:.5f} bound_chamfer: {ch_bound_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()
@@ -613,6 +771,163 @@ def optimize_boundaries(folder, frame):
     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
+
+
+def load_selected_saves():
+    registered = np.load('temp/matlab_match_registered.npy')
+    garment = np.load('temp/matlab_match_garment.npy')
+    mask = np.load('temp/matlab_match_mask.npy')
+    mask_bool = np.load('temp/matlab_match_mask.npy')
+    weights_faces = np.load('temp/matlab_match_weights_faces.npy')
+    weights_weights = np.load('temp/matlab_match_weights_weights.npy')
+    faces = np.load('temp/matlab_match_faces.npy')
+    boundaries = np.load('temp/matlab_match_boundaries.npy')
+
+    return registered, garment, mask, mask_bool, weights_faces, weights_weights, faces, boundaries
+
+def optimize_boundaries_2(folder, frame):
+    registered, garment, mask, mask_bool, weights_faces, weights_weights, faces, boundaries = load_selected_saves()
+    id, frame_nr, garment_name = frame.split()
+
+    registered_cloth = registered
+    registered_faces = faces
+    registered_weights = weights_weights
+
+    _, template_faces, _, _ = readOBJ(os.path.join(folder, id, garment_name + '.obj'))
+    template_faces = quads2tris(template_faces)
+
+    template = np.load('temp/mapping_cloth.npy')
+    mapping_mask = mask
+    mapping_faces = weights_faces
+
+
+    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)
+    template_faces = torch.tensor(template_faces).to(device)
+    registered_faces = torch.tensor(registered_faces).to(device)
+    mapping_faces = torch.tensor(mapping_faces.astype(np.int64)).to(device)
+    registered_weights = torch.tensor(registered_weights).to(device).to(torch.float32)
+    registered_cloth = torch.tensor(registered_cloth).to(device).to(torch.float32)
+    registered_border = torch.tensor(registered_border).to(device).to(torch.int64)
+
+    mesh2 = Meshes(verts=template[None,:,:], faces=template_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(1000):
+        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.5 + 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} 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()
+    template_faces = template_faces.detach().cpu().numpy()
+    optim_mask = optim_mask.detach().cpu().numpy()
+    registered_faces = registered_faces.detach().cpu().numpy()
+    registered_border = registered_border.detach().cpu().numpy()
+    # faces, F = F, faces
+    # verts, star_body = star_body, verts
+
+    from scipy.spatial import KDTree
+    template_border_tree = KDTree(template[template_border,:])
+    for i, border in enumerate(registered_border):
+        dist, idx = template_border_tree.query(registered[border,:], k=2)
+        weights_faces[border,:] = np.concatenate((template_border[idx], [0]))
+        weights_weights[border,:] = np.concatenate((1 - dist / np.sum(dist), [0]))
+    registered = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            x=template[:,0], y=template[:,1], z=template[:,2],
+            i=template_faces[:,0], j=template_faces[:,1], k=template_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),   
+        go.Scatter3d(
+            x=registered[registered_border,0], y=registered[registered_border,1], z=registered[registered_border,2],
+            mode='markers', marker=dict(size=4, color='red')),
+        go.Scatter3d(
+            x=template[template_border,0], y=template[template_border,1], z=template[template_border,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()
 
     fig = go.Figure(data=[
         go.Scatter3d(
@@ -631,9 +946,98 @@ def optimize_boundaries(folder, frame):
 
     pass
 
+
+def flatten_boundary(folder, frame):
+    from scipy.spatial import KDTree
+    registered, garment, mask, mask_bool, weights_faces, weights_weights, faces, boundaries = load_selected_saves()
+    id, frame_nr, garment_name = frame.split()
+    registered_faces = faces
+    registered_weights = weights_weights
+
+    _, template_faces, _, _ = readOBJ(os.path.join(folder, id, garment_name + '.obj'))
+    template_faces = quads2tris(template_faces)
+
+    template = np.load('temp/mapping_cloth.npy')
+    registered_cloth = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+
+
+    template_border = np.load('temp/template_border.npy')
+    registered_border = np.load('temp/registered_border.npy')
+
+    template_border_tree = KDTree(template[template_border,:])
+
+    for i, border in enumerate(registered_border):
+        dist, idx = template_border_tree.query(registered_cloth[border,:], k=2)
+        weights_faces[border,:] = np.concatenate((template_border[idx], [0]))
+        weights_weights[border,:] = np.concatenate((1 - dist / np.sum(dist), [0]))
+
+    registered_cloth = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+
+
+    fig = go.Figure(data=[
+        go.Mesh3d(
+            x=template[:,0], y=template[:,1], z=template[:,2],
+            i=template_faces[:,0], j=template_faces[:,1], k=template_faces[:,2],
+            color='lightblue', opacity=0.80),
+        go.Mesh3d(
+            x=registered_cloth[:,0], y=registered_cloth[:,1], z=registered_cloth[:,2],
+            i=registered_faces[:,0], j=registered_faces[:,1], k=registered_faces[:,2],
+            color='lightpink', opacity=0.50),   
+        go.Scatter3d(
+            x=registered_cloth[registered_border,0], y=registered_cloth[registered_border,1], z=registered_cloth[registered_border,2],
+            mode='markers', marker=dict(size=4, color='red')),
+        go.Scatter3d(
+            x=template[template_border,0], y=template[template_border,1], z=template[template_border,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
+
+
+def fill_holes(folder, frame):
+    from scipy.spatial import KDTree
+    from pymeshfix import MeshFix
+    import pyvista as pv
+    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_border_tree = KDTree(template[template_border])
+    registered_border_tree = KDTree(registered_cloth[registered_border])
+
+    registered_meshfix = MeshFix(registered_cloth, registered_faces)
+    registered_holes = registered_meshfix.extract_holes()
+
+    template_meshfix = MeshFix(template, faces)
+    template_holes = template_meshfix.extract_holes()
+
+    for border in np.array(registered_holes.points):
+        q = registered_border_tree.query(border)
+
+
 if __name__ == "__main__":
-    folder = '/data-net/code/cloth3d/train'
+    # folder = '/data-net/code/cloth3d/train'
-    # folder = r'D:\Work\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')
@@ -659,7 +1063,9 @@ if __name__ == "__main__":
         
 
         # register(folder, frame)
-        # get_boundaries(folder, frame)
+        get_boundaries(folder, frame)
-        optimize_boundaries(folder, frame)
+        # optimize_boundaries_2(folder, frame)
+        # flatten_boundary(folder, frame)
+        # fill_holes(folder, frame)
         break
         pass