File size: 1,769 Bytes
35f60b9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | # Create a script to split the dataset into train, validation, and test sets
# The filenames are from BASE_PATH/cfs/signals_MAE/EEG/C3-M2/*.npy
import os
import random
import numpy as np
BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'datasets')
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
def split_data(data_dir, output_dir, train_ratio=0.7, val_num=50, test_ratio=0.3):
# Get all the filenames in the data directory
all_files = []
for root, dirs, files in os.walk(data_dir):
for file in files:
if file.endswith('.npy'):
pid = file.split('.')[0] # Assuming filename format is PID_XXX.npy
all_files.append(pid)
# Shuffle the filenames
random.shuffle(all_files)
# Split the data into train, validation, and test sets
total_files = len(all_files)
train_end = int(total_files * train_ratio) - val_num
val_end = train_end + val_num
train_files = all_files[:train_end]
val_files = all_files[train_end:val_end]
test_files = all_files[val_end:]
# Save the splits to output directory
os.makedirs(output_dir, exist_ok=True)
with open(os.path.join(output_dir, 'train.txt'), 'w') as f:
for file in train_files:
f.write(file + '\n')
with open(os.path.join(output_dir, 'val.txt'), 'w') as f:
for file in val_files:
f.write(file + '\n')
with open(os.path.join(output_dir, 'test.txt'), 'w') as f:
for file in test_files:
f.write(file + '\n')
if __name__ == "__main__":
data_dir = os.path.join(BASE_PATH, 'cfs', 'signals_MAE', 'EEG', 'C3-M2')
output_dir = '.'
split_data(data_dir, output_dir) |