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

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

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

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

38
    elif config['model'] == 'deepeye':
Ard Kastrati's avatar
Ard Kastrati committed
39
        logging.info("Started running DeepEye. If you want to run other methods please choose another model in the config.py file.")
40
41
42
        deepeye_x = np.transpose(trainX, (0, 2, 1))
        logging.info(deepeye_x.shape)
        deepEye.run(trainX=deepeye_x, trainY=trainY)
zpgeng's avatar
zpgeng committed
43
44
45
46
		
	elif config['model'] == 'deepeye-lstm':
        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)
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
if __name__=='__main__':
Ard Kastrati's avatar
Ard Kastrati committed
54
    main()