add fix_boundaries

data/fix_boundaries.py

@@ -0,0 +1,642 @@
+import os
+import sys
+import random
+import numpy as np
+import plotly.graph_objects as go
+from pymeshfix import MeshFix
+import scipy.spatial as spatial
+from glob import glob
+import time
+from tqdm import tqdm
+
+# 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
+)
+
+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 = []
+    sub_folders = sorted(os.listdir(folder))
+    for sub_folder in sub_folders:
+        if not os.path.isdir(os.path.join(folder, sub_folder)):
+            continue
+        sample_folder = os.path.join(folder, sub_folder)
+        files = os.listdir(sample_folder)
+        garments = [file.split('.')[0] for file in files if file.endswith('.obj')]
+        for i in range(len(garments)):
+            frames.append([sub_folder, garments[i]])
+
+    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
+from pymeshfix import MeshFix
+from scipy.spatial import KDTree
+
+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 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 get_end(verts, edges):
+    vertex_counts = np.zeros(verts.shape[0], dtype=int)
+    for i in range(edges.shape[0]):
+        vertex_counts[edges[i, 0]] += 1
+        vertex_counts[edges[i, 1]] += 1
+
+    # Find the index of the first vertex that appears only once in the edges array
+    ends = np.where(vertex_counts == 1)
+    if ends[0].shape[0] == 0:
+        end_vertex_index = 0
+    else:
+        end_vertex_index = ends[0][0]
+
+    return end_vertex_index
+
+def order(verts, edges):
+    # Choose a starting vertex
+    start_vertex = get_end(verts, edges)
+
+    # 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])
+        if len(ordered_edges) == edges.shape[0] - 1:
+            # check if exist and edge with edges[ordered_edges[-1], 1] and edges[ordered_edges[0], 0] and if so add its index
+            for i in range(edges.shape[0]):
+                if edges[i, 0] == edges[ordered_edges[-1], 1] and edges[i, 1] == edges[ordered_edges[0], 0]:
+                    ordered_edges.append(i)
+
+    # Perform depth-first search starting from the chosen vertex
+    dfs(start_vertex)
+
+    assert len(ordered_vertices) == verts.shape[0], 'Ordering failed'
+    if edges.max() > verts.shape[0]:
+        assert len(ordered_edges) == edges.shape[0], 'Ordering failed'
+
+    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(1, 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()[0] in registered_border and ind2.ravel()[0] in registered_border:
+            lines_filtered.append(lines[i:i+2])
+    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, directed=False)
+
+    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
+
+
+_F_skirt = np.load('smpl\smpl_faces_skirt.npy')
+_F = np.load('smpl\smpl_faces.npy')
+def get_F(garment):
+    if garment == 'Dress' or garment == 'Skirt':
+        return _F_skirt
+    else:
+        return _F
+
+
+def get_boundaries(folder, folder_reg, frame):
+    id, garment = frame
+    registration_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights.mat'))
+    registration_mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'] - 1
+    registration_mask = np.zeros(14475, dtype=bool); registration_mask[registration_mask_idx] = True
+    registered_weight_faces = registration_data['F']
+    registered_weights = registration_data['W']
+    _, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj'))
+    faces = quads2tris(faces)
+    _, registered_faces = _maskFaces(np.zeros((14475, 3)), get_F(garment), registration_mask)
+    template = sio.loadmat(os.path.join(folder_reg, id, garment + '.mat'))['V']
+
+    registered_cloth = (template[registered_weight_faces] * registered_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)
+    template_holes = template_meshfix.extract_holes()
+    template_border = []
+    for border in np.array(template_holes.points):
+        # find the index of the vertex in the original mesh
+        ind = np.argwhere(np.all(template == border, axis=1))
+        template_border.extend(ind.ravel())
+
+    # template_border_tree = KDTree(template[template_border])
+
+    registered_border = []
+    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())
+
+
+    
+    lines = generate_edges(registered_cloth, registered_holes, registered_border)
+    lines_new = lines
+    bounds, bounds_edges_global, bounds_edges_local, labels = get_boundary_components(registered_border, lines_new)
+    
+    
+    # assert labels.max() <= 6, 'More than 6 boundaries found'
+
+    # 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()
+
+    # crotch = 2256
+    # crotch_idx = np.zeros((14475))
+    # crotch_idx[crotch] = 1
+    # crotch_idx = crotch_idx[registration_mask]
+    # crotch_idx = np.argwhere(crotch_idx)
+    # if crotch_idx.shape[0] == 0:
+    #     crotch_idx = crotch_idx.ravel()[0]
+    # else:
+    #     crotch_idx = None
+
+
+    bounds_filtered = []
+    bounds_edges_global_filtered = []
+    for k in range(labels.max()+1):
+        if bounds[k].shape[0] < 15:
+            continue
+
+        bounds_edges_reindexed = reindex_edges(bounds_edges_local[k].copy(), labels == k)
+        verts_order, edges_order = order(bounds[k], bounds_edges_reindexed)
+        # bounds_edges_global[k] = bounds_edges_global[k][edges_order]
+        bounds_filtered.append(bounds[k][verts_order])
+        bounds_edges_global_filtered.append(bounds_edges_global[k][edges_order])
+        # bounds[k] = bounds[k][verts_order]
+        # bounds_edges_global[k] = bounds_edges_global[k][edges_order]
+
+        # fixed_bounds = []
+        # for i in range(bounds_edges_global[k].shape[0] - 1):
+        #     dist_1 = np.linalg.norm(registered_cloth[bounds_edges_global[k][i,0]] - registered_cloth[bounds_edges_global[k][i,1]])
+        #     dist_2 = np.linalg.norm(registered_cloth[bounds_edges_global[k][i+1,0]] - registered_cloth[bounds_edges_global[k][i+1,1]])
+        #     if dist_1 > 0.1 or dist_2 > 0.1:
+        #         fixed_bounds.append(bounds_edges_global[k][i,0])
+        #         # fixed_bounds.append(bounds_edges_global[k][i+1,1])
+        #     else:
+        #         fixed_bounds.append(bounds_edges_global[k][i,0])
+        #         fixed_bounds.append(bounds_edges_global[k][i,1])
+        # bounds[k] = fixed_bounds
+
+        # if crotch_idx is not None:
+        #     if np.linalg.norm(registered_cloth[bounds[k]].mean() - registered_cloth[crotch_idx]) < 0.04:
+        #         continue
+        # traces = []
+        # traces.append(go.Scatter3d(
+        #     x=registered_cloth[:,0], y=registered_cloth[:,1], z=registered_cloth[:,2],
+        #     mode='markers', marker=dict(color='red', size=1)))
+        # traces.append(go.Scatter3d(
+        #     x=registered_cloth[bounds[k],0], y=registered_cloth[bounds[k],1], z=registered_cloth[bounds[k],2],
+        #     mode='markers', marker=dict(color='blue', size=2)))
+        # traces.append(go.Scatter3d(
+        #     x=registered_cloth[crotch_idx,0], y=registered_cloth[crotch_idx,1], z=registered_cloth[crotch_idx,2],
+        #     mode='markers', marker=dict(color='green', size=5)))
+        # 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()
+
+        # pass
+
+        # traces = []
+        # for i in range(bounds_edges_global[k].shape[0]):
+        #     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='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()
+
+        # pass
+
+    # np.savez('temp/registered_border.npz', {f'bound_{i}': bounds[i] for i in range(len(bounds))})
+    # np.savez('temp/registered_border_edges.npz', {f'bound_edges_{i}': bounds_edges_global[i] for i in range(len(bounds_edges_global))})
+    # np.save('temp/template_border.npy', template_border)
+
+
+    # plot registered cloth and tempalte as scatterplots and add another two scatter plots for the two boundary sets
+    # fig = go.Figure(data=[
+    #     go.Scatter3d(
+    #         x=registered_cloth[:,0], y=registered_cloth[:,1], z=registered_cloth[:,2],
+    #         mode='markers', marker=dict(size=2, color='red')),
+    #     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='green')),
+    # ])
+    # 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=template[:,0], y=template[:,1], z=template[:,2],
+    #         mode='markers', marker=dict(size=2, color='blue')),
+    #     go.Scatter3d(
+    #         x=template[template_border,0], y=template[template_border,1], z=template[template_border,2],
+    #         mode='markers', marker=dict(size=4, color='green')),
+    # ])
+    # 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()
+
+    return bounds_filtered, bounds_edges_global_filtered, np.asarray(template_border), template, faces, registered_weight_faces, registered_weights, registered_faces
+
+
+def flatten_boundary(folder, frame, registered_border, registered_border_edges, template_border, template, template_faces, registered_weight_faces, registered_weights, registered_faces):
+    weights_faces, weights_weights =  registered_weight_faces, registered_weights
+
+    registered_cloth = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+
+    registered_borders = registered_border
+    registered_border = np.concatenate(registered_border)
+
+    template_border_tree = KDTree(template[template_border,:])
+    # registered_border_tree = KDTree(registered_cloth[registered_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()
+
+    # fig = go.Figure(data=[
+    #     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=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()
+
+    # save registered_cloth as obj
+    # writeOBJ('temp/test_registered.obj', registered_cloth, registered_faces.astype(np.int32))
+
+ 
+
+    for k in range(len(registered_borders)):
+        border = registered_borders[k]
+        for i in range(border.shape[0]):
+            v_1 = registered_cloth[border[i],:]
+            v_2 = registered_cloth[border[(i+1)%border.shape[0]],:]
+            v_3 = registered_cloth[border[(i+2)%border.shape[0]],:]
+
+            dist_1 = np.linalg.norm(v_1 - v_2, axis=0)
+            dist_2 = np.linalg.norm(v_1 - v_3, axis=0)
+
+            if dist_1 > dist_2:
+                weights_weights[border[[(i+1)%border.shape[0],(i+2)%border.shape[0]]],:] = weights_weights[border[[(i+2)%border.shape[0],(i+1)%border.shape[0]]],:]
+                registered_cloth = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+
+    return weights_weights
+    # registered_cloth = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+
+    # for k in range(len(registered_borders)):
+    #     border = registered_borders[k]
+    #     for i in range(border.shape[0]):
+    #         v_1 = registered_cloth[border[i],:]
+    #         v_2 = registered_cloth[border[(i+1)%border.shape[0]],:]
+    #         v_3 = registered_cloth[border[(i+2)%border.shape[0]],:]
+    #         v_4 = registered_cloth[border[(i+3)%border.shape[0]],:]
+
+    #         dist_1 = np.linalg.norm(v_3 - v_2, axis=0)
+    #         dist_2 = np.linalg.norm(v_3 - v_1, axis=0)
+    #         dist_3 = np.linalg.norm(v_1 - v_4, axis=0)
+
+    #         if dist_1 > dist_2:
+    #             weights_weights[border[[(i+1)%border.shape[0],i]],:] = weights_weights[border[[i,(i+1)%border.shape[0]]],:]
+           
+
+    # registered_cloth = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+    # border_starters = [registered_borders[k][0] for k in range(len(registered_borders))]
+    # writeOBJ('temp/test_registered.obj', registered_cloth, registered_faces.astype(np.int32))
+    # traces = [
+    #     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=1),
+    # ]
+    # # for i in range(len(registered_borders[-1])):
+    # #     traces.append(go.Scatter3d(
+    # #         x=registered_cloth[registered_borders[-1][i]:registered_borders[-1][i]+1,0], y=registered_cloth[registered_borders[-1][i]:registered_borders[-1][i]+1,1], z=registered_cloth[registered_borders[-1][i]:registered_borders[-1][i]+1,2],
+    # #         mode='markers', marker=dict(size=3, color='red'), text=str(i)))
+    # 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()
+
+    pass
+
+
+def fill_holes(verts, faces):
+    import pyvista as pv
+    
+    mesh = pv.PolyData(verts, faces.astype(np.int32))
+
+    # Fill the holes in the mesh
+    filled_mesh = mesh.fill_holes(hole_size=100, progress_bar=True)
+
+    # Get the vertices and faces of the filled mesh
+    filled_vertices = filled_mesh.points
+    filled_faces = filled_mesh.faces.reshape(-1, 4)[:, 1:]
+
+
+    traces = [
+        go.Mesh3d(
+            x=filled_vertices[:,0], y=filled_vertices[:,1], z=filled_vertices[:,2],
+            i=filled_faces[:,0], j=filled_faces[:,1], k=filled_faces[:,2],
+            color='lightpink', opacity=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()
+
+    return filled_vertices, filled_faces
+
+if __name__ == "__main__":
+    # folder = '/data-net/code/cloth3d/train'
+    folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
+    registration_folder = r'D:\Work\Dataset\cloth3d_registration\train'
+    # folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train'
+    frames = get_frames(folder)
+    # frames = [['05740', 'Jumpsuit']] ['03102', 'Jumpsuit']
+    # frames = [['03102', 'Jumpsuit']] 
+    # frames = [['05060', 'Jumpsuit']] 
+
+    root = folder[:folder.rfind(os.path.basename(folder))]
+    type = os.path.basename(folder)
+    # random.shuffle(frames)
+
+    for frame in tqdm(frames):
+        # frame = ['03555', 'Dress']
+        # frame = ['03555', 'Dress']
+        # frame = ['06196', 'Skirt']
+        # frame = ['05060', 'Jumpsuit']
+        frame = ['01467', 'Jumpsuit']
+        # if os.path.exists(os.path.join(registration_folder, frame[0], frame[1] + '_weights_fixed.mat')):
+        #     continue
+        start = time.time()
+        registered_border, registered_border_edges, template_border, template, faces, registered_weight_faces, registered_weights, registered_faces = get_boundaries(folder, registration_folder, frame)
+        registered_weights_updated = flatten_boundary(folder, frame, registered_border, registered_border_edges, template_border, template, faces, registered_weight_faces, registered_weights, registered_faces)
+        sio.savemat(os.path.join(registration_folder, frame[0], frame[1] + '_weights_fixed.mat'), {'W': registered_weights_updated, 'F': registered_weight_faces })
+        # print('Sample: ' + frame[0] + ' ' + frame[1] + ': ' + str(time.time() - start))
+        # registered = (template[registered_weight_faces] * registered_weights_updated[:,:,None]).sum(1)
+        # fill_holes(registered, registered_faces)
+        # 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 = 'cpu'
+device = 'cuda'
 smpl = SMPL(device, 'smpl')
 
 
@@ -555,12 +555,12 @@ def get_boundaries(folder, frame):
     registered_cloth = np.load('temp/matlab_match_registered.npy')
     registered_faces = np.load('temp/matlab_match_faces.npy')
     registered_weight_faces = np.load('temp/matlab_match_weights_faces.npy')
-    registred_weights = np.load('temp/matlab_match_weights_weights.npy')
+    registered_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_cloth = (template[registered_weight_faces] * registred_weights[:,:,None]).sum(1)
+    registered_cloth = (template[registered_weight_faces] * registered_weights[:,:,None]).sum(1)
 
     registered_meshfix = MeshFix(registered_cloth, registered_faces.astype(np.int32))
     registered_holes = registered_meshfix.extract_holes()
@@ -586,46 +586,51 @@ def get_boundaries(folder, frame):
         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()
+
     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()
+        bounds_edges_global[k] = bounds_edges_global[k][edges_order]
+        bounds[k] = bounds[k][verts_order]
+
+    np.savez('temp/registered_border.npz', {f'bound_{i}': bounds[i] for i in range(len(bounds))})
+    np.savez('temp/registered_border_edges.npz', {f'bound_edges_{i}': bounds_edges_global[i] for i in range(len(bounds_edges_global))})
+
+
+    # for k in range(labels.max()+1):
+    #     for j in range(bounds[k].shape[0]):
+    #         w_1 = registered_weights[bounds[k][j]]
+    #         w_2 = registered_weights[bounds[k][(j+1)%bounds[k].shape[0]]]
+
+            # pass
 
 
 
-    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=[
@@ -962,7 +967,11 @@ def flatten_boundary(folder, frame):
 
 
     template_border = np.load('temp/template_border.npy')
-    registered_border = np.load('temp/registered_border.npy')
+    registered_border = np.load('temp/registered_border.npz', allow_pickle=True)['arr_0'].item()
+    registered_borders = [registered_border[k] for k in registered_border.keys()]
+    registered_border_edges = np.load('temp/registered_border_edges.npz', allow_pickle=True)['arr_0'].item()
+    registered_border_edges = [registered_border_edges[k] for k in registered_border_edges.keys()]
+    registered_border = np.concatenate(registered_borders)
 
     template_border_tree = KDTree(template[template_border,:])
 
@@ -996,6 +1005,63 @@ def flatten_boundary(folder, frame):
         yaxis = dict(range=[-1,1],),
         zaxis = dict(range=[-1,1],),),)
     fig.update_layout(scene_aspectmode='cube')
+    # fig.show()
+
+    fig = go.Figure(data=[
+        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=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()
+
+    # save registered_cloth as obj
+    # writeOBJ('temp/test_registered.obj', registered_cloth, registered_faces.astype(np.int32))
+
+
+    for k in range(len(registered_borders)):
+        border = registered_borders[k]
+        edges = registered_border_edges[k]
+        for i in range(border.shape[0]):
+            w_1 = weights_weights[border[i],:]
+            w_2 = weights_weights[border[(i+1)%border.shape[0]],:]
+            f_1 = weights_faces[border[i],:]
+            f_2 = weights_faces[border[(i+1)%border.shape[0]],:]
+
+            if np.all(f_1 == f_2):
+                reference = registered_cloth[f_1[0],:]
+                dist_1 = np.linalg.norm(registered_cloth[border[i],:] - reference, axis=0)
+                dist_2 = np.linalg.norm(registered_cloth[border[(i+1)%border.shape[0]],:] - reference, axis=0)
+                if (dist_1 < dist_2 and w_1[0] < w_2[0]) or (dist_1 > dist_2 and w_1[0] > w_2[0]):
+                    weights_weights[border[(i+1)%border.shape[0]],:] = weights_weights[border[i],:]
+
+                reference = registered_cloth[f_1[1],:]
+                dist_1 = np.linalg.norm(registered_cloth[border[i],:] - reference, axis=0)
+                dist_2 = np.linalg.norm(registered_cloth[border[(i+1)%border.shape[0]],:] - reference, axis=0)
+                if (dist_1 < dist_2 and w_1[1] < w_2[1]) or (dist_1 > dist_2 and w_1[1] > w_2[1]):
+                    weights_weights[border[(i+1)%border.shape[0]],:] = weights_weights[border[i],:]
+                # pass
+            # registered_cloth
+
+            # pass
+
+    registered_cloth = (template[weights_faces] * weights_weights[:,:,None]).sum(1)
+    fig = go.Figure(data=[
+        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=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
@@ -1036,8 +1102,8 @@ def fill_holes(folder, frame):
 
 
 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')
@@ -1063,9 +1129,9 @@ if __name__ == "__main__":
         
 
         # register(folder, frame)
-        get_boundaries(folder, frame)
+        # get_boundaries(folder, frame)
         # optimize_boundaries_2(folder, frame)
-        # flatten_boundary(folder, frame)
+        flatten_boundary(folder, frame)
         # fill_holes(folder, frame)
         break
         pass