diff --git a/mcc-flow/merge_records.py b/mcc-flow/merge_records.py index dea8c296b74600a5617855da748bb0be0382a910..a0403520aca561886f3ebd8084466877d245a6a7 100644 --- a/mcc-flow/merge_records.py +++ b/mcc-flow/merge_records.py @@ -3,9 +3,11 @@ from datetime import datetime, timedelta from pathlib import Path from struct import unpack +import mne import numpy as np from matplotlib import pyplot as plt -from pyedflib import EdfWriter, FILETYPE_EDFPLUS +from pyedflib import EdfWriter, FILETYPE_BDFPLUS +from scipy import signal HEADER_FS = 1 EEG_EMG_FS = 250 @@ -92,17 +94,16 @@ class FlowPacket: def to_edf(self, filename): to_save = [self.eeg, self.eeg_giro, self.eeg_acc, self.emg, self.emg_giro, self.emg_acc, - self.flow, self.hr, self.spo2, - self.sound, self.gsr] + self.flow, self.hr, self.spo2, self.gsr] - with EdfWriter(filename, sum(pkt.chs for pkt in to_save), FILETYPE_EDFPLUS) as edf: + with EdfWriter(filename, sum(pkt.chs for pkt in to_save), FILETYPE_BDFPLUS) as edf: data_buffer = [] sig = 0 for pkt in to_save: for i, ch_dat in enumerate(pkt.data): # print(sig) - edf.setPhysicalMaximum(sig, 3000 * EEG_EMG_SCALE) - edf.setPhysicalMinimum(sig, -3000 * EEG_EMG_SCALE) + edf.setPhysicalMaximum(sig, 3000) + edf.setPhysicalMinimum(sig, -3000) # edf.setDigitalMaximum(sig, 2**7) # edf.setDigitalMinimum(sig, -2**7) edf.setPhysicalDimension(sig, pkt.unit) @@ -112,6 +113,28 @@ class FlowPacket: sig += 1 edf.writeSamples(data_buffer) + def to_mne(self, filename): + to_save = [self.eeg, self.eeg_giro, self.eeg_acc, + self.emg, self.emg_giro, self.emg_acc, + self.flow, self.hr, self.spo2, self.gsr] + + data_buffer = [] + ch_types = [] + ch_names = [] + for pkt in to_save: + if pkt.fs == EEG_EMG_FS: + data_buffer.extend(pkt.data * 1e-6) + ch_types.extend([pkt.type] * pkt.chs) + else: + # https://docs.scipy.org/doc/scipy-1.9.3/reference/generated/scipy.interpolate.interp1d.html#scipy.interpolate.interp1d + dat = np.array([signal.resample(t, self.eeg.data.shape[-1]) for t in pkt.data]) + data_buffer.extend(dat) + ch_types.extend(['misc'] * pkt.chs) + ch_names.extend([f'{pkt.type}{i}' for i in range(pkt.chs)]) + info = mne.create_info(ch_names, EEG_EMG_FS, ch_types) + raw = mne.io.RawArray(np.array(data_buffer), info) + raw.save(filename) + def _read_block_data(file, shape, format_, format_size): sh = np.prod(shape) @@ -195,16 +218,18 @@ def read_data(path): if dat.date - data.date < timedelta(seconds=data.duration + 10): data += dat else: - data.plot() - # data.to_edf(f'exp{exp:03}-rec{rec:03}.edf') + # data.plot() + # data.to_edf(f'exp{exp:03}-rec{rec:03}.bdf') + data.to_mne(f'exp{exp:03}-rec{rec:03}_raw.fif') if dat.date - data.date < timedelta(minutes=10): rec += 1 else: rec = 1 exp += 1 data = dat - data.plot() - # data.to_edf(f'exp{exp:03}-rec{rec:03}.edf') + # data.plot() + # data.to_edf(f'exp{exp:03}-rec{rec:03}.bdf') + data.to_mne(f'exp{exp:03}-rec{rec:03}_raw.fif') if __name__ == '__main__':