add new volume mapping

data/calculate_normals.py

@@ -0,0 +1,112 @@
+import os
+import sys
+import numpy as np
+import trimesh as tm
+import scipy.io as sio
+import plotly.graph_objects as go
+import plotly.express as px
+from tqdm import tqdm
+import torch
+from scipy.spatial.transform import Rotation as R
+import argparse
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)) + '/..')
+
+from data.cloth3d_io import readOBJ, readPC2
+from data.utils import Frame
+from smpl.smpl_torch import SMPL
+
+
+def calculate_normals(vertices, mesh_file, poses):
+    _, faces, _, _ = readOBJ(mesh_file)
+    rodr = smpl.rodrigues(torch.tensor(poses.T.reshape(-1, 1, 3), dtype=torch.float32)).numpy().reshape(vertices.shape[0], -1, 3, 3)
+
+    normals = []
+    for i in range(vertices.shape[0]):
+        r = R.from_matrix(rodr[i, 0, :, :])
+        # assert r.as_matrix() == rodr[frame, 0, :, :]
+        out = r.apply(vertices[i], inverse=True)
+
+        # create a plotly mesh plot of out and vertices [i]
+        # fig = go.Figure(data=[
+        #     go.Mesh3d(
+        #         x=out[:, 0], y=out[:, 2], z=out[:, 1],
+        #         # i, j and k give the vertices of triangles
+        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
+        #         name='y',
+        #         showscale=False
+        #     ),
+        #     go.Mesh3d(
+        #         x=vertices[i][:, 0], y=vertices[i][:, 1], z=vertices[i][:, 2],
+        #         # i, j and k give the vertices of triangles
+        #         i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
+        #         name='x',
+        #         showscale=False
+        #     )
+        # ])
+        # fig.show()
+    
+        mesh = tm.Trimesh(vertices=out, faces=faces, process=False, validate=False)
+        normals.append(mesh.vertex_normals)
+    return np.stack(normals)
+
+
+def get_frames(folder):
+    uv_static_folder = os.path.join(folder, 'pose_shape') 
+    frames = sorted(os.listdir(uv_static_folder))
+
+    frames = [frame.split('.')[0] for frame in frames] # .split('_')[0]
+    frames = [Frame(frame.split('_')[0], frame.split('_')[1]) for frame in frames]
+    return frames
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--folder", help = "end of range", default='')
+    parser.add_argument("-f_obj", "--folder_object", help = "end of range", default='/')
+    parser.add_argument("-s", "--start", help = "start of range", default=0)
+    parser.add_argument("-e", "--end", help = "end of range", default=10)
+    parser.add_argument("-n", "--name", help = "end of range", default='lhunor')
+    args = parser.parse_args()
+
+    obj_folder = args.folder_object # '/data-net/dataset/cloth3d/train'
+    folder = args.folder #'/data-local/127/train'
+
+    os.makedirs(os.path.join(folder, 'mesh_normals'), exist_ok=True)
+    frames = get_frames(folder)
+    smpl = SMPL('cpu', 'smpl')
+
+    ids = {frame.id for frame in frames}
+    ids = sorted(list(ids))[int(args.start):int(args.end)]
+
+    for id in tqdm(ids):
+        # find obj files in the folder
+        garment_obj_folder = os.path.join(obj_folder, id[:5])
+        garments = [file.split('.')[0] for file in os.listdir(garment_obj_folder) if file.endswith('.obj')]
+
+        if len(id) == 5:
+            garment = garments[0]
+        elif len(id) == 6:
+            garment = garments[1]
+        else:
+            raise ValueError('Invalid frame id')
+
+        mesh_file = os.path.join(garment_obj_folder, f'{garment}.obj')
+        vertices =readPC2(os.path.join(garment_obj_folder, f'{garment}.pc16'), True)['V']
+        poses = sio.loadmat(os.path.join(garment_obj_folder, f'info.mat'))['poses']
+        normals = calculate_normals(vertices, mesh_file, poses)
+
+        map = sio.loadmat(os.path.join(garment_obj_folder, f'{garment}_map.mat'))['map'][0]
+        mask = np.load(os.path.join(folder, 'mesh_mask', f'{id}.npy'))
+
+        mesh_normals = normals[:,map,:]
+        mesh_normals = mesh_normals * mask[None,:, None]
+
+        for i in range(vertices.shape[0]):
+            # np.savez(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals)
+            np.savez_compressed(os.path.join(folder, 'mesh_normals', f'{id}_{i}.npz'), mesh_normals=mesh_normals[i])
+
+        # pass
+
+        # print(frame.id, frame.frame_nr)
+

docker/docker-compose.yml

@@ -34,6 +34,7 @@ services:
       - '/data-net/guest/lhunor/data:/data-net'
       - '/data-fast/126-data1/lhunor/data:/data-local/126'
       - '/data-fast/127-data1/lhunor/data:/data-local/127'
+      - '/data-fast/127-data2/lhunor/data:/data-local/127_2'
       - '/data-fast/128-data1/lhunor/data:/data-local/128'
     deploy:
       resources:

run_parallel.py

@@ -4,24 +4,27 @@ import os
 
 if __name__ == '__main__':
 
-    folder = '/home/workspace/data/dataset/cloth3d/train'
+    # folder = '/home/workspace/data/dataset/cloth3d/train'
     # folder = '/home/workspace/local-data/cape_processed/train/male'
     # folder = '/home/workspace/local-data/valid_skirt'
     # folder = '/home/workspace/local-data/valid'
     # ids = os.listdir(folder)
     # ids = os.listdir(folder + '/mesh_posed_no_orient')#[280000:]
     skirt = 0
+    obj_folder = '/data-net/dataset/cloth3d/train'
+    folder = '/data-local/127/train'
 
-    count = int(500 * 1e3)
+    count = 0
-    step = 10000
+    step = 100
 
     # !!!!!!! change skirt parameter in script
 
-    while count < 510 * 1e3:
+    while count <1200:
         # os.system('nohup python convert_to_uv_cape.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_2_cape_train_male_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
-        os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
+        # os.system('nohup python data/convert_to_uv.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_2_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv/train_2_' + str(count) + '_' + str(count+step) + '_again.out' + ' 2>&1 &')
+        os.system(f'nohup python data/calculate_normals.py -f {folder} -f_obj {obj_folder}  -s {count} -e {(count+step)} -n normals_train_calc_{count}_{(count+step)} > nohup/normals/normals_train_calc_{count}_{(count+step)}.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_train_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python convert_to_uv_local_single.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'to_uv_valid_skirt_local_unposed_onehot_3_' + str(count) + '_' + str(count+step) + ' > nohup/to_uv_valid_skirt_local_unposed_onehot_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')
         # os.system('nohup python downscale.py -f ' + folder + ' -s ' + str(count) + ' -e ' + str(count + step) + ' -n ' + 'downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + ' > nohup/downscale_valid_posed_no_orient_' + str(count) + '_' + str(count+step) + '.out' + ' 2>&1 &')