try registration

data/registration.py

@@ -0,0 +1,162 @@
+import os
+import sys
+import random
+import numpy as np
+import plotly.graph_objects as go
+from pymeshfix import MeshFix
+import trimesh
+import open3d as o3d
+
+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 custom_utils.custom_utils import _maskFaces
+
+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
+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 = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = scipy.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 = scipy.spatial.transform.Rotation.from_euler('xyz', pose[:3]).inv()
+        # R = scipy.spatial.transform.Rotation.from_euler([0,0,-zrot])
+        verts = R.apply(verts)
+
+    return verts
+
+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)
+    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'))
+    faces = quads2tris(faces)
+
+    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)
+    # v = mesh_posed[mask,:]
+
+    pass
+
+if __name__ == "__main__":
+    # folder = '/data-local/127/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'
+    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)
+
+        pass