main.py 2.66 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

19
    if config['model'] == 'cnn':
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

23
    elif config['model'] == 'inception':
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

27
    elif config['model'] == 'eegnet':
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

33
    elif config['model'] == 'deepeye':
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
39
40
41

    elif config['model'] == 'deepeye2':
        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

43
44
45
46
    elif config['model'] == 'deepeye3':
        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)

47
48
49
    elif config['model'] == 'xception':
        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
51
    
    elif config['model'] == 'deepeye-lstm':
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()