kraken/backend/core/sig_chain/sig_save.py

"""保存数据为bdf格式"""
import logging

import pyedflib

logger = logging.getLogger(__name__)

class SigSaveHigh():
    """数据保存类:使用highlevel API，简化了保存代码，
    一次保存一个文件"""

    def __init__(self, mne_raw_data):
        """初始化SigSave

        Args:
            mne_raw_data (class): mne.io.array.array.RawArray
        """
        self.raw_data = mne_raw_data


    def save(self, file_name, name, gender):
        """保存

        Args:
            file_name (str): 保存的文件名
            name (Optional[str], optional): bdf头信息中的姓名. Defaults to None.
            gender (Optional[str], optional): bdf头信息中的性别. Defaults to None.
        """
        signals = self.raw_data.get_data()
        channel_names = self.raw_data.info.get_montage().ch_names
        signal_headers = pyedflib.highlevel.make_signal_headers(
            channel_names, sample_frequency=self.raw_data.info["sfreq"])
        header = pyedflib.highlevel.make_header(patientname=name, gender=gender)
        pyedflib.highlevel.write_edf(file_name, signals, signal_headers, header)


class SigSave():
    """数据保存类: 使用基本的保存API,可以持续将数据写入一个文件"""


    def __init__(self, channel_labels, sample_rate, physical_max, physical_min):
        """初始化保存数据类

        Args:
            channel_labels (list): 导联标签，例如：['C3','C4']
            sample_rate (int): 采样率
            physical_max (int): 最大物理值，与设备相关, 例如pony:375000 neo:200000
            physical_min (int): 最小物理值，与设备相关, 例如pony:-375000 neo:-200000
        """
        self.channel_labels = channel_labels
        self.sample_rate = sample_rate
        self.physical_max = physical_max
        self.physical_min = physical_min
        self.is_ready = False
        self.is_first = False


    def set_edf_header(self, subject, filename, task_per_run, path):
        """ 用于设置EDF头部信息

        Args:
            subject (class): 受试数据库实体
            path (str): 存储数据路径
        """
        channel_info = []
        channel_count = len(self.channel_labels)
        self.path = path + "/" + filename
        self.edf_w = pyedflib.EdfWriter(self.path,
                                        channel_count,
                                        file_type=pyedflib.FILETYPE_BDFPLUS)
        for label_num in range(channel_count):
            ch_dict = {
                "label": self.channel_labels[label_num],
                "dimension": "uV",
                "sample_frequency": self.sample_rate,
                "physical_max": self.physical_max,
                "physical_min": self.physical_min,
                "digital_max": 8388607,
                "digital_min": -8388608,
            }
            channel_info.append(ch_dict)
        self.edf_w.setSignalHeaders(channel_info)
        self.edf_w.setPatientName(subject.name)
        if subject.gender == "男":
            self.edf_w.setGender(1)
        elif subject.gender == "女":
            self.edf_w.setGender(0)
        self.edf_w.setBirthdate(subject.birthday)
        self.edf_w.setPatientCode(subject.id_card)
        self.edf_w.setRecordingAdditional(str(task_per_run))
        self.is_ready = True
        self.is_first = True
        self.start_record_timestamp = 0


    def save_raw_data(self, signals, timestamp=None):
        """向打开的文件写入数据，可以持续写入，直到调用close_edf_file时则无法写入

        Args:
            signals (np array): 需要保存的数据，channels*samples
        """
        if self.is_ready:
            if self.is_first and timestamp:
                self.start_record_timestamp = timestamp
                self.is_first = False
            self.edf_w.writeSamples(signals)
        else:
            logger.info(
                "not ready for save, maybe edf/bdf header has not been set")



    def edf_data_mark(self, timestamp, mark: str):
        """给数据打标记

        Args:
            timestamp (int): 标记的时间点
            mark (str): 标记的信息
        """
        if self.is_ready:
            time_seconds = (timestamp - self.start_record_timestamp) / 1000
            self.edf_w.writeAnnotation(time_seconds, -1, mark)
        else:
            logger.info(
                "not ready for save, maybe edf/bdf header has not been set")


    def close_edf_file(self):
        """关闭BDF文件
        """
        self.edf_w.close()
        self.is_ready = False