diff --git a/data/registration.py b/data/registration.py index 229c969..569ccce 100644 --- a/data/registration.py +++ b/data/registration.py @@ -3,9 +3,10 @@ import sys import random import numpy as np import plotly.graph_objects as go -from pymeshfix import MeshFix -import trimesh -import open3d as o3d +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 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..') from data.cloth3d_io import ( @@ -22,8 +23,12 @@ from data.cloth3d_io import ( readOBJ, quads2tris ) +from smpl.smpl_torch import SMPL # from custom_utils.custom_utils import _maskFaces +smpl = SMPL('cuda', 'smpl') + + def _maskFaces(V, F, mask, return_unfixed=False): V_mask = mask F_mask = np.take(V_mask, F.reshape(-1), axis=0) @@ -117,25 +122,65 @@ def register(folder, frame): 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) - v, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj')) + template, faces, _, _ = readOBJ(os.path.join(folder, id, garment + '.obj')) faces = quads2tris(faces) + pose = torch.zeros(72) + shape = torch.zeros(11) + shape[-1] = 1 + # 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() 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(v, faces) - plot(mesh_posed, faces) - plot(mesh_posed, F) + plot(template, faces) # v = mesh_posed[mask,:] + # Define the meshes to optimize + mesh1 = Meshes(verts=template[None,:,:], faces=faces[None,:,:]) + mesh2 = Meshes(verts=star_body, faces=F[None,:,:]) + + # Define the optimizer and learning rate + optimizer = torch.optim.Adam([mesh1.verts_packed()], lr=0.01) + + # Define the loss functions + def laplacian_smoothing_loss(mesh): + return mesh_laplacian_smoothing(mesh, method="uniform") + + def edge_loss(mesh1, mesh2): + return mesh_edge_loss(mesh1, mesh2) + + def normal_consistency_loss(mesh1, mesh2): + return mesh_normal_consistency(mesh1, mesh2) + + # Optimization loop + for i in range(100): + optimizer.zero_grad() + + # 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) + + # Backward pass + loss.backward() + optimizer.step() + + # Print the loss + print(f"Loss: {loss.item()}") + pass if __name__ == "__main__": - # folder = '/data-local/127/train' + folder = '/data-net/dataset/cloth3d/train' # folder = r'D:\Work\Dataset\cloth3d\train' # frames = [f'00761_{i}' for i in range(300)] - folder = r'Z:\Work\Huni\Barcelona\Dataset\cloth3d\train' + # 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]