-
Notifications
You must be signed in to change notification settings - Fork 1
/
run.py
33 lines (22 loc) · 824 Bytes
/
run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
# import os
# import sys
# sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from eye_movement.demo import eye_movement_process
from EEG.inference import inference
from EEG.train import load_model
from EEG.model import SleepModel
from EEG.config import config as cfg
import torch.backends.cudnn as cudnn
def prepare_model(test_model):
model = SleepModel(5, is_training=True)
model = model.to(cfg.device)
if cfg.cuda:
cudnn.benchmark = True
if cfg.resume:
load_model(model, test_model)
return model
if __name__ == "__main__":
test_model = "./EEG/model/FC_best.pth"
dataset_path = "./EEG/data/"
model = prepare_model(test_model)
eye_movement_process(inference_func=inference, dataset_path=dataset_path, model=model, outcall=True)