calculate average garments

data/preprocessing.py

@@ -679,7 +679,8 @@ def check_irregular_poses(folder, folder_reg, out_folder):
     # mesh_static = np.load(f'{out_folder}/mesh_static/{frame.id}_{frame.garment}.npy')
 
     mesh_mask = np.load(f'{out_folder}/mesh_mask/{frame.id}_{frame.garment}.npy')
-    mesh_body_posed = readPC2Frame(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), frame.frame_nr, True)
+    mesh_body_posed = readPC2(os.path.join(out_folder, 'mesh_body_posed_no_orient', f'{frame.id}.pc16'), True)['V']
+    mesh_body_posed = mesh_body_posed[int(mesh_body_posed.shape[0]/2)+frame.frame_nr]
 
     F = np.load('smpl/smpl_faces_skirt.npy')
     _, faces = _maskFaces(np.zeros((14475, 3)), F, mesh_mask)
@@ -694,6 +695,46 @@ def check_irregular_poses(folder, folder_reg, out_folder):
     fig.show()
     pass
 
+
+def calculate_average(folder, out_folder, frames_file='frames_train_skirt.txt'):
+    from custom_utils.custom_utils import UV2Mesh
+    import plotly.express as ex
+    import plotly.graph_objects as go
+
+    frames = []
+    with open(os.path.join(out_folder, frames_file), 'r') as f:
+        for line in f:
+            line = line.strip()
+            id, frame_nr, garment = line.split(' ')
+            frames.append(Frame(id, int(frame_nr), garment))
+
+    uv_sum = np.zeros((512, 256, 3), dtype=np.float32)
+    mask_sum = np.zeros((512, 256), dtype=np.float32)
+
+    for frame in tqdm(frames):
+        uv = np.load(os.path.join(out_folder, 'uv_unposed_no_shape', f'{frame.id}_{frame.frame_nr}_{frame.garment}.npz'))['uv_unposed']
+        mask = np.load(os.path.join(out_folder, 'uv_mask', f'{frame.id}_{frame.garment}.npz'))['uv_mask']
+        uv_sum += uv
+        mask_sum += mask
+
+    mask_sum[mask_sum == 0] = 1
+    avg_garment = uv_sum / mask_sum[:, :, None]
+
+    avg_mesh = UV2Mesh()(torch.from_numpy(avg_garment.transpose(2,0,1))[None,:,:,:]).numpy()[0]
+    F = np.load('smpl/smpl_faces.npy')
+
+    fig = go.Figure(data=[go.Mesh3d(x=avg_mesh[:,0], y=avg_mesh[:,2], z=avg_mesh[:,1], i=F[:,0], j=F[:,1], k=F[:,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()
+
+    np.save('custom_data/avg/train_skirt_no_shape_avg_uv.npy', avg_garment)
+
+
+
 if __name__ == '__main__':
 
 
@@ -721,9 +762,10 @@ if __name__ == '__main__':
     # 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)
+    calculate_average(folder, out_folder)
 
     # create_frames_file(folder, out_folder)
-    create_frames_file_partial(folder, out_folder, part=0.5)
+    # create_frames_file_partial(folder, out_folder, part=0.5)
     
     # folder = '/data-local/127_2/cape_processed/test/female'
     # create_mesh_normals_cape(folder, None)