Sleep Stage Detection
- TYPE
- EEG deep learning
- STACK
- 07 COMPONENTS
- STATUS
- SOURCE AVAILABLE
- FEATURED
- —
- LINKS
- SOURCE ↗
01 / OVERVIEW
A sleep-stage classifier for EEG data that fights class imbalance with a VAE-based data generator and squeezes more signal out of unlabeled recordings with self-supervised pretraining.
Built as an AFML course project, Sleep Stage Detection tackles automatic sleep stage classification on the Sleep-EDF (Sleep Cassette) dataset — EEG recordings from 20 healthy adults across 39 nights, labeled into five stages (W, N1, N2, N3, REM). The N1 and N3 stages are badly underrepresented in the raw data, which cripples a naive classifier. We started from a 1D-CNN baseline trained with class-weighted and focal loss, then layered on two complementary fixes: a Variational Autoencoder that learns the distribution of the minority classes and generates synthetic EEG segments to rebalance the training set, and a SimCLR-based semi-supervised pipeline that pretrains on unlabeled EEG before fine-tuning on as little as 5% labeled data. The full pipeline — preprocessing, k-fold splitting, VAE training, synthetic augmentation, and SSL pretraining/fine-tuning — is scripted end-to-end so any of the four experimental configurations (baseline, VAE-augmented, SSL, SSL+VAE) can be reproduced with a single command per stage.
02/ STACK & ARCHITECTURE
BACKEND
- Python
- PyTorch
- MNE
- NumPy
- Pandas
- scikit-learn
- einops
03 / WHAT WAS ACTUALLY HARD
- 01
1D-CNN baseline with imbalance-aware losses
A supervised 1D-CNN trained directly on preprocessed EEG epochs, using class-weighted and focal loss so the heavily underrepresented N1 and N3 stages aren't drowned out by the dominant N2 and W classes.
- 02
VAE-based synthetic minority-class augmentation
A Variational Autoencoder trained specifically on target classes (N1 and N3) generates realistic synthetic EEG segments to rebalance the dataset toward a configurable target ratio, then the augmented set is fed back into the same 1D-CNN training pipeline for a direct before/after comparison.
- 03
SimCLR semi-supervised pretraining
A contrastive self-supervised pretraining stage (SimCLR) learns representations from unlabeled EEG, then fine-tunes on small labeled slices (as low as 5%) to test how much labeling effort can be saved without sacrificing accuracy.
- 04
Combined SSL + VAE pipeline and full experiment harness
SSL pretraining and fine-tuning can also run on the VAE-augmented dataset, giving a four-way comparison (baseline, VAE-only, SSL-only, SSL+VAE) across identical k-fold splits, driven by a consistent set of preprocessing, splitting, and training scripts for reproducibility.