File size: 2,519 Bytes
4fcc185 | 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 34 35 36 37 38 39 40 41 42 43 44 45 46 | """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}
|