As part of an university project I am currently working on an eeg-biosignal classifier. While the project itself doesn't really focus on design ("anything that works") I am trying to learn and apply some best practices.
Some Background
The experimental setup is essentially as follows: A subject sits in front of a screen on which an optical stimulus is shown. EEG measurements are used to observe the triggered response to this stimulus. Each stimulus is assigned to an event code. The goal is to build a neural network that can classify a stimulus for a particular EEG measurement.
Thus, we have several datasets in the general data format (.gdf) containing eeg measurements labeled with event codes.
My goal now is to load these datasets and extract the information I need to create a CustomDataset suitable for PyTorch.
This is my Dataset class:
import torch
from torch.utils.data.dataset import Dataset as TorchDataset
class Dataset(TorchDataset):
class_mapping: dict[str, int]
labels: torch.Tensor
samples: torch.Tensor
def __init__(
self,
class_mapping: dict[str, int],
labels: torch.Tensor,
samples: torch.Tensor
) -> None:
self.class_mapping = class_mapping
self.labels = labels
self.samples = samples
def __len__(self) -> int:
return len(self.labels)
def __getitem__(self, item) -> tuple[torch.Tensor, torch.Tensor]:
return self.labels[item], self.samples[item]
The class_mapping variable is a dict that maps the event code labels as key to class labels used in the PyTorch framework.
The process of obtaining the samples and labels is as follows:
- load raw files
- extract data fragments in a specified time window around the event markers (aka. epochs)
- extract the data and labels from the epochs objects
As these steps are quite something to do I decided to put them into a simple factory:
from datasets.loading import LoadingStrategy
from datasets.data import Dataset
from typing import Iterable, Optional, Callable
from pathlib import Path
import numpy as np
import torch
import mne
from mne.io import Raw
from mne import Epochs
class DatasetFactoryV2:
@staticmethod
def create_dataset(
path_to_raw_files: Iterable[Path],
file_loading_strategy: LoadingStrategy,
considered_events: Iterable[str],
time_window_min: float,
time_window_max: float
) -> Dataset:
class_mapping = _get_class_mapping(considered_events)
raw_files = file_loading_strategy(path_to_raw_files)
labels, samples = _get_labels_and_samples(
raw_files=raw_files,
class_mapping=class_mapping,
time_window_min=time_window_min,
time_window_max=time_window_max
)
return Dataset(
class_mapping=class_mapping,
labels=labels,
samples=samples
)
def _get_class_mapping(events_to_consider: Iterable[str]) -> dict[str, int]:
return {event: idx for idx, event in enumerate(events_to_consider)}
def _get_labels_and_samples(
raw_files: Iterable[Raw],
class_mapping: dict[str, int],
time_window_min: float,
time_window_max: float,
) -> tuple[torch.Tensor, torch.Tensor]:
labels = []
samples = []
for raw in raw_files:
events_from_annot, event_dict = mne.events_from_annotations(raw=raw, event_id=class_mapping)
channel_picks = mne.pick_types(raw.info, eeg=True, exclude='bads')
epochs = Epochs(
raw=raw,
events=events_from_annot,
on_missing='warn',
tmin=time_window_min,
tmax=time_window_max,
baseline=(None, None),
event_repeated='drop',
preload=True,
picks=channel_picks,
reject_by_annotation=True # exclude Epochs fully/partially overlapping with 'bad'-labelled data spans
)
labels.extend(epochs.events[:, -1])
samples.extend(epochs.get_data())
labels = torch.tensor(data=np.array(labels), dtype=torch.int64)
samples = torch.tensor(data=np.array(samples), dtype=torch.float32)
return labels, samples
And only for the completeness the LoadingStrategy:
from typing import Protocol, Iterable
from pathlib import Path
from mne.io import Raw, read_raw_gdf
class LoadingStrategy(Protocol):
def __call__(self, path_to_files: Iterable[Path]) -> list[Raw]:
...
class GDFLoadingStrategy:
def __call__(self, path_to_files: Iterable[Path]) -> list[Raw]:
return [read_raw_gdf(str(filepath)) for filepath in path_to_files]
This is basically the only part I expect to change in the scope of that educational project. So I decided to go with a strategy pattern that can be extended for any of the many possible data formats.
Problem / Question
In several lectures I have now heard that coupling is a bad thing and should be avoided as much as possible. If I recognize it correctly, I must admit that my DatasetFactory is strongly coupled with the creation of mne.io.Raw and mne.Epochs and a number of different functions from this library as well.
Thinking in terms of tests for example it seems quite obvious that I will have a hard time testing without having real data to load.
This brings me to my question:
How could I decouple a library if I need a lot of logic from it?
Should I really create e.g. Protocols that mimic the mne.io.Raw and mne.Epochs classes and use them instead of the mne ones? It would certainly make mocking and thus testing easier, but it seems like a lot of work relative to the benefits.
Thanks for any help!
Julian
