forked from ChinaYi/ASFormer
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
20 lines (19 loc) · 788 Bytes
/
utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
import os
def get_train_val_lists(split_num, fold_path, features_path):
files = os.listdir(fold_path)
train_list = []
val_list = []
test_list = []
for file in files:
file_path = os.path.join(fold_path, file)
with open(file_path) as f:
lines = f.read().splitlines()
lines = [line[:line.index('.')] for line in lines]
if file.startswith(f'valid {split_num}'):
val_list += lines
elif file.startswith(f'test {split_num}'):
test_list += lines
files_in_fold = os.listdir(features_path)
files_in_fold = [file[:file.index('.')] for file in files_in_fold]
train_list = [video for video in files_in_fold if video not in val_list + test_list]
return train_list, val_list, test_list