New Feature: saving to mne, DBF

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__':