kraken/backend/core/sig_chain/sig_reader.py

"""读取数据文件
"""
from typing import List

import mne
import numpy as np


class Reader:
    """读取bdf文件
    """
    def __init__(self) -> None:
        self._montage = mne.channels.make_standard_montage('standard_1020')

    def read(self, filename: str, ch_names: List[str]):
        raw = mne.io.read_raw_bdf(filename, preload=True)
        raw.set_montage(self._montage)
        raw.pick_channels(ch_names=ch_names)

        return raw

    def fix_annotation(self, raw:mne.io.Raw):
        """在线数据按秒打标签，这里将相同秒标签合并

        Args:
            raw (mne.io.Raw): eeg data
        """
        annotations = raw.annotations
        for item in ['miFailed', 'miSuccess']:
            if item in set(annotations.description):
                annotations.rename({item: 'mi'})
        all_idxes = np.arange(0, len(annotations))
        valid_idxes = []
        last_label = None
        for ii, annot in enumerate(annotations):
            if last_label != annot['description']:
                last_label = annot['description']
                valid_idxes.append(ii)
        valid_idxes = np.array(valid_idxes)
        delete_mask = ~np.isin(all_idxes, valid_idxes)
        delete_idxes = all_idxes[delete_mask]
        annotations.delete(delete_idxes)
        raw.set_annotations(annotations)