Alaska2

example_pipeline/00_truncate_edf_annotations.py

@@ -51,7 +51,8 @@
         file_path_raw_signals = dataset['file_path_raw_signals']
         with EdfReader(file_path_raw_signals) as f:
             total_time_in_seconds = f.file_duration
-
+            edf_header = f.getHeader()
+
         file_path_manual_state_annotation = dataset['file_path_manual_state_annotation']
         states, intervals = _load_edf_hypnogram(file_path_manual_state_annotation)
 
@@ -63,5 +64,5 @@
             elif (start < total_time_in_seconds) and (stop > total_time_in_seconds):
                 intervals[-1] = (start, total_time_in_seconds)
             start, stop = intervals[-1]
-
-        _export_edf_hypnogram(file_path_manual_state_annotation, states, intervals)
+          
+        _export_edf_hypnogram(file_path_manual_state_annotation, states, intervals, edf_header)

example_pipeline/04_run_state_annotation.py

@@ -11,6 +11,7 @@
 import numpy as np
 import matplotlib.pyplot as plt
 
+from pyedflib import EdfReader, EdfWriter
 from somnotate._automated_state_annotation import StateAnnotator
 from somnotate._utils import (
     convert_state_vector_to_state_intervals,
@@ -89,11 +90,13 @@ def _get_score(vec):
     # check contents of spreadsheet
     check_dataframe(datasets,
                     columns = [
+                        'file_path_raw_signals',                   # ? get name of original raw file from spreadsheet
                         'file_path_preprocessed_signals',
                         'file_path_automated_state_annotation',
                         'file_path_review_intervals',
                     ],
                     column_to_dtype = {
+                        'file_path_raw_signals' : str,             # ? get name of original raw file from spreadsheet
                         'file_path_preprocessed_signals' : str,
                         'file_path_automated_state_annotation' : str,
                         'file_path_review_intervals' : str,
@@ -114,7 +117,11 @@ def _get_score(vec):
         predicted_state_vector = annotator.predict(signal_array)
         predicted_states, predicted_intervals = convert_state_vector_to_state_intervals(
             predicted_state_vector, mapping=int_to_state, time_resolution=time_resolution)
-        export_hypnogram(dataset['file_path_automated_state_annotation'], predicted_states, predicted_intervals)
+
+        with EdfReader(dataset['file_path_raw_signals']) as f:
+            edf_header = f.getHeader() # ?read header from the original file
+
+        export_hypnogram(dataset['file_path_automated_state_annotation'], predicted_states, predicted_intervals, edf_header)
 
         # compute intervals for manual review
         state_probability = annotator.predict_proba(signal_array)

example_pipeline/data_io.py

@@ -9,7 +9,12 @@
 import pandas
 
 from argparse import ArgumentParser
-from collections import Iterable
+try:
+    # Python <= 3.9
+    from collections import Iterable
+except ImportError:
+    # Python > 3.9
+    from collections.abc import Iterable
 from pyedflib import EdfReader, EdfWriter
 from six import ensure_str
 
@@ -338,7 +343,7 @@ def _load_edf_hypnogram(file_path):
     return states, intervals
 
 
-def _export_edf_hypnogram(file_path, states, intervals):
+def _export_edf_hypnogram(file_path, states, intervals, header=None):
     """Export hypnogram as "Time-stamped Annotations Lists (TALs)" in
     EDF annotations.
 
@@ -361,6 +366,8 @@ def _export_edf_hypnogram(file_path, states, intervals):
     with EdfWriter(file_path, 0) as writer:
         for (start, stop), state in zip(intervals, states):
             writer.writeAnnotation(start, stop-start, state)
+        if header:
+            writer.setHeader(header)
         writer.close()