main.py 2.4 KB
Newer Older
1
from config import config
2
import time
3
4
from CNN import CNN
from utils import IOHelper
5
6
7
from DeepEye import deepeye
from DeepEyeRNN import deepeyeRNN
from Xception import xception
8
9
from InceptionTime import inception
from EEGNet import eegNet
10
import numpy as np
Ard Kastrati's avatar
Ard Kastrati committed
11
12
import logging

13

14
def main():
Ard Kastrati's avatar
Ard Kastrati committed
15
16
    logging.basicConfig(filename=config['info_log'], level=logging.INFO)
    logging.info('Started the Logging')
17
    start_time = time.time()
18
19
    # try:
    trainX, trainY = IOHelper.get_mat_data(config['data_dir'], verbose=True)
20

zpgeng's avatar
zpgeng committed
21
    if config['model'] == 'cnn' or config['model'] == 'cnn_cluster':
22
        logging.info("Started running CNN. If you want to run other methods please choose another model in the config.py file.")
23
        CNN.run(trainX, trainY)
24

zpgeng's avatar
zpgeng committed
25
    elif config['model'] == 'eegnet' or config['model'] == 'eegnet_cluster':
26
27
        logging.info(
            "Started running EEGNet. If you want to run other methods please choose another model in the config.py file.")
28
        eegnet_x = np.transpose(trainX, (0, 2, 1))
Ard Kastrati's avatar
Ard Kastrati committed
29
        logging.info(eegnet_x.shape)
30
        eegNet.run(trainX=eegnet_x, trainY=trainY)
31

32
33
34
    elif config['model'] == 'inception' or config['model'] == 'inception_cluster':
        logging.info("Started running InceptionTime. If you want to run other methods please choose another model in the config.py file.")
        inception.run(trainX=trainX, trainY=trainY)
35

zpgeng's avatar
zpgeng committed
36
    elif config['model'] == 'xception' or config['model'] == 'xception_cluster':
37
        logging.info("Started running XceptionTime. If you want to run other methods please choose another model in the config.py file.")
38
39
40
41
42
43
44
45
46
        xception.run(trainX=trainX, trainY=trainY)

    elif config['model'] == 'deepeye' or config['model'] == 'deepeye_cluster':
        logging.info("Started running DeepEye. If you want to run other methods please choose another model in the config.py file.")
        deepeye.run(trainX=trainX, trainY=trainY)

    elif config['model'] == 'deepeye-rnn' or config['model'] == 'deepeye-rnn_cluster':
        logging.info("Started running deepeye-rnn. If you want to run other methods please choose another model in the config.py file.")
        deepeyeRNN.run(trainX=trainX, trainY=trainY)
47

48
    else:
Ard Kastrati's avatar
Ard Kastrati committed
49
        logging.info('Cannot start the program. Please choose one model in the config.py file')
50

Ard Kastrati's avatar
Ard Kastrati committed
51
52
    logging.info("--- Runtime: %s seconds ---" % (time.time() - start_time))
    logging.info('Finished Logging')
53

54
if __name__=='__main__':
zigeng's avatar
zigeng committed
55
    main()