diff --git a/data/preprocessing.py b/data/preprocessing.py index 595dcab..f48b859 100644 --- a/data/preprocessing.py +++ b/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') \ No newline at end of file + # generate_frames_file(folder, 'frames_all.txt') + split_data(folder, 0.9) \ No newline at end of file