| """Inference accepts measured EEG only; labels never enter the prediction method.""" |
| import os,time |
| os.environ.setdefault('MPLCONFIGDIR',os.path.abspath('.cache/matplotlib')) |
| from pathlib import Path |
| import numpy as np,joblib |
| from bci_data import preprocess,LABELS |
| from classical import features |
| class Predictor: |
| def __init__(self,path): |
| self.path=str(path);self.neural=Path(path).suffix=='.pt';self.backbone=None |
| if self.neural: |
| import torch |
| from neural import build |
| torch.set_num_threads(2);self.state=torch.load(path,map_location='cpu',weights_only=False);self.model=build(self.state['kind'],self.state['channels'],[self.state['names'][i] for i in self.state['indices']]);self.model.load_state_dict(self.state['model']);self.model.eval() |
| else: |
| self.state=joblib.load(path);self.model=self.state['model'] |
| def predict(self,raw): |
| t=time.perf_counter();idx=self.state['indices'];x=preprocess(np.asarray(raw)[None],idx) |
| if self.neural: |
| import torch |
| x=np.clip(x/self.state['scale'],-20,20) |
| if self.state['kind'].startswith('cbramod_'):x=x.reshape(-1,len(idx),2,200) |
| with torch.inference_mode():prob=self.model(torch.tensor(x)).softmax(-1).numpy()[0] |
| label=int(prob.argmax()) |
| else: |
| kind=self.state['kind'] |
| if kind in ['reve','csbrain']: |
| import torch |
| from foundation_models import Foundation |
| if self.backbone is None:self.backbone=Foundation(kind,[self.state['names'][i] for i in idx]).eval();torch.set_num_threads(2) |
| with torch.inference_mode():f=self.backbone(torch.tensor(x/self.state['scale'])).float().numpy() |
| elif kind=='cbramod': |
| import torch |
| from cbramod_probe import backbone |
| if self.backbone is None:self.backbone=backbone().eval();torch.set_num_threads(2) |
| with torch.inference_mode():f=self.backbone(torch.tensor(x/100).reshape(-1,len(idx),2,200)).flatten(1).numpy() |
| elif self.state.get('mode')=='late_matched': |
| from late_control import late |
| f=features(late(np.asarray(raw)[None].copy()),'bandpower') |
| elif self.state.get('mode')=='precue': |
| from controls import precue |
| f=features(precue(np.asarray(raw)[None]),'bandpower') |
| else:f=x.astype(float) if kind in ['fbcsp','fbts'] else features(x,kind) |
| prob=self.model.predict_proba(f)[0];label=int(self.model.predict(f)[0]) |
| return {'prediction':label,'command':LABELS[label],'confidence':float(prob[label]),'probabilities':dict(zip(LABELS,map(float,prob))),'latency_ms':(time.perf_counter()-t)*1000,'measured_channels':len(idx),'model':Path(self.path).name} |
|
|
