split dataset

data/preprocessing.py

@@ -1,6 +1,7 @@
 import os
 import sys
 
+import math
 import shutil
 import numpy as np
 import scipy.io as sio
@@ -50,6 +51,31 @@ def generate_frames_file(folder, file_name):
             f.write(str(frame.id) + ' ' + str(frame.frame_nr) + ' ' + str(frame.garment) + '\n')
 
 
+def split_data(folder, ratio):
+    with open(os.path.join(folder, 'frames_all.txt'), "r") as f:
+        frames = f.readlines()
+
+    ids = sorted(list(set([frame.split()[0] for frame in frames])))
+    train_last_id = round(len(ids) * ratio)
+
+    train_frames = []
+    valid_frames = []
+
+    for frame in tqdm(frames):
+        id = frame.split()[0]
+        if id in ids[:train_last_id]:
+            train_frames.append(frame)
+        else:
+            valid_frames.append(frame)
+
+    with open(os.path.join(folder, f'frames_train_{ratio}.txt'), 'w') as f:
+        f.writelines(train_frames)
+
+    with open(os.path.join(folder, f'frames_valid_{round(1-ratio,1)}.txt'), 'w') as f:
+        f.writelines(valid_frames)
+
+
+    pass
 
 
 
@@ -58,4 +84,5 @@ if __name__ == '__main__':
     folder = '/home/workspace/data/dataset/cloth3d/train'
 
     # clean_dataset(folder)
-    generate_frames_file(folder, 'frames_all.txt')
+    # generate_frames_file(folder, 'frames_all.txt')
+    split_data(folder, 0.9)