Merge pull request #71 from peeplika/WID-69

[WID-69] Support for EDF File Data in TSWidget

tvbwidgets/tests/ts/test_data_wrappers.py

@@ -9,9 +9,11 @@
 import numpy as np
 import pytest
 
+from pathlib import Path
 from tvbwidgets.tests.ts.ts_generator import generate_ts_with_mode_and_sv
 from tvbwidgets.ui.ts.data_wrappers.numpy_data_wrapper import WrapperNumpy
 from tvbwidgets.ui.ts.data_wrappers.tvb_data_wrapper import WrapperTVB
+from tvbwidgets.ui.ts.data_wrappers.edf_data_wrapper import WrapperEDF
 
 
 # =========================================== TEST WRAPPER NUMPY =======================================================
@@ -172,3 +174,18 @@ def test_get_hover_channel_value_tvb(wrapper_tvb):
     val = wrapper_tvb.data.data[x_int, sel1, ch_index, sel2]
     val = round(val, 4)
     assert ch_value == val
+
+
+# ============================================ TEST WRAPPER EDF ========================================================
+file_path = Path(__file__).parent /"test_file.edf"
+
+@pytest.fixture
+def edf_wrapper():
+    """Returns an initialized WrapperEDF"""
+    return WrapperEDF(file_path)
+
+def test_edf_wrapper(edf_wrapper):
+    assert edf_wrapper.ch_idx == 1
+    assert edf_wrapper.sample_rate == 400.0
+    assert isinstance(edf_wrapper.data, np.ndarray)
+

tvbwidgets/ui/ts/base_ts_widget.py

@@ -14,6 +14,7 @@
 from tvbwidgets.ui.base_widget import TVBWidget
 from tvbwidgets.ui.ts.data_wrappers.tvb_data_wrapper import WrapperTVB
 from tvbwidgets.ui.ts.data_wrappers.numpy_data_wrapper import WrapperNumpy
+from tvbwidgets.ui.ts.data_wrappers.edf_data_wrapper import WrapperEDF
 
 class TimeSeriesWidgetBase(widgets.VBox, TVBWidget):
     # =========================================== SETUP ================================================================
@@ -27,10 +28,16 @@ def add_data_array(self, numpy_array, sample_freq, ch_idx):
         # type: (np.array, float, int) -> None
         data_wrapper = WrapperNumpy(numpy_array, sample_freq, ch_idx=ch_idx)
         self._populate_from_data_wrapper(data_wrapper)
+
+    def add_edf_data(self,filepath):
+        data_wrapper = WrapperEDF(filepath)
+        self._populate_from_data_wrapper(data_wrapper)
 
     def add_data(self, data, sample_freq=None, ch_idx=None):
         if isinstance(data, TimeSeries):
             self.add_datatype(data)
+        elif isinstance(data, str) and data.endswith(".edf"):
+            self.add_edf_data(data)
         else:
             self.add_data_array(data, sample_freq, ch_idx)
 
@@ -112,4 +119,4 @@ def _dimensions_selection_update(self, _):
         # update self.raw
         sel1, sel2 = self._get_selection_values()
         new_slice = self.data.get_update_slice(sel1, sel2)
-        self.raw = self.data.build_raw(new_slice)
+        self.raw = self.data.build_raw(new_slice)

tvbwidgets/ui/ts/data_wrappers/edf_data_wrapper.py

@@ -0,0 +1,16 @@
+import mne
+import numpy as np
+from tvbwidgets.ui.ts.data_wrappers.numpy_data_wrapper import WrapperNumpy
+
+class WrapperEDF(WrapperNumpy):
+    """Wraps EDF data for tswidget"""
+
+    def __init__(self, filepath):
+        # type: (str) -> None 
+
+        self.raw = mne.io.read_raw_edf(filepath)
+        data = self.raw.get_data()
+        self.data = np.transpose(data)
+        self.ch_idx = len(data.shape)-1
+        self.sample_rate = self.raw.info['sfreq']
+        super().__init__(self.data,self.sample_rate,ch_idx= self.ch_idx)