From 9a7aaefb8fe67b93a064a00ccbc7a1ea31b4157d Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Sun, 25 Jun 2023 23:16:48 +0200 Subject: [PATCH] registration progress save --- data/registration.py | 310 ++++++++++++++++++++-- data/registration_post.py | 482 +++++++++++++++++++++++++++++++++++ data/registration_weights.py | 440 ++++++++++++++++++++++++++++++-- 3 files changed, 1187 insertions(+), 45 deletions(-) create mode 100644 data/registration_post.py diff --git a/data/registration.py b/data/registration.py index 569ccce..e135374 100644 --- a/data/registration.py +++ b/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 = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot]) + 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 @@ -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 = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot]) + 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 + + 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) - shape[-1] = 1 + + # boundaries = raw_mask == 2 + # # 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_body = smpl(pose.unsqueeze(0), shape.unsqueeze(0))[0][0].cpu().numpy() + # 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) - 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): - return mesh_edge_loss(mesh1, mesh2) + def edge_loss(mesh): + return mesh_edge_loss(mesh, target_length=avg_edge_length) - def normal_consistency_loss(mesh1, mesh2): - return mesh_normal_consistency(mesh1, mesh2) + def normal_consistency_loss(mesh): + 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) - loss += laplacian_smoothing_loss(mesh1_out) - loss += edge_loss(mesh1_out) - loss += normal_consistency_loss(mesh1_out) + 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 * 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 = r'D:\Work\Dataset\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') @@ -203,5 +457,5 @@ if __name__ == "__main__": register(folder, frame) - + break pass diff --git a/data/registration_post.py b/data/registration_post.py new file mode 100644 index 0000000..e035deb --- /dev/null +++ b/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 diff --git a/data/registration_weights.py b/data/registration_weights.py index 121b5d1..6f3f326 100644 --- a/data/registration_weights.py +++ b/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_faces = np.load('temp/registered_faces.npy') - registred_weights = np.load('temp/registered_weights.npy') + 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_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} 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}") + 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 = 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') @@ -659,7 +1063,9 @@ if __name__ == "__main__": # register(folder, frame) - # get_boundaries(folder, frame) - optimize_boundaries(folder, frame) + get_boundaries(folder, frame) + # optimize_boundaries_2(folder, frame) + # flatten_boundary(folder, frame) + # fill_holes(folder, frame) break pass