main.py 2.98 KB
Newer Older
1
from config import config
2
import time
3
4
from CNN import CNN
from utils import IOHelper
5
from DeepEye import deepEye, RNNdeep, deepEye2, deepEye3, Xception
6
7
from InceptionTime import inception
from EEGNet import eegNet
8
import numpy as np
Ard Kastrati's avatar
Ard Kastrati committed
9
10
import logging

11

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

zpgeng's avatar
zpgeng committed
19
    if config['model'] == 'cnn' or config['model'] == 'cnn_cluster':
Ard Kastrati's avatar
Ard Kastrati committed
20
        logging.info("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

zpgeng's avatar
zpgeng committed
23
    elif config['model'] == 'inception' or config['model'] == 'inception_cluster':
Ard Kastrati's avatar
Ard Kastrati committed
24
        logging.info("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

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

zpgeng's avatar
zpgeng committed
33
    elif config['model'] == 'deepeye' or config['model'] == 'deepeye_cluster':
Ard Kastrati's avatar
Ard Kastrati committed
34
        logging.info("Started running DeepEye. If you want to run other methods please choose another model in the config.py file.")
35
36
37
        deepeye_x = np.transpose(trainX, (0, 2, 1))
        logging.info(deepeye_x.shape)
        deepEye.run(trainX=deepeye_x, trainY=trainY)
38

zpgeng's avatar
zpgeng committed
39
    elif config['model'] == 'deepeye2' or config['model'] == 'deepeye2_cluster':
40
41
        logging.info("Started running DeepEye2. If you want to run other methods please choose another model in the config.py file.")
        deepEye2.run(trainX=trainX, trainY=trainY)
42

zpgeng's avatar
zpgeng committed
43
    elif config['model'] == 'deepeye3' or config['model'] == 'deepeye3_cluster':
44
45
46
        logging.info("Started running DeepEye3. If you want to run other methods please choose another model in the config.py file.")
        deepEye3.run(trainX=trainX, trainY=trainY)

zpgeng's avatar
zpgeng committed
47
    elif config['model'] == 'xception' or config['model'] == 'xception_cluster':
48
49
        logging.info("Started running XceptionTime. If you want to run other methods please choose another model in the config.py file.")
        Xception.run(trainX=trainX, trainY=trainY)
zigeng's avatar
zigeng committed
50
    
zpgeng's avatar
zpgeng committed
51
    elif config['model'] == 'deepeye-lstm' or config['model'] == 'deepeye-lstm_cluster':
zpgeng's avatar
zpgeng committed
52
53
        logging.info("Started running deepeye-lstm. If you want to run other methods please choose another model in the config.py file.")
        RNNdeep.run(trainX=trainX, trainY=trainY)
54

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

Ard Kastrati's avatar
Ard Kastrati committed
58
59
    logging.info("--- Runtime: %s seconds ---" % (time.time() - start_time))
    logging.info('Finished Logging')
60

61
if __name__=='__main__':
zigeng's avatar
zigeng committed
62
    main()