main.py 2.42 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
import logging
okiss's avatar
okiss committed
12
13
import scipy
import h5py
14

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

22

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

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

34
35
36
    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)
37

zpgeng's avatar
zpgeng committed
38
    elif config['model'] == 'xception' or config['model'] == 'xception_cluster':
39
        logging.info("Started running XceptionTime. If you want to run other methods please choose another model in the config.py file.")
40
41
42
43
44
        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)
okiss's avatar
okiss committed
45

46
47
48
    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)
49

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

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

56
if __name__=='__main__':
zigeng's avatar
zigeng committed
57
    main()