diff --git a/mcc-flow/merge_records.py b/mcc-flow/merge_records.py index 54a3c9b373a9cfa8c296433965c26b68309b5ebc..6916fb8179495845833f505a28df537db53a940c 100644 --- a/mcc-flow/merge_records.py +++ b/mcc-flow/merge_records.py @@ -1,70 +1,116 @@ +from dataclasses import dataclass, field from datetime import datetime from pathlib import Path from struct import unpack import numpy as np from matplotlib import pyplot as plt +from pyedflib import EdfWriter, FILETYPE_EDFPLUS +HEADER_FS = 1 +EEG_EMG_FS = 250 +GYRO_ACC_FS = 50 +SOUND_FS = 6000 +GSR_FS = 10 +EEG_EMG_SCALE = 0.04470348358154 # (mV) +GYRO_SCALE = 0.0152587890625 # (deg/sec) +ACC_SCALE = 0.00006103515625 # (g) + + +@dataclass +class PacketData: + data: np.ndarray + type: str + unit: str + fs: int + + @property + def chs(self): + return self.data.shape[0] + + def __add__(self, other): + assert self.type == other.type, 'Cannot combine PacketDatas with different signal types' + assert self.unit == other.unit, 'Cannot combine PacketDatas with different units' + assert self.fs == other.fs, 'Cannot combine PacketDatas with different sampling frequency' + assert self.chs == other.chs, 'Cannot combine PacketDatas with different number of channels' + self.data = np.concatenate((self.data, other.data), axis=-1) + return self + + +@dataclass class FlowPacket: - header_fs = 1 - signal_fs = 250 - gyro_acc_fs = 50 - sound_fs = 6000 - grs_fs = 10 - - def __init__(self, date_, flow, hr, spo2, - # vbat_eeg, vbat_arm, - eeg, eeg_giro, eeg_acc, - emg, emg_giro, emg_acc, - sound, gsr): - self.date = date_ - self.duration = 1 - self.flow = [flow] - self.hr = [hr] - self.spo2 = [spo2] - - self.eeg = eeg - self.eeg_giro = eeg_giro - self.eeg_acc = eeg_acc - self.emg = emg - self.emg_giro = emg_giro - self.emg_acc = emg_acc - self.sound = sound - self.gsr = gsr + date: datetime + duration: int = field(init=False, default=1) + flow: PacketData + hr: PacketData + spo2: PacketData + eeg: PacketData + eeg_giro: PacketData + eeg_acc: PacketData + emg: PacketData + emg_giro: PacketData + emg_acc: PacketData + sound: PacketData + gsr: PacketData def __add__(self, other): - self.flow.extend(other.flow) - self.hr.extend(other.hr) - self.spo2.extend(other.spo2) - self.eeg = np.concatenate((self.eeg, other.eeg), axis=-1) - self.eeg_giro = np.concatenate((self.eeg_giro, other.eeg_giro), axis=-1) - self.eeg_acc = np.concatenate((self.eeg_acc, other.eeg_acc), axis=-1) - self.emg = np.concatenate((self.emg, other.emg), axis=-1) - self.emg_giro = np.concatenate((self.emg_giro, other.emg_giro), axis=-1) - self.emg_acc = np.concatenate((self.emg_acc, other.emg_acc), axis=-1) - self.sound = np.concatenate((self.sound, other.sound), axis=-1) - self.gsr = np.concatenate((self.gsr, other.gsr), axis=-1) + self.flow += other.flow + self.hr += other.hr + self.spo2 += other.spo2 + self.eeg += other.eeg + self.eeg_giro += other.eeg_giro + self.eeg_acc += other.eeg_acc + self.emg += other.emg + self.emg_giro += other.emg_giro + self.emg_acc += other.emg_acc + self.sound += other.sound + self.gsr += other.gsr self.duration += other.duration return self - def _plot(self, data, title): + def _plot(self, pkt_data, title): + data = pkt_data.data plt.figure() plt.suptitle(title) chs = data.shape[0] for i in range(chs): plt.subplot(chs, 1, i + 1) plt.plot(np.linspace(0, self.duration, data.shape[1]), data[i, :]) - plt.title(f'ch{i + 1}') + if chs > 1: + plt.title(f'ch{i + 1}') plt.xlabel('time (s)') def plot(self): self._plot(self.eeg, 'EEG') self._plot(self.eeg_giro, 'EEG Gyro') self._plot(self.emg, 'EMG') - + self._plot(self.spo2, 'SPO2') plt.show() + 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] + + with EdfWriter(filename, sum(pkt.chs for pkt in to_save), FILETYPE_EDFPLUS) as edf: + data_buffer = [] + sig = 0 + for pkt in to_save: + for i, ch_dat in enumerate(pkt.data): + print(sig) + # edf.setPhysicalMaximum(sig, 2**7 * EEG_EMG_SCALE) + # edf.setPhysicalMinimum(sig, -2**7 * EEG_EMG_SCALE) + # edf.setDigitalMaximum(sig, 2**7) + # edf.setDigitalMinimum(sig, -2**7) + edf.setPhysicalDimension(sig, pkt.unit) + edf.setSamplefrequency(sig, pkt.fs) + edf.setLabel(sig, f'{pkt.type}{i}') + data_buffer.append(ch_dat) + sig += 1 + edf.writeSamples(data_buffer) + def _read_block_data(file, shape, format_, format_size): sh = np.prod(shape) @@ -82,26 +128,37 @@ def read_packet(file): flow = int.from_bytes(flow, 'big') date_ = datetime.fromisoformat(f'{date_[:4]}-{date_[4:6]}-{date_[6:]}T{time_[:2]}:{time_[2:4]}:{time_[4:]}') spo2 /= 10. - print(date_, flow, hr, spo2, vbat_eeg, vbat_arm) + flow = PacketData(np.array([[flow]]), 'flow', 'unit', HEADER_FS) + hr = PacketData(np.array([[hr]]), 'hr', 'unit', HEADER_FS) + spo2 = PacketData(np.array([[spo2]]), 'spo2', 'unit', HEADER_FS) + # print(date_, flow, hr, spo2, vbat_eeg, vbat_arm) # eeg - eeg = _read_block_data(file, (250, 4), 'i', 4).T - eeg_giro = _read_block_data(file, (50, 3), 'h', 2).T - eeg_acc = _read_block_data(file, (50, 3), 'h', 2).T - print(eeg.shape, eeg_giro.shape, eeg_acc.shape) + eeg = _read_block_data(file, (250, 4), 'i', 4).T * EEG_EMG_SCALE + eeg_giro = _read_block_data(file, (50, 3), 'h', 2).T * GYRO_SCALE + eeg_acc = _read_block_data(file, (50, 3), 'h', 2).T * ACC_SCALE + eeg = PacketData(eeg, 'eeg', 'uV', EEG_EMG_FS) + eeg_giro = PacketData(eeg_giro, 'eeg_giro', 'unit', GYRO_ACC_FS) + eeg_acc = PacketData(eeg_acc, 'eeg_acc', 'unit', GYRO_ACC_FS) + # print(eeg, eeg_giro, eeg_acc) # emg - emg = _read_block_data(file, (250, 6), 'i', 4).T - emg_giro = _read_block_data(file, (50, 3), 'h', 2).T - emg_acc = _read_block_data(file, (50, 3), 'h', 2).T - print(emg.shape, emg_giro.shape, emg_acc.shape) + emg = _read_block_data(file, (250, 6), 'i', 4).T * EEG_EMG_SCALE + emg_giro = _read_block_data(file, (50, 3), 'h', 2).T * GYRO_SCALE + emg_acc = _read_block_data(file, (50, 3), 'h', 2).T * ACC_SCALE + emg = PacketData(emg, 'emg', 'uV', EEG_EMG_FS) + emg_giro = PacketData(emg_giro, 'emg_giro', 'unit', GYRO_ACC_FS) + emg_acc = PacketData(emg_acc, 'emg_acc', 'unit', GYRO_ACC_FS) + # print(emg, emg_giro, emg_acc) # extra sound = _read_block_data(file, (6000, 1), 'h', 2).T gsr = _read_block_data(file, (10, 1), 'h', 2).T - print(sound.shape, gsr.shape) + sound = PacketData(sound, 'sound', 'unit', SOUND_FS) + gsr = PacketData(gsr, 'gsr', 'unit', GSR_FS) + # print(sound, gsr) - packet = FlowPacket(date_, flow, hr, spo2, # vbat_eeg, vbat_arm, + packet = FlowPacket(date_, flow, hr, spo2, eeg, eeg_giro, eeg_acc, emg, emg_giro, emg_acc, sound, gsr) @@ -125,7 +182,8 @@ def read_data(path): for file in files: with open(file, 'rb') as f: data = read_dat_file(f) - data.plot() + # data.plot() + data.to_edf('test.edf') exit(12)