accusleepy 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

accusleepy/classification.py

@@ -11,6 +11,7 @@ from tqdm import trange
 
 import accusleepy.constants as c
 from accusleepy.brain_state_set import BrainStateSet
+from accusleepy.fileio import EMGFilter, Hyperparameters
 from accusleepy.models import SSANN
 from accusleepy.signal_processing import (
     create_eeg_emg_image,
@@ -19,11 +20,6 @@ from accusleepy.signal_processing import (
     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"""
@@ -61,11 +57,36 @@ def get_device():
     )
 
 
+def create_dataloader(
+    annotations_file: str,
+    img_dir: str,
+    hyperparameters: Hyperparameters,
+    shuffle: bool = True,
+) -> DataLoader:
+    """Create DataLoader for a dataset of training or calibration images
+
+    :param annotations_file: file with information on each training image
+    :param img_dir: training image location
+    :param hyperparameters: model training hyperparameters
+    :param shuffle: reshuffle data for every epoch
+    :return: DataLoader for the data
+
+    """
+    image_dataset = AccuSleepImageDataset(
+        annotations_file=annotations_file,
+        img_dir=img_dir,
+    )
+    return DataLoader(
+        image_dataset, batch_size=hyperparameters.batch_size, shuffle=shuffle
+    )
+
+
 def train_ssann(
     annotations_file: str,
     img_dir: str,
     mixture_weights: np.array,
     n_classes: int,
+    hyperparameters: Hyperparameters,
 ) -> SSANN:
     """Train a SSANN classification model for sleep scoring
 
@@ -73,13 +94,14 @@ def train_ssann(
     :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
+    :param hyperparameters: model training hyperparameters
     :return: trained Sleep Scoring Artificial Neural Network model
     """
-    training_data = AccuSleepImageDataset(
+    train_dataloader = create_dataloader(
         annotations_file=annotations_file,
         img_dir=img_dir,
+        hyperparameters=hyperparameters,
     )
-    train_dataloader = DataLoader(training_data, batch_size=BATCH_SIZE, shuffle=True)
 
     device = get_device()
     model = SSANN(n_classes=n_classes)
@@ -90,9 +112,13 @@ def train_ssann(
     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)
+    optimizer = optim.SGD(
+        model.parameters(),
+        lr=hyperparameters.learning_rate,
+        momentum=hyperparameters.momentum,
+    )
 
-    for _ in trange(TRAINING_EPOCHS):
+    for _ in trange(hyperparameters.training_epochs):
         for data in train_dataloader:
             inputs, labels = data
             (inputs, labels) = (inputs.to(device), labels.to(device))
@@ -115,6 +141,7 @@ def score_recording(
     epoch_length: int | float,
     epochs_per_img: int,
     brain_state_set: BrainStateSet,
+    emg_filter: EMGFilter,
 ) -> np.array:
     """Use classification model to get brain state labels for a recording
 
@@ -130,7 +157,8 @@ def score_recording(
     :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
+    :param emg_filter: EMG filter parameters
+    :return: brain state labels, confidence scores
     """
     # prepare model
     device = get_device()
@@ -138,7 +166,7 @@ def score_recording(
     model.eval()
 
     # create and scale eeg+emg spectrogram
-    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length)
+    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length, emg_filter)
     img = mixture_z_score_img(
         img,
         mixture_means=mixture_means,
@@ -158,10 +186,12 @@ def score_recording(
     # perform classification
     with torch.no_grad():
         outputs = model(images)
-        _, predicted = torch.max(outputs, 1)
+        logits, predicted = torch.max(outputs, 1)
 
     labels = brain_state_set.convert_class_to_digit(predicted.cpu().numpy())
-    return labels
+    confidence_scores = 1 / (1 + np.e ** (-logits.cpu().numpy()))
+
+    return labels, confidence_scores
 
 
 def example_real_time_scoring_function(
@@ -174,6 +204,7 @@ def example_real_time_scoring_function(
     epoch_length: int | float,
     epochs_per_img: int,
     brain_state_set: BrainStateSet,
+    emg_filter: EMGFilter,
 ) -> int:
     """Example function that could be used for real-time scoring
 
@@ -202,6 +233,7 @@ def example_real_time_scoring_function(
     :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
+    :param emg_filter: EMG filter parameters
     :return: brain state label
     """
     # prepare model
@@ -211,7 +243,7 @@ def example_real_time_scoring_function(
     model.eval()
 
     # create and scale eeg+emg spectrogram
-    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length)
+    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length, emg_filter)
     img = mixture_z_score_img(
         img,
         mixture_means=mixture_means,
@@ -242,6 +274,7 @@ def create_calibration_file(
     sampling_rate: int | float,
     epoch_length: int | float,
     brain_state_set: BrainStateSet,
+    emg_filter: EMGFilter,
 ) -> None:
     """Create file of calibration data for a subject
 
@@ -255,8 +288,9 @@ def create_calibration_file(
     :param sampling_rate: sampling rate, in Hz
     :param epoch_length: epoch length, in seconds
     :param brain_state_set: set of brain state options
+    :param emg_filter: EMG filter parameters
     """
-    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length)
+    img = create_eeg_emg_image(eeg, emg, sampling_rate, epoch_length, emg_filter)
     mixture_means, mixture_sds = get_mixture_values(
         img=img,
         labels=brain_state_set.convert_digit_to_class(labels),

accusleepy/config.json

@@ -19,5 +19,19 @@
             "frequency": 0.55
         }
     ],
-    "default_epoch_length": 2.5
+    "default_epoch_length": 2.5,
+    "default_overwrite_setting": false,
+    "save_confidence_setting": true,
+    "default_min_bout_length": 5.0,
+    "emg_filter": {
+        "order": 8,
+        "bp_lower": 20.0,
+        "bp_upper": 50.0
+    },
+    "hyperparameters": {
+        "batch_size": 64,
+        "learning_rate": 0.001,
+        "momentum": 0.9,
+        "training_epochs": 6
+    }
 }

accusleepy/constants.py

@@ -8,6 +8,7 @@ EEG_COL = "eeg"
 EMG_COL = "emg"
 # label file columns
 BRAIN_STATE_COL = "brain_state"
+CONFIDENCE_SCORE_COL = "confidence_score"
 
 
 # really don't change these
@@ -35,7 +36,33 @@ LABEL_COL = "label"
 # recording list file header:
 RECORDING_LIST_NAME = "recording_list"
 RECORDING_LIST_FILE_TYPE = ".json"
-# key for default epoch length in config
-DEFAULT_EPOCH_LENGTH_KEY = "default_epoch_length"
 # filename used to store info about training image datasets
 ANNOTATIONS_FILENAME = "annotations.csv"
+# filename for annotation file for the calibration set
+CALIBRATION_ANNOTATION_FILENAME = "calibration_set.csv"
+
+# config file keys
+# ui setting keys
+DEFAULT_EPOCH_LENGTH_KEY = "default_epoch_length"
+DEFAULT_CONFIDENCE_SETTING_KEY = "save_confidence_setting"
+DEFAULT_MIN_BOUT_LENGTH_KEY = "default_min_bout_length"
+DEFAULT_OVERWRITE_KEY = "default_overwrite_setting"
+# EMG filter parameters key
+EMG_FILTER_KEY = "emg_filter"
+# model training hyperparameters key
+HYPERPARAMETERS_KEY = "hyperparameters"
+
+# default values
+# default UI settings
+DEFAULT_MIN_BOUT_LENGTH = 5.0
+DEFAULT_CONFIDENCE_SETTING = True
+DEFAULT_OVERWRITE_SETTING = False
+# default EMG filter parameters (order, bandpass frequencies)
+DEFAULT_EMG_FILTER_ORDER = 8
+DEFAULT_EMG_BP_LOWER = 20
+DEFAULT_EMG_BP_UPPER = 50
+# default hyperparameters
+DEFAULT_BATCH_SIZE = 64
+DEFAULT_LEARNING_RATE = 1e-3
+DEFAULT_MOMENTUM = 0.9
+DEFAULT_TRAINING_EPOCHS = 6

accusleepy/fileio.py

@@ -7,17 +7,26 @@ import pandas as pd
 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,
-    DEFAULT_EPOCH_LENGTH_KEY,
-    EEG_COL,
-    EMG_COL,
-    MIXTURE_MEAN_COL,
-    MIXTURE_SD_COL,
-    RECORDING_LIST_NAME,
-    UNDEFINED_LABEL,
-)
+import accusleepy.constants as c
+
+
+@dataclass
+class EMGFilter:
+    """Convenience class for a EMG filter parameters"""
+
+    order: int  # filter order
+    bp_lower: int | float  # lower bandpass frequency
+    bp_upper: int | float  # upper bandpass frequency
+
+
+@dataclass
+class Hyperparameters:
+    """Convenience class for model training hyperparameters"""
+
+    batch_size: int
+    learning_rate: float
+    momentum: float
+    training_epochs: int
 
 
 @dataclass
@@ -39,8 +48,8 @@ def load_calibration_file(filename: str) -> (np.array, np.array):
     :return: mixture means and SDs
     """
     df = pd.read_csv(filename)
-    mixture_means = df[MIXTURE_MEAN_COL].values
-    mixture_sds = df[MIXTURE_SD_COL].values
+    mixture_means = df[c.MIXTURE_MEAN_COL].values
+    mixture_sds = df[c.MIXTURE_SD_COL].values
     return mixture_means, mixture_sds
 
 
@@ -67,56 +76,121 @@ def load_recording(filename: str) -> (np.array, np.array):
     :return: arrays of EEG and EMG data
     """
     df = load_csv_or_parquet(filename)
-    eeg = df[EEG_COL].values
-    emg = df[EMG_COL].values
+    eeg = df[c.EEG_COL].values
+    emg = df[c.EMG_COL].values
     return eeg, emg
 
 
-def load_labels(filename: str) -> np.array:
-    """Load file of brain state labels
+def load_labels(filename: str) -> (np.array, np.array):
+    """Load file of brain state labels and confidence scores
 
     :param filename: filename
-    :return: array of brain state labels
+    :return: array of brain state labels and, optionally, array of confidence scores
     """
     df = load_csv_or_parquet(filename)
-    return df[BRAIN_STATE_COL].values
+    if c.CONFIDENCE_SCORE_COL in df.columns:
+        return df[c.BRAIN_STATE_COL].values, df[c.CONFIDENCE_SCORE_COL].values
+    else:
+        return df[c.BRAIN_STATE_COL].values, None
 
 
-def save_labels(labels: np.array, filename: str) -> None:
+def save_labels(
+    labels: np.array, filename: str, confidence_scores: np.array = None
+) -> None:
     """Save brain state labels to file
 
     :param labels: brain state labels
     :param filename: filename
+    :param confidence_scores: optional confidence scores
     """
-    pd.DataFrame({BRAIN_STATE_COL: labels}).to_csv(filename, index=False)
+    if confidence_scores is not None:
+        pd.DataFrame(
+            {c.BRAIN_STATE_COL: labels, c.CONFIDENCE_SCORE_COL: confidence_scores}
+        ).to_csv(filename, index=False)
+    else:
+        pd.DataFrame({c.BRAIN_STATE_COL: labels}).to_csv(filename, index=False)
 
 
-def load_config() -> tuple[BrainStateSet, int | float]:
+def load_config() -> tuple[
+    BrainStateSet, int | float, bool, bool, int | float, EMGFilter, Hyperparameters
+]:
     """Load configuration file with brain state options
 
-    :return: set of brain state options and default epoch length
+    :return: set of brain state options,
+        default epoch length,
+        default overwrite setting,
+        default confidence score output setting,
+        default minimum bout length,
+        EMG filter parameters,
+        model training hyperparameters
     """
     with open(
-        os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG_FILE), "r"
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), c.CONFIG_FILE), "r"
     ) as f:
         data = json.load(f)
-    return BrainStateSet(
-        [BrainState(**b) for b in data[BRAIN_STATES_KEY]], UNDEFINED_LABEL
-    ), data[DEFAULT_EPOCH_LENGTH_KEY]
+
+    return (
+        BrainStateSet(
+            [BrainState(**b) for b in data[BRAIN_STATES_KEY]], c.UNDEFINED_LABEL
+        ),
+        data[c.DEFAULT_EPOCH_LENGTH_KEY],
+        data.get(c.DEFAULT_OVERWRITE_KEY, c.DEFAULT_OVERWRITE_SETTING),
+        data.get(c.DEFAULT_CONFIDENCE_SETTING_KEY, c.DEFAULT_CONFIDENCE_SETTING),
+        data.get(c.DEFAULT_MIN_BOUT_LENGTH_KEY, c.DEFAULT_MIN_BOUT_LENGTH),
+        EMGFilter(
+            **data.get(
+                c.EMG_FILTER_KEY,
+                {
+                    "order": c.DEFAULT_EMG_FILTER_ORDER,
+                    "bp_lower": c.DEFAULT_EMG_BP_LOWER,
+                    "bp_upper": c.DEFAULT_EMG_BP_UPPER,
+                },
+            )
+        ),
+        Hyperparameters(
+            **data.get(
+                c.HYPERPARAMETERS_KEY,
+                {
+                    "batch_size": c.DEFAULT_BATCH_SIZE,
+                    "learning_rate": c.DEFAULT_LEARNING_RATE,
+                    "momentum": c.DEFAULT_MOMENTUM,
+                    "training_epochs": c.DEFAULT_TRAINING_EPOCHS,
+                },
+            )
+        ),
+    )
 
 
 def save_config(
-    brain_state_set: BrainStateSet, default_epoch_length: int | float
+    brain_state_set: BrainStateSet,
+    default_epoch_length: int | float,
+    overwrite_setting: bool,
+    save_confidence_setting: bool,
+    min_bout_length: int | float,
+    emg_filter: EMGFilter,
+    hyperparameters: Hyperparameters,
 ) -> None:
     """Save configuration of brain state options to json file
 
     :param brain_state_set: set of brain state options
-    :param default_epoch_length: epoch length to use when the GUI starts
+    :param default_epoch_length: default epoch length
+    :param save_confidence_setting: default setting for
+        saving confidence scores
+    :param emg_filter: EMG filter parameters
+    :param min_bout_length: default minimum bout length
+    :param overwrite_setting: default setting for overwriting
+        existing labels
+    :param hyperparameters: model training hyperparameters
     """
     output_dict = brain_state_set.to_output_dict()
-    output_dict.update({DEFAULT_EPOCH_LENGTH_KEY: default_epoch_length})
+    output_dict.update({c.DEFAULT_EPOCH_LENGTH_KEY: default_epoch_length})
+    output_dict.update({c.DEFAULT_OVERWRITE_KEY: overwrite_setting})
+    output_dict.update({c.DEFAULT_CONFIDENCE_SETTING_KEY: save_confidence_setting})
+    output_dict.update({c.DEFAULT_MIN_BOUT_LENGTH_KEY: min_bout_length})
+    output_dict.update({c.EMG_FILTER_KEY: emg_filter.__dict__})
+    output_dict.update({c.HYPERPARAMETERS_KEY: hyperparameters.__dict__})
     with open(
-        os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG_FILE), "w"
+        os.path.join(os.path.dirname(os.path.abspath(__file__)), c.CONFIG_FILE), "w"
     ) as f:
         json.dump(output_dict, f, indent=4)
 
@@ -129,7 +203,7 @@ def load_recording_list(filename: str) -> list[Recording]:
     """
     with open(filename, "r") as f:
         data = json.load(f)
-    recording_list = [Recording(**r) for r in data[RECORDING_LIST_NAME]]
+    recording_list = [Recording(**r) for r in data[c.RECORDING_LIST_NAME]]
     for i, r in enumerate(recording_list):
         r.name = i + 1
     return recording_list
@@ -142,7 +216,7 @@ def save_recording_list(filename: str, recordings: list[Recording]) -> None:
     :param recordings: list of recordings to export
     """
     recording_dict = {
-        RECORDING_LIST_NAME: [
+        c.RECORDING_LIST_NAME: [
             {
                 "recording_file": r.recording_file,
                 "label_file": r.label_file,