save

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]