Read MIT Annotations

doc/changes/devel/13030.newfeature.rst

@@ -0,0 +1 @@
+Adds the ability to read WaveForm Database annotations in :func:`mne.read_annotations()` by `Jacob Woessner`_.

environment.yml

@@ -60,4 +60,5 @@ dependencies:
   - trame-vtk
   - trame-vuetify
   - vtk =9.3.1=qt_*
+  - wfdb
   - xlrd

mne/annotations.py

@@ -36,6 +36,7 @@
     _check_option,
     _check_pandas_installed,
     _check_time_format,
+    _check_wfdb_installed,
     _convert_times,
     _DefaultEventParser,
     _dt_to_stamp,
@@ -1140,13 +1141,21 @@ def _write_annotations_txt(fname, annot):
 
 @fill_doc
 def read_annotations(
-    fname, sfreq="auto", uint16_codec=None, encoding="utf8", ignore_marker_types=False
+    fname,
+    sfreq="auto",
+    *,
+    uint16_codec=None,
+    encoding="utf8",
+    ignore_marker_types=False,
+    fmt="auto",
+    suffix=None,
 ) -> Annotations:
     r"""Read annotations from a file.
 
     This function reads a ``.fif``, ``.fif.gz``, ``.vmrk``, ``.amrk``,
-    ``.edf``, ``.bdf``, ``.gdf``, ``.txt``, ``.csv``, ``.cnt``, ``.cef``, or
-    ``.set`` file and makes an :class:`mne.Annotations` object.
+    ``.edf``, ``.bdf``, ``.gdf``, ``.txt``, ``.csv``, ``.cnt``, ``.cef``,
+    ``.set``, or ``.seizures`` file and makes an :class:`mne.Annotations`
+    object.
 
     Parameters
     ----------
@@ -1174,6 +1183,14 @@ def read_annotations(
     ignore_marker_types : bool
         If ``True``, ignore marker types in BrainVision files (and only use their
         descriptions). Defaults to ``False``.
+    fmt : str | None
+        Used to manually specify the format of the annotations file. If
+        ``None`` (default), the format is inferred from the file extension.
+        Currently only supports ``'wfdb'``.
+    suffix : str | None
+        Used to manually specify the suffix of the annotations file for WFDB
+        files.
+        If ``None`` (default), the suffix is inferred from the file extension.
 
     Returns
     -------
@@ -1212,6 +1229,7 @@ def read_annotations(
         ".vmrk": _read_annotations_brainvision,
         ".amrk": _read_annotations_brainvision,
         ".txt": _read_annotations_txt,
+        ".seizures": _read_wfdb_annotations,
     }
     kwargs = {
         ".vmrk": {"sfreq": sfreq, "ignore_marker_types": ignore_marker_types},
@@ -1231,6 +1249,8 @@ def read_annotations(
             annotations = _read_annotations_fif(fid, tree)
     elif fname.name.startswith("events_") and fname.suffix == ".mat":
         annotations = _read_brainstorm_annotations(fname)
+    elif fmt == "wfdb":
+        annotations = _read_wfdb_annotations(fname, suffix=suffix)
     else:
         raise OSError(f'Unknown annotation file format "{fname}"')
 
@@ -1513,6 +1533,20 @@ def _check_event_description(event_desc, events):
     return event_desc
 
 
+def _read_wfdb_annotations(fname, suffix=None, sfreq="auto"):
+    """Read annotations from wfdb format."""
+    if suffix is None:
+        suffix = fname.suffix[1:]
+    wfdb = _check_wfdb_installed(strict=True)
+    anno = wfdb.io.rdann(fname.stem, extension=suffix)
+    anno_dict = anno.__dict__
+    sfreq = anno_dict["fs"] if sfreq == "auto" else sfreq
+    onset = anno_dict["sample"] / sfreq
+    duration = np.zeros_like(onset)
+    description = anno_dict["symbol"]
+    return Annotations(onset, duration, description)
+
+
 @verbose
 def events_from_annotations(
     raw,

mne/tests/test_annotations.py

@@ -51,6 +51,7 @@
 first_samps = pytest.mark.parametrize("first_samp", (0, 10000))
 edf_reduced = data_path / "EDF" / "test_reduced.edf"
 edf_annot_only = data_path / "EDF" / "SC4001EC-Hypnogram.edf"
+wfdb_dir = data_path / "wfbd"
 
 
 needs_pandas = pytest.mark.skipif(not check_version("pandas"), reason="Needs pandas")
@@ -416,6 +417,16 @@ def test_read_edf_annotations(fname, n_annot):
     assert len(annot) == n_annot
 
 
+@testing.requires_testing_data
+def test_read_wfdb_annotations():
+    """Test reading WFDB annotations."""
+    fname = wfdb_dir / "chb06_04.edf.seizures"
+    annot = read_annotations(fname)
+    # Verify that [ and ] are in description
+    assert "[" in annot.description
+    assert "]" in annot.description
+
+
 @first_samps
 def test_raw_reject(first_samp):
     """Test raw data getter with annotation reject."""

mne/utils/__init__.pyi

@@ -48,6 +48,7 @@ __all__ = [
     "_check_preload",
     "_check_pybv_installed",
     "_check_pymatreader_installed",
+    "_check_wfdb_installed",
     "_check_qt_version",
     "_check_range",
     "_check_rank",
@@ -254,6 +255,7 @@ from .check import (
     _check_stc_units,
     _check_subject,
     _check_time_format,
+    _check_wfdb_installed,
     _ensure_events,
     _ensure_int,
     _import_h5io_funcs,

mne/utils/check.py

@@ -451,6 +451,11 @@ def _check_edfio_installed(strict=True):
     return _soft_import("edfio", "exporting to EDF", strict=strict)
 
 
+def _check_wfdb_installed(strict=True):
+    """Aux function."""
+    return _soft_import("wfdb", "MIT WFDB IO", strict=strict)
+
+
 def _check_pybv_installed(strict=True):
     """Aux function."""
     return _soft_import("pybv", "exporting to BrainVision", strict=strict)

pyproject.toml

@@ -132,6 +132,7 @@ full-no-qt = [
   "trame-vtk",
   "trame-vuetify",
   "vtk >= 9.2",
+  "wfdb",
   "xlrd",
 ]
 full-pyqt6 = ["mne[full]"]