Automated Sleep Stage Classification using EEG — Dreem Headband Dataset

An end-to-end pipeline for automated sleep-stage classification from wearable EEG and accelerometer signals, comparing classical machine learning, deep learning, and transfer learning approaches on the Dreem Headband dataset.

Full reports included in this repo:

• Sleep_Stage_Classification_Report1.pdf — Dataset analysis, feature extraction, classical ML
• Sleep_Stage_Classification_Report2.pdf — Deep learning, transfer learning, LOEO/LOSO evaluation

---

Overview

Sleep stage classification is traditionally done via manual polysomnography (PSG), which is costly and time-consuming. This project builds a scalable, automated pipeline using the Dreem Headband — a wearable EEG device — classifying each 30-second epoch into one of five stages:

Label	Stage	Description
0	Wake	High-frequency, low-amplitude EEG
1	N1	Light sleep, reduced alpha activity
2	N2	Stable sleep, spindles and K-complexes
3	N3	Deep sleep, slow high-amplitude waves
4	REM	Mixed-frequency, low muscle tone

---

Dataset

Dreem Automated Sleep Staging Dataset

• 7 participants, ~6,175 labeled epochs
• 5 EEG channels @ 250 Hz + 3 accelerometer axes @ 50 Hz
• Each epoch: 30-second window → 42,001 features per row
• Significant class imbalance: N2 = 37.5%, N1 = 4.0%

Dataset not included in repo — download from Kaggle and structure as:

data/
├── record_0.npy ... record_6.npy
└── targets_train.csv

---

Pipeline

Raw EEG/ACC Data
      ↓
Preprocessing (Resample 250Hz→128Hz, Bandpass Filter 0.5–45Hz, Z-score Normalization)
      ↓
Feature Extraction (Time-domain: Mean, Std, RMS, ZCR | Frequency-domain: Dominant Freq, Spectral Energy, Entropy)
      ↓
Model Training (Classical ML / Deep Learning / Transfer Learning)
      ↓
Evaluation (Accuracy, Macro-F1, Cohen's κ | LOEO / LOSO)

---

Models & Results

Classical Machine Learning (56 features × 8 channels)

Model	Accuracy	Macro F1
XGBoost	0.8197	0.6735
Random Forest	0.8087	0.6096
AdaBoost	0.7322	0.5734
Decision Tree	0.7377	0.5669
SVM	0.4098	0.1423
Naïve Bayes	0.2022	0.1040
Logistic Regression	0.1808	0.1447

Deep Learning (CNN & RNN — trained from scratch)

Model	Accuracy	Macro F1
CNN	0.3806	0.1381
RNN	0.3814	0.1829

Both suffered from severe overfitting and class bias toward N2.

Transfer Learning (MobileNetV2 on Log-Mel Spectrograms)

Model	Accuracy	Macro F1	Cohen's κ
Pre-trained CNN (Frozen)	0.589	0.523	0.463
Pre-trained CNN (Fine-tuned)	0.712	0.623	0.618
CNN→BiLSTM (Frozen)	0.581	0.510	0.427
CNN→BiLSTM (Fine-tuned)	0.662	0.582	0.528

Best overall: Fine-tuned MobileNetV2 CNN — 71.2% accuracy, macro-F1 = 0.623, κ = 0.618

---

Key Findings

• XGBoost dominated classical ML — 82% accuracy, F1 = 0.67, leveraging engineered EEG spectral features
• Transfer learning dramatically outperformed scratch-trained deep models — 3× higher macro-F1
• CNN→BiLSTM added limited gains — single-epoch spectrograms already capture full stage patterns; BiLSTM increased parameters without proportional benefit
• N1 remains the hardest stage to classify across all models due to low sample count (4%)
• Optimal window size: 2–5 seconds — short windows capture localized EEG events (spindles, K-complexes) better than standard 30s epochs
• EEG features dominate — RMS, spectral energy, and entropy (EEG ch1/ch2) were top predictors; accelerometer contributed less

---

Evaluation Protocols

• Hold-out: 64% train / 16% val / 20% test, stratified splits
• LOEO (Leave-One-Event-Out): Train on 6 nights, test on 1, rotated across all 7 — accuracy range: 29.3%–46.2%
• LOSO (Leave-One-Subject-Out): Equivalent to LOEO here (1 night per subject)

---

Project Structure

CS_289_Project/
├── SleepStageClassification.ipynb      # Full pipeline notebook
├── Sleep_Stage_Classification_Report1.pdf  # ML analysis report
└── Sleep_Stage_Classification_Report2.pdf  # Deep learning report

---

Setup

git clone https://github.com/19ska/Sleep-Stage-Classification.git
cd Sleep-Stage-Classification

python -m venv venv
source venv/bin/activate
pip install -r requirements.txt

Key dependencies: numpy, scipy, pandas, scikit-learn, xgboost, torch, tensorflow, librosa

---

Future Work

• Focal loss or class-balanced sampling to handle N1/Wake imbalance
• Multi-epoch context window (±2 neighboring epochs) for LSTM/Transformer
• Multimodal fusion of EEG + accelerometer for improved REM/Wake discrimination
• Data augmentation: time masking, frequency dropout, Gaussian noise injection
• β-VAE or Conditional VAE for latent space exploration