An end-to-end pipeline that scores 30-second epochs of polysomnography into the five American Academy of Sleep Medicine (AASM) sleep stages, Wake, N1, N2, N3, and REM, from EEG. It extracts interpretable spectral and time-domain features, trains a tree-ensemble classifier, refines the per-epoch predictions with a hidden Markov model, and evaluates the whole system with subject-wise cross-validation on the Sleep-EDF Expanded database from PhysioNet.
The project is built around one idea that is easy to state and easy to get wrong: a sleep classifier must be tested on people it has never seen, and it must respect the fact that sleep unfolds in time. Getting those two things right, along with the class imbalance and the choice of metric, is what separates a number that means something from a number that does not.
Author: Rohan Shirur
Sleep is scored by hand. A trained technologist looks at each 30-second window of an overnight recording and assigns a stage, a process that takes hours per night and varies between scorers. Automating it well would lower the cost of sleep studies and make large-scale sleep research tractable, which matters because sleep architecture is tightly coupled to metabolic and cardiovascular health: fragmented and shortened slow-wave sleep is associated with impaired glucose regulation and insulin resistance. A reliable, transparent automatic scorer is a useful building block for that kind of work.
The task is a five-class classification problem with three properties that shape every design choice here:
- The classes are very imbalanced. N2 typically makes up close to half of a night; N1 is scarce. A model that ignores this can post a high accuracy while being useless on the rare stages.
- The stages are genuinely confusable. N1, the transitional stage, has low agreement even between expert human scorers. Any honest result shows N1 as the weakest class, and any honest evaluation reports per-stage behaviour rather than hiding it inside an average.
- The sequence is not random. Stages persist for minutes and follow an ordered progression. A frame-by-frame classifier that ignores this produces physiologically impossible flips (a single N3 epoch marooned in REM) that are almost always errors.
The figures and metrics below come from the synthetic demonstration that ships with the repository (see Note on data). They are not the Sleep-EDF benchmark; they exist so the pipeline can be run and inspected without downloading the recordings. Running on real data is a single command, described under Usage.
The table reports the per-epoch model and the same model after temporal smoothing, on pooled out-of-fold predictions:
| Metric | Per-epoch | + Temporal smoothing |
|---|---|---|
| Accuracy | 0.865 | 0.908 |
| Cohen's kappa | 0.805 | 0.866 |
| Macro F1 | 0.854 | 0.896 |
| Weighted F1 | 0.862 | 0.906 |
Per-stage recall shows the expected pattern, with N1 the hardest stage and Wake, N2 and REM the easiest. This is the confusion matrix of the final (smoothed) system:
The most direct view of quality is the hypnogram: the expert scoring and the model's predictions plotted over the night. Where the two step functions track, the model agrees with the scorer.
Because the features are named quantities rather than opaque activations, the model can be audited. The most important features are the ones a sleep technologist would point to: slow-wave (delta) power, spindle-band (sigma) energy, and the spectral edge frequency, which is exactly how the stages are distinguished by eye.
For reference, classical feature-based pipelines of this kind on Sleep-EDF commonly report accuracy in the low-to-mid 0.80s and Cohen's kappa around 0.75 to 0.80, with an HMM or similar sequence step adding a few points and N1 as the limiting class throughout. Confirm the current numbers against the literature when you write up your own run.
The pipeline is five stages, each in its own module, connected by plain NumPy arrays so that real and synthetic data flow through identical code.
1. Data (data.py). Each recording is a pair of EDF files, one of signals
and one of expert annotations. The loader attaches the annotations to the raw
recording, splits them into consecutive 30-second events, and cuts fixed-length
epochs with MNE, which keeps the signal and the scoring aligned through their
shared recording clock rather than by fragile manual indexing. Annotations are
mapped onto the five-stage taxonomy (merging the historical R&K stages 3 and 4
into N3, the modern AASM convention), movement and unscored epochs are dropped,
and long stretches of wakefulness at the start and end of the night are trimmed
to a fixed margin so that Wake does not swamp the sleep stages. Signals are
returned in microvolts.
2. Features (features.py). For each epoch and channel the code computes:
- absolute power in the delta, theta, alpha, sigma and beta bands, via Welch's method and Simpson integration;
- the same powers as fractions of total power, which are robust to per-recording amplitude scaling;
- spectral entropy and the 95% spectral edge frequency, two single-number summaries of the spectrum's shape;
- the three Hjorth parameters (activity, mobility, complexity); and
- waveform statistics: standard deviation, peak-to-peak, zero-crossing rate, skewness and kurtosis.
Every feature maps onto a quantity with a physiological meaning, which is what makes the model interpretable.
3. Model (model.py). A StandardScaler feeds a random forest (a gradient
boosting option is included). Class imbalance is handled with balanced class
weights so the objective does not simply learn to predict N2. Evaluation is done
with grouped k-fold cross-validation that splits on subject, and the out-of-fold
predictions and probabilities are pooled so that every epoch is scored exactly
once by a model that never saw its subject.
4. Temporal smoothing (smoothing.py). The per-epoch probabilities are the
emission model of a hidden Markov model whose transition matrix is estimated
from the training hypnograms. Viterbi decoding then recovers the single most
probable stage sequence for each night, which removes isolated one-epoch flips
and sharpens stage boundaries. The transition matrix is estimated per fold from
training subjects only, so the smoothed numbers are honest out-of-sample
estimates rather than a leak.
5. Evaluation (evaluate.py). Metrics lead with Cohen's kappa, the
field-standard measure that discounts chance agreement, alongside accuracy and
macro F1, and always include the per-stage precision, recall and F1 and the
confusion matrix. The per-epoch and smoothed results are reported side by side
so the effect of the HMM is explicit. Figures are the confusion matrix, the
expert-versus-predicted hypnogram, the feature-importance ranking, and the class
distribution.
Four decisions do most of the work of making the results trustworthy.
Subject-wise cross-validation. Consecutive epochs within one night are
strongly autocorrelated, and every person has an idiosyncratic EEG signature. If
epochs are split randomly, the model is tested on epochs from subjects it also
trained on, and it will look far better than it actually is. Every split in this
project is by subject. The test suite encodes this as an executable check
(test_cross_validation_partitions_subjects) that fails if a future edit ever
reintroduces a naive split.
Leakage-free temporal smoothing. Sequence post-processing is where many sleep-staging projects accidentally cheat, by fitting the temporal model on the same nights they evaluate. Here the transition matrix is estimated inside each fold from training subjects only and applied to the held-out subjects, so the reported gain is real. On the demonstration data it adds about four points of accuracy and six of kappa.
Class imbalance is handled and reported, not hidden. Balanced class weights keep the rare stages in the objective, and the reporting never leans on overall accuracy alone. The confusion matrix and per-stage F1 make it obvious where the model is strong and where it is not.
Cohen's kappa is the headline metric. Accuracy is inflated by the dominant N2 class and is hard to compare across datasets with different stage proportions. Kappa corrects for chance agreement and is directly comparable to the inter-scorer reliability of human experts, which is the right yardstick for this task.
git clone https://github.com/rohans11/sleep-stage-classifier.git
cd sleep-stage-classifier
python -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"Python 3.10 or newer is required.
Synthetic demonstration (no download, runs anywhere, reproduces the figures above):
python scripts/run_pipeline.py --synthetic --subjects 15Real Sleep-EDF data. The recordings download automatically on first run via MNE's dataset fetcher:
# optionally pre-download
python scripts/download_data.py --subjects 0 1 2 3 4 --data-dir data/
# run the full pipeline on real recordings
python scripts/run_pipeline.py --subjects 0 1 2 3 4 --classifier random_forestUseful flags: --classifier {random_forest,hist_gradient_boosting},
--n-splits, --data-dir, --output-dir, --seed. Figures and a
metrics.json are written to the output directory.
Using the package directly:
from sleepstager import (
extract_features, subject_cross_validate, compute_metrics, format_metrics,
)
from sleepstager.data import load_dataset
dataset = load_dataset(subjects=[0, 1, 2, 3, 4])
X, names = extract_features(dataset.epochs, dataset.channel_names)
cv = subject_cross_validate(X, dataset.labels, dataset.groups)
print("Per-epoch:")
print(format_metrics(compute_metrics(cv.y_true, cv.y_pred)))
print("\nAfter temporal smoothing:")
print(format_metrics(compute_metrics(cv.y_true, cv.y_pred_smoothed)))sleep-stage-classifier/
src/sleepstager/
config.py constants, frequency bands, stage taxonomy
data.py real Sleep-EDF loading + synthetic generator
features.py spectral and time-domain feature extraction
model.py pipeline + subject-wise cross-validation
smoothing.py HMM / Viterbi temporal smoothing
evaluate.py metrics and figures
scripts/
run_pipeline.py end-to-end pipeline (CLI)
download_data.py pre-fetch Sleep-EDF recordings
tests/ 48 tests covering every module
figures/ committed demo figures + metrics.json (shown above)
.github/workflows/ci.yml lint, test matrix, smoke run
pytest # 48 tests
ruff check src scripts testsThe tests check real properties rather than just running the code: that a pure sine wave deposits its power in the correct frequency band, that Hjorth activity equals the signal variance, that cross-validation truly partitions subjects, that the EDF loading path recovers the correct labels and epoch shapes on an in-memory recording, that Viterbi decoding corrects an isolated flip while leaving a clean sequence untouched, and that the metrics are correct on inputs with a known answer. GitHub Actions runs the linter, the test suite on Python 3.10 through 3.12, and an end-to-end synthetic run on every push.
The committed figures are generated from synthetic signals, for two practical
reasons: the Sleep-EDF recordings are large and are not mine to redistribute,
and continuous integration needs to run without network access. The generator
in data.py builds EEG-like signals whose spectral content depends on an
assigned stage (delta-dominant N3, spindle-band N2, alpha and beta in Wake, and
so on), and it deliberately injects epoch-level ambiguity, above all for N1, so
that the demonstration is realistically difficult rather than trivial.
The synthetic data is a way to exercise and verify the code, not a scientific
result. Because its hypnograms come from a clean Markov process, the temporal
smoother helps a little more on synthetic data than it would on real
recordings, so treat the smoothing gain here as an upper end. The real numbers
come from load_dataset, which reads the actual PhysioNet recordings, and the
two paths share all downstream code, so what the tests verify on synthetic data
is exactly what runs on real data.
- Single montage. The features use the two Sleep-EDF EEG derivations. EOG and EMG carry complementary information, particularly for REM, and are the natural next channels to add.
- Classical model. Deep architectures on raw signals or spectrograms can edge out feature-based models, at the cost of interpretability and reproducibility. The clean feature interface here makes such a model a natural drop-in comparison.
- Fixed HMM transitions. The smoother uses a first-order transition matrix. Duration-aware (semi-Markov) models capture the long, stable bouts of deep sleep more faithfully.
- Single cohort. Reported numbers are within Sleep-EDF. Cross-dataset validation (for example on the MASS or SHHS cohorts) is the real test of generalisation.
- Kemp, B., Zwinderman, A. H., Tuk, B., Kamphuisen, H. A. C., & Oberye, J. J. L. (2000). Analysis of a sleep-dependent neuronal feedback loop: the slow-wave microcontinuity of the EEG. IEEE Transactions on Biomedical Engineering. (The Sleep-EDF database.)
- Goldberger, A. L., et al. (2000). PhysioBank, PhysioToolkit, and PhysioNet. Circulation, 101(23), e215-e220.
- Iber, C., Ancoli-Israel, S., Chesson, A. L., & Quan, S. F. (2007). The AASM Manual for the Scoring of Sleep and Associated Events. American Academy of Sleep Medicine.
- Gramfort, A., et al. (2013). MEG and EEG data analysis with MNE-Python. Frontiers in Neuroscience, 7, 267.
- Rabiner, L. R. (1989). A tutorial on hidden Markov models and selected applications in speech recognition. Proceedings of the IEEE, 77(2), 257-286.
MIT. See LICENSE.