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 import scipy.spatial as spatial # 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 = 'cpu' 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 _, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj')) faces = quads2tris(faces) template = np.load('temp/mapping_cloth.npy') mapping_faces = np.load('temp/mapping_faces.npy') mapping_weights = np.load('temp/mapping_weights.npy') mapping_mask = np.load('temp/mapping_mask.npy') mapping_faces, mapping_weights = mapping_faces[mapping_mask], mapping_weights[mapping_mask] 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) _, registered_faces = _maskFaces(np.zeros((14475, 3)), F, mapping_mask) template = torch.tensor(template).to(device).to(torch.float32) template_orig = template.clone() faces = torch.tensor(faces).to(device) mapping_faces = torch.tensor(mapping_faces).to(device) mapping_weights = torch.tensor(mapping_weights).to(device).to(torch.float32) F = torch.tensor(F).to(device) registered_faces = torch.tensor(registered_faces).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() mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:]) avg_edge_length = mesh_edge_lengths(mesh2).mean() # Define the meshes to optimize registered = (template[mapping_faces,:] * mapping_weights[:,:,None]).sum(1) optim_mask = torch.ones(registered.shape[0]).to(device).to(torch.bool) # registered_tree = spatial.KDTree(template.detach().cpu().numpy()) # for i in range(registered.shape[0]): # # if distance between registerd and template is too big set it to require grad that weight # res = registered_tree.query(registered[i].detach().cpu().numpy(), k=1) # if res[0] > avg_edge_length / 4: # optim_mask[i] = True mapping_weights_masked = mapping_weights[optim_mask] mapping_weights_nonmasked = mapping_weights[~optim_mask] mapping_weights_masked.requires_grad = True # mesh1 = Meshes(verts=registered[None,:,:], faces=registered_faces[None,:,:]) # 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(1): 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(mapping_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()), ) loss: torch.Tensor = ch_loss * 20 #+ ch_norm_loss * 10 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 * 20 + 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} 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() optim_mask = optim_mask.detach().cpu().numpy() registered_faces = registered_faces.detach().cpu().numpy() # faces, F = F, faces # verts, star_body = star_body, verts fig = go.Figure(data=[ go.Mesh3d( x=template[:,0], y=template[:,1], z=template[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=0.80), go.Mesh3d( x=registered[:,0], y=registered[:,1], z=registered[:,2], i=registered_faces[:,0], j=registered_faces[:,1], k=registered_faces[:,2], color='lightpink', opacity=0.50), # go.Scatter3d( # x=template[:,0], y=template[:,1], z=template[:,2], # mode='markers', marker=dict(size=2, color='red')), # go.Scatter3d( # x=registered[:,0], y=registered[:,1], z=registered[:,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() 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')), ]) # 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() # fig = go.Figure(data=[ # go.Scatter3d( # x=registered[:optim_mask.sum(),0], y=registered[:optim_mask.sum(),1], z=registered[:optim_mask.sum(),2], # mode='markers', marker=dict(size=4, color='red')), # go.Scatter3d( # x=registered[:,0], y=registered[:,1], z=registered[:,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 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/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') _, 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_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 = [] 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( 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() pass def optimize_boundaries(folder, frame): id, frame_nr, garment = frame.split() frame_nr = int(frame_nr) registered_cloth = np.load('temp/registered_cloth.npy') registered_faces = np.load('temp/registered_faces.npy') registered_weights = np.load('temp/registered_weights.npy') _, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj')) faces = quads2tris(faces) template = np.load('temp/mapping_cloth.npy') mapping_mask = np.load('temp/mapping_mask.npy') mapping_faces = np.load('temp/mapping_faces.npy') mapping_weights = np.load('temp/mapping_weights.npy') mapping_faces, mapping_weights = mapping_faces[mapping_mask], mapping_weights[mapping_mask] template_border = np.load('temp/template_border.npy') registered_border = np.load('temp/registered_border.npy') template = torch.tensor(template).to(device).to(torch.float32) faces = torch.tensor(faces).to(device) registered_faces = torch.tensor(registered_faces).to(device) mapping_faces = torch.tensor(mapping_faces).to(device) registered_weights = torch.tensor(registered_weights).to(device).to(torch.float32) registered_cloth = torch.tensor(registered_cloth).to(device).to(torch.float32) mesh2 = Meshes(verts=template[None,:,:], faces=faces[None,:,:]) avg_edge_length = mesh_edge_lengths(mesh2).mean() optim_mask = torch.zeros(registered_cloth.shape[0]).to(device).to(torch.bool) optim_mask[registered_border] = True mapping_weights_masked = registered_weights[optim_mask] mapping_weights_nonmasked = registered_weights[~optim_mask] mapping_weights_masked.requires_grad = True # Define the optimizer and learning rate optimizer = torch.optim.Adam([mapping_weights_masked], lr=0.001) # optimizer = torch.optim.SGD([mesh1.verts_packed()], lr=0.001) # Define the loss functions def laplacian_smoothing_loss(mesh): return mesh_laplacian_smoothing(mesh, method="uniform") # cotcurv uniform cot def edge_loss(mesh): return mesh_edge_loss(mesh, target_length=avg_edge_length) def normal_consistency_loss(mesh): return mesh_normal_consistency(mesh) # Optimization loop for i in range(1): 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() faces = faces.detach().cpu().numpy() optim_mask = optim_mask.detach().cpu().numpy() registered_faces = registered_faces.detach().cpu().numpy() # faces, F = F, faces # verts, star_body = star_body, verts fig = go.Figure(data=[ go.Mesh3d( x=template[:,0], y=template[:,1], z=template[:,2], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightblue', opacity=0.80), go.Mesh3d( x=registered[:,0], y=registered[:,1], z=registered[:,2], i=registered_faces[:,0], j=registered_faces[:,1], k=registered_faces[:,2], color='lightpink', opacity=0.50), ]) fig.update_layout(scene = dict( xaxis = dict(range=[-1,1],), yaxis = dict(range=[-1,1],), zaxis = dict(range=[-1,1],),),) fig.update_layout(scene_aspectmode='cube') fig.show() fig = go.Figure(data=[ go.Scatter3d( x=registered[:,0], y=registered[:,1], z=registered[:,2], mode='markers', marker=dict(size=4, color='red')), go.Scatter3d( x=template[:,0], y=template[:,1], z=template[:,2], mode='markers', marker=dict(size=2, color='blue')), ]) fig.update_layout(scene = dict( xaxis = dict(range=[-1,1],), yaxis = dict(range=[-1,1],), zaxis = dict(range=[-1,1],),),) fig.update_layout(scene_aspectmode='cube') fig.show() pass 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( 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 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 = 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' frame = '05060 1 Jumpsuit' # register(folder, frame) get_boundaries(folder, frame) # optimize_boundaries_2(folder, frame) # flatten_boundary(folder, frame) # fill_holes(folder, frame) break pass