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