main.py 1.79 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
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
    try:
Ard Kastrati's avatar
Ard Kastrati committed
17
18
        trainX, trainY = IOHelper.get_mat_data(config['data_dir'], verbose=True)
        # trainX, trainY = IOHelper.get_pickle_data(config['data_dir'], verbose=True)
19
        # IOHelper.store(trainX, trainY)
20
21
22
    except:
        return

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

27
    elif config['model'] == 'inception':
Ard Kastrati's avatar
Ard Kastrati committed
28
        logging.info("Started running InceptionTime. If you want to run other methods please choose another model in the config.py file.")
29
        inception.run(trainX=trainX, trainY=trainY)
30

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

37
    elif config['model'] == 'deepeye':
Ard Kastrati's avatar
Ard Kastrati committed
38
        logging.info("Started running DeepEye. If you want to run other methods please choose another model in the config.py file.")
39
        deepEye.run(trainX=trainX, trainY=trainY)
40

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

Ard Kastrati's avatar
Ard Kastrati committed
44
45
    logging.info("--- Runtime: %s seconds ---" % (time.time() - start_time))
    logging.info('Finished Logging')
46
if __name__=='__main__':
Ard Kastrati's avatar
Ard Kastrati committed
47
    main()