EEG Attention Classification

This project uses the user's EEG data to classify their attention into three states:

• distracted
• moderate
• focused

It has two parts:

1. Train and export the classifier in Google Colab.
2. Run local real-time inference from a live EEG stream.

YouTube demo: https://youtu.be/j-fiuQFRf-I

Hardware and input data

The hardware is a Muse 2 EEG headband. This project uses its four EEG electrodes/channels:

Model channel	Muse 2 electrode
eeg1	TP9
eeg2	AF7
eeg3	AF8
eeg4	TP10

The model expects four raw EEG values for every sample, in this channel order.

Project files

• ML_Attention_classification_using_EEG.ipynb - Google Colab notebook for data preparation, feature extraction, training, evaluation, and model export.
• EEG_Attention_prediction.py - local inference module. It loads the trained model assets and converts each 100-sample EEG window into attention probabilities.
• EEG_Attention_main.py - local WebSocket server. It receives live EEG messages, maintains the raw EEG buffer, calls the prediction module, prints results, and sends results back to the client.
• model/ - folder containing the required exported model assets:
  • attention_classifier_eeg_model.keras
  • feature_scaler_time_domain.pkl
  • feature_scaler_freq_domain.pkl
  • label_encoder.pkl

1. Train the model in Google Colab

Open ML_Attention_classification_using_EEG.ipynb in Google Colab and run its cells in order.

Load labelled recordings

Load four-channel EEG recordings collected for the distracted, moderate, and focused attention states. Labels tell the model which attention level each recording represents.

EEG signals were recorded from a subject under three cognitive states: distracted, moderately focused, and focused. Three experimental trials were conducted for each state.

Clean and resample the EEG

The notebook removes missing values, orders samples by timestamp, and resamples recordings to 100 Hz. This makes recordings comparable even if their original timestamps are uneven.

Segment the signal

Each resampled recording is split into non-overlapping windows of 100 samples. At 100 Hz, each window represents approximately one second of EEG data and becomes one training sample.

Extract features

The notebook calculates features independently for each EEG channel, then combines them into one feature row per window. Features include:

• Time-domain statistics such as mean, variance, minimum, maximum, RMS, skewness, and kurtosis.
• Hjorth activity, mobility, and complexity.
• Spectral entropy, spectral centroid, and peak frequency.
• Delta, theta, alpha, beta, and gamma band powers and relative band powers.

These features describe the EEG window in a compact form for the classifier.

Prepare and train the classifier

The feature data is combined, shuffled, and split into training, validation, and test sets. Time-domain and frequency-domain features are standardised with separate scalers. Attention labels are encoded before training the two-input neural-network classifier. The model learns the time-domain and frequency-domain features separately.

Export model assets

After training, export the Keras model, two feature scalers, and label encoder. Place the following files in the local model/ directory before running real-time prediction:

attention_classifier_eeg_model.keras
feature_scaler_time_domain.pkl
feature_scaler_freq_domain.pkl
label_encoder.pkl

The local scripts use the model, scalers, and encoder generated by the same training run.

2. Run real-time prediction locally

Install the required packages in the Python environment used to run the local scripts:

pip install numpy pandas scipy scikit-learn joblib websockets
pip install tensorflow==2.20.0
pip install keras==3.13.2

Start the WebSocket server:

python EEG_Attention_main.py

The server listens at:

ws://localhost:8001

Connect an EEG data streaming application to this address and start the live Muse 2 stream.

Live message format

Each incoming JSON message must provide raw EEG values in an eeg array.

{
  "timestamp": 0,
  "eeg": [eeg1, eeg2, eeg3, eeg4]
}

Additional values and fields are allowed; the local server uses the first four EEG values. It can also accept a batch of four-channel samples.

Real-time prediction flow

1. EEG_Attention_main.py receives live JSON EEG messages over WebSocket.
2. It validates the first four EEG values and stores them in a rolling buffer for the four channels.
3. When the buffer reaches 100 valid samples, it represents the current EEG window.
4. Every 10 new samples, the window is passed to EEG_Attention_prediction.py.
5. The prediction module extracts the same features used during training, applies the saved scalers (time-domain and frequency-domain scalers), and runs the Keras model.
6. The server smooths the probabilities, prints the attention percentages, and sends a JSON prediction response to the connected client.

Example console output:

Received 100 valid EEG samples (100/100 buffered).

ATTENTION (%)
  distracted: 12.4%
  moderate: 31.7%
  focused: 55.9%

Use Ctrl+C to stop the local server cleanly.

Notes

• Use the same four-channel order for recorded training data and live data.
• The local predictor uses the latest 100 raw valid samples, so it does not rely on repeated stream timestamps.
• ATTENTION LABELS ARE ML MODEL ESTIMATES, NOT MEDICAL OR DIAGNOSTIC RESULTS.