From 1122a88e293eb29d01ee76bfccdb15744582801d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?K=C3=B6ll=C5=91d=20Csaba?= <kollod.csaba@itk.ppke.hu>
Date: Thu, 5 Jan 2023 17:55:15 +0100
Subject: [PATCH] Release: file converter

---
 mcc-flow/merge_records.py | 63 ++++++++++++++++++++++++++-------------
 1 file changed, 43 insertions(+), 20 deletions(-)

diff --git a/mcc-flow/merge_records.py b/mcc-flow/merge_records.py
index a040352..353d2f4 100644
--- a/mcc-flow/merge_records.py
+++ b/mcc-flow/merge_records.py
@@ -6,8 +6,12 @@ from struct import unpack
 import mne
 import numpy as np
 from matplotlib import pyplot as plt
-from pyedflib import EdfWriter, FILETYPE_BDFPLUS
+from pyedflib import EdfWriter, FILETYPE_EDFPLUS
 from scipy import signal
+from scipy.interpolate import interp1d
+
+TASK_SPLIT_TIME = 30  # (sec)
+EXPERIMENT_SPLIT_TIME = 10  # (min)
 
 HEADER_FS = 1
 EEG_EMG_FS = 250
@@ -87,7 +91,7 @@ class FlowPacket:
         self._plot(self.eeg, 'EEG')
         self._plot(self.eeg_giro, 'EEG Gyroscope')
         self._plot(self.eeg_acc, 'EEG Accelerometer')
-        # self._plot(self.emg, 'EMG')
+        self._plot(self.emg, 'EMG')
         # self._plot(self.spo2, 'SPO2')
         plt.show()
 
@@ -96,14 +100,14 @@ class FlowPacket:
                    self.emg, self.emg_giro, self.emg_acc,
                    self.flow, self.hr, self.spo2, self.gsr]
 
-        with EdfWriter(filename, sum(pkt.chs for pkt in to_save), FILETYPE_BDFPLUS) as edf:
+        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, 3000)
-                    edf.setPhysicalMinimum(sig, -3000)
+                    edf.setPhysicalMaximum(sig, max(ch_dat.max(), 1))
+                    edf.setPhysicalMinimum(sig, min(ch_dat.min(), -1))
                     # edf.setDigitalMaximum(sig, 2**7)
                     # edf.setDigitalMinimum(sig, -2**7)
                     edf.setPhysicalDimension(sig, pkt.unit)
@@ -114,7 +118,8 @@ class FlowPacket:
             edf.writeSamples(data_buffer)
 
     def to_mne(self, filename):
-        to_save = [self.eeg, self.eeg_giro, self.eeg_acc,
+        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]
 
@@ -126,14 +131,21 @@ class FlowPacket:
                 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])
+                if pkt.fs < EEG_EMG_FS:
+                    interp = interp1d(np.arange(pkt.data.shape[-1]), pkt.data, kind='linear', axis=-1)
+                    dat = interp(np.linspace(0, pkt.data.shape[-1] - 1, self.eeg.data.shape[-1]))
+                    # # Spline interpolation
+                    # t0 = np.arange(pkt.data.shape[-1])
+                    # tn = np.linspace(0, pkt.data.shape[-1] - 1, self.eeg.data.shape[-1])
+                    # dat = np.array([splev(tn, splrep(t0, d)) for d in pkt.data])
+                else:
+                    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)
+        raw.save(filename, overwrite=True)
 
 
 def _read_block_data(file, shape, format_, format_size):
@@ -203,9 +215,24 @@ def read_dat_file(file):
     return packet_holder
 
 
-def read_data(path):
+def process_data_files_in_dir(path, out_file_type='edf'):
     path = Path(path)
-    files = list(path.rglob('*.dat'))
+    files = list(sorted(path.rglob('*.dat')))
+
+    def save():
+        save_path = path.joinpath(f'exp{exp:03}')
+        if out_file_type in ['edf', 'fif', 'mne']:
+            save_path.mkdir(parents=True, exist_ok=True)
+            base_file = str(save_path.joinpath(f'exp{exp:03}-rec{rec:03}'))
+            if out_file_type == 'edf':
+                data.to_edf(base_file + '.edf')
+            else:
+                data.to_mne(base_file + '_raw.fif')
+        else:
+            if out_file_type != 'plot':
+                print(f'Supported out_file_types are [edf, fif], got {out_file_type} instead. \n'
+                      f'Files are not saved, but displayed with matplotlib.')
+            data.plot()
 
     with open(files[0], 'rb') as f:
         data = read_dat_file(f)
@@ -215,23 +242,19 @@ def read_data(path):
     for i in range(1, len(files)):
         with open(files[i], 'rb') as f:
             dat = read_dat_file(f)
-            if dat.date - data.date < timedelta(seconds=data.duration + 10):
+            if dat.date - data.date < timedelta(seconds=data.duration + TASK_SPLIT_TIME):
                 data += dat
             else:
-                # 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):
+                save()
+                if dat.date - data.date < timedelta(minutes=EXPERIMENT_SPLIT_TIME):
                     rec += 1
                 else:
                     rec = 1
                     exp += 1
                 data = dat
-    # 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')
+    save()
 
 
 if __name__ == '__main__':
     path = r'D:\Users\Csabi\OneDrive - Pázmány Péter Katolikus Egyetem\MCC Flow\database'
-    read_data(path)
+    process_data_files_in_dir(path)
-- 
GitLab