accusleepy 0.1.0__py3-none-any.whl

accusleepy/__main__.py

@@ -0,0 +1,4 @@
+from accusleepy.gui.main import run_primary_window
+
+if __name__ == "__main__":
+    run_primary_window()

accusleepy/brain_state_set.py

@@ -0,0 +1,89 @@
+from dataclasses import dataclass
+
+import numpy as np
+
+BRAIN_STATES_KEY = "brain_states"
+
+
+@dataclass
+class BrainState:
+    """Convenience class for a brain state and its attributes"""
+
+    name: str  # friendly name
+    digit: int  # number 0-9 - used as keyboard shortcut and in label files
+    is_scored: bool  # whether a classification model should score this state
+    frequency: int | float  # typical relative frequency, between 0 and 1
+
+
+class BrainStateSet:
+    def __init__(self, brain_states: list[BrainState], undefined_label: int):
+        """Initialize set of brain states
+
+        :param brain_states: list of BrainState objects
+        :param undefined_label: label for undefined epochs
+        """
+        self.brain_states = brain_states
+
+        # The user can choose any subset of the digits 0-9 to represent
+        # brain states, but not all of them are necessarily intended to be
+        # scored by a classifier, and pytorch requires that all input
+        # labels are in the 0-n range for training and inference.
+        # So, we have to have a distinction between "brain states" (as
+        # represented in label files and keyboard inputs) and "classes"
+        # (AccuSleep's internal representation).
+
+        # map digits to classes, and vice versa
+        self.digit_to_class = {undefined_label: None}
+        self.class_to_digit = dict()
+        # relative frequencies of each class
+        self.mixture_weights = list()
+
+        i = 0
+        for brain_state in self.brain_states:
+            if brain_state.digit == undefined_label:
+                raise Exception(
+                    f"Digit for {brain_state.name} matches 'undefined' label"
+                )
+            if brain_state.is_scored:
+                self.digit_to_class[brain_state.digit] = i
+                self.class_to_digit[i] = brain_state.digit
+                self.mixture_weights.append(brain_state.frequency)
+                i += 1
+            else:
+                self.digit_to_class[brain_state.digit] = None
+
+        self.n_classes = i
+
+        self.mixture_weights = np.array(self.mixture_weights)
+        if np.sum(self.mixture_weights) != 1:
+            raise Exception("Typical frequencies for scored brain states must sum to 1")
+
+    def convert_digit_to_class(self, digits: np.array) -> np.array:
+        """Convert array of digits to their corresponding classes
+
+        :param digits: array of digits
+        :return: array of classes
+        """
+        return np.array([self.digit_to_class[i] for i in digits])
+
+    def convert_class_to_digit(self, classes: np.array) -> np.array:
+        """Convert array of classes to their corresponding digits
+
+        :param classes: array of classes
+        :return: array of digits
+        """
+        return np.array([self.class_to_digit[i] for i in classes])
+
+    def to_output_dict(self) -> dict:
+        """Return dictionary of brain states"""
+        return {
+            BRAIN_STATES_KEY: [
+                {
+                    "name": b.name,
+                    "digit": b.digit,
+                    "is_scored": b.is_scored,
+                    "frequency": b.frequency,
+                }
+                for b in self.brain_states
+            ]
+        }

accusleepy/classification.py

@@ -0,0 +1,267 @@
+import os
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.optim as optim
+from torch import nn
+from torch.utils.data import DataLoader, Dataset
+from torchvision.io import read_image
+from tqdm import trange
+
+import accusleepy.constants as c
+from accusleepy.brain_state_set import BrainStateSet
+from accusleepy.models import SSANN
+from accusleepy.signal_processing import (
+    create_eeg_emg_image,
+    format_img,
+    get_mixture_values,
+    mixture_z_score_img,
+)
+
+BATCH_SIZE = 64
+LEARNING_RATE = 1e-3
+MOMENTUM = 0.9
+TRAINING_EPOCHS = 6
+
+
+class AccuSleepImageDataset(Dataset):
+    """Dataset for loading AccuSleep training images"""
+
+    def __init__(
+        self, annotations_file, img_dir, transform=None, target_transform=None
+    ):
+        self.img_labels = pd.read_csv(annotations_file)
+        self.img_dir = img_dir
+        self.transform = transform
+        self.target_transform = target_transform
+
+    def __len__(self):
+        return len(self.img_labels)
+
+    def __getitem__(self, idx):
+        img_path = str(
+            os.path.join(self.img_dir, self.img_labels.at[idx, c.FILENAME_COL])
+        )
+        image = read_image(img_path)
+        label = self.img_labels.at[idx, c.LABEL_COL]
+        if self.transform:
+            image = self.transform(image)
+        if self.target_transform:
+            label = self.target_transform(label)
+        return image, label
+
+
+def get_device():
+    """Get accelerator, if one is available"""
+    return (
+        torch.accelerator.current_accelerator().type
+        if torch.accelerator.is_available()
+        else "cpu"
+    )
+
+
+def train_model(
+    annotations_file: str,
+    img_dir: str,
+    mixture_weights: np.array,
+    n_classes: int,
+) -> SSANN:
+    """Train a classification model for sleep scoring
+
+    :param annotations_file: file with information on each training image
+    :param img_dir: training image location
+    :param mixture_weights: typical relative frequencies of brain states
+    :param n_classes: number of classes the model will learn
+    :return: trained Sleep Scoring Artificial Neural Network model
+    """
+    training_data = AccuSleepImageDataset(
+        annotations_file=annotations_file,
+        img_dir=img_dir,
+    )
+    train_dataloader = DataLoader(training_data, batch_size=BATCH_SIZE, shuffle=True)
+
+    device = get_device()
+    model = SSANN(n_classes=n_classes)
+    model.to(device)
+    model.train()
+
+    # correct for class imbalance
+    weight = torch.tensor((mixture_weights**-1).astype("float32")).to(device)
+
+    criterion = nn.CrossEntropyLoss(weight=weight)
+    optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, momentum=MOMENTUM)
+
+    for _ in trange(TRAINING_EPOCHS):
+        for data in train_dataloader:
+            inputs, labels = data
+            (inputs, labels) = (inputs.to(device), labels.to(device))
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, labels)
+            loss.backward()
+            optimizer.step()
+
+    return model
+
+
+def score_recording(
+    model: SSANN,
+    eeg: np.array,
+    emg: np.array,
+    mixture_means: np.array,
+    mixture_sds: np.array,
+    sampling_rate: int | float,
+    epoch_length: int | float,
+    epochs_per_img: int,
+    brain_state_set: BrainStateSet,
+) -> np.array:
+    """Use classification model to get brain state labels for a recording
+
+    This assumes signals have been preprocessed to contain an integer
+    number of epochs.
+
+    :param model: classification model
+    :param eeg: EEG signal
+    :param emg: EMG signal
+    :param mixture_means: mixture means, for calibration
+    :param mixture_sds: mixture standard deviations, for calibration
+    :param sampling_rate: sampling rate, in Hz
+    :param epoch_length: epoch length, in seconds
+    :param epochs_per_img: number of epochs for the model to consider
+    :param brain_state_set: set of brain state options
+    :return: brain state labels
+    """
+    # prepare model
+    device = get_device()
+    model = model.to(device)
+    model.eval()
+
+    # create and scale eeg+emg spectrogram
+    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length)
+    img = mixture_z_score_img(
+        img,
+        mixture_means=mixture_means,
+        mixture_sds=mixture_sds,
+        brain_state_set=brain_state_set,
+    )
+    img = format_img(img=img, epochs_per_img=epochs_per_img, add_padding=True)
+
+    # create dataset for inference
+    images = []
+    for i in range(img.shape[1] - epochs_per_img + 1):
+        images.append(img[:, i : (i + epochs_per_img)].astype("float32"))
+    images = torch.from_numpy(np.array(images))
+    images = images[:, None, :, :]  # add channel
+    images = images.to(device)
+
+    # perform classification
+    with torch.no_grad():
+        outputs = model(images)
+        _, predicted = torch.max(outputs, 1)
+
+    labels = brain_state_set.convert_class_to_digit(predicted.cpu().numpy())
+    return labels
+
+
+def example_real_time_scoring_function(
+    model: SSANN,
+    eeg: np.array,
+    emg: np.array,
+    mixture_means: np.array,
+    mixture_sds: np.array,
+    sampling_rate: int | float,
+    epoch_length: int | float,
+    epochs_per_img: int,
+    brain_state_set: BrainStateSet,
+) -> int:
+    """Example function that could be used for real-time scoring
+
+    This function demonstrates how you could use a model trained in
+    "real-time" mode (current epoch on the right side of each image)
+    to score incoming data. By passing a segment of EEG/EMG data
+    into this function, the most recent epoch will be scored. For
+    example, if the model expects 9 epochs worth of data and the
+    epoch length is 5 seconds, you would pass in 45 seconds of data
+    and would obtain the brain state of the most recent 5 seconds.
+
+    Note:
+    - The EEG and EMG signals must have length equal to
+        sampling_rate * epoch_length * <number of epochs per image>.
+    - The number of samples per epoch must be an integer.
+    - This is just a demonstration, you should customize this for
+        your application and there are probably ways to make it
+        run faster.
+
+    :param model: classification model
+    :param eeg: EEG signal
+    :param emg: EMG signal
+    :param mixture_means: mixture means, for calibration
+    :param mixture_sds: mixture standard deviations, for calibration
+    :param sampling_rate: sampling rate, in Hz
+    :param epoch_length: epoch length, in seconds
+    :param epochs_per_img: number of epochs shown to the model at once
+    :param brain_state_set: set of brain state options
+    :return: brain state label
+    """
+    # prepare model
+    # this could be done outside the function
+    device = get_device()
+    model = model.to(device)
+    model.eval()
+
+    # create and scale eeg+emg spectrogram
+    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length)
+    img = mixture_z_score_img(
+        img,
+        mixture_means=mixture_means,
+        mixture_sds=mixture_sds,
+        brain_state_set=brain_state_set,
+    )
+    img = format_img(img=img, epochs_per_img=epochs_per_img, add_padding=False)
+
+    # create dataset for inference
+    images = torch.from_numpy(np.array([img.astype("float32")]))
+    images = images[:, None, :, :]  # add channel
+    images = images.to(device)
+
+    # perform classification
+    with torch.no_grad():
+        outputs = model(images)
+        _, predicted = torch.max(outputs, 1)
+
+    label = int(brain_state_set.convert_class_to_digit(predicted.cpu().numpy())[0])
+    return label
+
+
+def create_calibration_file(
+    filename: str,
+    eeg: np.array,
+    emg: np.array,
+    labels: np.array,
+    sampling_rate: int | float,
+    epoch_length: int | float,
+    brain_state_set: BrainStateSet,
+) -> None:
+    """Create file of calibration data for a subject
+
+    This assumes signals have been preprocessed to contain an integer
+    number of epochs.
+
+    :param filename: filename for the calibration file
+    :param eeg: EEG signal
+    :param emg: EMG signal
+    :param labels: brain state labels, as digits
+    :param sampling_rate: sampling rate, in Hz
+    :param epoch_length: epoch length, in seconds
+    :param brain_state_set: set of brain state options
+    """
+    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length)
+    mixture_means, mixture_sds = get_mixture_values(
+        img=img,
+        labels=brain_state_set.convert_digit_to_class(labels),
+        brain_state_set=brain_state_set,
+    )
+    pd.DataFrame(
+        {c.MIXTURE_MEAN_COL: mixture_means, c.MIXTURE_SD_COL: mixture_sds}
+    ).to_csv(filename, index=False)

accusleepy/config.json

@@ -0,0 +1,22 @@
+{
+    "brain_states": [
+        {
+            "name": "REM",
+            "digit": 1,
+            "is_scored": true,
+            "frequency": 0.1
+        },
+        {
+            "name": "Wake",
+            "digit": 2,
+            "is_scored": true,
+            "frequency": 0.35
+        },
+        {
+            "name": "NREM",
+            "digit": 3,
+            "is_scored": true,
+            "frequency": 0.55
+        }
+    ]
+}

accusleepy/constants.py

@@ -0,0 +1,37 @@
+# probably don't change these unless you really need to
+UNDEFINED_LABEL = -1  # can't be the same as a brain state's digit, must be an integer
+# calibration file columns
+MIXTURE_MEAN_COL = "mixture_mean"
+MIXTURE_SD_COL = "mixture_sd"
+# recording file columns
+EEG_COL = "eeg"
+EMG_COL = "emg"
+# label file columns
+BRAIN_STATE_COL = "brain_state"
+
+
+# really don't change these
+# config file location
+CONFIG_FILE = "config.json"
+# number of times to include the EMG power in a training image
+EMG_COPIES = 9
+# minimum spectrogram window length, in seconds
+MIN_WINDOW_LEN = 5
+# frequency above which to downsample EEG spectrograms
+DOWNSAMPLING_START_FREQ = 20
+# upper frequency cutoff for EEG spectrograms
+UPPER_FREQ = 50
+# classification model types
+DEFAULT_MODEL_TYPE = "default"  # current epoch is centered
+REAL_TIME_MODEL_TYPE = "real-time"  # current epoch on the right
+# valid filetypes
+RECORDING_FILE_TYPES = [".parquet", ".csv"]
+LABEL_FILE_TYPE = ".csv"
+CALIBRATION_FILE_TYPE = ".csv"
+MODEL_FILE_TYPE = ".pth"
+# annotation file columns
+FILENAME_COL = "filename"
+LABEL_COL = "label"
+# recording list file header:
+RECORDING_LIST_NAME = "recording_list"
+RECORDING_LIST_FILE_TYPE = ".json"

accusleepy/fileio.py

@@ -0,0 +1,201 @@
+import json
+import os
+from dataclasses import dataclass
+
+import numpy as np
+import pandas as pd
+import torch
+from PySide6.QtWidgets import QListWidgetItem
+
+from accusleepy.brain_state_set import BRAIN_STATES_KEY, BrainState, BrainStateSet
+from accusleepy.constants import (
+    BRAIN_STATE_COL,
+    CONFIG_FILE,
+    EEG_COL,
+    EMG_COL,
+    MIXTURE_MEAN_COL,
+    MIXTURE_SD_COL,
+    RECORDING_LIST_NAME,
+    UNDEFINED_LABEL,
+)
+from accusleepy.models import SSANN
+
+
+@dataclass
+class Recording:
+    """Store information about a recording"""
+
+    name: int = 1  # name to show in the GUI
+    recording_file: str = ""  # path to recording file
+    label_file: str = ""  # path to label file
+    calibration_file: str = ""  # path to calibration file
+    sampling_rate: int | float = 0.0  # sampling rate, in Hz
+    widget: QListWidgetItem = None  # list item widget shown in the GUI
+
+
+def load_calibration_file(filename: str) -> (np.array, np.array):
+    """Load a calibration file
+
+    :param filename: filename
+    :return: mixture means and SDs
+    """
+    df = pd.read_csv(filename)
+    mixture_means = df[MIXTURE_MEAN_COL].values
+    mixture_sds = df[MIXTURE_SD_COL].values
+    return mixture_means, mixture_sds
+
+
+def save_model(
+    model: SSANN,
+    filename: str,
+    epoch_length: int | float,
+    epochs_per_img: int,
+    model_type: str,
+    brain_state_set: BrainStateSet,
+) -> None:
+    """Save classification model and its metadata
+
+    :param model: classification model
+    :param epoch_length: epoch length used when training the model
+    :param epochs_per_img: number of epochs in each model input
+    :param model_type: default or real-time
+    :param brain_state_set: set of brain state options
+    :param filename: filename
+    """
+    state_dict = model.state_dict()
+    state_dict.update({"epoch_length": epoch_length})
+    state_dict.update({"epochs_per_img": epochs_per_img})
+    state_dict.update({"model_type": model_type})
+    state_dict.update(
+        {BRAIN_STATES_KEY: brain_state_set.to_output_dict()[BRAIN_STATES_KEY]}
+    )
+
+    torch.save(state_dict, filename)
+
+
+def load_model(filename: str) -> tuple[SSANN, int | float, int, str, dict]:
+    """Load classification model and its metadata
+
+    :param filename: filename
+    :return: model, epoch length used when training the model,
+        number of epochs in each model input, model type
+        (default or real-time), set of brain state options
+        used when training the model
+    """
+    state_dict = torch.load(filename, weights_only=True)
+    epoch_length = state_dict.pop("epoch_length")
+    epochs_per_img = state_dict.pop("epochs_per_img")
+    model_type = state_dict.pop("model_type")
+    brain_states = state_dict.pop(BRAIN_STATES_KEY)
+    n_classes = len([b for b in brain_states if b["is_scored"]])
+
+    model = SSANN(n_classes=n_classes)
+    model.load_state_dict(state_dict)
+    return model, epoch_length, epochs_per_img, model_type, brain_states
+
+
+def load_csv_or_parquet(filename: str) -> pd.DataFrame:
+    """Load a csv or parquet file as a dataframe
+
+    :param filename: filename
+    :return: dataframe of file contents
+    """
+    extension = os.path.splitext(filename)[1]
+    if extension == ".csv":
+        df = pd.read_csv(filename)
+    elif extension == ".parquet":
+        df = pd.read_parquet(filename)
+    else:
+        raise Exception("file must be csv or parquet")
+    return df
+
+
+def load_recording(filename: str) -> (np.array, np.array):
+    """Load recording of EEG and EMG time series data
+
+    :param filename: filename
+    :return: arrays of EEG and EMG data
+    """
+    df = load_csv_or_parquet(filename)
+    eeg = df[EEG_COL].values
+    emg = df[EMG_COL].values
+    return eeg, emg
+
+
+def load_labels(filename: str) -> np.array:
+    """Load file of brain state labels
+
+    :param filename: filename
+    :return: array of brain state labels
+    """
+    df = load_csv_or_parquet(filename)
+    return df[BRAIN_STATE_COL].values
+
+
+def save_labels(labels: np.array, filename: str) -> None:
+    """Save brain state labels to file
+
+    :param labels: brain state labels
+    :param filename: filename
+    """
+    pd.DataFrame({BRAIN_STATE_COL: labels}).to_csv(filename, index=False)
+
+
+def load_config() -> BrainStateSet:
+    """Load configuration file with brain state options
+
+    :return: set of brain state options
+    """
+    with open(
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG_FILE), "r"
+    ) as f:
+        data = json.load(f)
+    return BrainStateSet(
+        [BrainState(**b) for b in data[BRAIN_STATES_KEY]], UNDEFINED_LABEL
+    )
+
+
+def save_config(brain_state_set: BrainStateSet) -> None:
+    """Save configuration of brain state options to json file
+
+    :param brain_state_set: set of brain state options
+    """
+    with open(
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG_FILE), "w"
+    ) as f:
+        json.dump(brain_state_set.to_output_dict(), f, indent=4)
+
+
+def load_recording_list(filename: str) -> list[Recording]:
+    """Load list of recordings from file
+
+    :param filename: filename of list of recordings
+    :return: list of recordings
+    """
+    with open(filename, "r") as f:
+        data = json.load(f)
+    recording_list = [Recording(**r) for r in data[RECORDING_LIST_NAME]]
+    for i, r in enumerate(recording_list):
+        r.name = i + 1
+    return recording_list
+
+
+def save_recording_list(filename: str, recordings: list[Recording]) -> None:
+    """Save list of recordings to file
+
+    :param filename: where to save the list
+    :param recordings: list of recordings to export
+    """
+    recording_dict = {
+        RECORDING_LIST_NAME: [
+            {
+                "recording_file": r.recording_file,
+                "label_file": r.label_file,
+                "calibration_file": r.calibration_file,
+                "sampling_rate": r.sampling_rate,
+            }
+            for r in recordings
+        ]
+    }
+    with open(filename, "w") as f:
+        json.dump(recording_dict, f, indent=4)