updated dataset

data/preprocessing.py

@@ -510,6 +510,62 @@ def create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=0, end=1)
             # fig.show()
             pass
 
+
+def create_frames_file(folder, out_folder, split=0.9):
+    ids = os.listdir(folder)
+
+    frames_non_skirt = []
+    frames_skirt = []
+    ids_non_skirt = set()
+    ids_skirt = set()
+
+    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'))
+        garments = [f.replace('.obj','') for f in os.listdir(garment_dir) if f.endswith('.obj')]
+        frames = info['poses'].shape[1]
+
+        assert info['poses'].shape[0] == 72
+
+        for garment in garments:
+            for i in range(frames):
+                if garment == 'Skirt' or garment == 'Dress':
+                    frames_skirt.append(Frame(id, i, garment))
+                    ids_skirt.add(id)
+                else:
+                    frames_non_skirt.append(Frame(id, i, garment))
+                    ids_non_skirt.add(id)
+
+    nr_train_non_skirt = int(len(ids_non_skirt) * split)
+    nr_train_skirt = int(len(ids_skirt) * split)
+
+    ids_skirt = list(ids_skirt)
+    ids_non_skirt = list(ids_non_skirt)
+
+    with open(os.path.join(out_folder, 'frames_train.txt'), 'w') as f:
+        for frame in frames_non_skirt:
+            if frame.id in ids_non_skirt[:nr_train_non_skirt]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, 'frames_train_skirt.txt'), 'w') as f:
+        for frame in frames_skirt:
+            if frame.id in ids_skirt[:nr_train_skirt]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, 'frames_valid.txt'), 'w') as f:
+        for frame in frames_non_skirt:
+            if frame.id in ids_non_skirt[nr_train_non_skirt:]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+    with open(os.path.join(out_folder, 'frames_valid_skirt.txt'), 'w') as f:
+        for frame in frames_skirt:
+            if frame.id in ids_skirt[nr_train_skirt:]:
+                f.write(f'{frame.id} {frame.frame_nr} {frame.garment}\n')
+
+
 if __name__ == '__main__':
 
 
@@ -530,4 +586,6 @@ if __name__ == '__main__':
     # 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)
+    # create_mesh_unposed_no_shape(folder, folder_reg, out_folder, start=args.start, end=args.end)
+
+    create_frames_file(folder, out_folder)