final smpl fix

data/convert_to_uv.py

@@ -12,7 +12,7 @@ import pickle
 import numpy as np
 import scipy as sp
 import trimesh as tm
-import cmapy as cm
+# import cmapy as cm
 import tables as tb
 import scipy.io as sio
 import scipy.interpolate as spi
@@ -24,7 +24,7 @@ import plotly.graph_objects as go
 
 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
-from data.utils import save_to_h5file
+from data.cloth3d_io import readPC2
 
 
 os.environ["CUDA_VISIBLE_DEVICES"] = "6"
@@ -34,26 +34,26 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
     uv_values = np.zeros((uv_map.shape[0], 3))
     uv_values[m2uv[:, 0]] = verts
     uv_values[m2uv[:, 1]] += verts
-    uv_values = uv_values[uv_map_mask, :] #   add garment mask
-    uv_values /= uv_counter[:, np.newaxis]
+    uv_values = uv_values[uv_map_mask[type], :] #   add garment mask
+    uv_values /= uv_counter[type][:, np.newaxis]
     return_garment_uv_mask = False
 
     if mask is not None:
         garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
         garment_mask[m2uv[:, 0]] = mask
         garment_mask[m2uv[:, 1]] += mask
-        garment_mask = garment_mask[uv_map_mask] #   add garment mask
+        garment_mask = garment_mask[uv_map_mask[type]] #   add garment mask
     else:
         garment_mask = np.sum(uv_values, axis=-1) != 0
 
     # test
-    verts_uv[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    verts_uv[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values
     # verts_uv_save = verts_uv.copy()
 
     # garment_mask = np.sum(uv_values, axis=-1) != 0
     unique_uv_values = uv_values[garment_mask, :]
-    uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
-    rbf_dest_current = rbf_dest
+    uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :]
+    rbf_dest_current = rbf_dest[type]
 
     if garment_uv_mask is None and mask is not None:
         garment_uv_mask = np.zeros((res * 2, res))
@@ -62,8 +62,8 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
         garment_uv_mask_border = garment_uv_mask.astype(bool)
         garment_uv_mask = cv2.dilate(garment_uv_mask, cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)
         garment_uv_mask_border = ~garment_uv_mask_border & garment_uv_mask
-        garment_uv_mask = garment_uv_mask * uv_mask[:,:,0]
-        rbf_dest_current = rbf_dest * garment_uv_mask
+        garment_uv_mask = garment_uv_mask * uv_mask[type][:,:,0]
+        rbf_dest_current = rbf_dest[type] * garment_uv_mask
         return_garment_uv_mask = True
 
 
@@ -86,7 +86,7 @@ def convert_to_uv(res, m2uv, uv_map, verts, neighbours=10, mask=None, method='rb
     # rbf_dest[:, [0,1]] = rbf_dest[:, [1, 0]]
 
     result[rbf_dest_current[:, 0], rbf_dest_current[:, 1], :] = final
-    result[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    result[uv_map_discreet_unique[type][:, 0], uv_map_discreet_unique[type][:, 1], :] = uv_values
     if garment_uv_mask is not None:
         result = result * garment_uv_mask[:,:,np.newaxis]
 
@@ -102,26 +102,26 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
     uv_values = np.zeros((batch, uv_map.shape[0], 3))
     uv_values[:, m2uv[:, 0]] = verts
     uv_values[:, m2uv[:, 1]] += verts
-    uv_values = uv_values[:, uv_map_mask, :]
-    uv_values /= uv_counter[np.newaxis, :, np.newaxis]
+    uv_values = uv_values[:, uv_map_mask['non-skirt'], :]
+    uv_values /= uv_counter['non-skirt'][np.newaxis, :, np.newaxis]
     return_garment_uv_mask = False
     # test
-    verts_uv[:, uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1], :] = uv_values
+    verts_uv[:, uv_map_discreet_unique['non-skirt'][:, 0], uv_map_discreet_unique['non-skirt'][:, 1], :] = uv_values
 
     if mask is not None:
         garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
         garment_mask[m2uv[:, 0]] = mask
         garment_mask[m2uv[:, 1]] += mask
-        garment_mask = garment_mask[uv_map_mask] #   add garment mask
+        garment_mask = garment_mask[uv_map_mask['non-skirt']] #   add garment mask
     else:
         garment_mask = np.sum(uv_values[0,:,:], axis=-1) != 0
 
     unique_uv_values = uv_values[:, garment_mask, :]
-    uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
-    rbf_dest_current = rbf_dest
+    uv_map_discreet_unique_unique = uv_map_discreet_unique['non-skirt'][garment_mask, :]
+    rbf_dest_current = rbf_dest['non-skirt']
 
     # add back mesh vertices that are not part of the garment
-    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
+    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique['non-skirt'][np.invert(garment_mask), :]
     rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
     rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
     
@@ -133,7 +133,7 @@ def convert_to_uv_batch_body(res, m2uv, uv_map, verts, neighbours=15, mask=None,
         return_garment_uv_mask = True
 
     if garment_uv_mask is not None:
-        rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+        rbf_dest_current = (rbf_dest['non-skirt'] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
 
     rbf_dest_current = np.nonzero(rbf_dest_current)
     rbf_dest_current = np.array(rbf_dest_current).T
@@ -172,18 +172,18 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
     garment_mask = np.zeros((uv_map.shape[0]), dtype=bool)
     garment_mask[m2uv[:, 0]] = mask
     garment_mask[m2uv[:, 1]] += mask
-    garment_mask = garment_mask[uv_map_mask] #   add garment mask
+    garment_mask = garment_mask[uv_map_mask[type]] #   add garment mask
 
-    rbf_dest_current = rbf_dest
+    rbf_dest_current = rbf_dest[type]
 
     # add back mesh vertices that are not part of the garment
-    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[np.invert(garment_mask), :]
+    uv_map_discreet_unique_unique_inv = uv_map_discreet_unique[type][np.invert(garment_mask), :]
     rbf_dest_fill = np.zeros((res * 2, res), dtype=bool)
     rbf_dest_fill[uv_map_discreet_unique_unique_inv[:, 0], uv_map_discreet_unique_unique_inv[:, 1]] = True
 
     assert not (garment_uv_mask is None and mask is not None)
     assert garment_uv_mask is not None
-    rbf_dest_current = (rbf_dest + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
+    rbf_dest_current = (rbf_dest[type] + rbf_dest_fill) * garment_uv_mask + garment_uv_mask_border
 
     rbf_dest_current = np.nonzero(rbf_dest_current)
     rbf_dest_current = np.array(rbf_dest_current).T
@@ -196,11 +196,11 @@ def convert_to_uv_batch(res, m2uv, uv_map, verts, neighbours=15, mask=None, meth
         uv_values = np.zeros((uv_map.shape[0], 3))
         uv_values[m2uv[:, 0]] = verts[i]
         uv_values[m2uv[:, 1]] += verts[i]
-        uv_values = uv_values[uv_map_mask, :]
-        uv_values /= uv_counter[:, np.newaxis]
+        uv_values = uv_values[uv_map_mask[type], :]
+        uv_values /= uv_counter[type][:, np.newaxis]
 
         unique_uv_values = uv_values[garment_mask, :]
-        uv_map_discreet_unique_unique = uv_map_discreet_unique[garment_mask, :]
+        uv_map_discreet_unique_unique = uv_map_discreet_unique[type][garment_mask, :]
 
         interpolator =  spi.RBFInterpolator(y=uv_map_discreet_unique_unique, d=unique_uv_values[:,:], neighbors=neighbours, kernel='linear')
         result = np.zeros((res * 2, res, 3), np.float32)
@@ -386,12 +386,12 @@ def test(verts, faces, values):
 
 
 def get_shapedirs():
-    to_keep = sio.loadmat(os.path.join('custom_utils/keep_remove.mat'))['keep'].reshape(-1)
+    to_keep = sio.loadmat(os.path.join('smpl/keep_remove.mat'))['keep'].reshape(-1)
     shapedirs = []
-    with open(os.path.join('custom_utils/super_model_f.pkl'), 'rb') as f:
+    with open(os.path.join('smpl/super_model_f.pkl'), 'rb') as f:
         female = pickle.load(f, encoding="latin1")
         shapedirs.append(female['shapedirs'][to_keep, :])
-    with open(os.path.join('custom_utils/super_model_m.pkl'), 'rb') as f:
+    with open(os.path.join('smpl/super_model_m.pkl'), 'rb') as f:
         male = pickle.load(f, encoding="latin1")
         shapedirs.append(male['shapedirs'][to_keep, :])
     return shapedirs
@@ -399,162 +399,139 @@ def get_shapedirs():
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser()
-    parser.add_argument("-f", "--folder", help = "end of range", default='/home/workspace/data/dataset/cloth3d/train')
-    parser.add_argument("-s", "--start", help = "start of range", default=70500)
-    parser.add_argument("-e", "--end", help = "end of range", default=90000)
+    parser.add_argument("-f", "--folder", help = "end of range", default='/data-local/127/cloth3d_processed')
+    parser.add_argument("-s", "--start", help = "start of range", default=0, type=int)
+    parser.add_argument("-e", "--end", help = "end of range", default=100, type=int)
     parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
-    parser.add_argument("-sk", "--skirt", help = "end of range", default=0)
     args = parser.parse_args()
 
-    F = np.load("custom_utils/smpl_faces" + ("_skirt" if args.skirt == 1 else "") + ".npy")
-    uv_map = np.load('custom_utils/stretched_uv_map_' + ("skirt" if args.skirt == 1 else "3") + '.npy')
-    m2uv = np.load('custom_utils/m2uv' + ("_skirt" if args.skirt == 1 else "") + '.npy', allow_pickle=True).item()
-    m2uv = m2uv_array(m2uv)
-    uv_mask = np.load('custom_utils/mask_256x512' + ("_skirt" if args.skirt == 1 else "") + '.npy')
-    curvature_label_type = '6_25 '
+    F = {'non-skirt': np.load("smpl/smpl_faces.npy"), 'skirt': np.load("smpl/smpl_faces_skirt.npy")}
+    uv_map = {'non-skirt': np.load('custom_utils/stretched_uv_map_3.npy'), 'skirt': np.load('custom_utils/stretched_uv_map_skirt.npy')}
+    m2uv = {'non-skirt': np.load('custom_utils/m2uv.npy', allow_pickle=True).item(), 'skirt': np.load('custom_utils/m2uv_skirt.npy', allow_pickle=True).item()}
+    m2uv = {k: m2uv_array(v) for k, v in m2uv.items()}
+    uv_mask = {'non-skirt': np.load('custom_utils/mask_256x512.npy'), 'skirt': np.load('custom_utils/mask_256x512_skirt.npy')}
+
 
     uv_res = 256
-    uv_mask = (cv2.resize(uv_mask * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
-    # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis]
-    # uv_mask_border = uv_mask * np.invert(uv_mask_dilate)
-    uv_mask = cv2.dilate(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
+    uv_map_discreet = dict()
+    uv_map_discreet_unique = dict()
+    uv_counter = dict()
+    rbf_dest = dict()
+    uv_map_mask = dict()
+    for k,v in uv_mask.items():
+        uv_mask[k] = (cv2.resize(v * 255.0, (uv_res, uv_res * 2)) / 255).astype(bool)[:, :, np.newaxis]
+        # uv_mask_dilate = cv2.erode(uv_mask.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_CROSS, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=1).astype(bool)[:, :, np.newaxis]
+        # uv_mask_border = uv_mask * np.invert(uv_mask_dilate)
+        uv_mask[k] = cv2.dilate(v.astype(np.float32), cv2.getStructuringElement(cv2.MORPH_RECT, (2 * 2 + 1, 2 * 2 + 1), (2, 2)), iterations=2).astype(bool)[:, :, np.newaxis]
 
-    uv_map_discreet = np.floor(uv_map * uv_res).astype(np.int32)
-    uv_map_mask = unique_uv_map_mask(uv_map_discreet)
+        uv_map_discreet[k] = np.floor(uv_map[k] * uv_res).astype(np.int32)
+        uv_map_mask[k] = unique_uv_map_mask(uv_map_discreet[k])
 
-    uv_counter = np.zeros((uv_map.shape[0]))
-    uv_counter[m2uv[:, 0]] = 1
-    uv_counter[m2uv[:, 1]] += 1
-    uv_counter = uv_counter[uv_map_mask]
+        uv_counter[k] = np.zeros((uv_map[k].shape[0]))
+        uv_counter[k][m2uv[k][:, 0]] = 1
+        uv_counter[k][m2uv[k][:, 1]] += 1
+        uv_counter[k] = uv_counter[k][uv_map_mask[k]]
 
-    uv_map_discreet_unique = uv_map_discreet[uv_map_mask, :]
+        uv_map_discreet_unique[k] = uv_map_discreet[k][uv_map_mask[k], :]
 
-    rbf_dest = np.ones((uv_res * 2, uv_res))
-    rbf_dest[uv_map_discreet_unique[:, 0], uv_map_discreet_unique[:, 1]] = 0
-    rbf_dest[np.invert(uv_mask[:,:,0])] = 0
+        rbf_dest[k] = np.ones((uv_res * 2, uv_res))
+        rbf_dest[k][uv_map_discreet_unique[k][:, 0], uv_map_discreet_unique[k][:, 1]] = 0
+        rbf_dest[k][np.invert(uv_mask[k][:,:,0])] = 0
 
 
     #########################################################################################
 
-    use_no_orient = False
-    no_shape = False
-    is_skirt = True if args.skirt == 1 else False
     folder = args.folder
-    neighbourhood = 25
-    frames_file = 'frames_all.txt'
-    data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_curvature', 'body_posed'] # 'mesh_unposed',
+    data_to_convert = ['mesh_posed', 'mesh_static', 'mesh_normals', 'body_posed', 'mesh_unposed'] # 'mesh_unposed',
 
-    print("skirt" if args.skirt == 1 else "not skirt")
-    print(f'params: -f {args.folder} -s {args.start} -e {args.end} -n {args.name} -skirt {args.skirt}')
-
-    skirt_modifier = '_skirt' if is_skirt else ''
-    orient_mod = '_no_orient' if use_no_orient else ''
-    shape_mod = '_no_shape' if no_shape else ''
-
-    frames = []
-    with open(os.path.join(folder, frames_file), 'r') as f:
-        lines = f.readlines()
-        lines = [line.strip() for line in lines]
-    frames = lines[int(args.start):int(args.end)]
-    ids = [frame.split()[0] for frame in frames]
-    ids = sorted(list(set(ids)))
+    ids = sorted(os.listdir(folder + '/mesh_body_posed_no_orient'))[args.start:args.end]
+    ids = [id.split('.')[0].split('_')[0] for id in ids]
 
     shapedirs = get_shapedirs()
 
     root = '/'.join(folder.split('/')[:-1])
     type = folder.split('/')[-1]
-    # h5file_curv = tb.open_file(os.path.join(root, f'{type}_extras.h5'), mode='r')
-    # h5file = tb.open_file(os.path.join(root, f'{type}_extras_uv.h5'), mode='a')
 
-    extras_folder = os.path.join(root, f'{type}_extras')
-    os.makedirs(extras_folder, exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_posed' + orient_mod), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_static'), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_curvature'), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_unposed'), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_body_posed' + orient_mod), exist_ok=True)
-    os.makedirs(os.path.join(extras_folder, 'uv_mask'), exist_ok=True)
+    os.makedirs(folder, exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_posed_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_static'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_normals'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_unposed_no_shape'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_body_posed_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(folder, 'uv_mask'), exist_ok=True)
 
 
     for id in tqdm.tqdm(ids):
+        garments = glob.glob(os.path.join(folder, 'mesh_mask', f'{id}_*'))
+        garments = [garment.split('/')[-1].split('.')[0].split('_')[-1] for garment in garments]
 
+        # if os.path.exists(os.path.join(folder, 'uv_static',  f'{id}_{garment}.npz')):
+        #     continue
 
-        info = loadInfo(os.path.join(folder, id, 'info.mat'))
-        garments = list(info['outfit'].keys())
-        assert info['poses'].shape[1] <= 300
-
-        # put all this in a function and call for all garments
-
-        # also masking is not the same as just filling the empty values with neighbouring
-        # values, at least the plots are different
-
-        for garment in garments:
-        # garment = garments[0]
-
-            # if all the uv maps have been generated, skip this id
-            # if all([os.path.exists(os.path.join(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz')) for i in range(info['poses'].shape[1])]):
+        for k, garment in enumerate(garments):
+            # if not ('Dress' in garments or 'Skirt' in garments):
+            #     if k != 1:
+            #         files = glob.glob(os.path.join(folder, 'uv_body_posed_no_orient',  f'{id}_*_{garment}.npz'))
+            #         for file in files:
+            #             shutil.move(file, '_'.join(file.split('_')[:-1]) + '.npz')
             #     continue
 
-            mask = read_garment_mask(folder, id, garment)
-
-            if 'mesh_posed' in data_to_convert:
-                # mesh_posed_frame = read_posed_garment_frame(folder, id, garment, 0, mask=mask, shape=info['shape'], pose=info['poses'][:,0], deorient=use_no_orient, unshape=no_shape)
-                mesh_posed = read_posed_garment(folder, id, garment, mask=mask, shape=info['shape'], poses=info['poses'], deorient=use_no_orient, unshape=no_shape)
+            mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}_{garment}.npy'))
+            x = readPC2(os.path.join(folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
+            mesh_posed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
+            mesh_posed[:,mask, :] = x
             if 'mesh_unposed' in data_to_convert:
-                raise NotImplementedError()
-                mesh_unposed = read_unposed_garment(folder, id, garment, mask=mask, shape=info['shape'], unshape=no_shape)
+                mesh_unposed = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
+                mesh_unposed[:,mask, :] = readPC2(os.path.join(folder, 'mesh_unposed_no_shape' ,f'{id}_{garment}.pc16'), True)['V']
             if 'mesh_static' in data_to_convert:
-                mesh_static = read_garment_template(folder, id, garment, mask)
-            if 'mesh_curvature' in data_to_convert:
-                mesh_curvature = read_curvatures(folder + '_extras', id, garment, count=info['poses'].shape[1])
-            if 'body_posed' in data_to_convert:
-                body_posed = read_posed_body(folder, id, info['poses'], use_no_orient)
-
+                mesh_static = np.zeros((14475, 3), dtype=np.float32)
+                mesh_static[mask, :] = np.load(os.path.join(folder, 'mesh_static', f'{id}_{garment}.npy'))
+            if 'mesh_normals' in data_to_convert:
+                mesh_normals = np.zeros((x.shape[0], 14475, 3), dtype=np.float32)
+                mesh_normals[:,mask, :] = readPC2(os.path.join(folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
+            if 'body_posed' in data_to_convert and k != 1:
+                body_posed = readPC2(os.path.join(folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
 
             start = time.perf_counter()
+            type = 'skirt' if garment == 'Dress' or garment == 'Skirt' else 'non-skirt'
 
-            uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv, uv_map, mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
-            uv_static = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-            uv_static = uv_static.astype(np.float32)
+            uv_static, uv_mask_local, garment_uv_mask_border = convert_to_uv(uv_res, m2uv[type], uv_map[type], mesh_static, neighbours=3, mask=mask, method='rbf', garment_uv_mask=None)
+            uv_static = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_static[None, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+            uv_static = uv_static[0].astype(np.float32)
 
             if 'mesh_posed' in data_to_convert:
-                uv_posed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+                uv_posed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_posed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
             if 'mesh_unposed' in data_to_convert:
-                raise NotImplementedError()
-                uv_unposed = convert_to_uv_batch(uv_res, m2uv, uv_map, mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-            if 'mesh_curvature' in data_to_convert:
-                uv_curvature = convert_to_uv_batch_label(uv_res, m2uv, uv_map, mesh_curvature[:, :], neighbours=3, mask=mask, method='nearest', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
-            if 'body_posed' in data_to_convert:
-                uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv, uv_map, body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
-
-
+                uv_unposed = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_unposed[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+            if 'mesh_normals' in data_to_convert:
+                uv_normals = convert_to_uv_batch(uv_res, m2uv[type], uv_map[type], mesh_normals[:, :, :], neighbours=3, mask=mask, method='rbf', garment_uv_mask=uv_mask_local, garment_uv_mask_border=garment_uv_mask_border)
+            if 'body_posed' in data_to_convert and k != 1:
+                uv_body_posed = convert_to_uv_batch_body(uv_res, m2uv['non-skirt'], uv_map['non-skirt'], body_posed[:, :, :], neighbours=3, mask=None, method='rbf', garment_uv_mask=None, garment_uv_mask_border=garment_uv_mask_border)
 
             print(time.perf_counter() - start)
 
-            # mesh_static
-            # save_to_h5file(h5file, id, 0, garment, uv_static, 'uv_static', 'float32')
-            # save_to_h5file(h5file, id, 0, garment, uv_mask, 'uv_mask', 'float32')
-            np.savez_compressed(os.path.join(extras_folder, 'uv_static',  f'{id}_{garment}.npz'), uv_static=uv_static)
-            np.savez_compressed(os.path.join(extras_folder, 'uv_mask',  f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
+            # ex.imshow(uv_mask_local).show()
+            # ex.imshow((uv_static+1)/2).show()
+            # ex.imshow((uv_posed[0]+1)/2).show()
+            # ex.imshow((uv_unposed[0]+1)/2).show()
+            # ex.imshow((uv_normals[0]+1.5)/3).show()
+            # ex.imshow((uv_body_posed[0]+1)/2).show()
+
+            # pass
+            np.savez_compressed(os.path.join(folder, 'uv_static',  f'{id}_{garment}.npz'), uv_static=uv_static)
+            np.savez_compressed(os.path.join(folder, 'uv_mask',  f'{id}_{garment}.npz'), uv_mask=uv_mask_local)
 
             if 'mesh_posed' in data_to_convert:
                 for i in range(len(uv_posed)):
-                    np.savez_compressed(os.path.join(extras_folder, 'uv_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
-                    # save_to_h5file(h5file, id, i, garment, uv_posed[i, :, :, :], 'uv_posed' + orient_mod + shape_mod, 'float32')
+                    np.savez_compressed(os.path.join(folder, 'uv_posed_no_orient',  f'{id}_{i}_{garment}.npz'), uv_posed=uv_posed[i])
             if 'mesh_unposed' in data_to_convert:
-                raise NotImplementedError()
-                for i in range(uv_unposed.shape[0]):
-                    save_to_h5file(h5file, id, i, garment, uv_unposed[i, :, :, :], 'uv_unposed' + orient_mod + shape_mod, 'float32')
-            if 'mesh_curvature' in data_to_convert:
-                for i in range(len(uv_curvature)):
-                    np.savez_compressed(os.path.join(extras_folder, 'uv_curvature',  f'{id}_{i}_{garment}.npz'), uv_curvature=uv_curvature[i].astype(np.uint8))
-                    # save_to_h5file(h5file, id, i, garment, uv_curvature[i, :, :, :], 'uv_curvature', 'uint8')
-            if 'body_posed' in data_to_convert:
+                for i in range(len(uv_unposed)):
+                    np.savez_compressed(os.path.join(folder, 'uv_unposed_no_shape',  f'{id}_{i}_{garment}.npz'), uv_unposed=uv_unposed[i])
+            if 'mesh_normals' in data_to_convert:
+                for i in range(len(uv_normals)):
+                    np.savez_compressed(os.path.join(folder, 'uv_normals',  f'{id}_{i}_{garment}.npz'), uv_normals=uv_normals[i])
+            if 'body_posed' in data_to_convert and k != 1:
                 for i in range(len(uv_body_posed)):
-                    np.savez_compressed(os.path.join(extras_folder, 'uv_body_posed' + orient_mod,  f'{id}_{i}_{garment}.npz'), uv_body_posed=uv_body_posed[i])
-                    # save_to_h5file(h5file, id, i, garment, uv_body_posed[i, :, :, :], 'uv_body_posed', 'float32')
+                    np.savez_compressed(os.path.join(folder, 'uv_body_posed_no_orient',  f'{id}_{i}.npz'), uv_body_posed=uv_body_posed[i])
             
-            pass
-
-
-    h5file_curv.close()
+            # pass

data/preprocessing.py

@@ -1,22 +1,39 @@
 import os
 import sys
 
+import argparse
+parser = argparse.ArgumentParser()
+# parser.add_argument("-f", "--folder", help = "end of range")
+parser.add_argument("-s", "--start", help = "start of range", type=int, default=0)
+parser.add_argument("-e", "--end", help = "end of range", type=int, default=1)
+parser.add_argument("-g", "--gpu", help = "gpu id")
+# parser.add_argument("-n", "--name", help = "end of range")
+args = parser.parse_args()
+
+os.environ["CUDA_DEVICE_ORDER"]="PCI_BUS_ID"
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
+
 import math
 import shutil
 import warnings
 import tables as tb
 import numpy as np
 import scipy.io as sio
-
+import trimesh as tm
 from tqdm import tqdm
 
 from tables import NaturalNameWarning
 warnings.filterwarnings('ignore', category=NaturalNameWarning)
+import plotly.graph_objects as go
+
+import torch
+from scipy.spatial.transform import Rotation as R
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
-from data.cloth3d_io import readPC2, loadInfo
+from data.cloth3d_io import readPC2, loadInfo, writePC2, readOBJ, quads2tris
 from data.utils import Frame
-
+from smpl.smpl_torch import SMPL
+from custom_utils.custom_utils import _maskFaces
 
 
 
@@ -132,12 +149,385 @@ def create_extras_file(folder):
     h5file.close()
 
 
+def deorient(verts, pose):
+    rodr = smpl.rodrigues(torch.tensor(pose.T, device=smpl.device).reshape(-1, 1, 3)).reshape(verts.shape[0], -1, 3, 3).cpu().numpy()
+    rotated = np.zeros_like(verts)
+    for frame in range(verts.shape[0]):
+        r = R.from_matrix(rodr[frame, 0, :, :])
+        rotated[frame] = r.apply(verts[frame], inverse=True)
+
+    return rotated
+
+
+def create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder):
+    ids = os.listdir(folder)
+    os.makedirs(os.path.join(out_folder, 'mesh_posed_no_orient'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'mesh_mask'), exist_ok=True)
+
+    for id in tqdm(ids):
+        # id = '01467'
+
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        info = loadInfo(os.path.join(folder, id, 'info.mat'))
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue
+        
+            verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
+            reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
+            verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
+            mask_idx = sio.loadmat(os.path.join(folder_reg, id, garment + '_mask.mat'))['mask_idx'].reshape((-1)) - 1
+            mask = np.zeros((14475), dtype=bool)
+            mask[mask_idx] = True
+
+            verts = deorient(verts, info['poses'])
+
+            # # plot as scatter plot with plotly
+            # import plotly.graph_objects as go
+            # frame = 5
+            # fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))])
+            # fig.show()
+
+
+            writePC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), verts, True)
+            np.save(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'), mask)
+
+
+def create_pose_shape(folder, out_folder):
+    ids = os.listdir(folder)
+    os.makedirs(os.path.join(out_folder, 'pose'), exist_ok=True)
+    os.makedirs(os.path.join(out_folder, 'shape'), exist_ok=True)
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        info = loadInfo(os.path.join(folder, id, 'info.mat'))
+        
+        shape = np.concatenate([info['shape'], np.asarray([info['gender']])]).astype(np.float32)
+        np.save(os.path.join(out_folder, 'pose', f'{id}.npy'), info['poses'].T.astype(np.float32))
+        np.save(os.path.join(out_folder, 'shape', f'{id}.npy'), shape)
+
+
+def create_mesh_static(folder, out_folder):
+    ids = os.listdir(folder)
+    os.makedirs(os.path.join(out_folder, 'mesh_static'), exist_ok=True)
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue
+
+            verts_static, _, _, _ = readOBJ(os.path.join(garment_dir, f'{garment}.obj'))
+            reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
+            verts_static = (verts_static[reg_data['F'], :] * reg_data['W'][:, :, None]).sum(axis=1).astype(np.float32)
+
+            # rotate 90 degrees around x axis
+            verts_static = np.matmul(verts_static, np.array([[1,0,0],[0,0,-1],[0,1,0]]))
+
+            # fig = go.Figure(data=[go.Scatter3d(x=verts_static[:,0], y=verts_static[:,2], z=verts_static[:,1], mode='markers', marker=dict(size=1))])
+            # fig.update_layout(scene_aspectmode='cube')
+            # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1])))
+            # fig.show()
+
+            np.save(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'), verts_static.astype(np.float32))
+
+
+
+def create_mesh_body_posed_no_orient(folder, out_folder):
+    ids = os.listdir(folder)
+    os.makedirs(os.path.join(out_folder, 'mesh_body_posed_no_orient'), exist_ok=True)
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+    
+
+        pose, shape, gender = info['poses'], info['shape'], info['gender']
+        pose = np.transpose(pose, [1, 0]).astype(np.float32)
+
+        pose[:, :3] = 0
+        beta = np.squeeze(np.concatenate((shape, gender), axis=1)).astype(np.float32)
+
+        T = smpl(torch.tensor(pose, device=smpl.device), torch.tensor(np.tile(beta[np.newaxis, :], [pose.shape[0], 1]), device=smpl.device))[0]
+        T = T.cpu().numpy()
+
+        writePC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), T, True)
+
+        # verts = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{id}.pc16'), True)['V']
+        # # plot as scatter plot with plotly
+        # import plotly.graph_objects as go
+        # frame = 5
+        # fig = go.Figure(data=[go.Scatter3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], mode='markers', marker=dict(size=1))])
+        # fig.show()
+
+        pass
+
+
+def calculate_normals(vertices, faces):
+    normals = []
+    for i in range(vertices.shape[0]):
+        mesh = tm.Trimesh(vertices=vertices[i], faces=faces, process=False, validate=False)
+        normals.append(mesh.vertex_normals)
+    return np.stack(normals)
+
+
+def create_mesh_normals(folder, out_folder):
+    ids = os.listdir(folder)
+    os.makedirs(os.path.join(out_folder, 'mesh_normals'), exist_ok=True)
+    # F = np.load('smpl/smpl_faces.npy')
+    # F_skirt = np.load('smpl/smpl_faces_skirt.npy')
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue
+
+            verts = readPC2(os.path.join(garment_dir, garment + '.pc16'), True)['V']
+            verts = deorient(verts, info['poses'])
+            _, faces, _, _ = readOBJ(os.path.join(garment_dir, garment + '.obj'))
+            faces = quads2tris(faces)
+            normals = calculate_normals(verts, faces.astype(np.int32))
+
+            reg_data = sio.loadmat(os.path.join(folder_reg, id, garment + '_weights_fixed.mat'))
+            normals = (normals[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
+
+            writePC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), normals, True)
+
+            # normals = readPC2(os.path.join(out_folder, 'mesh_normals', f'{id}_{garment}.pc16'), True)['V']
+            # colors = np.clip(np.abs(normals), 0, 1)
+            # verts = (verts[:, reg_data['F'], :] * reg_data['W'][None, :, :, None]).sum(axis=2)
+            # # plot as scatter plot with plotly
+            # import plotly.graph_objects as go
+            # frame = 5
+            # fig = go.Figure(data=[go.Scatter3d(
+            #     x=verts[frame,:,0],
+            #     y=verts[frame,:,2],
+            #     z=verts[frame,:,1],
+            #     mode='markers',
+            #     marker=dict(
+            #         size=1,
+            #         color=normals[frame,:,:],
+            #         colorscale='Viridis',
+            #         opacity=0.8
+            #     )
+            # )])
+            # moved_verts = verts + 0.05 * normals
+
+            # # Create the scatter plot for the moved vertices
+            # fig.add_trace(go.Scatter3d(
+            #     x=moved_verts[frame,:,0],
+            #     y=moved_verts[frame,:,2],
+            #     z=moved_verts[frame,:,1],
+            #     mode='markers',
+            #     marker=dict(
+            #         size=1,
+            #         color='red',
+            #         opacity=0.8
+            #     )
+            # ))
+            # fig.show()
+
+            # # create a plotly mesh plot for verts and faces
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1)])
+            # fig.show()
+
+            pass
+
+
+def verify_registration_ans_mask(folder, out_folder):
+    ids = os.listdir(folder)
+
+    for id in ids:
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        
+        for garment in garments:
+            vertices = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
+            mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
+
+            if mask.sum() == vertices.shape[1]:
+                print(id, garment)
+
+            pass
+
+
+def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1):
+    ids = sorted(os.listdir(folder))[start:end]
+    os.makedirs(os.path.join(out_folder, 'mesh_unposed_no_shape'), exist_ok=True)
+    F = np.load('smpl/smpl_faces.npy')
+    # F_skirt = np.load('smpl/smpl_faces_skirt.npy')
+
+    for id in tqdm(ids):
+        garment_dir = os.path.join(folder, id)
+        if not os.path.isdir(garment_dir):
+            continue
+
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        # info = sio.loadmat(os.path.join(folder, id, 'info.mat'))
+        pose = np.load(os.path.join(out_folder, 'pose', f'{id}.npy'))
+        shape = np.tile(np.load(os.path.join(out_folder, 'shape', f'{id}.npy'))[None,:], [pose.shape[0],1])
+
+        pose = torch.tensor(pose, dtype=torch.float32, device='cuda')
+        shape = torch.tensor(shape, dtype=torch.float32, device='cuda') 
+        
+        for garment in garments:
+            if os.path.exists(os.path.join(out_folder, f'{id}_{garment}.pc16')):
+                continue    
+
+            mesh_static = np.load(os.path.join(out_folder, 'mesh_static', f'{id}_{garment}.npy'))
+            mesh_posed = readPC2(os.path.join(out_folder, 'mesh_posed_no_orient', f'{id}_{garment}.pc16'), True)['V']
+            mask = np.load(os.path.join(out_folder, 'mesh_mask', f'{id}_{garment}.npy'))
+            # mask_idx = np.where(mask)[0]
+
+            pose[:,:3] = 0
+            # pose = pose[:10]
+            # shape = shape[:10]
+            # mesh_posed = mesh_posed[:10]
+
+            unposed = np.zeros((14475, 3), dtype=np.float32)
+            unposed[mask] = mesh_static
+
+            unposed = torch.tensor(np.tile(unposed[None,:,:], [mesh_posed.shape[0],1,1]), dtype=torch.float32, device='cuda')
+            unposed.requires_grad = True
+            mesh_posed = torch.tensor(mesh_posed, dtype=torch.float32, device='cuda')
+
+            optimizer = torch.optim.Adam([unposed], lr=0.01)
+
+            # verts_1 = unposed[0].cpu().detach().numpy()[mask]
+            # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
+            # fig = go.Figure(data=[go.Mesh3d(x=verts_1[:,0], y=verts_1[:,2], z=verts_1[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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()
+
+            # verts_2 = mesh_posed[0].cpu().detach().numpy()
+            # fig = go.Figure(data=[go.Mesh3d(x=verts_2[:,0], y=verts_2[:,2], z=verts_2[:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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()
+
+            # zero_body = smpl(torch.zeros_like(pose), torch.zeros_like(shape))[0]
+
+            # verts_3 = smpl(pose, shape, zero_body)[0][0].cpu().detach().numpy()
+            # # verts_3 = zero_body[0].cpu().detach().numpy()
+            # faces_3 = F
+            # fig = go.Figure(data=[go.Mesh3d(x=verts_3[:,0], y=verts_3[:,2], z=verts_3[:,1], i=faces_3[:,0], j=faces_3[:,1], k=faces_3[:,2], color='lightpink', opacity=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()
+
+            for i in range(5000):
+                    optimizer.zero_grad()
+
+                    pred = smpl(pose, shape, unposed)[0]
+                    loss : torch.Tensor = ((pred[:,mask,:] - mesh_posed)**2).sum()
+
+                    # Backward pass
+                    loss.backward()
+                    optimizer.step()
+
+                    with torch.inference_mode():
+                        euclidean_error = torch.sqrt(((pred[:,mask,:] - mesh_posed)**2).sum(-1)).mean()
+
+                    # Print the loss
+                    if i % 50 == 0:
+                        print(f"it: {i} loss: {loss:.8f} euclidean_error: {euclidean_error:.8f}")
+ 
+                    if euclidean_error < 0.0000001:
+                        break
+
+                    if euclidean_error < 0.0000005 and i > 2000:
+                        break
+
+                    if euclidean_error < 0.000001 and i > 4000:
+                        break
+
+            unposed = unposed.detach().cpu().numpy()
+
+            gender = int(shape[0].cpu().numpy()[-1] > 0)
+            shape_offsets = torch.tensordot(shape[:,:10], smpl.shapedirs[gender], dims=[[1], [2]])
+            shape_offsets = shape_offsets.cpu().numpy()[:,smpl.to_keep,:]
+
+            unshaped = (unposed - shape_offsets)[:,mask,:]
+
+            writePC2(os.path.join(out_folder, 'mesh_unposed_no_shape', f'{id}_{garment}.pc16'), unshaped, True)
+
+            # verts = unposed[:,mask,:]
+            # verts = unshaped
+            # _, faces = _maskFaces(np.zeros((14475, 3)), F, mask)
+            # frame = 9
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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()
+
+            # verts = pred.detach().cpu().numpy()[:,mask,:]
+            # frame = 9
+            # fig = go.Figure(data=[go.Mesh3d(x=verts[frame,:,0], y=verts[frame,:,2], z=verts[frame,:,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=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 = '/home/workspace/data/dataset/cloth3d/train'
+
+    folder = '/data-net/dataset/cloth3d/train'
+    folder_reg = '/data-local/127/cloth3d_registration/train'
+    out_folder = '/data-local/127/cloth3d_processed'
+
+    smpl = SMPL(path='smpl', device='cuda')
 
     # clean_dataset(folder)
     # generate_frames_file(folder, 'frames_all.txt')
     # split_data(folder, 0.9)
-    create_extras_file(folder)
+    # create_extras_file(folder)
+
+    # create_mesh_posed_no_orient_and_mask(folder, folder_reg, out_folder)
+    # verify_registration_ans_mask(folder, out_folder)
+    # create_pose_shape(folder, out_folder)
+    # create_mesh_static(folder, out_folder)
+    # create_mesh_body_posed_no_orient(folder, out_folder)
+    # create_mesh_normals(folder, out_folder)
+    create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)

smpl/smpl_torch.py

@@ -163,9 +163,12 @@ class SMPL(Module):
     v_shaped = torch.stack(v_shaped, axis=-1)
     
     # J = torch.matmul(self.J_regressor, v_shaped)
-    J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
-    J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
-    J = J - J0
+    if verts is None:
+      J = torch.stack([torch.matmul(self.J_regressor[g.to(torch.long)], v_shaped[i, :, :, g.to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+    else:
+      J = torch.stack([torch.matmul(self.J_regressor[torch.greater(g, 0.5).to(torch.long)], v_shaped[i, :, :, torch.greater(g, 0.5).to(torch.long)]) for i, g in enumerate(betas[:,-1])])
+      J0 = torch.reshape(J[:,0,:], (-1, 1, 3))
+      J = J - J0
     R_cube_big = self.rodrigues(pose.reshape(-1, 1, 3)).reshape(batch_num, -1, 3, 3)
 
     if verts is not None:
@@ -208,7 +211,8 @@ class SMPL(Module):
     
     joints = torch.matmul(results, torch.transpose(torch.concat([J, torch.ones((J.shape[0], 24, 1), dtype=torch.float32, device=J.get_device())], axis=-1)[:,:,None,:],2,3))[:,:,:3,0]
             
-    T = torch.stack([torch.tensordot(results[i], self.weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])])
+    weights = self.weights if verts is not None else self.weights_all 
+    T = torch.stack([torch.tensordot(results[i], weights[torch.greater(g, 0.5).to(torch.long)], dims=([0], [1])).permute(2, 0, 1) for i, g in enumerate(betas[:,-1])]) 
     rest_shape_h = torch.cat(
       (v_posed, torch.ones((batch_num, v_posed.shape[1], 1), dtype=torch.float32).to(self.device)), dim=2
     )