Repository: torcheeg/torcheeg (language: python)

Issue:
Add SEED examples and fix typos



Discussion / hints:

Fix patch:
diff --git a/torcheeg/datasets/functional/__init__.py b/torcheeg/datasets/functional/__init__.py
index 6a8dfd5..fb7a57c 100644
--- a/torcheeg/datasets/functional/__init__.py
+++ b/torcheeg/datasets/functional/__init__.py
@@ -1,6 +1,7 @@
 from .emotion_recognition import *
 from .mne import mne_constructor
 from .numpy import numpy_constructor
+from .folder import folder_constructor
 from .personal_identification import *
 from .ssvep import *
 from .hooks import *
\ No newline at end of file
diff --git a/torcheeg/datasets/functional/folder.py b/torcheeg/datasets/functional/folder.py
new file mode 100644
index 0000000..00a3ac4
--- /dev/null
+++ b/torcheeg/datasets/functional/folder.py
@@ -0,0 +1,192 @@
+import os
+import re
+from functools import partial
+from typing import Tuple, Callable
+from multiprocessing import Manager, Pool, Process, Queue
+from typing import Callable, Union
+from pathlib import Path
+import scipy.io as scio
+from torcheeg.io import EEGSignalIO, MetaInfoIO
+from tqdm import tqdm
+import mne
+import numpy as np
+mne.set_log_level(40)#ERROR
+MAX_QUEUE_SIZE = 1024
+
+
+def transform_producer(file_path_id: Tuple[str, int], read_file: Callable, chunk_size: int, overlap: int,
+                       num_channel: int, before_trial: Union[None, Callable],
+                       transform: Union[None, Callable],
+                       after_trial: Union[Callable, None], queue: Queue):
+    
+    file_path = file_path_id[0]
+    subject_id = file_path_id[1]
+    file_label= file_path.split(os.sep)
+    label, file_name  = file_label[-2], file_label[-1]
+    trial_samples = read_file(file_path, chunk_size)
+    events = [i[0] for i in trial_samples.events]
+    events.append(events[-1] + np.diff(events)[0])# time interval between all events are same
+    if before_trial:
+        raise NotImplementedError
+ 
+    trial_queue = []
+    write_pointer = 0
+    for i, trial_signal in enumerate(trial_samples.get_data()):
+        t_eeg = trial_signal[:num_channel, :]
+        if not transform is None:
+            t = transform(eeg=trial_signal[:num_channel, :])
+            t_eeg = t['eeg']
+
+        clip_id = f'{file_name}_{write_pointer}'
+        write_pointer += 1
+
+
+        record_info = {
+            'subject_id': subject_id,
+            'trial_id': i,
+            'file_id': file_name,
+            'start_at': events[i],
+            'end_at': events[i+1],
+            'clip_id': clip_id,
+            'label': label
+        }
+
+        if after_trial:
+            trial_queue.append({
+                'eeg': t_eeg,
+                'key': clip_id,
+                'info': record_info
+            })
+        else:
+            queue.put({'eeg': t_eeg, 'key': clip_id, 'info': record_info})
+    
+    if len(trial_queue) and after_trial:
+        trial_queue = after_trial(trial_queue)
+        for obj in trial_queue:
+            assert 'eeg' in obj and 'key' in obj and 'info' in obj, 'after_trial must return a list of dictionaries, where each dictionary corresponds to an EEG sample, containing `eeg`, `key` and `info` as keys.'
+            queue.put(obj)
+
+def io_consumer(write_eeg_fn: Callable, write_info_fn: Callable, queue: Queue):
+    while True:
+        item = queue.get()
+        if not item is None:
+            eeg = item['eeg']
+            key = item['key']
+            write_eeg_fn(eeg, key)
+            if 'info' in item:
+                info = item['info']
+                write_info_fn(info)
+        else:
+            break
+
+
+class SingleProcessingQueue:
+    def __init__(self, write_eeg_fn: Callable, write_info_fn: Callable):
+        self.write_eeg_fn = write_eeg_fn
+        self.write_info_fn = write_info_fn
+
+    def put(self, item):
+        eeg = item['eeg']
+        key = item['key']
+        self.write_eeg_fn(eeg, key)
+        if 'info' in item:
+            info = item['info']
+            self.write_info_fn(info)
+
+
+def folder_constructor(
+    root_path: str = './eeg_raw_data',
+    chunk_size: int = -1,
+    overlap: int = 0,
+    num_channel: int = -1,
+    before_trial: Union[None, Callable] = None,
+    transform: Union[None, Callable] = None,
+    after_trial: Union[Callable, None] = None,
+    io_path: str = './io/seed_iv',
+    io_size: int = 10485760,
+    io_mode: str = 'lmdb',
+    num_worker: int = 0,
+    verbose: bool = True,
+    read_func =None
+    
+) -> None:
+    # init IO
+
+    meta_info_io_path = os.path.join(io_path, 'info.csv')
+    eeg_signal_io_path = os.path.join(io_path, 'eeg')
+
+    if os.path.exists(io_path) and not os.path.getsize(meta_info_io_path) == 0:
+        print(
+            f'The target folder already exists, if you need to regenerate the database IO, please delete the path {io_path}.'
+        )
+        return
+
+    os.makedirs(io_path, exist_ok=True)
+
+    meta_info_io_path = os.path.join(io_path, 'info.csv')
+    eeg_signal_io_path = os.path.join(io_path, 'eeg')
+
+    info_io = MetaInfoIO(meta_info_io_path)
+    eeg_io = EEGSignalIO(eeg_signal_io_path, io_size=io_size, io_mode=io_mode)
+
+    # loop to access the dataset files
+    
+    folder_path = Path(root_path)
+    file_paths = folder_path.glob('**/*.*')
+    file_path_list = [str(i).replace('\\','/') for i in file_paths]   
+    subjects = list(range(len(file_path_list))) 
+    file_path_list_subject = zip(file_path_list,subjects)
+    if verbose:
+        # show process bar
+        pbar = tqdm(total=len(file_path_list))
+        pbar.set_description("[Folder Data]")
+
+    if num_worker < 0:
+        num_worker = os.cpu_count() + num_worker +1
+
+    if num_worker > 1:
+        manager = Manager()
+        queue = manager.Queue(maxsize=MAX_QUEUE_SIZE)
+        io_consumer_process = Process(target=io_consumer,
+                                      args=(eeg_io.write_eeg,
+                                            info_io.write_info, queue),
+                                      daemon=True)
+        io_consumer_process.start()
+
+        partial_mp_fn = partial(transform_producer,
+                                chunk_size=chunk_size,
+                                read_file=read_func,
+                                overlap=overlap,
+                                num_channel=num_channel,
+                                before_trial=before_trial,
+                                transform=transform,
+                                after_trial=after_trial,
+                                queue=queue)
+
+        for _ in Pool(num_worker).imap(partial_mp_fn, file_path_list_subject):
+            if verbose:
+                pbar.update(1)
+
+        queue.put(None)
+
+        io_consumer_process.join()
+        io_consumer_process.close()
+
+    else:
+        for file_path in file_path_list_subject:
+            transform_producer(file_path_id=file_path,
+                               chunk_size=chunk_size,
+                               read_file=read_func,
+                               overlap=overlap,
+                               num_channel=num_channel,
+                               before_trial=before_trial,
+                               transform=transform,
+                               after_trial=after_trial,
+                               queue=SingleProcessingQueue(
+                                   eeg_io.write_eeg, info_io.write_info))
+            if verbose:
+                pbar.update(1)
+
+    if verbose:
+        pbar.close()
+        print('Please wait for the writing process to complete...')
diff --git a/torcheeg/datasets/module/__init__.py b/torcheeg/datasets/module/__init__.py
index ab8c3df..60a40fc 100644
--- a/torcheeg/datasets/module/__init__.py
+++ b/torcheeg/datasets/module/__init__.py
@@ -1,6 +1,7 @@
 from .base_dataset import BaseDataset
 from .emotion_recognition import *
 from .mne_dataset import MNEDataset
+from .folder_dataset import FolderDataset
 from .numpy_dataset import NumpyDataset
 from .personal_identification import *
 from .sspev import *
diff --git a/torcheeg/datasets/module/folder_dataset.py b/torcheeg/datasets/module/folder_dataset.py
new file mode 100644
index 0000000..c1501d2
--- /dev/null
+++ b/torcheeg/datasets/module/folder_dataset.py
@@ -0,0 +1,89 @@
+from typing import Callable, Dict, Tuple, Union
+from ..functional.folder import folder_constructor
+from .base_dataset import BaseDataset
+
+
+class FolderDataset(BaseDataset):
+    def __init__(self,
+                 root_path: str = './eeg_raw_data',
+                 chunk_size: int = -1,
+                 overlap: int = 0,
+                 num_channel: int = -1,
+                 online_transform: Union[None, Callable] = None,
+                 offline_transform: Union[None, Callable] = None,
+                 label_transform: Union[None, Callable] = None,
+                 before_trial: Union[None, Callable] = None,
+                 after_trial: Union[Callable, None] = None,
+                 io_path: str = './io/seed',
+                 io_size: int = 10485760,
+                 io_mode: str = 'lmdb',
+                 num_worker: int = 0,
+                 verbose: bool = True,
+                 in_memory: bool = False,
+                 read_func = None):
+        folder_constructor(
+            root_path=root_path,
+            chunk_size=chunk_size,
+            overlap=overlap,
+            num_channel=num_channel,
+            before_trial=before_trial,
+            transform=offline_transform,
+            after_trial=after_trial,
+            io_path=io_path,
+            io_size=io_size,
+            io_mode=io_mode,
+            num_worker=num_worker,
+            verbose=verbose,
+            read_func = read_func
+        )
+        super().__init__(io_path=io_path,
+                         io_size=io_size,
+                         io_mode=io_mode,
+                         in_memory=in_memory)
+
+        self.root_path = root_path
+        self.chunk_size = chunk_size
+        self.overlap = overlap
+        self.num_channel = num_channel
+        self.online_transform = online_transform
+        self.offline_transform = offline_transform
+        self.label_transform = label_transform
+        self.before_trial = before_trial
+        self.after_trial = after_trial
+        self.num_worker = num_worker
+        self.verbose = verbose
+
+    def __getitem__(self, index: int) -> Tuple[any, any, int, int, int]:
+        info = self.read_info(index)
+
+        eeg_index = str(info['clip_id'])
+        eeg = self.read_eeg(eeg_index)
+
+        signal = eeg
+        label = info
+
+        if self.online_transform:
+            signal = self.online_transform(eeg=eeg)['eeg']
+
+        if self.label_transform:
+            label = self.label_transform(y=info)['y']
+
+        return signal, label
+
+    @property
+    def repr_body(self) -> Dict:
+        return dict(
+            super().repr_body, **{
+                'root_path': self.root_path,
+                'chunk_size': self.chunk_size,
+                'overlap': self.overlap,
+                'num_channel': self.num_channel,
+                'online_transform': self.online_transform,
+                'offline_transform': self.offline_transform,
+                'label_transform': self.label_transform,
+                'before_trial': self.before_trial,
+                'after_trial': self.after_trial,
+                'num_worker': self.num_worker,
+                'verbose': self.verbose,
+                'io_size': self.io_size
+            })


Test patch:
diff --git a/test/datasets/test_folder.py b/test/datasets/test_folder.py
new file mode 100644
index 0000000..055f98c
--- /dev/null
+++ b/test/datasets/test_folder.py
@@ -0,0 +1,117 @@
+
+import sys
+sys.path.insert(0, "../torcheeg")
+
+import os
+import random
+import shutil
+import unittest
+import numpy as np
+from torcheeg import transforms
+from torcheeg.datasets import FolderDataset
+from torcheeg.datasets.functional import folder_constructor
+import mne
+
+class TestFolderDataset(unittest.TestCase):
+    def setUp(self):
+        if  os.path.isdir("./tmp_out/"):
+            shutil.rmtree('./tmp_out/')
+            os.mkdir('./tmp_out/')
+            
+        if  os.path.isdir("./tmp_in/"):
+            shutil.rmtree('./tmp_in/')
+            os.mkdir('./tmp_in/')
+
+    def generate_dummy_eeg_files(self,num_files, folder_path):
+        # Define EEG file specifications
+        sfreq = 128  # Sampling rate
+        n_channels = 14  # Number of channels
+        duration = 5  # Data collected for 5 seconds
+        for i in range(num_files):
+            # Generate dummy EEG data
+            n_samples = sfreq * duration
+            data = np.random.randn(n_channels, n_samples)
+
+            # Create MNE Raw object
+            ch_names = [f'ch_{i+1:03}' for i in range(n_channels)]
+            ch_types = ['eeg'] * n_channels
+            info = mne.create_info(ch_names, sfreq, ch_types)
+            raw = mne.io.RawArray(data, info)
+
+            # Save to file
+            file_name = f'dummy_eeg_{i+1}.fif'
+            file_path = os.path.join(folder_path, file_name)
+            raw.save(file_path)
+            # print(f'Saved EEG file {file_path}')
+
+    def read_eeg_file(self, file_path, chunk_size = 100):
+        # 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
+
+    def test_folder_constructor(self):
+        # Generate dummy EEG files and save them to folders
+        folder1 = './tmp_in/data/folder1'
+        folder2 = './tmp_in/data/folder2'
+        if not os.path.exists(folder1):
+            os.makedirs(folder1)
+        if not os.path.exists(folder2):
+            os.makedirs(folder2)
+
+        self.generate_dummy_eeg_files(2, folder1)
+        self.generate_dummy_eeg_files(2, folder2)
+
+        # Define input and output paths for folder_constructor
+        io_path = './tmp_out/eeg_folder'
+        root_path = './tmp_in/data'
+
+        # Call folder_constructor and pass read_eeg_file as argument
+        folder_constructor(io_path=io_path,
+                            root_path=root_path,
+                            num_channel=14,
+                            read_func=self.read_eeg_file,
+                            num_worker=0)
+
+
+    def test_folder_dataset(self):
+        # Generate dummy EEG files and save them to folders
+        folder1 = './tmp_in/data/folder1'
+        folder2 = './tmp_in/data/folder2'
+        if not os.path.exists(folder1):
+            os.makedirs(folder1)
+        if not os.path.exists(folder2):
+            os.makedirs(folder2)
+
+        self.generate_dummy_eeg_files(2, folder1)
+        self.generate_dummy_eeg_files(2, folder2)
+
+        # Define input and output paths for folder_constructor
+        io_path = './tmp_out/eeg_folder'
+        root_path = './tmp_in/data'
+
+        label_map = {'folder1':0,'folder2':1}
+        dataset = FolderDataset(
+            io_path=io_path,
+            root_path=root_path,
+            num_channel=14,
+            read_func = self.read_eeg_file,
+            online_transform=transforms.ToTensor(),
+            label_transform=transforms.Compose([
+                transforms.Select('label'),
+                transforms.Lambda(lambda x: label_map[x] )
+            ]),
+            num_worker=0)
+        
+        self.assertEqual(len(dataset), 20)
+        first_item = dataset[0]
+        self.assertEqual(first_item[0].shape, (14, 128))
+        self.assertEqual(first_item[1], 0)
+        last_item = dataset[-1]
+        self.assertEqual(last_item[0].shape, (14, 128))
+        self.assertEqual(last_item[1], 1)
+
+if __name__ == '__main__':
+    unittest.main()