main.py 1.47 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:
13
        trainX, trainY = IOHelper.get_data(verbose=True)
14
        # IOHelper.store(trainX, trainY)
15
16
17
    except:
        return

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

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

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

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

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

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

if __name__=='__main__':
    main()