Note
Go to the end to download the full example code
Introduction to the datasets Module¶
Welcome to the guide on TorchEEG’s datasets
Module! This module provides you with various benchmark datasets for EEG-based emotion recognition, such as DEAP, DREAMER, SEED, MAHNOB, AMIGOS, and MPED. These datasets use a range of stimuli like music and videos to trigger emotional responses. Once the emotional experiences are recorded, participants label them using methods like valence-arousal dimensions or discrete emotion categories.
This guide will help you understand how to easily load and manipulate these datasets for your training needs, including signal segmentation, applying transformations, and more.
An Example of Emotion Recognition Dataset¶
To begin, you’ll need to download the DEAP dataset from DEAP download link. Once you have the dataset, use the following code to load it. This will read EEG signals and labels, apply offline transformations, and save them for easy access later on.
from torcheeg.datasets import DEAPDataset
from torcheeg import transforms
from torcheeg.datasets.constants.emotion_recognition.deap import \
DEAP_CHANNEL_LOCATION_DICT
dataset = DEAPDataset(
io_path=f'./deap',
root_path='./data_preprocessed_python',
offline_transform=transforms.Compose([
transforms.BandDifferentialEntropy(apply_to_baseline=True),
transforms.ToGrid(DEAP_CHANNEL_LOCATION_DICT, apply_to_baseline=True)
]),
online_transform=transforms.Compose(
[transforms.BaselineRemoval(),
transforms.ToTensor()]),
label_transform=transforms.Compose([
transforms.Select('valence'),
transforms.Binary(5.0),
]),
num_worker=8)
You can get individual samples from the dataset like so:
print(dataset[0])
This will output a tuple, where the first element is the EEG signal and the second is its label.
If you want to visualize the EEG signals, you can use the helper function as shown below:
import torch
from torcheeg.utils import plot_3d_tensor
img = plot_3d_tensor(torch.tensor(dataset[0][0]))
You can also refer to the document for more information on visualizing EEG signals
via the plot_3d_tensor
, plot_feature_topomap
, plot_raw_topomap
, plot_signal
, and plot_adj_connectivity
function.
Important Notes¶
If you make changes to the dataset’s configuration, remember to clear the cache or specify a new io_path. Otherwise, only online transformations (online_transform and label_transform) will take effect:
!rm -rf ./deap
If you remove all transforms and declare a dataset, you’ll find that it returns segmented EEG signals that haven’t undergone any preprocessing. Additionally, there is a dictionary representing all the meta-information corresponding to the EEG signal samples.
from torcheeg.datasets import DEAPDataset
dataset = DEAPDataset(io_path=f'./deap', root_path='./data_preprocessed_python')
print(dataset[0])
Advanced Usage¶
Most emotion analysis datasets can be found in TorchEEG. Moreover,
TorchEEG provides support for MoABB
, allowing access to
motor-imagery-related datasets with the help of MoABB
.
import torcheeg.datasets.moabb as moabb_dataset
from moabb.datasets import BNCI2014001
from moabb.paradigms import LeftRightImagery
dataset = BNCI2014001()
dataset.subject_list = [1, 2, 3]
paradigm = LeftRightImagery()
dataset = moabb_dataset.MOABBDataset(
dataset=dataset,
paradigm=paradigm,
io_path='./moabb',
offline_transform=transforms.Compose([transforms.BandDifferentialEntropy()
]),
online_transform=transforms.ToTensor(),
label_transform=transforms.Compose([transforms.Select('label')]))
Custom Datasets¶
TorchEEG also supports custom datasets. You can place recorded EEG signal files in a folder following a specific rule, such as:
label01
|- sub01.edf
|- sub02.edf
label02
|- sub01.edf
|- sub02.edf
Then, you can use FolderDataset
to automatically access the
corresponding EEG signal samples.
from torcheeg.datasets import FolderDataset
label_map = {'label01': 0, 'label02': 1}
dataset = FolderDataset(io_path='./folder',
root_path='./root_folder',
structure='subject_in_label',
num_channel=14,
online_transform=transforms.ToTensor(),
label_transform=transforms.Compose([
transforms.Select('label'),
transforms.Lambda(lambda x: label_map[x])
]),
num_worker=4)
Alternatively, you can use a CSV file to specify more detailed meta-information for reading:
| subject_id | trial_id | label | file_path |
| ---------- | ------- | ----- | ------------------------- |
| sub1 | 0 | 0 | './data/label1/sub1.fif' |
| sub1 | 1 | 1 | './data/label2/sub1.fif' |
| sub1 | 2 | 2 | './data/label3/sub1.fif' |
| sub2 | 0 | 0 | './data/label1/sub2.fif' |
| sub2 | 1 | 1 | './data/label2/sub2.fif' |
| sub2 | 2 | 2 | './data/label3/sub2.fif' |
from torcheeg.datasets import CSVFolderDataset
dataset = CSVFolderDataset(csv_path='./data.csv',
online_transform=transforms.ToTensor(),
label_transform=transforms.Select('label'),
num_worker=4)
By default, TorchEEG uses mne to read recorded EEG signals, but you can
also specify your own file reading logic through read_fn
.
import mne
def default_read_fn(file_path, **kwargs):
# Load EEG file
raw = mne.io.read_raw(file_path)
# Convert raw to epochs
epochs = mne.make_fixed_length_epochs(raw, duration=1)
# Return EEG data
return epochs
dataset = CSVFolderDataset(io_path='./csv_folder',
csv_path='./data.csv',
read_fn=default_read_fn,
online_transform=transforms.ToTensor(),
label_transform=transforms.Select('label'),
num_worker=4)