kraken/backend/bci_core/online.py

import joblib
import numpy as np
import random
import logging
import os
from scipy import signal
from .utils import parse_model_type, reref
from .pipeline import data_evaluation


logger = logging.getLogger(__name__)


class Controller:
    """在线控制接口
    运行时主要调用decision方法，
    每次气动手反馈后调用reset_buffer方法，用以跳过气动手不应期
    Args:
        virtual_feedback_rate (float): 0-1之间浮点数，控制假反馈占比
        model_path (string): 模型文件路径
        buffer_steps (int): 
    """
    def __init__(self,
                 virtual_feedback_rate=1., 
                 real_feedback_model=None,
                 reref_method='monopolar'):
        
        self.real_feedback_model = real_feedback_model
        self.virtual_feedback_rate = virtual_feedback_rate
        self.reref_method = reref_method

    def step_decision(self, data, true_label=None):
        """抓握训练调用接口，只进行单次判决，不涉及马尔可夫过程，
        假反馈的错误反馈默认输出为10000
        Args: 
            data (mne.io.RawArray): 数据
            true_label (None or int): 训练时假反馈的真实标签
        Return:
            int: 统一化标签 (-1: keep, 0: rest, 1: cylinder, 2: ball, 3: flex, 4: double, 5: treble)
        """
        virtual_feedback = self.virtual_feedback(true_label)
        logger.debug('step_decision: virtual feedback: {}'.format(virtual_feedback))
        if virtual_feedback is not None:
            return virtual_feedback

        if self.real_feedback_model is not None:
            fs, data = self.parse_data(data)
            p = self.real_feedback_model.step_probability(fs, data)
            logger.debug('step_decison: model probability: {}'.format(str(p)))
            pred = np.argmax(p)
            real_decision = self.real_feedback_model.model.classes_[pred]
            return real_decision
        else:
            raise ValueError('Neither decision model nor true label are given')
    
    def decision(self, data, true_label=None):
        """决策主要方法，输出逻辑如下：
            如果有决策模型，无论是否有true_label，都会使用模型进行一步决策计算并填入buffer（不一定返回）

            如果有true_label（训练模式），产生一个随机数确定本trial是否为假反馈，
                是假反馈，产生一个随机数确定本trial产生正确or错误的假反馈，假反馈的标签为10000
                不是假反馈，使用模型决策
            如果没有true_label（测试模式），直接使用模型决策

            模型决策逻辑：
                根据模型记录的last_state，
                    如果当前state和last_state相同，输出-1
                    如果当前state和last_state不同，输出当前state

        Args: 
            data (mne.io.RawArray): 数据
            true_label (None or int): 训练时假反馈的真实标签
        Return:
            int: 统一化标签 (-1: keep, 0: rest, 1: cylinder, 2: ball, 3: flex, 4: double, 5: treble)
        """
        if self.real_feedback_model is not None:
            fs, data = self.parse_data(data)
            real_decision = self.real_feedback_model.viterbi(fs, data)
            # map to unified label
            if real_decision != -1:
                real_decision = self.real_feedback_model.model.classes_[real_decision]
        
        virtual_feedback = self.virtual_feedback(true_label)
        if virtual_feedback is not None:
            return virtual_feedback
        
        # true_label is None or not running virtual feedback in this trial
        # if no real model, raise ValueError
        if self.real_feedback_model is None:
            raise ValueError('Neither decision model nor true label are given')
        return real_decision

    def virtual_feedback(self, true_label=None):
        if true_label is not None:
            p = random.random()
            if p < self.virtual_feedback_rate:  # virtual feedback (error rate 0.2)
                p_correct = random.random()
                if p_correct < 0.8:
                    return true_label
                else:
                    return 10000
        return None
    
    def parse_data(self, data):
        fs, event, data_array = data
        # do preprocessing
        data_array = reref(data_array, self.reref_method)
        return fs, data_array


class HMMModel:
    """HMMModel 是一个基于隐马尔可夫模型（Hidden Markov Model, HMM）的框架，用于建模状态转移和更新。"""
    def __init__(self, 
                 transmat=None, 
                 n_classes=2, 
                 state_trans_prob=0.6, 
                 state_change_threshold=0.5,
                 momentum=0.5):
        """
        初始化HMM模型。
            transmat: 状态转移矩阵，如果为 None，则自动生成一个简单的转移矩阵。
            n_classes: 状态的数量。
            state_trans_prob: 状态保持不变的概率。
            state_change_threshold: 状态改变的阈值。
            momentum: 用于更新状态概率的动量因子。
        """
        self.n_classes = n_classes
        self.set_current_state(0)

        self.state_change_threshold = state_change_threshold

        self.hold_state = False

        if transmat is None:
            # build state transition matrix
            self.state_trans_matrix = np.zeros((n_classes, n_classes))
            # fill diagonal
            np.fill_diagonal(self.state_trans_matrix, state_trans_prob)
            # fill 0 -> each state, 
            self.state_trans_matrix[0, 1:] = (1 - state_trans_prob) / (n_classes - 1)
            self.state_trans_matrix[1:, 0] = 1 - state_trans_prob
        else:
            if isinstance(transmat, str):
                transmat = np.loadtxt(transmat)
            self.state_trans_matrix = transmat

        # momentum factor
        self.momentum = momentum
    
    def set_current_state(self, current_state):
        self._last_state = current_state
        self._probability = np.zeros(self.n_classes)
        self._probability[current_state] = 1.
    
    def step_probability(self, fs, data):
        raise NotImplementedError
    
    def offset_updater(self, decision):
        if decision == 0:
            if self.hold_state:
                self.hold_state = False
                return 0
            else:
                self.hold_state = True
                return -1
        else:
            return decision
            
    
    def viterbi(self, fs, data, return_step_p=False):
        """
            Interface for class decision

        """
        p = self.step_probability(fs, data)
        if return_step_p:
            return p, self.update_state(p)
        else:
            return self.update_state(p)
    
    def update_state(self, current_p):
        # veterbi algorithm
        prob = (self.state_trans_matrix * self._probability.T).sum(axis=1) * current_p
        # normalize
        prob /= np.sum(prob)
        # momentum
        self._probability = self.momentum * self._probability + (1 - self.momentum) * prob

        logger.debug("viterbi probability, {}".format(str(self._probability)))

        current_state = np.argmax(self._probability)
        if current_state == self._last_state:
            return -1
        else:
            if self._probability[current_state] > self.state_change_threshold:
                self.set_current_state(current_state)
                return self.offset_updater(current_state)
            else:
                return -1
    
    @property
    def probability(self):
        return self._probability.copy()


class ClfEmissionHMM(HMMModel):
    """
    ClfEmissionHMM 则是 HMMModel 的一个扩展，结合了分类模型的输出作为HMM的发射概率。
    """
    def __init__(self, model, **kwargs):
        """
        初始化分类器发射的HMM模型。
            model: 包含特征提取器、嵌入器和分类模型的元组或模型文件路径。
        """
        if isinstance(model, str):
            model = joblib.load(model)
        self.feat_extractor, self.embedder, self.model = model

        super(ClfEmissionHMM, self).__init__(n_classes=len(self.model.classes_), **kwargs)
    
    def step_probability(self, fs, data):
        p = data_evaluation([self.feat_extractor, self.embedder, self.model], data, fs, None, None, False).squeeze()
        return p


def model_loader(model_path, **kwargs):
    """
    模型如果存在训练好的transmat，会直接load
    """
    model_root, model_filename = os.path.dirname(model_path), os.path.basename(model_path)
    model_name = model_filename.split('.')[0]
    transmat_path = os.path.join(model_root, model_name + '_transmat.txt')
    if os.path.isfile(transmat_path):
        transmat = np.loadtxt(transmat_path)
    else:
        transmat = None
    kwargs['transmat'] = transmat

    return ClfEmissionHMM(model_path, **kwargs)