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

accusleepy/gui/main.py

@@ -36,8 +36,16 @@ from accusleepy.bouts import enforce_min_bout_length
 from accusleepy.brain_state_set import BRAIN_STATES_KEY, BrainState, BrainStateSet
 from accusleepy.constants import (
     ANNOTATIONS_FILENAME,
+    CALIBRATION_ANNOTATION_FILENAME,
     CALIBRATION_FILE_TYPE,
     DEFAULT_MODEL_TYPE,
+    DEFAULT_EMG_FILTER_ORDER,
+    DEFAULT_EMG_BP_LOWER,
+    DEFAULT_EMG_BP_UPPER,
+    DEFAULT_BATCH_SIZE,
+    DEFAULT_LEARNING_RATE,
+    DEFAULT_MOMENTUM,
+    DEFAULT_TRAINING_EPOCHS,
     LABEL_FILE_TYPE,
     MODEL_FILE_TYPE,
     REAL_TIME_MODEL_TYPE,
@@ -55,6 +63,8 @@ from accusleepy.fileio import (
     save_config,
     save_labels,
     save_recording_list,
+    EMGFilter,
+    Hyperparameters,
 )
 from accusleepy.gui.manual_scoring import ManualScoringWindow
 from accusleepy.gui.primary_window import Ui_PrimaryWindow
@@ -96,21 +106,31 @@ class AccuSleepWindow(QMainWindow):
         self.setWindowTitle("AccuSleePy")
 
         # fill in settings tab
-        self.brain_state_set, self.epoch_length = load_config()
+        (
+            self.brain_state_set,
+            self.epoch_length,
+            self.only_overwrite_undefined,
+            self.save_confidence_scores,
+            self.min_bout_length,
+            self.emg_filter,
+            self.hyperparameters,
+        ) = load_config()
+
         self.settings_widgets = None
         self.initialize_settings_tab()
 
         # initialize info about the recordings, classification data / settings
         self.ui.epoch_length_input.setValue(self.epoch_length)
+        self.ui.overwritecheckbox.setChecked(self.only_overwrite_undefined)
+        self.ui.save_confidence_checkbox.setChecked(self.save_confidence_scores)
+        self.ui.bout_length_input.setValue(self.min_bout_length)
         self.model = None
-        self.only_overwrite_undefined = False
-        self.min_bout_length = 5
 
         # initialize model training variables
         self.training_epochs_per_img = 9
         self.delete_training_images = True
-        self.training_image_dir = ""
         self.model_type = DEFAULT_MODEL_TYPE
+        self.calibrate_trained_model = True
 
         # metadata for the currently loaded classification model
         self.model_epoch_length = None
@@ -166,16 +186,25 @@ class AccuSleepWindow(QMainWindow):
         self.ui.load_model_button.clicked.connect(partial(self.load_model, None))
         self.ui.score_all_button.clicked.connect(self.score_all)
         self.ui.overwritecheckbox.stateChanged.connect(self.update_overwrite_policy)
+        self.ui.save_confidence_checkbox.stateChanged.connect(
+            self.update_confidence_policy
+        )
         self.ui.bout_length_input.valueChanged.connect(self.update_min_bout_length)
         self.ui.user_manual_button.clicked.connect(self.show_user_manual)
         self.ui.image_number_input.valueChanged.connect(self.update_epochs_per_img)
         self.ui.delete_image_box.stateChanged.connect(self.update_image_deletion)
-        self.ui.training_folder_button.clicked.connect(self.set_training_folder)
+        self.ui.calibrate_checkbox.stateChanged.connect(
+            self.update_training_calibration
+        )
         self.ui.train_model_button.clicked.connect(self.train_model)
         self.ui.save_config_button.clicked.connect(self.save_brain_state_config)
         self.ui.export_button.clicked.connect(self.export_recording_list)
         self.ui.import_button.clicked.connect(self.import_recording_list)
         self.ui.default_type_button.toggled.connect(self.model_type_radio_buttons)
+        self.ui.reset_emg_params_button.clicked.connect(self.reset_emg_filter_settings)
+        self.ui.reset_hyperparams_button.clicked.connect(
+            self.reset_hyperparams_settings
+        )
 
         # user input: drag and drop
         self.ui.recording_file_label.installEventFilter(self)
@@ -294,11 +323,12 @@ class AccuSleepWindow(QMainWindow):
                 )
             )
             return
-        if self.training_image_dir == "":
-            self.show_message(
-                ("ERROR: no output location selected for training images.")
-            )
-            return
+
+        # determine fraction of training data to use for calibration
+        if self.calibrate_trained_model:
+            calibration_fraction = self.ui.calibration_spinbox.value() / 100
+        else:
+            calibration_fraction = 0
 
         # check some inputs for each recording
         for recording_index in range(len(self.recordings)):
@@ -320,9 +350,10 @@ class AccuSleepWindow(QMainWindow):
             return
         model_filename = os.path.normpath(model_filename)
 
-        # create (probably temporary) image folder
+        # create (probably temporary) image folder in
+        # the same folder as the trained model
         temp_image_dir = os.path.join(
-            self.training_image_dir,
+            os.path.dirname(model_filename),
             "images_" + datetime.datetime.now().strftime("%Y%m%d%H%M"),
         )
 
@@ -334,7 +365,12 @@ class AccuSleepWindow(QMainWindow):
 
         # create training images
         self.show_message("Training, please wait. See console for progress updates.")
-        self.show_message((f"Creating training images in {temp_image_dir}"))
+        if not self.delete_training_images:
+            self.show_message((f"Creating training images in {temp_image_dir}"))
+        else:
+            self.show_message(
+                (f"Creating temporary folder of training images: {temp_image_dir}")
+            )
         self.ui.message_area.repaint()
         QApplication.processEvents()
         print("Creating training images")
@@ -345,14 +381,17 @@ class AccuSleepWindow(QMainWindow):
             epochs_per_img=self.training_epochs_per_img,
             brain_state_set=self.brain_state_set,
             model_type=self.model_type,
+            calibration_fraction=calibration_fraction,
+            emg_filter=self.emg_filter,
         )
         if len(failed_recordings) > 0:
             if len(failed_recordings) == len(self.recordings):
                 self.show_message("ERROR: no recordings were valid!")
+                return
             else:
                 self.show_message(
                     (
-                        "WARNING: the following recordings could not be"
+                        "WARNING: the following recordings could not be "
                         "loaded and will not be used for training: "
                         f"{', '.join([str(r) for r in failed_recordings])}"
                     )
@@ -363,16 +402,32 @@ class AccuSleepWindow(QMainWindow):
         self.ui.message_area.repaint()
         QApplication.processEvents()
         print("Training model")
-        from accusleepy.classification import train_ssann
+        from accusleepy.classification import create_dataloader, train_ssann
         from accusleepy.models import save_model
+        from accusleepy.temperature_scaling import ModelWithTemperature
 
         model = train_ssann(
             annotations_file=os.path.join(temp_image_dir, ANNOTATIONS_FILENAME),
             img_dir=temp_image_dir,
             mixture_weights=self.brain_state_set.mixture_weights,
             n_classes=self.brain_state_set.n_classes,
+            hyperparameters=self.hyperparameters,
         )
 
+        # calibrate the model
+        if self.calibrate_trained_model:
+            calibration_annotation_file = os.path.join(
+                temp_image_dir, CALIBRATION_ANNOTATION_FILENAME
+            )
+            calibration_dataloader = create_dataloader(
+                annotations_file=calibration_annotation_file,
+                img_dir=temp_image_dir,
+                hyperparameters=self.hyperparameters,
+            )
+            model = ModelWithTemperature(model)
+            print("Calibrating model")
+            model.set_temperature(calibration_dataloader)
+
         # save model
         save_model(
             model=model,
@@ -381,29 +436,26 @@ class AccuSleepWindow(QMainWindow):
             epochs_per_img=self.training_epochs_per_img,
             model_type=self.model_type,
             brain_state_set=self.brain_state_set,
+            is_calibrated=self.calibrate_trained_model,
         )
 
         # optionally delete images
         if self.delete_training_images:
+            print("Cleaning up training image folder")
             shutil.rmtree(temp_image_dir)
 
         self.show_message(f"Training complete. Saved model to {model_filename}")
         print("Training complete.")
 
-    def set_training_folder(self) -> None:
-        """Select location in which to create a folder for training images"""
-        training_folder_parent = QFileDialog.getExistingDirectory(
-            self, "Select directory for training images"
-        )
-        if training_folder_parent:
-            training_folder_parent = os.path.normpath(training_folder_parent)
-            self.training_image_dir = training_folder_parent
-            self.ui.image_folder_label.setText(training_folder_parent)
-
     def update_image_deletion(self) -> None:
         """Update choice of whether to delete images after training"""
         self.delete_training_images = self.ui.delete_image_box.isChecked()
 
+    def update_training_calibration(self) -> None:
+        """Update choice of whether to calibrate model after training"""
+        self.calibrate_trained_model = self.ui.calibrate_checkbox.isChecked()
+        self.ui.calibration_spinbox.setEnabled(self.calibrate_trained_model)
+
     def update_epochs_per_img(self, new_value) -> None:
         """Update number of epochs per image
 
@@ -491,7 +543,8 @@ class AccuSleepWindow(QMainWindow):
             label_file = self.recordings[recording_index].label_file
             if os.path.isfile(label_file):
                 try:
-                    existing_labels = load_labels(label_file)
+                    # ignore any existing confidence scores; they will all be overwritten
+                    existing_labels, _ = load_labels(label_file)
                 except Exception:
                     self.show_message(
                         (
@@ -544,7 +597,7 @@ class AccuSleepWindow(QMainWindow):
                 )
                 continue
 
-            labels = score_recording(
+            labels, confidence_scores = score_recording(
                 model=self.model,
                 eeg=eeg,
                 emg=emg,
@@ -554,6 +607,7 @@ class AccuSleepWindow(QMainWindow):
                 epoch_length=self.epoch_length,
                 epochs_per_img=self.model_epochs_per_img,
                 brain_state_set=self.brain_state_set,
+                emg_filter=self.emg_filter,
             )
 
             # overwrite as needed
@@ -569,8 +623,14 @@ class AccuSleepWindow(QMainWindow):
                 min_bout_length=self.min_bout_length,
             )
 
+            # ignore confidence scores if desired
+            if not self.save_confidence_scores:
+                confidence_scores = None
+
             # save results
-            save_labels(labels, label_file)
+            save_labels(
+                labels=labels, filename=label_file, confidence_scores=confidence_scores
+            )
             self.show_message(
                 (
                     "Saved labels for recording "
@@ -586,8 +646,6 @@ class AccuSleepWindow(QMainWindow):
 
         :param filename: model filename, if it's known
         """
-        from accusleepy.models import load_model
-
         if filename is None:
             file_dialog = QFileDialog(self)
             file_dialog.setWindowTitle("Select classification model")
@@ -606,6 +664,12 @@ class AccuSleepWindow(QMainWindow):
             self.show_message("ERROR: model file does not exist")
             return
 
+        self.show_message("Loading classification model")
+        self.ui.message_area.repaint()
+        QApplication.processEvents()
+
+        from accusleepy.models import load_model
+
         try:
             model, epoch_length, epochs_per_img, model_type, brain_states = load_model(
                 filename=filename
@@ -648,6 +712,8 @@ class AccuSleepWindow(QMainWindow):
         if len(config_warnings) > 0:
             for w in config_warnings:
                 self.show_message(w)
+        else:
+            self.show_message(f"Loaded classification model from {filename}")
 
         self.ui.model_label.setText(filename)
 
@@ -716,7 +782,7 @@ class AccuSleepWindow(QMainWindow):
             self.show_message("ERROR: label file does not exist")
             return
         try:
-            labels = load_labels(label_file)
+            labels, _ = load_labels(label_file)
         except Exception:
             self.ui.calibration_status.setText("could not load labels")
             self.show_message(
@@ -728,6 +794,7 @@ class AccuSleepWindow(QMainWindow):
             return
         label_error_message = check_label_validity(
             labels=labels,
+            confidence_scores=None,
             samples_in_recording=eeg.size,
             sampling_rate=sampling_rate,
             epoch_length=self.epoch_length,
@@ -758,6 +825,7 @@ class AccuSleepWindow(QMainWindow):
             sampling_rate=sampling_rate,
             epoch_length=self.epoch_length,
             brain_state_set=self.brain_state_set,
+            emg_filter=self.emg_filter,
         )
 
         self.ui.calibration_status.setText("")
@@ -814,6 +882,15 @@ class AccuSleepWindow(QMainWindow):
         """
         self.only_overwrite_undefined = checked
 
+    def update_confidence_policy(self, checked) -> None:
+        """Toggle policy for saving confidence scores
+
+        If the checkbox is enabled, confidence scores will be saved to the label files.
+
+        :param checked: state of the checkbox
+        """
+        self.save_confidence_scores = checked
+
     def manual_scoring(self) -> None:
         """View the selected recording for manual scoring"""
         # immediately display a status message
@@ -833,7 +910,7 @@ class AccuSleepWindow(QMainWindow):
         label_file = self.recordings[self.recording_index].label_file
         if os.path.isfile(label_file):
             try:
-                labels = load_labels(label_file)
+                labels, confidence_scores = load_labels(label_file)
             except Exception:
                 self.ui.manual_scoring_status.setText("could not load labels")
                 self.show_message(
@@ -848,10 +925,14 @@ class AccuSleepWindow(QMainWindow):
                 np.ones(int(eeg.size / (sampling_rate * self.epoch_length)))
                 * UNDEFINED_LABEL
             ).astype(int)
+            # manual scoring will not add a new confidence score column
+            # to a label file that does not have one
+            confidence_scores = None
 
         # check that all labels are valid
         label_error = check_label_validity(
             labels=labels,
+            confidence_scores=confidence_scores,
             samples_in_recording=eeg.size,
             sampling_rate=sampling_rate,
             epoch_length=self.epoch_length,
@@ -866,6 +947,10 @@ class AccuSleepWindow(QMainWindow):
                 epochs_in_recording = round(eeg.size / samples_per_epoch)
                 if epochs_in_recording - labels.size == 1:
                     labels = np.concatenate((labels, np.array([UNDEFINED_LABEL])))
+                    if confidence_scores is not None:
+                        confidence_scores = np.concatenate(
+                            (confidence_scores, np.array([0]))
+                        )
                     self.show_message(
                         (
                             "WARNING: an undefined epoch was added to "
@@ -874,6 +959,8 @@ class AccuSleepWindow(QMainWindow):
                     )
                 elif labels.size - epochs_in_recording == 1:
                     labels = labels[:-1]
+                    if confidence_scores is not None:
+                        confidence_scores = confidence_scores[:-1]
                     self.show_message(
                         (
                             "WARNING: the last epoch was removed from "
@@ -900,8 +987,10 @@ class AccuSleepWindow(QMainWindow):
             emg=emg,
             label_file=label_file,
             labels=labels,
+            confidence_scores=confidence_scores,
             sampling_rate=sampling_rate,
             epoch_length=self.epoch_length,
+            emg_filter=self.emg_filter,
         )
         manual_scoring_window.setWindowTitle(f"AccuSleePy viewer: {label_file}")
         manual_scoring_window.exec()
@@ -1067,15 +1156,6 @@ class AccuSleepWindow(QMainWindow):
 
     def initialize_settings_tab(self):
         """Populate settings tab and assign its callbacks"""
-        # show information about the settings tab
-        config_guide_file = open(
-            os.path.join(os.path.dirname(os.path.abspath(__file__)), CONFIG_GUIDE_FILE),
-            "r",
-        )
-        config_guide_text = config_guide_file.read()
-        config_guide_file.close()
-        self.ui.settings_text.setText(config_guide_text)
-
         # store dictionary that maps digits to rows of widgets
         # in the settings tab
         self.settings_widgets = {
@@ -1152,7 +1232,21 @@ class AccuSleepWindow(QMainWindow):
         }
 
         # update widget state to display current config
+        # UI defaults
         self.ui.default_epoch_input.setValue(self.epoch_length)
+        self.ui.overwrite_default_checkbox.setChecked(self.only_overwrite_undefined)
+        self.ui.confidence_setting_checkbox.setChecked(self.save_confidence_scores)
+        self.ui.default_min_bout_length_spinbox.setValue(self.min_bout_length)
+        # EMG filter
+        self.ui.emg_order_spinbox.setValue(self.emg_filter.order)
+        self.ui.bp_lower_spinbox.setValue(self.emg_filter.bp_lower)
+        self.ui.bp_upper_spinbox.setValue(self.emg_filter.bp_upper)
+        # model training hyperparameters
+        self.ui.batch_size_spinbox.setValue(self.hyperparameters.batch_size)
+        self.ui.learning_rate_spinbox.setValue(self.hyperparameters.learning_rate)
+        self.ui.momentum_spinbox.setValue(self.hyperparameters.momentum)
+        self.ui.training_epochs_spinbox.setValue(self.hyperparameters.training_epochs)
+        # brain states
         states = {b.digit: b for b in self.brain_state_set.brain_states}
         for digit in range(10):
             if digit in states.keys():
@@ -1171,6 +1265,15 @@ class AccuSleepWindow(QMainWindow):
                 self.settings_widgets[digit].frequency_widget.setEnabled(False)
 
         # set callbacks
+        self.ui.emg_order_spinbox.valueChanged.connect(self.emg_filter_order_changed)
+        self.ui.bp_lower_spinbox.valueChanged.connect(self.emg_filter_bp_lower_changed)
+        self.ui.bp_upper_spinbox.valueChanged.connect(self.emg_filter_bp_upper_changed)
+        self.ui.batch_size_spinbox.valueChanged.connect(self.hyperparameters_changed)
+        self.ui.learning_rate_spinbox.valueChanged.connect(self.hyperparameters_changed)
+        self.ui.momentum_spinbox.valueChanged.connect(self.hyperparameters_changed)
+        self.ui.training_epochs_spinbox.valueChanged.connect(
+            self.hyperparameters_changed
+        )
         for digit in range(10):
             state = self.settings_widgets[digit]
             state.enabled_widget.stateChanged.connect(
@@ -1233,6 +1336,41 @@ class AccuSleepWindow(QMainWindow):
         # check that configuration is valid
         _ = self.check_config_validity()
 
+    def emg_filter_order_changed(self, new_value: int) -> None:
+        """Called when user modifies EMG filter order
+
+        :param new_value: new EMG filter order
+        """
+        self.emg_filter.order = new_value
+
+    def emg_filter_bp_lower_changed(self, new_value: int | float) -> None:
+        """Called when user modifies EMG filter lower cutoff
+
+        :param new_value: new lower bandpass cutoff frequency
+        """
+        self.emg_filter.bp_lower = new_value
+        _ = self.check_config_validity()
+
+    def emg_filter_bp_upper_changed(self, new_value: int | float) -> None:
+        """Called when user modifies EMG filter upper cutoff
+
+        :param new_value: new upper bandpass cutoff frequency
+        """
+        self.emg_filter.bp_upper = new_value
+        _ = self.check_config_validity()
+
+    def hyperparameters_changed(self, new_value) -> None:
+        """Called when user modifies model training hyperparameters
+
+        :param new_value: unused
+        """
+        self.hyperparameters = Hyperparameters(
+            batch_size=self.ui.batch_size_spinbox.value(),
+            learning_rate=self.ui.learning_rate_spinbox.value(),
+            momentum=self.ui.momentum_spinbox.value(),
+            training_epochs=self.ui.training_epochs_spinbox.value(),
+        )
+
     def check_config_validity(self) -> str:
         """Check if brain state configuration on screen is valid"""
         # error message, if we get one
@@ -1259,6 +1397,10 @@ class AccuSleepWindow(QMainWindow):
         if sum(frequencies) != 1:
             message = "Error: sum(frequencies) != 1"
 
+        # check validity of EMG filter settings
+        if self.emg_filter.bp_lower >= self.emg_filter.bp_upper:
+            message = "Error: EMG filter cutoff frequencies are invalid"
+
         if message is not None:
             self.ui.save_config_status.setText(message)
             self.ui.save_config_button.setEnabled(False)
@@ -1290,12 +1432,41 @@ class AccuSleepWindow(QMainWindow):
         self.brain_state_set = BrainStateSet(brain_states, UNDEFINED_LABEL)
 
         # save to file
-        save_config(self.brain_state_set, self.ui.default_epoch_input.value())
+        save_config(
+            brain_state_set=self.brain_state_set,
+            default_epoch_length=self.ui.default_epoch_input.value(),
+            overwrite_setting=self.ui.overwrite_default_checkbox.isChecked(),
+            save_confidence_setting=self.ui.confidence_setting_checkbox.isChecked(),
+            min_bout_length=self.ui.default_min_bout_length_spinbox.value(),
+            emg_filter=EMGFilter(
+                order=self.emg_filter.order,
+                bp_lower=self.emg_filter.bp_lower,
+                bp_upper=self.emg_filter.bp_upper,
+            ),
+            hyperparameters=Hyperparameters(
+                batch_size=self.hyperparameters.batch_size,
+                learning_rate=self.hyperparameters.learning_rate,
+                momentum=self.hyperparameters.momentum,
+                training_epochs=self.hyperparameters.training_epochs,
+            ),
+        )
         self.ui.save_config_status.setText("configuration saved")
 
+    def reset_emg_filter_settings(self) -> None:
+        self.ui.emg_order_spinbox.setValue(DEFAULT_EMG_FILTER_ORDER)
+        self.ui.bp_lower_spinbox.setValue(DEFAULT_EMG_BP_LOWER)
+        self.ui.bp_upper_spinbox.setValue(DEFAULT_EMG_BP_UPPER)
+
+    def reset_hyperparams_settings(self):
+        self.ui.batch_size_spinbox.setValue(DEFAULT_BATCH_SIZE)
+        self.ui.learning_rate_spinbox.setValue(DEFAULT_LEARNING_RATE)
+        self.ui.momentum_spinbox.setValue(DEFAULT_MOMENTUM)
+        self.ui.training_epochs_spinbox.setValue(DEFAULT_TRAINING_EPOCHS)
+
 
 def check_label_validity(
     labels: np.array,
+    confidence_scores: np.array,
     samples_in_recording: int,
     sampling_rate: int | float,
     epoch_length: int | float,
@@ -1307,6 +1478,7 @@ def check_label_validity(
     brain state labels.
 
     :param labels: brain state labels
+    :param confidence_scores: confidence scores
     :param samples_in_recording: number of samples in the recording
     :param sampling_rate: sampling rate, in Hz
     :param epoch_length: epoch length, in seconds
@@ -1325,6 +1497,12 @@ def check_label_validity(
     ):
         return "label file contains invalid entries"
 
+    if confidence_scores is not None:
+        if np.min(confidence_scores) < 0 or np.max(confidence_scores) > 1:
+            return "label file contains invalid confidence scores"
+
+    return None
+
 
 def check_config_consistency(
     current_brain_states: dict,