main.py 1.58 KB
Newer Older
1
from config import config
2
import time
3
4
5
6
7
from CNN import CNN
from utils import IOHelper
from DeepEye import deepEye
from InceptionTime import inception
from EEGNet import eegNet
8
import numpy as np
9

10
11
def main():
    start_time = time.time()
12
    try:
Ard Kastrati's avatar
Ard Kastrati committed
13
14
        trainX, trainY = IOHelper.get_mat_data(config['data_dir'], verbose=True)
        # trainX, trainY = IOHelper.get_pickle_data(config['data_dir'], verbose=True)
15
        # IOHelper.store(trainX, trainY)
16
17
18
    except:
        return

19
    if config['model'] == 'cnn':
20
        print("Started running CNN-1. If you want to run other methods please choose another model in the config.py file.")
21
        CNN.run(trainX, trainY)
22

23
    elif config['model'] == 'inception':
24
        print("Started running InceptionTime. If you want to run other methods please choose another model in the config.py file.")
25
        inception.run(trainX=trainX, trainY=trainY)
26

27
    elif config['model'] == 'eegnet':
28
        print("Started running EEGNet. If you want to run other methods please choose another model in the config.py file.")
29
30
31
        eegnet_x = np.transpose(trainX, (0, 2, 1))
        print(eegnet_x.shape)
        eegNet.run(trainX=eegnet_x, trainY=trainY)
32

33
    elif config['model'] == 'deepeye':
34
        print("Started running DeepEye. If you want to run other methods please choose another model in the config.py file.")
35
        deepEye.run(trainX=trainX, trainY=trainY)
36

37
    else:
Ard Kastrati's avatar
Ard Kastrati committed
38
        print('Cannot start the program. Please choose one model in the config.py file')
39
40
41
42

    print("--- Runtime: %s seconds ---" % (time.time() - start_time))

if __name__=='__main__':
Ard Kastrati's avatar
Ard Kastrati committed
43
    main()