File size: 868 Bytes
4943752 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
import os
from findfile import find_files, find_dir
filter_key_words = ['.py', '.md', 'readme', 'log', 'result', 'zip',
'.state_dict', '.model', '.png', 'acc_', 'f1_', '.backup', '.bak']
def detect_infer_dataset(dataset_path, task='apc'):
dataset_file = []
if isinstance(dataset_path, str) and os.path.isfile(dataset_path):
dataset_file.append(dataset_path)
return dataset_file
for d in dataset_path:
if not os.path.exists(d):
search_path = find_dir(os.getcwd(), [d, task, 'dataset'], exclude_key=filter_key_words, disable_alert=False)
dataset_file += find_files(search_path, ['.inference', d], exclude_key=['train.'] + filter_key_words)
else:
dataset_file += find_files(d, ['.inference', task], exclude_key=['train.'] + filter_key_words)
return dataset_file
|