Initial commit for public release

.gitattributes

@@ -0,0 +1 @@
+*.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,22 @@
+*__pycache__*
+*.ipynb
+*.pyc
+*.pyo
+*.pyd
+*.pkl
+*.pkl.gz
+*.h5
+*.hdf5
+*.csv
+*.log
+experiments/
+wandb/
+ckpt/
+.ipynb_checkpoints/
+lightning_logs/*
+lightning_logs/
+datasets
+Notebooks/
+embeddings/
+figures/
+ft_results/

README.md

@@ -0,0 +1,69 @@
+# SleepMaMi
+This is the code respository for **SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures**
+
+## Abstract
+>While the shift toward unified foundation models has revolutionized many deep learning domains, sleep medicine remains largely restricted to task-specific models that focus on localized micro-structure features. These approaches often neglect the rich, multi-modal context of Polysomnography (PSG) and fail to capture the global macro-structure of a full night's sleep. To address this, we introduce SleepMaMi, a Sleep Foundation Model engineered to master both hour-long sleep architectures and fine-grained signal morphologies. Our framework utilizes a hierarchical dual-encoder design: a Macro-Encoder to model full-night temporal dependencies and a Micro-Encoder to capture short-term characteristics from biosignals. Macro-Encoder is trained via Demographic-Guided Contrastive Learning, which aligns overnight sleep patterns with objective subject metadata, such as age and sex, to refine global representations. Micro-Encoder is optimized via a hybrid Masked Autoencoder (MAE) and multi-modal contrastive objective. Pre-trained on a massive corpus of >20,000 PSG recordings (15.8K hours), SleepMaMi outperforms existing foundation models across a diverse suite of downstream tasks, demonstrating superior generalizability and label-efficient adaptation for clinical sleep analysis.
+
+
+## Prerequisites
+```
+conda create -n sleepmami python==3.9
+conda activate sleepmami
+pip install -r requirements.txt
+
+# conda install -c "nvidia/label/cuda-12.1.1" cuda-nvcc cuda-cudart-dev libcublas-dev libcurand-dev # if necessary
+pip install flash-attn==2.7.4.post1 --no-build-isolation --cache-dir /tmp/pip_cache
+apt update
+apt install libgl1-mesa-glx
+pip install mamba-ssm==2.2.3.post2 --no-build-isolation
+```
+
+
+## Data preprocessing
+Please prepare pretraining and finetuning data, following [README](data/README.md) in `data` directory.
+
+## Pretraining
+### Micro-Encoder pretraining
+After pretraining dataset is prepared, prepare a config file(`*.yaml`) in `configs` directory.
+
+Example config files are provided in `configs` directory.
+
+Use the following command to start pretraining:
+```
+python pretrain.py --config PT_6mod_9ds_CL
+```
+
+
+### Macro-Encoder pretraining
+Similar to Micro-Encoder pretraining, prepare a config file in `configs` directory.
+
+Then use the following command to start Macro-Encoder pretraining
+```
+python pretrain_seq.py --config PTseq_6mod_9ds_CL_full_night
+```
+
+
+### Downstream tasks
+After pretraining is completed, finetuning on downstream tasks can be done with the following command:
+```
+python finetune.py --config FT_6mod_CL
+```
+The config file should be located in `configs` directory. `ft_task` variable should be set to specify the finetuning task (e.g. sleep_staging, apnea_seg, diag, etc.)
+
+
+### Pretrained checkpoints and Inference
+To be completed.
+
+
+## BibTeX
+Please use the below bibtex when you reference our work!
+```
+@inproceedings{
+park2026sleepmami,
+title={SleepMaMi: A Universal Sleep Foundation Model for Integrating Macro- and Micro-structures},
+author={Park, Keondo and Na, Younghoon and Choi, You Rim and Ryu, Hyunwoo and Shin, Hyun-Woo and Kim, Hyung-Sin},
+booktitle={Forty-third International Conference on Machine Learning},
+year={2026},
+url={https://openreview.net/forum?id=e9hvlHlUas}
+}
+```

configs/Eval_6mod_CL_full_night.yaml

@@ -0,0 +1,34 @@
+project_name: ConvMAE1d_FT
+experiment_name: Eval_PTseq_p50_6mod_9ds_kissFix_diag_micro
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: diagnosis
+patch_size: 50
+num_classes: 6
+num_workers: 6
+loss: coxph
+batch_size_v: 32
+proportion_v: 1.0
+gpus: [0]
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix_diag_micro-final.ckpt  # PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP","EMG","SaO2"
+]
+sig_len_target: 30
+use_revin: True
+full_night: True
+from_seq_pretrain: True
+seq_truncate: False
+use_lstm: False
+use_micro: True
+save_results: True
+save_emb: False

configs/Eval_6mod_CL_full_night_measure.yaml

@@ -0,0 +1,34 @@
+project_name: ConvMAE1d_FT
+experiment_name: Eval_PTseq_p50_6mod_9ds_kissFix
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+loss: wce
+batch_size_v: 8
+proportion_v: 1.0
+gpus: [0]
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix-final.ckpt  # PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP","EMG","SaO2"
+]
+sig_len_target: 30
+use_revin: True
+full_night: True
+from_seq_pretrain: True
+seq_truncate: False
+use_lstm: False
+use_micro: False
+save_results: False
+save_emb: False

configs/FT_6mod.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p100_10_6mod_cat
+data_root: datasets
+dataset: ['phy']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: sleep_staging
+patch_sizes: [10,100]
+num_classes: 5
+num_workers: 6
+lr: 0.002
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 3
+loss: wce
+batch_size_t: 512
+batch_size_v: 512
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p100_10_6mod-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","EMG","ECG","AIRFLOW","ABD"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p50_6mod_CL_shhs_CLonly_cat
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 512
+batch_size_v: 512
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_CL_shhs_CLonly_cat-final.ckpt #p50_6mod_6ds_CL5-final.ckpt #p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix5-final.ckpt # p25_6mod_hier_2_2_4_lgonly_5ds_CL_nonmix6-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL_abl1.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p50_6mod_noshared_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 512
+batch_size_v: 512
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_noshared_shhs-final.ckpt #p50_6mod_6ds_CL5-final.ckpt #p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix5-final.ckpt # p25_6mod_hier_2_2_4_lgonly_5ds_CL_nonmix6-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL_abl2.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p50_6mod_MAEonly_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 512
+batch_size_v: 512
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_MAEonly_shhs-final.ckpt #p50_6mod_6ds_CL5-final.ckpt #p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix5-final.ckpt # p25_6mod_hier_2_2_4_lgonly_5ds_CL_nonmix6-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL_abl3.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p50_6mod_MAEonly_noMoE_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 512
+batch_size_v: 512
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_MAEonly_noMoE_shhs-final.ckpt #p50_6mod_6ds_CL5-final.ckpt #p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix5-final.ckpt # p25_6mod_hier_2_2_4_lgonly_5ds_CL_nonmix6-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL_abl4.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p50_6mod_CL_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 512
+batch_size_v: 512
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_CL_shhs-final.ckpt #p50_6mod_6ds_CL5-final.ckpt #p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix5-final.ckpt # p25_6mod_hier_2_2_4_lgonly_5ds_CL_nonmix6-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL_ah.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p50_6mod_9ds_kissFix_ah_sof
+data_root: datasets
+datasets: ['sof']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: apnea_seg
+patch_size: 50
+num_classes: 2
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 1024
+batch_size_v: 1024
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_kissFix-final.ckpt # p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix2-final.ckpt # p25_6mod_hier_2_2_4_lgonly_5ds_CL_nonmix6-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG", "SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL_fewshot_ah.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_p50_6mod_9ds_kissFix_ah_cfs_1_e20
+data_root: datasets
+datasets: ['cfs_1'] #['isruc_sg1_1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi
+ft_task: apnea_seg
+patch_size: 50
+num_classes: 2
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 20
+loss: wce
+batch_size_t: 1024
+batch_size_v: 1024
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_kissFix-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True

configs/FT_6mod_CL_full_night.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_kissFix
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_micro: False
+use_lstm: False

configs/FT_6mod_CL_full_night_abl1.yaml

@@ -0,0 +1,39 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_CL_w_1_improve_shhs_bmi2
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_CL_w_1_improve_shhs_bmi2-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False

configs/FT_6mod_CL_full_night_abl2.yaml

@@ -0,0 +1,39 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_CL_w_1_improve_shhs_age2
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_CL_w_1_improve_shhs_age2-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False

configs/FT_6mod_CL_full_night_abl3.yaml

@@ -0,0 +1,39 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_CL_w_1_shhs_sex
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_CL_w_1_shhs_sex-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False

configs/FT_6mod_CL_full_night_abl4.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_kissFix_shhs_mamba_uni
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 32
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix_shhs_mamba_uni-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_transformer: False
+bidi: False

configs/FT_6mod_CL_full_night_age.yaml

@@ -0,0 +1,40 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_CL_w_1_improve2_age
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: age
+patch_size: 50
+num_classes: 1
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: mae
+batch_size_t: 16
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_CL_w_1_improve2-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_micro: False

configs/FT_6mod_CL_full_night_ahi.yaml

@@ -0,0 +1,40 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_CL_w_1_improve2_ahi_micro
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: ahi
+patch_size: 50
+num_classes: 1
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: mae
+batch_size_t: 16
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_CL_w_1_improve2-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_micro: False

configs/FT_6mod_CL_full_night_dino.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_kissFix_dino
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix_dino-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_micro: False
+use_lstm: False

configs/FT_6mod_CL_full_night_fewshot.yaml

@@ -0,0 +1,39 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_kissFix_cfs_100_e20
+data_root: datasets
+datasets: ['cfs_100'] #['isruc_sg1_1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 20
+loss: wce
+batch_size_t: 5
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False

configs/FT_6mod_CL_full_night_lstm.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_kissFix_cfs_lstm
+data_root: datasets
+datasets: ['cfs']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.001
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 20
+loss: wce
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_micro: False
+use_lstm: True

configs/FT_6mod_CL_full_night_sex.yaml

@@ -0,0 +1,40 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_CL_w_1_improve2_sex
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sex
+patch_size: 50
+num_classes: 2
+num_workers: 6
+lr: 0.004
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: ce
+batch_size_t: 8
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_CL_w_1_improve2-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_micro: False

configs/FT_6mod_CL_full_night_simclr.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_FT
+experiment_name: FT_PTseq_p50_6mod_9ds_kissFix_simclr2
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: sleep_staging
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: wce
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_kissFix_simclr2-final.ckpt  # PTseq_p50_6mod_6ds_CL5_5ds-epoch=00-train_loss=10.93.ckpt ## PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True 
+seq_truncate: False
+use_micro: False
+use_lstm: False

configs/FT_6mod_diag.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_FT
+experiment_name: PTseq_p50_6mod_9ds_CL_w_1_improve_diag
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+ft_task: diagnosis
+patch_size: 50
+num_classes: 6
+num_workers: 6
+lr: 0.01
+layer_decay: 1.0
+weight_decay: 0.0001
+num_epochs: 3
+loss: coxph
+batch_size_t: 4
+batch_size_v: 32
+proportion_t: 1.0
+proportion_v: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: PTseq_p50_6mod_9ds_CL_w_1_improve-final.ckpt  # PTseq_p50_4mod_hier_2_2_lgonly_4ds_CL_koelo_full_night-final.ckpt  # p50_4mod_hier_2_2_lgonly_4ds_CL_koelo-final.ckpt # 
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP","EMG","SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+from_seq_pretrain: True
+seq_truncate: False
+use_lstm: False
+use_micro: False

configs/PT_6mod_5ds_CL_test.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_kiss_CL_w_1_test
+data_root: datasets
+datasets: ['kiss'] #['phy-tr','phy-te','shhs1','shhs2','kiss','kvss','mros1','mros1-val','mros1-te','mros2','mros2-val','mros2-te','mesa','mesa-val','mesa-te']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 1
+batch_size_t: 1
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: True
+use_moe: True
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 1.0,
+  'moe': 0.001,
+  'ortho': 0.0,
+  'contrastive': 0.1,
+  'koleo': 0.01
+}

configs/PT_6mod_9ds_CL.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_9ds_kissFix
+data_root: datasets
+datasets: ['phy-tr','phy-te','shhs1','shhs2','kiss','kvss','mros1','mros1-val','mros1-te','mros2','mros2-val','mros2-te','mesa','mesa-val','mesa-te']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 2
+batch_size_t: 512
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: True
+use_moe: True
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 1.0,
+  'moe': 0.001,
+  'ortho': 0.0,
+  'contrastive': 0.1,
+  'koleo': 0.01
+}

configs/PT_6mod_9ds_CL_abl1.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_noshared_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 2
+batch_size_t: 512
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: False
+use_moe: False
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 1.0,
+  'moe': 0.0,
+  'ortho': 0.0,
+  'contrastive': 0.0,
+  'koleo': 0.0
+}

configs/PT_6mod_9ds_CL_abl2.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_MAEonly_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 2
+batch_size_t: 512
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: True
+use_moe: True
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 1.0,
+  'moe': 0.001,
+  'ortho': 0.0,
+  'contrastive': 0.0,
+  'koleo': 0.0
+}

configs/PT_6mod_9ds_CL_abl3.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_MAEonly_noMoE_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 2
+batch_size_t: 512
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: True
+use_moe: False
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 1.0,
+  'moe': 0.0,
+  'ortho': 0.0,
+  'contrastive': 0.0,
+  'koleo': 0.0
+}

configs/PT_6mod_9ds_CL_abl4.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_CL_shhs
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 2
+batch_size_t: 512
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: True
+use_moe: True
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 1.0,
+  'moe': 0.001,
+  'ortho': 0.0,
+  'contrastive': 0.1,
+  'koleo': 0.01
+}

configs/PT_6mod_9ds_CL_abl5.yaml

@@ -0,0 +1,41 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_CL_shhs_w3
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 2
+batch_size_t: 512
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: True
+use_moe: True
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 1.0,
+  'moe': 0.001,
+  'ortho': 0.0,
+  'contrastive': 0.3,
+  'koleo': 0.01
+}

configs/PT_6mod_9ds_CL_abl6.yaml

@@ -0,0 +1,42 @@
+project_name: ConvMAE1d_3
+experiment_name: p50_6mod_CL_shhs_CLonly_cat
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: convmae1d_multi
+patch_size: 50
+num_classes: 5
+num_workers: 6
+lr: 0.0005
+weight_decay: 0.05
+num_epochs: 2
+batch_size_t: 512
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+previous_run_id: null
+previous_checkpoint: null
+pretrained: null #experiments/p50_6mod_hier_2_2_lgonly_5ds_CL_nonmix8-final.ckpt 
+scheduler_name: CosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+lg_ratio: 1.0
+sig_len_target: 60
+use_shared: True
+use_moe: True
+use_revin: True
+loss: MSE
+loss_weights: {
+  'mae': 0.0,
+  'moe': 0.001,
+  'ortho': 0.0,
+  'contrastive': 1.0,
+  'koleo': 0.1
+}
+mask_ratio: 0.0

configs/PTseq_6mod_CL_full_night.yaml

@@ -0,0 +1,37 @@
+project_name: ConvMAE1d_PTseq
+experiment_name: PTseq_p50_6mod_9ds_kissFix
+data_root: datasets
+datasets: ['shhs1','shhs2','kiss','kvss','mesa','mesa-val','mesa-te','mros1','mros1-val','mros1-te','mros2','mros2-val','mros2-te'] #'phy-te','mros1'] # 
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+patch_size: 50
+num_classes: 5
+num_workers: 8
+lr: 0.0001
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 4
+batch_size_t: 40
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_kissFix-final.ckpt #p50_6mod_6ds_CL_w_1-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: WarmupCosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+seq_truncate: True
+demo_keys: ['age','sex','bmi']
+cycle_len: 180
+use_transformer: False

configs/PTseq_6mod_CL_full_night_abl1.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_PTseq
+experiment_name: PTseq_p50_6mod_9ds_CL_w_1_improve_shhs_bmi2
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+patch_size: 50
+num_classes: 5
+num_workers: 8
+lr: 0.0001
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 4
+batch_size_t: 40
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_CL_w_1-final.ckpt #p50_6mod_6ds_CL_w_1-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: WarmupCosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+seq_truncate: True
+demo_keys: ['age','sex','bmi']
+cycle_len: 180

configs/PTseq_6mod_CL_full_night_abl2.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_PTseq
+experiment_name: PTseq_p50_6mod_9ds_CL_w_1_improve_shhs_age2
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+patch_size: 50
+num_classes: 5
+num_workers: 8
+lr: 0.0001
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 4
+batch_size_t: 40
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_CL_w_1-final.ckpt #p50_6mod_6ds_CL_w_1-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: WarmupCosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+seq_truncate: True
+demo_keys: ['age','sex','bmi']
+cycle_len: 180

configs/PTseq_6mod_CL_full_night_abl3.yaml

@@ -0,0 +1,36 @@
+project_name: ConvMAE1d_PTseq
+experiment_name: PTseq_p50_6mod_9ds_kissFix_shhs_cycle1080
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+patch_size: 50
+num_classes: 5
+num_workers: 8
+lr: 0.0001
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 4
+batch_size_t: 40
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_kissFix-final.ckpt #p50_6mod_6ds_CL_w_1-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: WarmupCosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+seq_truncate: True
+demo_keys: ['age','sex','bmi']
+cycle_len: 1080

configs/PTseq_6mod_CL_full_night_abl4.yaml

@@ -0,0 +1,38 @@
+project_name: ConvMAE1d_PTseq
+experiment_name: PTseq_p50_6mod_9ds_kissFix_shhs_mamba_uni
+data_root: datasets
+datasets: ['shhs1']
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+patch_size: 50
+num_classes: 5
+num_workers: 8
+lr: 0.0001
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 4
+batch_size_t: 32
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_kissFix-final.ckpt #p50_6mod_6ds_CL_w_1-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: WarmupCosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+seq_truncate: True
+demo_keys: ['age','sex','bmi']
+cycle_len: 180
+use_transformer: False
+bidi: False

configs/PTseq_6mod_CL_full_night_dino.yaml

@@ -0,0 +1,39 @@
+project_name: ConvMAE1d_PTseq
+experiment_name: PTseq_p50_6mod_9ds_kissFix_dino
+data_root: datasets
+datasets: ['shhs1','shhs2','kiss','kvss','mesa','mesa-val','mesa-te','mros1','mros1-val','mros1-te','mros2','mros2-val','mros2-te'] #'phy-te','mros1'] # 
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+patch_size: 50
+num_classes: 5
+num_workers: 8
+lr: 0.0001
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 4
+batch_size_t: 24
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_kissFix-final.ckpt #p50_6mod_6ds_CL_w_1-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: WarmupCosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+seq_truncate: True
+demo_keys: ['age','sex','bmi']
+cycle_len: 180
+use_transformer: False
+simclr: False
+dino: True

configs/PTseq_6mod_CL_full_night_simclr.yaml

@@ -0,0 +1,38 @@
+project_name: ConvMAE1d_PTseq
+experiment_name: PTseq_p50_6mod_9ds_kissFix_simclr2
+data_root: datasets
+datasets: ['shhs1','shhs2','kiss','kvss','mesa','mesa-val','mesa-te','mros1','mros1-val','mros1-te','mros2','mros2-val','mros2-te'] #'phy-te','mros1'] # 
+config_save_path: experiments
+results_save_path: experiments
+ckpt_save_path: experiments
+model_name: classification_multi_full_night
+patch_size: 50
+num_classes: 5
+num_workers: 8
+lr: 0.0001
+layer_decay: 1.0
+weight_decay: 0.05
+num_epochs: 4
+batch_size_t: 24
+proportion_t: 1.0
+optimizer_name: AdamW
+gpus: [0,1]
+accumulate_grad_batches: 1
+num_experts: 4
+num_experts_per: 2
+pretrain_checkpoint: p50_6mod_9ds_kissFix-final.ckpt #p50_6mod_6ds_CL_w_1-final.ckpt
+previous_run_id: null
+previous_checkpoint: null
+scheduler_name: WarmupCosineAnnealingLR
+modalities: [
+  "EEG","EOG","ECG","RESP", "EMG" ,"SaO2"
+]
+sig_len_target: 30
+full_finetune: False
+use_revin: True
+full_night: True
+seq_truncate: True
+demo_keys: ['age','sex','bmi']
+cycle_len: 180
+use_transformer: False
+simclr: True

configs/global_cfg.py

@@ -0,0 +1,73 @@
+TARGET_RESAMPLE=100
+EPOCH_LEN_MAX=1080
+EPOCH_LEN_MIN=360
+DATASET_PATH = {
+    'phy-tr': 'physionet_2018/mmap_signals_MAE_processed',
+    'phy-te': 'physionet_test/mmap_signals_MAE_processed',
+    'shhs1': 'shhs1/mmap_signals_MAE_processed',
+    'shhs2': 'shhs2/mmap_signals_MAE_processed',
+    'kiss': 'kiss/mmap_signals_MAE_processed',
+    'kvss': 'kvss/mmap_signals_MAE_processed',
+    'mros1': 'mros1/mmap_signals_MAE_processed',    
+    'mros2': 'mros2/mmap_signals_MAE_processed',    
+    'mesa': 'mesa/mmap_signals_MAE_processed',
+    'mros1-val': 'mros1_val',
+    'mros1-te': 'mros1_te',
+    'mros2-val': 'mros2_val',
+    'mros2-te': 'mros2_te',
+    'mesa-val': 'mesa_val',
+    'mesa-te': 'mesa_te',
+    'isruc_sg1': 'isruc_all/mmap_signals_MAE_processed',
+    'isruc_sg1_1': 'isruc_all/mmap_signals_MAE_processed_train1',
+    'isruc_sg1_5': 'isruc_all/mmap_signals_MAE_processed_train5',
+    'isruc_sg1_10': 'isruc_all/mmap_signals_MAE_processed_train10',
+    'isruc_sg1_20': 'isruc_all/mmap_signals_MAE_processed_train20',
+    'isruc_sg1_30': 'isruc_all/mmap_signals_MAE_processed_train30',
+    'isruc_sg1_50': 'isruc_all/mmap_signals_MAE_processed_train50',
+    'cfs': 'cfs/mmap_signals_MAE_processed',
+    'cfs_1': 'cfs/mmap_signals_MAE_processed_train1',
+    'cfs_5': 'cfs/mmap_signals_MAE_processed_train5',
+    'cfs_10': 'cfs/mmap_signals_MAE_processed_train10',
+    'cfs_30': 'cfs/mmap_signals_MAE_processed_train30',
+    'cfs_100': 'cfs/mmap_signals_MAE_processed_train100',
+    'sof': 'sof/mmap_signals_MAE_processed',
+    
+}
+LABELS_PATH = {
+    'sleep_staging': {
+        'phy-tr': 'physionet_2018/annotations',
+        'shhs1': 'shhs1/annotations_split',
+        'kiss': 'kiss/annotations_split',
+        'isruc_sg1': 'isruc_all/annotations',       
+        'isruc_sg1_1': 'isruc_all/annotations',
+        'isruc_sg1_5': 'isruc_all/annotations',
+        'isruc_sg1_10': 'isruc_all/annotations',
+        'isruc_sg1_20': 'isruc_all/annotations',
+        'isruc_sg1_30': 'isruc_all/annotations',
+        'isruc_sg1_50': 'isruc_all/annotations',
+        'cfs': 'cfs/annotations_split',
+        'cfs_1': 'cfs/annotations',
+        'cfs_5': 'cfs/annotations',
+        'cfs_10': 'cfs/annotations',
+        'cfs_30': 'cfs/annotations',
+        'cfs_100': 'cfs/annotations',
+        'sof': 'sof/annotations',
+    },
+    'apnea_seg': {
+        'phy-tr': 'physionet_2018/ah_labels',
+        'kiss': 'kiss/ah_labels_split',
+        'shhs1': 'shhs1/ah_labels_split',
+        'cfs': 'cfs/ah_labels_split',
+        'cfs_1': 'cfs/ah_labels',
+        'cfs_5': 'cfs/ah_labels',
+        'cfs_10': 'cfs/ah_labels',
+        'cfs_30': 'cfs/ah_labels',
+        'cfs_100': 'cfs/ah_labels',
+        'sof': 'sof/ah_labels',
+
+    }
+}
+
+
+CLASSIFICATION_TASKS = ['sleep_staging', 'apnea_seg', 'sex']
+REGRESSION_TASKS = ['age', 'ahi']

data/README.md

@@ -0,0 +1,63 @@
+# Data Preprocessing
+
+This directory contains the preprocessing pipeline for several sleep study datasets, including **Physionet 2018, SHHS1/2, KISS/KVSS, MrOS1/2, and MESA**.
+
+---
+
+## 1. Raw Data Extraction
+### `{dataset_name}_process_rawdata.py`
+
+This script extracts signals and labels from raw data (e.g., `.edf` files) and saves them in `.npy` format.
+
+* **Channel Selection:** Extracts specific required channels: `EEG`, `EOG`, `EMG`, `ECG`, `ABD`, `CHEST`, `AIRFLOW`, and `SaO2`.
+* **Label Alignment:** * Synchronizes sleep staging labels with signal timeframes.
+    * Adjusts staging formats (e.g., converting **R&K** to **AASM** standards).
+    * Excludes error cases or undefined stages (e.g., "Unknown").
+* **Respiratory Events:** Extracts Apnea-Hypopnea labels where applicable (Apnea: 2, Hypopnea: 1).
+* **Wake Trimming:** If excessive "Wake" periods exist at the beginning or end of a recording, they are trimmed to a maximum of 30 minutes.
+* **Sampling Rate:** The original sampling rate is preserved during this step.
+* **Note:** For KISS/KVSS, the extraction logic is stored in a dedicated sub-folder.
+
+**Output Paths:**
+* `datasets/{dataset_name}/signals_MAE/{CHN_GRP}/{CHN}`: Raw biosignals.
+* `datasets/{dataset_name}/annotations`: Sleep staging labels.
+* `datasets/{dataset_name}/ah_labels`: Apnea-Hypopnea labels.
+
+---
+
+## 2. Signal Refinement & Memory Mapping
+### `{dataset_name}_preprocess.py`
+
+This script processes the `.npy` files (filtering, resampling) and saves them as memory-mapped (`.mmap`) files for efficient data loading.
+
+### Filtering Specifications
+
+| Signal Type | Filtering Parameters |
+| :--- | :--- |
+| **EEG, EOG** | Notch: 60 Hz \| Bandpass: 0.3–40 Hz |
+| **ECG** | Notch: 60 Hz \| Bandpass: 0.3–60 Hz \| Peak-based inversion correction |
+| **EMG** | Notch: 60 Hz \| Highpass: 10 Hz |
+| **AIRFLOW, ABD, CHEST** | Lowpass: 15 Hz |
+
+* **Resampling:** All signals are unified to a sampling rate of **100 Hz**.
+* **Metadata Generation:** A `.pickle` file is created for each biosignal to manage the global `.mmap` file. It includes:
+    * **Indices:** Start and end positions for each case within the combined `.mmap`.
+    * **Statistics:** Sum and square sum per case (used later for **Standardization**).
+    * **Shape:** The total dimensions of the `.mmap` file.
+* **Data Splitting:** Automatically performs the **Train/Val/Test** split (Only for finetuning datasets - SHHS1 and KISS).
+
+**Output Paths:**
+* `datasets/{dataset_name}/mmap_signals_MAE_processed/{split}/{CHN_GRP}/{CHN}`: `.mmap` and `.pickle` files.
+* `datasets/{dataset_name}/annotations_split/{split}`: Split sleep staging labels.
+* `datasets/{dataset_name}/ah_labels_split/{split}`: Split Apnea-Hypopnea labels.
+
+---
+
+## 3. Patient Demographics
+### `get_demographics.py`
+
+This script extracts and aggregates patient-level information into a structured format.
+
+* **Data Points:** Collects **Sex, Age, BMI, and AHI** (Apnea-Hypopnea Index).
+* **Handling Missing Values:** If a specific metric was not collected for a patient, it is ignored/skipped.
+* **Output Path:** `datasets/demographics/{dataset_name}_demographics.pkl`

data/cfs_meta_data/split_data.py

@@ -0,0 +1,51 @@
+# Create a script to split the dataset into train, validation, and test sets
+# The filenames are from BASE_PATH/cfs/signals_MAE/EEG/C3-M2/*.npy
+import os
+import random
+import numpy as np
+
+BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))), 'datasets')
+SEED = 42
+random.seed(SEED)
+np.random.seed(SEED)
+
+def split_data(data_dir, output_dir, train_ratio=0.7, val_num=50, test_ratio=0.3):
+    # Get all the filenames in the data directory
+    all_files = []
+    for root, dirs, files in os.walk(data_dir):
+        for file in files:
+            if file.endswith('.npy'):
+                pid = file.split('.')[0]  # Assuming filename format is PID_XXX.npy
+                all_files.append(pid)
+    
+    # Shuffle the filenames
+    random.shuffle(all_files)
+    
+    # Split the data into train, validation, and test sets
+    total_files = len(all_files)
+    train_end = int(total_files * train_ratio) - val_num
+    val_end = train_end + val_num
+    
+    train_files = all_files[:train_end]
+    val_files = all_files[train_end:val_end]
+    test_files = all_files[val_end:]
+    
+    # Save the splits to output directory
+    os.makedirs(output_dir, exist_ok=True)
+    
+    with open(os.path.join(output_dir, 'train.txt'), 'w') as f:
+        for file in train_files:
+            f.write(file + '\n')
+    
+    with open(os.path.join(output_dir, 'val.txt'), 'w') as f:
+        for file in val_files:
+            f.write(file + '\n')
+    
+    with open(os.path.join(output_dir, 'test.txt'), 'w') as f:
+        for file in test_files:
+            f.write(file + '\n')
+
+if __name__ == "__main__":
+    data_dir = os.path.join(BASE_PATH, 'cfs', 'signals_MAE', 'EEG', 'C3-M2')
+    output_dir = '.'
+    split_data(data_dir, output_dir)

data/cfs_meta_data/test.txt

@@ -0,0 +1,201 @@
+cfs-visit5-800900
+cfs-visit5-802451
+cfs-visit5-802745
+cfs-visit5-802201
+cfs-visit5-801900
+cfs-visit5-800048
+cfs-visit5-800891
+cfs-visit5-800195
+cfs-visit5-801220
+cfs-visit5-800680
+cfs-visit5-802110
+cfs-visit5-802180
+cfs-visit5-800412
+cfs-visit5-800688
+cfs-visit5-801645
+cfs-visit5-801942
+cfs-visit5-802065
+cfs-visit5-801782
+cfs-visit5-800916
+cfs-visit5-800690
+cfs-visit5-801677
+cfs-visit5-801945
+cfs-visit5-800038
+cfs-visit5-800544
+cfs-visit5-802104
+cfs-visit5-801602
+cfs-visit5-801773
+cfs-visit5-801814
+cfs-visit5-800375
+cfs-visit5-802242
+cfs-visit5-801035
+cfs-visit5-800968
+cfs-visit5-802385
+cfs-visit5-801323
+cfs-visit5-800604
+cfs-visit5-801651
+cfs-visit5-802380
+cfs-visit5-801126
+cfs-visit5-800780
+cfs-visit5-800113
+cfs-visit5-800170
+cfs-visit5-800365
+cfs-visit5-800908
+cfs-visit5-801770
+cfs-visit5-800705
+cfs-visit5-801085
+cfs-visit5-800697
+cfs-visit5-802335
+cfs-visit5-801337
+cfs-visit5-801788
+cfs-visit5-802449
+cfs-visit5-802303
+cfs-visit5-802228
+cfs-visit5-800861
+cfs-visit5-801162
+cfs-visit5-800194
+cfs-visit5-800985
+cfs-visit5-801735
+cfs-visit5-800021
+cfs-visit5-800980
+cfs-visit5-802482
+cfs-visit5-800078
+cfs-visit5-802185
+cfs-visit5-802171
+cfs-visit5-802691
+cfs-visit5-801396
+cfs-visit5-800769
+cfs-visit5-801576
+cfs-visit5-801794
+cfs-visit5-801762
+cfs-visit5-800926
+cfs-visit5-801432
+cfs-visit5-802038
+cfs-visit5-801058
+cfs-visit5-802460
+cfs-visit5-800247
+cfs-visit5-800271
+cfs-visit5-802075
+cfs-visit5-801454
+cfs-visit5-802010
+cfs-visit5-800226
+cfs-visit5-800672
+cfs-visit5-801445
+cfs-visit5-800197
+cfs-visit5-802533
+cfs-visit5-800313
+cfs-visit5-801225
+cfs-visit5-802716
+cfs-visit5-801299
+cfs-visit5-800847
+cfs-visit5-800122
+cfs-visit5-802593
+cfs-visit5-802573
+cfs-visit5-801037
+cfs-visit5-801638
+cfs-visit5-801719
+cfs-visit5-801246
+cfs-visit5-801325
+cfs-visit5-801587
+cfs-visit5-801438
+cfs-visit5-801291
+cfs-visit5-801393
+cfs-visit5-801540
+cfs-visit5-800716
+cfs-visit5-801152
+cfs-visit5-800229
+cfs-visit5-802612
+cfs-visit5-800490
+cfs-visit5-801311
+cfs-visit5-801528
+cfs-visit5-802072
+cfs-visit5-801675
+cfs-visit5-801101
+cfs-visit5-800617
+cfs-visit5-802005
+cfs-visit5-801080
+cfs-visit5-800203
+cfs-visit5-802252
+cfs-visit5-802549
+cfs-visit5-801352
+cfs-visit5-800665
+cfs-visit5-802658
+cfs-visit5-801144
+cfs-visit5-800434
+cfs-visit5-801902
+cfs-visit5-801806
+cfs-visit5-801559
+cfs-visit5-800841
+cfs-visit5-801150
+cfs-visit5-801685
+cfs-visit5-801335
+cfs-visit5-800658
+cfs-visit5-801159
+cfs-visit5-800136
+cfs-visit5-801497
+cfs-visit5-800719
+cfs-visit5-802724
+cfs-visit5-801458
+cfs-visit5-802678
+cfs-visit5-800105
+cfs-visit5-802475
+cfs-visit5-800449
+cfs-visit5-800326
+cfs-visit5-801650
+cfs-visit5-800380
+cfs-visit5-801446
+cfs-visit5-801859
+cfs-visit5-800236
+cfs-visit5-800953
+cfs-visit5-800049
+cfs-visit5-802637
+cfs-visit5-800971
+cfs-visit5-802487
+cfs-visit5-802648
+cfs-visit5-801833
+cfs-visit5-800321
+cfs-visit5-800494
+cfs-visit5-801381
+cfs-visit5-801264
+cfs-visit5-800395
+cfs-visit5-802717
+cfs-visit5-801270
+cfs-visit5-802368
+cfs-visit5-801747
+cfs-visit5-801166
+cfs-visit5-800440
+cfs-visit5-800311
+cfs-visit5-802123
+cfs-visit5-800811
+cfs-visit5-801137
+cfs-visit5-801294
+cfs-visit5-800545
+cfs-visit5-801095
+cfs-visit5-801724
+cfs-visit5-800512
+cfs-visit5-802498
+cfs-visit5-802125
+cfs-visit5-802753
+cfs-visit5-801278
+cfs-visit5-801501
+cfs-visit5-800594
+cfs-visit5-802154
+cfs-visit5-801662
+cfs-visit5-800011
+cfs-visit5-800905
+cfs-visit5-800017
+cfs-visit5-801706
+cfs-visit5-800329
+cfs-visit5-801855
+cfs-visit5-801231
+cfs-visit5-802295
+cfs-visit5-801732
+cfs-visit5-801019
+cfs-visit5-801526
+cfs-visit5-802268
+cfs-visit5-801110
+cfs-visit5-800211
+cfs-visit5-800851
+cfs-visit5-802537
+cfs-visit5-800471
+cfs-visit5-801809

data/cfs_meta_data/train.txt

@@ -0,0 +1,418 @@
+cfs-visit5-800238
+cfs-visit5-801487
+cfs-visit5-801341
+cfs-visit5-801427
+cfs-visit5-802462
+cfs-visit5-800733
+cfs-visit5-802346
+cfs-visit5-802467
+cfs-visit5-800914
+cfs-visit5-801181
+cfs-visit5-802435
+cfs-visit5-801850
+cfs-visit5-800325
+cfs-visit5-801660
+cfs-visit5-800099
+cfs-visit5-800801
+cfs-visit5-800805
+cfs-visit5-800854
+cfs-visit5-802247
+cfs-visit5-800565
+cfs-visit5-802273
+cfs-visit5-802398
+cfs-visit5-801005
+cfs-visit5-801173
+cfs-visit5-800173
+cfs-visit5-801267
+cfs-visit5-800522
+cfs-visit5-802408
+cfs-visit5-802711
+cfs-visit5-801928
+cfs-visit5-800946
+cfs-visit5-800860
+cfs-visit5-801300
+cfs-visit5-800220
+cfs-visit5-801457
+cfs-visit5-800722
+cfs-visit5-802183
+cfs-visit5-801350
+cfs-visit5-801452
+cfs-visit5-801282
+cfs-visit5-800640
+cfs-visit5-801800
+cfs-visit5-801252
+cfs-visit5-800151
+cfs-visit5-801122
+cfs-visit5-802233
+cfs-visit5-801917
+cfs-visit5-801263
+cfs-visit5-802302
+cfs-visit5-801312
+cfs-visit5-802078
+cfs-visit5-802272
+cfs-visit5-801027
+cfs-visit5-800316
+cfs-visit5-800341
+cfs-visit5-800346
+cfs-visit5-800047
+cfs-visit5-801734
+cfs-visit5-802195
+cfs-visit5-802676
+cfs-visit5-801222
+cfs-visit5-801052
+cfs-visit5-800299
+cfs-visit5-801521
+cfs-visit5-802294
+cfs-visit5-801183
+cfs-visit5-801125
+cfs-visit5-801867
+cfs-visit5-800086
+cfs-visit5-801346
+cfs-visit5-801377
+cfs-visit5-800344
+cfs-visit5-801631
+cfs-visit5-801416
+cfs-visit5-800492
+cfs-visit5-802491
+cfs-visit5-800371
+cfs-visit5-800109
+cfs-visit5-802003
+cfs-visit5-801013
+cfs-visit5-800937
+cfs-visit5-801449
+cfs-visit5-800198
+cfs-visit5-800285
+cfs-visit5-801656
+cfs-visit5-800427
+cfs-visit5-800714
+cfs-visit5-802235
+cfs-visit5-800390
+cfs-visit5-801488
+cfs-visit5-802308
+cfs-visit5-800359
+cfs-visit5-800161
+cfs-visit5-801609
+cfs-visit5-800426
+cfs-visit5-802509
+cfs-visit5-800177
+cfs-visit5-800939
+cfs-visit5-800575
+cfs-visit5-802400
+cfs-visit5-801212
+cfs-visit5-800506
+cfs-visit5-800138
+cfs-visit5-802466
+cfs-visit5-800118
+cfs-visit5-800973
+cfs-visit5-802340
+cfs-visit5-801790
+cfs-visit5-801366
+cfs-visit5-802366
+cfs-visit5-800110
+cfs-visit5-800799
+cfs-visit5-802726
+cfs-visit5-801139
+cfs-visit5-802094
+cfs-visit5-802042
+cfs-visit5-800757
+cfs-visit5-801453
+cfs-visit5-800280
+cfs-visit5-800187
+cfs-visit5-800509
+cfs-visit5-800576
+cfs-visit5-800081
+cfs-visit5-801158
+cfs-visit5-801289
+cfs-visit5-800216
+cfs-visit5-801849
+cfs-visit5-801780
+cfs-visit5-802474
+cfs-visit5-801530
+cfs-visit5-802386
+cfs-visit5-801610
+cfs-visit5-800347
+cfs-visit5-802182
+cfs-visit5-802073
+cfs-visit5-801813
+cfs-visit5-802106
+cfs-visit5-802500
+cfs-visit5-802459
+cfs-visit5-800228
+cfs-visit5-801909
+cfs-visit5-801199
+cfs-visit5-800659
+cfs-visit5-801848
+cfs-visit5-801012
+cfs-visit5-802529
+cfs-visit5-801448
+cfs-visit5-801206
+cfs-visit5-800421
+cfs-visit5-801474
+cfs-visit5-800618
+cfs-visit5-801998
+cfs-visit5-801912
+cfs-visit5-802211
+cfs-visit5-801585
+cfs-visit5-802060
+cfs-visit5-801357
+cfs-visit5-801792
+cfs-visit5-801785
+cfs-visit5-802503
+cfs-visit5-801021
+cfs-visit5-801355
+cfs-visit5-800227
+cfs-visit5-800032
+cfs-visit5-801590
+cfs-visit5-802731
+cfs-visit5-801919
+cfs-visit5-802555
+cfs-visit5-802103
+cfs-visit5-802205
+cfs-visit5-802734
+cfs-visit5-801591
+cfs-visit5-800364
+cfs-visit5-801319
+cfs-visit5-802004
+cfs-visit5-801112
+cfs-visit5-801376
+cfs-visit5-801825
+cfs-visit5-800072
+cfs-visit5-801491
+cfs-visit5-802424
+cfs-visit5-800159
+cfs-visit5-801707
+cfs-visit5-800404
+cfs-visit5-800630
+cfs-visit5-801038
+cfs-visit5-801858
+cfs-visit5-801936
+cfs-visit5-801064
+cfs-visit5-802109
+cfs-visit5-801340
+cfs-visit5-801864
+cfs-visit5-801373
+cfs-visit5-802659
+cfs-visit5-801389
+cfs-visit5-801444
+cfs-visit5-800400
+cfs-visit5-800911
+cfs-visit5-801133
+cfs-visit5-800666
+cfs-visit5-802626
+cfs-visit5-800505
+cfs-visit5-800393
+cfs-visit5-802619
+cfs-visit5-801331
+cfs-visit5-802427
+cfs-visit5-801326
+cfs-visit5-800883
+cfs-visit5-801569
+cfs-visit5-800148
+cfs-visit5-802206
+cfs-visit5-800527
+cfs-visit5-801731
+cfs-visit5-800928
+cfs-visit5-802011
+cfs-visit5-800114
+cfs-visit5-802634
+cfs-visit5-800949
+cfs-visit5-801213
+cfs-visit5-800643
+cfs-visit5-800024
+cfs-visit5-802218
+cfs-visit5-800737
+cfs-visit5-801666
+cfs-visit5-801370
+cfs-visit5-800962
+cfs-visit5-801233
+cfs-visit5-801983
+cfs-visit5-802322
+cfs-visit5-801421
+cfs-visit5-801771
+cfs-visit5-801982
+cfs-visit5-802523
+cfs-visit5-800917
+cfs-visit5-802495
+cfs-visit5-801633
+cfs-visit5-800349
+cfs-visit5-801380
+cfs-visit5-801078
+cfs-visit5-800821
+cfs-visit5-800557
+cfs-visit5-800551
+cfs-visit5-800657
+cfs-visit5-801838
+cfs-visit5-800996
+cfs-visit5-801461
+cfs-visit5-802578
+cfs-visit5-801963
+cfs-visit5-801514
+cfs-visit5-800461
+cfs-visit5-800340
+cfs-visit5-801177
+cfs-visit5-800312
+cfs-visit5-802718
+cfs-visit5-800528
+cfs-visit5-800221
+cfs-visit5-802465
+cfs-visit5-801916
+cfs-visit5-801757
+cfs-visit5-800531
+cfs-visit5-801805
+cfs-visit5-802300
+cfs-visit5-801248
+cfs-visit5-802522
+cfs-visit5-802018
+cfs-visit5-801517
+cfs-visit5-800839
+cfs-visit5-800798
+cfs-visit5-802315
+cfs-visit5-802257
+cfs-visit5-801420
+cfs-visit5-800002
+cfs-visit5-801310
+cfs-visit5-801243
+cfs-visit5-801831
+cfs-visit5-800616
+cfs-visit5-800410
+cfs-visit5-802421
+cfs-visit5-801075
+cfs-visit5-801494
+cfs-visit5-802113
+cfs-visit5-800129
+cfs-visit5-802610
+cfs-visit5-800890
+cfs-visit5-801554
+cfs-visit5-800411
+cfs-visit5-800392
+cfs-visit5-802382
+cfs-visit5-800613
+cfs-visit5-800693
+cfs-visit5-801996
+cfs-visit5-802720
+cfs-visit5-800183
+cfs-visit5-801502
+cfs-visit5-802587
+cfs-visit5-801001
+cfs-visit5-800037
+cfs-visit5-801933
+cfs-visit5-801463
+cfs-visit5-800279
+cfs-visit5-802099
+cfs-visit5-800923
+cfs-visit5-802721
+cfs-visit5-802101
+cfs-visit5-802196
+cfs-visit5-802132
+cfs-visit5-802595
+cfs-visit5-800853
+cfs-visit5-800825
+cfs-visit5-801196
+cfs-visit5-802021
+cfs-visit5-801841
+cfs-visit5-801425
+cfs-visit5-802144
+cfs-visit5-802081
+cfs-visit5-802635
+cfs-visit5-802041
+cfs-visit5-800363
+cfs-visit5-802654
+cfs-visit5-800912
+cfs-visit5-800498
+cfs-visit5-802438
+cfs-visit5-800451
+cfs-visit5-802298
+cfs-visit5-800597
+cfs-visit5-801307
+cfs-visit5-800027
+cfs-visit5-802355
+cfs-visit5-801193
+cfs-visit5-800104
+cfs-visit5-800537
+cfs-visit5-801658
+cfs-visit5-802213
+cfs-visit5-801966
+cfs-visit5-800115
+cfs-visit5-801835
+cfs-visit5-800718
+cfs-visit5-801668
+cfs-visit5-800789
+cfs-visit5-802502
+cfs-visit5-800625
+cfs-visit5-801147
+cfs-visit5-800040
+cfs-visit5-802086
+cfs-visit5-802330
+cfs-visit5-800356
+cfs-visit5-802076
+cfs-visit5-801157
+cfs-visit5-802056
+cfs-visit5-800249
+cfs-visit5-802401
+cfs-visit5-801674
+cfs-visit5-802388
+cfs-visit5-802739
+cfs-visit5-802301
+cfs-visit5-801413
+cfs-visit5-802709
+cfs-visit5-801034
+cfs-visit5-802579
+cfs-visit5-800334
+cfs-visit5-800960
+cfs-visit5-801811
+cfs-visit5-801563
+cfs-visit5-802675
+cfs-visit5-800212
+cfs-visit5-800385
+cfs-visit5-800230
+cfs-visit5-800179
+cfs-visit5-802348
+cfs-visit5-801224
+cfs-visit5-800092
+cfs-visit5-801033
+cfs-visit5-800975
+cfs-visit5-801387
+cfs-visit5-801121
+cfs-visit5-800045
+cfs-visit5-800535
+cfs-visit5-802219
+cfs-visit5-801041
+cfs-visit5-800820
+cfs-visit5-801907
+cfs-visit5-802627
+cfs-visit5-800927
+cfs-visit5-800199
+cfs-visit5-800331
+cfs-visit5-800232
+cfs-visit5-801228
+cfs-visit5-802223
+cfs-visit5-800243
+cfs-visit5-800176
+cfs-visit5-801532
+cfs-visit5-800647
+cfs-visit5-800272
+cfs-visit5-800990
+cfs-visit5-800715
+cfs-visit5-800606
+cfs-visit5-802483
+cfs-visit5-801414
+cfs-visit5-802748
+cfs-visit5-800035
+cfs-visit5-801226
+cfs-visit5-800501
+cfs-visit5-800389
+cfs-visit5-801922
+cfs-visit5-800481
+cfs-visit5-802176
+cfs-visit5-800464
+cfs-visit5-801827
+cfs-visit5-801596
+cfs-visit5-802645
+cfs-visit5-801161
+cfs-visit5-802283
+cfs-visit5-801713
+cfs-visit5-801873
+cfs-visit5-800554
+cfs-visit5-801044
+cfs-visit5-802643
+cfs-visit5-800276

data/cfs_meta_data/val.txt

@@ -0,0 +1,50 @@
+cfs-visit5-800202
+cfs-visit5-800184
+cfs-visit5-800702
+cfs-visit5-800071
+cfs-visit5-801071
+cfs-visit5-800367
+cfs-visit5-801960
+cfs-visit5-801764
+cfs-visit5-802304
+cfs-visit5-802001
+cfs-visit5-800130
+cfs-visit5-801320
+cfs-visit5-801689
+cfs-visit5-801131
+cfs-visit5-801149
+cfs-visit5-800337
+cfs-visit5-800246
+cfs-visit5-801273
+cfs-visit5-800174
+cfs-visit5-802628
+cfs-visit5-800474
+cfs-visit5-800945
+cfs-visit5-800010
+cfs-visit5-800182
+cfs-visit5-800667
+cfs-visit5-800785
+cfs-visit5-800967
+cfs-visit5-800497
+cfs-visit5-801077
+cfs-visit5-800065
+cfs-visit5-801760
+cfs-visit5-802454
+cfs-visit5-800355
+cfs-visit5-802332
+cfs-visit5-800051
+cfs-visit5-801988
+cfs-visit5-802471
+cfs-visit5-802279
+cfs-visit5-801949
+cfs-visit5-801908
+cfs-visit5-801553
+cfs-visit5-801435
+cfs-visit5-800596
+cfs-visit5-802664
+cfs-visit5-802177
+cfs-visit5-801868
+cfs-visit5-800407
+cfs-visit5-800901
+cfs-visit5-800577
+cfs-visit5-800373

data/cfs_preprocess.py

@@ -0,0 +1,372 @@
+import os
+import shutil
+import numpy as np
+import matplotlib.pyplot as plt
+from tqdm import tqdm
+import shutil
+import librosa
+
+from scipy.signal import iirnotch, lfilter
+from data_utils import butter_bandpass_filter, correct_ecg_inversion
+import pickle
+import pandas as pd
+
+BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'datasets')
+NOTCH_FREQ = 60
+EEG_LOWCUT = 0.3
+EEG_HIGHCUT = 40
+EMG_LOWCUT = 10
+EMG_HIGHCUT = 1000
+ECG_LOWCUT = 0.3
+ECG_HIGHCUT = 60
+FLOW_LOWCUT = 0
+FLOW_HIGHCUT = 15
+
+NUM_PROCESS = 4
+
+import argparse
+import multiprocessing
+
+TARGET_RESAMPLE = 100
+
+
+
+def process_ecg(args):     
+    sig, chn_grp, chn = args
+    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
+    data = np.load(sig_path)
+    sampling_rate = data.shape[1] // (30)
+                
+    data = data.reshape(-1, 1)
+    # Apply bandpass filter
+    filtered = butter_bandpass_filter(data, ECG_LOWCUT, ECG_HIGHCUT, sampling_rate)
+    b_notch, a_notch = iirnotch(NOTCH_FREQ, 30, sampling_rate)
+    filtered = lfilter(b_notch, a_notch, filtered)   
+
+    filtered = filtered.reshape(-1, sampling_rate*30)
+    # filtered = librosa.resample(filtered, orig_sr=SAMPLING_RATE, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq')
+    
+    filtered_epochs = []
+    
+    for epoch in filtered:
+        inv_epoch, _ = correct_ecg_inversion(epoch, TARGET_RESAMPLE)
+        filtered_epochs.append(inv_epoch)  
+
+    filtered = np.stack(filtered_epochs, axis=0)
+
+    if sampling_rate != TARGET_RESAMPLE:
+        filtered = np.reshape(filtered, (-1,1))
+        filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq')
+        filtered = filtered.reshape(-1, TARGET_RESAMPLE*30)                  
+    
+    if cfg.save_type == 'npy':
+        # Save processed data
+        output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig)
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        np.save(output_path, filtered.astype(np.float32))
+        return
+    else:
+        return filtered
+
+
+def process_eeg(args):
+    sig, chn_grp, chn = args
+    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
+    data = np.load(sig_path)
+    sampling_rate = data.shape[1] // (30)
+    data = data.reshape(-1, 1)
+    # Apply bandpass filter
+    filtered = butter_bandpass_filter(data, EEG_LOWCUT, EEG_HIGHCUT, sampling_rate)
+    b_notch, a_notch = iirnotch(NOTCH_FREQ, 30, sampling_rate)
+    filtered = lfilter(b_notch, a_notch, filtered)
+    if sampling_rate != TARGET_RESAMPLE:        
+        filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq')
+    filtered = filtered.reshape(-1, TARGET_RESAMPLE*30)
+    
+    if cfg.save_type == 'npy':
+        # Save processed data
+        output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig)
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        np.save(output_path, filtered.astype(np.float32))
+        return
+    else:
+        return filtered
+
+def process_emg(args):
+    sig, chn_grp, chn = args
+    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
+    data = np.load(sig_path)
+    sampling_rate = data.shape[1] // (30)
+    data = data.reshape(-1, 1)
+    # Apply bandpass filter
+    filtered = butter_bandpass_filter(data, EMG_LOWCUT, EMG_HIGHCUT, sampling_rate)
+    b_notch, a_notch = iirnotch(NOTCH_FREQ, 30, sampling_rate)
+    filtered = lfilter(b_notch, a_notch, filtered)    
+    if sampling_rate != TARGET_RESAMPLE:        
+        filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq')
+    filtered = filtered.reshape(-1, TARGET_RESAMPLE*30)
+    
+    if cfg.save_type == 'npy':
+        # Save processed data
+        output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig)
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        np.save(output_path, filtered.astype(np.float32))
+        return
+    else:
+        return filtered
+
+def process_flow(args):
+    sig, chn_grp, chn = args
+    sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
+    data = np.load(sig_path)
+    sampling_rate = data.shape[1] // (30)
+    data = data.reshape(-1, 1)
+    # Apply bandpass filter
+    filtered = butter_bandpass_filter(data, FLOW_LOWCUT, FLOW_HIGHCUT, sampling_rate)    
+    if sampling_rate != TARGET_RESAMPLE:        
+        filtered = librosa.resample(filtered, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq')
+    
+    # median filter?
+
+    filtered = filtered.reshape(-1, TARGET_RESAMPLE*30)
+    
+    if cfg.save_type == 'npy':
+        # Save processed data
+        output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig)
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+        np.save(output_path, filtered.astype(np.float32))
+        return
+    else:
+        return filtered
+
+def save_mmap(results, sig_list, chn_grp, chn, split='train'):
+    signame_list = [sig.split('.')[0] for sig in sig_list]
+    epoch_lengths = [sig.shape[0] for sig in results]
+    metadata = dict()
+    sig_info = dict()    
+    start, end = 0, 0    
+    for signame, sig in zip(signame_list, results):
+        length  = sig.shape[0]
+        end = start + length
+        info = dict()
+        info['pos'] = (start, end)
+        info['stats'] = {'sum': float(np.sum(sig)), 'sumsq': float(np.sum(sig ** 2))}
+        sig_info[signame] = info
+        start = end
+    
+
+    mmap_path = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.mmap')
+    signals_cat = np.concatenate(results, axis=0).astype(np.float32)    
+    os.makedirs(os.path.dirname(mmap_path), exist_ok=True)
+    data = np.memmap(mmap_path, dtype=np.float32, mode='w+', shape=signals_cat.shape)
+    data[:] = signals_cat
+    data.flush()
+
+    metadata['sig_info'] = sig_info
+    metadata['data_shape'] = signals_cat.shape
+    info_path = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.pickle')
+    with open(info_path, 'wb') as f:
+       pickle.dump(metadata, f)
+
+
+def preprocess(sig_list, chn_grp=None, chn=None, split='train'):
+    def _process(sig_list, chn_grp, chn, split):    
+        
+        print(f"Processing {chn_grp}/{chn}")
+        # if not (chn == "C4-A1" or chn == "AIRFLOW"): continue        
+        
+        if chn_grp in ['EEG','EOG']:
+            # sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))
+            input_list = [(sig, chn_grp, chn) for sig in sig_list]
+            with multiprocessing.Pool(NUM_PROCESS) as pool:
+                results = list(tqdm(pool.imap(process_eeg, input_list), total=len(input_list)))
+            if cfg.save_type == 'mmap':
+                save_mmap(results, sig_list, chn_grp, chn, split)
+
+        elif chn_grp in ['EMG']:
+            # sig_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn))
+            input_list = [(sig, chn_grp, chn) for sig in sig_list]
+            with multiprocessing.Pool(NUM_PROCESS) as pool:
+                results = list(tqdm(pool.imap(process_emg, input_list), total=len(input_list)))
+            if cfg.save_type == 'mmap':
+                save_mmap(results, sig_list, chn_grp, chn, split)
+
+        elif chn_grp in ['ECG']:            
+
+            # sig_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn))
+            input_list = [(sig, chn_grp, chn) for sig in sig_list]
+            with multiprocessing.Pool(NUM_PROCESS) as pool:
+                results = list(tqdm(pool.imap(process_ecg, input_list), total=len(input_list)))
+            if cfg.save_type == 'mmap':
+                save_mmap(results, sig_list, chn_grp, chn, split)
+
+        elif chn_grp in ['ABD', 'CHEST', 'AIRFLOW']:
+            # sig_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn))
+            input_list = [(sig, chn_grp, chn) for sig in sig_list]
+            with multiprocessing.Pool(NUM_PROCESS) as pool:
+                results = list(tqdm(pool.imap(process_flow, input_list), total=len(input_list)))
+            if cfg.save_type == 'mmap':
+                save_mmap(results, sig_list, chn_grp, chn, split)
+
+
+        else:
+            results = []
+            # for sig in tqdm(sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)))):
+            for sig in tqdm(sig_list):
+                sig_path = os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn, sig)
+                data = np.load(sig_path)
+                sampling_rate = data.shape[1] // (30)                
+                if sampling_rate != TARGET_RESAMPLE or chn == 'SaO2':
+                    # assuming 30s epochs
+                    data = data.reshape(-1, 1)
+                    if chn == 'SaO2':
+                        data = librosa.resample(data, orig_sr=sampling_rate, target_sr=1, axis=0, res_type='soxr_qq')
+                        print(np.min(data), np.max(data))
+                        data = data / 1e8 # data range 0 to 1e8
+                        if np.max(data) > 1 or np.min(data) < 0:
+                            print(f"Warning: SaO2 values out of range: {sig} with min {np.min(data)} and max {np.max(data)}")
+                        data = data.reshape(-1, 1*30)
+                    else:
+                        data = librosa.resample(data, orig_sr=sampling_rate, target_sr=TARGET_RESAMPLE, axis=0, res_type='soxr_qq')
+                        data = data.reshape(-1, TARGET_RESAMPLE*30)
+
+                if cfg.save_type == 'npy':
+                    # Save processed data
+                    output_path = os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn, sig)
+                    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+                    np.save(output_path, data)
+                elif cfg.save_type == 'mmap':
+                    results.append(data)
+
+            save_mmap(results, sig_list, chn_grp, chn, split)
+            
+
+    if chn_grp is not None and chn is not None:
+        _process(sig_list, chn_grp, chn, split)
+        return
+    else:
+        for chn_grp in os.listdir(os.path.join(BASE_PATH, cfg.data_dir)):
+            if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp)):
+                continue
+            for chn in os.listdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp)):  
+                if not os.path.isdir(os.path.join(BASE_PATH, cfg.data_dir, chn_grp, chn)):
+                    continue
+                _process(sig_list, chn_grp, chn, split) 
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--data_dir', default='cfs/signals_MAE', help='SHHS raw data location')
+parser.add_argument('--output_dir', default='cfs/mmap_signals_MAE_processed', help='SHHS processed data location')
+parser.add_argument('--ann_split', action='store_true', help='To split annotation files into train/val/test')
+parser.add_argument('--save_type', default='mmap', help='Which data type to save. "npy" or "mmap"')
+parser.add_argument('--chn_grp', default=None, help='If specified, only preprocess this channel group')
+parser.add_argument('--chn', default=None, help='If specified, only preprocess this channel')
+parser.add_argument('--num_train', type=int, default=50, help='Number of training samples')
+
+cfg = parser.parse_args()
+
+if __name__ == '__main__':
+
+    skip_val, skip_test = False, False
+    if cfg.num_train <= 100:        
+        if os.path.exists(os.path.join(BASE_PATH, cfg.output_dir, 'val')):
+            skip_val = True
+        if os.path.exists(os.path.join(BASE_PATH, cfg.output_dir, 'test')):
+            skip_test = True
+        cfg.output_dir += f"_train{cfg.num_train}"
+
+
+    def get_cfs_lists(cfs_split_file):
+        ids = pd.read_csv(os.path.join("cfs_meta_data", cfs_split_file), header=None, names=['a'])
+        ids = ids['a'].values.tolist()                
+        return [id + ".npy" for id in ids]
+
+    list_files_train = get_cfs_lists("train.txt")
+    list_files_val = get_cfs_lists("val.txt")
+    list_files_test = get_cfs_lists("test.txt")
+
+    sig_list = sorted(os.listdir(os.path.join(BASE_PATH, cfg.data_dir, "EEG", "C3-M2")))    
+    
+    list_files_train = sorted(list(set(list_files_train).intersection(set(sig_list))))
+    list_files_val = list(set(list_files_val).intersection(set(sig_list)))
+    list_files_test = list(set(list_files_test).intersection(set(sig_list)))
+
+    if cfg.num_train <= 100:
+        list_files_train = list_files_train[:cfg.num_train]
+
+    print("Number of records in training set:", len(list_files_train)) 
+    print("Number of records in validation set:", len(list_files_val)) 
+    print("Number of records in test set:", len(list_files_test)) 
+
+    print("Processing training set...")
+    preprocess(list_files_train, cfg.chn_grp, cfg.chn, split='train')
+    
+    if skip_val:
+        val_output_dir = cfg.output_dir
+        val_output_dir = val_output_dir.replace(f"_train{cfg.num_train}","")
+        os.symlink(os.path.join(BASE_PATH, val_output_dir, 'val'), os.path.join(BASE_PATH, cfg.output_dir, 'val'))
+    else:
+        print("Processing evaluation set...")
+        preprocess(list_files_val, cfg.chn_grp, cfg.chn, split='val')
+    
+    if skip_test:
+        test_output_dir = cfg.output_dir
+        test_output_dir = test_output_dir.replace(f"_train{cfg.num_train}","")
+        os.symlink(os.path.join(BASE_PATH, test_output_dir, 'test'), os.path.join(BASE_PATH, cfg.output_dir, 'test'))
+    else:
+        print("Processing test set...")
+        preprocess(list_files_test, cfg.chn_grp, cfg.chn, split='test')
+
+
+    # Reorganize annotations files
+    # ANN_PATH = os.path.join(BASE_PATH, 'shhs1/annotations')
+    if cfg.ann_split:
+        ANN_PATH = os.path.join(os.path.dirname(os.path.join(BASE_PATH, cfg.data_dir)), 'annotations')
+        NEW_ANN_PATH = ANN_PATH.replace('annotations', 'annotations_split')
+        os.makedirs(NEW_ANN_PATH, exist_ok=True)
+        os.makedirs(os.path.join(NEW_ANN_PATH,'train'), exist_ok=True)
+        os.makedirs(os.path.join(NEW_ANN_PATH,'val'), exist_ok=True)
+        os.makedirs(os.path.join(NEW_ANN_PATH,'test'), exist_ok=True)
+
+        for fname in tqdm(list_files_train):
+            ann_fname = fname.replace('.npy', '_ann.npy')
+            shutil.copy(os.path.join(ANN_PATH, ann_fname), os.path.join(NEW_ANN_PATH, 'train', ann_fname))
+        for fname in tqdm(list_files_val):
+            ann_fname = fname.replace('.npy', '_ann.npy')
+            shutil.copy(os.path.join(ANN_PATH, ann_fname), os.path.join(NEW_ANN_PATH, 'val', ann_fname))
+        for fname in tqdm(list_files_test):
+            ann_fname = fname.replace('.npy', '_ann.npy')
+            shutil.copy(os.path.join(ANN_PATH, ann_fname), os.path.join(NEW_ANN_PATH, 'test', ann_fname))
+
+        AH_PATH = os.path.join(os.path.dirname(os.path.join(BASE_PATH, cfg.data_dir)), 'ah_labels')        
+        NEW_AH_PATH = AH_PATH.replace('ah_labels', 'ah_labels_split')
+        os.makedirs(NEW_AH_PATH, exist_ok=True)
+        os.makedirs(os.path.join(NEW_AH_PATH,'train'), exist_ok=True)
+        os.makedirs(os.path.join(NEW_AH_PATH,'val'), exist_ok=True)
+        os.makedirs(os.path.join(NEW_AH_PATH,'test'), exist_ok=True)
+        for fname in tqdm(list_files_train):
+            ann_fname = fname.replace('.npy', '_ah.npy')
+            shutil.copy(os.path.join(AH_PATH, ann_fname), os.path.join(NEW_AH_PATH, 'train', ann_fname))
+        for fname in tqdm(list_files_val):
+            ann_fname = fname.replace('.npy', '_ah.npy')
+            shutil.copy(os.path.join(AH_PATH, ann_fname), os.path.join(NEW_AH_PATH, 'val', ann_fname))
+        for fname in tqdm(list_files_test):
+            ann_fname = fname.replace('.npy', '_ah.npy')
+            shutil.copy(os.path.join(AH_PATH, ann_fname), os.path.join(NEW_AH_PATH, 'test', ann_fname))
+
+
+    for split in ['train', 'val', 'test']:
+        #  create link files to respiratory channels
+        os.makedirs(os.path.join(BASE_PATH, cfg.output_dir, split, 'RESP'), exist_ok=True)
+        for chn_grp in ['ABD', 'CHEST', 'AIRFLOW']:
+            chns = [f.split('.')[0] for f in os.listdir(os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp))]
+            for chn in chns:
+                src_mmap = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.mmap')
+                src_pickle = os.path.join(BASE_PATH, cfg.output_dir, split, chn_grp, f'{chn}.pickle')
+                dst_mmap = os.path.join(BASE_PATH, cfg.output_dir, split, 'RESP', f'{chn}.mmap')
+                dst_pickle = os.path.join(BASE_PATH, cfg.output_dir, split, 'RESP', f'{chn}.pickle')
+                if not os.path.exists(dst_mmap):
+                    os.symlink(src_mmap, dst_mmap)
+                if not os.path.exists(dst_pickle):
+                    os.symlink(src_pickle, dst_pickle)
+        

data/cfs_process_rawdata.py

@@ -0,0 +1,291 @@
+import os
+import scipy
+import numpy as np
+#!pip install wfdb
+import pandas as pd
+from mne.io import concatenate_raws, read_raw_edf
+import xml.etree.ElementTree as ET
+import multiprocessing
+
+BASE_PATH = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'datasets')
+# BASE_PATH = 'datasets'
+EPOCH_SEC_SIZE = 30
+
+# AIRFLOW_NAMES = ["NEW AIR", "NEWAIR", "AIRFLOW", "AIRFLOW-0", "airflow", "new A/F", "New Air", "AUX"]
+# EEG2_NAMES = ["EEG2", "EEG 2", "EEG(sec)", "EEG sec", "EEG(SEC)"]
+SELECT_CHNS = ['SaO2', 'SpO2','ECG1', 'EMG1', 'LOC', 'ROC', 'C3', 'C4', 'M1', 'M2', 'A1','A2', 'ABDO EFFORT', 'THOR EFFORT', 'AIRFLOW']
+CHN_SAVE_DICT = {
+    'SaO2': ('SaO2','SaO2'),
+    'SpO2': ('SaO2','SaO2'),
+    'ECG1': ('ECG', 'ECG'),
+    'EMG1': ('EMG', 'EMG'),
+    'LOC': ('EOG', 'EOG-Diff'),
+    'ROC': ('EOG', 'EOG-Diff'),
+    'C3': ('EEG', 'C3-M2'),
+    'C4': ('EEG', 'C4-M1'),
+    'M1': ('EEG', 'C4-M1'),
+    'M2': ('EEG', 'C3-M2'),
+    'A1': ('EEG', 'A1-M2'),
+    'A2': ('EEG', 'A2-M1'),
+    'ABDO EFFORT': ('ABD', 'ABD'),
+    'THOR EFFORT': ('CHEST', 'CHEST'),    
+    'AIRFLOW': ('AIRFLOW', 'AIRFLOW')
+}
+
+
+
+import argparse
+parser = argparse.ArgumentParser()
+parser.add_argument('--data_dir', default="CFS_raw/polysomnography/edfs", help='CFS raw data location')
+parser.add_argument('--output_dir', default='cfs/signals_MAE', help='CFS processed data location')
+parser.add_argument('--annotation_dir', default='CFS_raw/polysomnography/annotations-events-profusion', help='CFS annotations location')
+parser.add_argument('--ah_dir', default='CFS_raw/polysomnography/annotations-events-nsrr', help='CFS apnea/hypopnea labels location')
+parser.add_argument('--demographics_file', default='CFS_raw/datasets/cfs-visit5-harmonized-dataset-0.7.0.csv', help='CFS demographics file location')
+
+cfg = parser.parse_args()
+
+from tqdm.auto import tqdm
+
+def process_signal(raw, select_idx, chns, fname):
+    # sampling rate from instance
+    raw_sampling_rate = int(raw.info['sfreq'])
+    raw_df = raw.to_data_frame()
+    # print(raw_df.head())
+    print(raw_df.columns)
+    for chn in chns:    
+        if chn == "LOC":
+            raw_df['EOG-Diff'] = raw_df['LOC'] - raw_df['ROC']
+            raw_ch_df = raw_df['EOG-Diff']
+        elif chn == "ROC":
+            continue
+        elif chn == 'C3':
+            if 'M2' in raw_df.columns:
+                raw_df['C3-M2'] = raw_df['C3'] - raw_df['M2']
+            else:
+                raw_df['C3-M2'] = raw_df['C3'] - raw_df['A2']
+            raw_ch_df = raw_df['C3-M2']
+        elif chn == 'C4':
+            if 'M1' in raw_df.columns:
+                raw_df['C4-M1'] = raw_df['C4'] - raw_df['M1']
+            else:
+                raw_df['C4-M1'] = raw_df['C4'] - raw_df['A1']
+            raw_ch_df = raw_df['C4-M1']
+        elif chn in ['M1', 'M2', 'A1', 'A2']:
+            continue
+        elif chn in ['SaO2', 'SpO2']:
+            if 'SaO2' in raw_df.columns:                
+                raw_ch_df = raw_df['SaO2']
+            else:
+                raw_ch_df = raw_df['SpO2']
+        else:        
+            raw_ch_df = raw_df[chn]
+        raw_ch_df = raw_ch_df.to_frame()
+        raw_ch_df.set_index(np.arange(len(raw_ch_df)))
+
+        # Remove movement and unknown stages if any
+        raw_ch = raw_ch_df.values
+
+        # Verify that we can split into 30-s epochs
+        if len(raw_ch) % (EPOCH_SEC_SIZE * raw_sampling_rate) != 0:
+            print(f"Length of raw channel ({len(raw_ch)}) \
+                is not multiple of raw sampling rate and epoch size ({EPOCH_SEC_SIZE} * {raw_sampling_rate})")
+            return np.array([])                    
+
+        # get (# of epochs) = (# of datapoints) / (30sec * 100Hz)
+        n_epochs = len(raw_ch) / (EPOCH_SEC_SIZE * raw_sampling_rate)
+        
+        # Get epochs and their corresponding labels
+        x = np.asarray(np.split(raw_ch, n_epochs)).astype(np.float32)   # split epochs
+
+        x = x[select_idx]
+        x = x.reshape((-1, EPOCH_SEC_SIZE * raw_sampling_rate))
+
+        chn_grp, chn_rename = CHN_SAVE_DICT[chn]
+        # chn_grp = CHN_SAVE_DICT[chn][0] if chn not in AIRFLOW_NAMES else 'AIRFLOW'
+        # chn_rename = CHN_SAVE_DICT[chn][1] if chn not in AIRFLOW_NAMES else 'AIRFLOW'
+        # print(chn_grp, chn_rename)
+        os.makedirs(os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn_rename), exist_ok=True)
+        np.save(os.path.join(BASE_PATH, cfg.output_dir, chn_grp, chn_rename, f"{fname}.npy"), x)
+
+def process_data(args):
+
+    ann_fname, ah_fname, edf_fname, file_id = args
+    print(args)
+
+    ann_filename = os.path.basename(edf_fname).replace(".edf", "_ann.npy")
+    # if os.path.exists(os.path.join(BASE_PATH, cfg.output_dir, "annotations", ann_filename)):
+    #     return
+
+    labels = []
+
+    # Read annotation and its header
+    t = ET.parse(ann_fname)
+
+    r = t.getroot()
+    faulty_File = 0
+
+    for i in range(len(r[4])):
+        lbl = int(r[4][i].text)
+        # print(lbl)
+        if lbl == 4:  # make stages N3, N4 same as N3
+            labels.append(3)
+        elif lbl == 5:  # Assign label 4 for REM stage
+            labels.append(4)
+        else:
+            labels.append(lbl)
+        if lbl > 5:  # some files may contain labels > 5 BUT not the selected ones.
+            faulty_File = 1
+
+    if faulty_File == 1:
+        print("Faulty labels found!!")
+        return -1
+
+    labels = np.asarray(labels)     # label list to npaprray
+
+    # corresponding labels
+    y = labels.astype(np.int32)
+
+    # Select on sleep periods
+    w_edge_mins = 30
+    # get index of first not WAKE
+    nw_idx = np.where(y != 0)[0]
+    # start from 30min before first WAKE
+    start_idx = nw_idx[0] - (w_edge_mins * 2)
+    # end at 30min after first WAKE
+    end_idx = nw_idx[-1] + (w_edge_mins * 2)
+    if start_idx < 0:
+        start_idx = 0
+    if end_idx >= len(y):
+        end_idx = len(y) - 1
+
+    select_idx = np.arange(start_idx, end_idx + 1)
+
+    # only selected labels
+    y = y[select_idx]
+
+    # get RawEDF instance
+    # !!!!!!!!!!!!!!!!!! 12~13 sec    
+    raw = read_raw_edf(edf_fname, preload=True, stim_channel=None, verbose=None)        
+    process_signal(raw, select_idx, SELECT_CHNS, os.path.basename(edf_fname).split('.')[0])
+
+    ann_dir = os.path.join(BASE_PATH, 'cfs', 'annotations')
+    os.makedirs(ann_dir, exist_ok=True)
+    np.save(os.path.join(ann_dir, ann_filename), arr=y)
+
+    tree = ET.parse(ah_fname)
+    root = tree.getroot()
+    events = root.findall('.//ScoredEvent')
+    
+    total_seconds = 0
+    
+    for event in events:
+        if event.find('EventConcept').text == 'Recording Start Time':
+            duration = float(event.find('Duration').text)
+            total_seconds = int(np.ceil(duration))
+            break
+            
+    # 만약 Recording Start Time을 못 찾았을 경우 (예외 처리), 기존 방식대로 마지막 이벤트 시간으로 계산
+    if total_seconds == 0:
+        print("예외처리로 인한 Start time ~ End time 계산")
+        max_time = 0
+        for event in events:
+            try:
+                start = float(event.find('Start').text)
+                dur = float(event.find('Duration').text)
+                if start + dur > max_time:
+                    max_time = start + dur
+            except:
+                continue
+        total_seconds = int(np.ceil(max_time))
+
+    # 2. 라벨 어레이 생성
+    ah_labels = np.zeros(total_seconds, dtype=int)    
+
+    
+    # 3. 호흡 이벤트 파싱 및 라벨링
+    for event in events:
+        try:
+            event_type = event.find('EventType').text
+            
+            if event_type != 'Respiratory|Respiratory': continue             # 호흡 이벤트가 아니면 스킵
+
+            concept = event.find('EventConcept').text
+            start = float(event.find('Start').text)
+            duration = float(event.find('Duration').text)
+            
+            start_idx = int(round(start))
+            end_idx = start_idx + int(round(duration))
+            
+            # 인덱스 범위 보정
+            if end_idx > total_seconds:
+                end_idx = total_seconds
+            
+            # 라벨링 (Hypopnea=1, Apnea=2)
+            if 'Hypopnea' in concept:
+                mask = ah_labels[start_idx:end_idx] != 2
+                ah_labels[start_idx:end_idx][mask] = 1
+                
+            elif any(x in concept for x in ['Obstructive', 'Central', 'Mixed']):
+                ah_labels[start_idx:end_idx] = 2
+                
+        except AttributeError:
+            continue
+
+    ah_labels = np.reshape(ah_labels, (-1, 30))
+    ah_labels = ah_labels[select_idx]    
+    assert len(ah_labels) == len(y), f"Length of ah_labels ({len(ah_labels)}) does not match length of sleep stage labels ({len(y)})"
+
+    ah_labels = np.reshape(ah_labels, (-1,))    
+
+    ah_dir = os.path.join(BASE_PATH, 'cfs', 'ah_labels')
+    ah_filename = os.path.basename(edf_fname).replace(".edf", "_ah.npy")
+    os.makedirs(ah_dir, exist_ok=True)
+    np.save(os.path.join(ah_dir, ah_filename), arr=ah_labels)
+
+
+if __name__ == '__main__':
+
+    # data_dir = '/tf/00_data/kdpark/physionet_2018/physionet_rawdata'
+    file_list = os.listdir(os.path.join(BASE_PATH, cfg.data_dir))    
+    
+    demo = pd.read_csv(os.path.join(BASE_PATH, cfg.demographics_file))
+    def get_nsrrid(filename):
+        nsrrid = filename.split('-')[-1].replace('.edf', '')
+        return int(nsrrid)
+    
+    age_filtered_df = demo[demo['nsrr_age'] > 13]    
+    
+    file_list = [filename for filename in file_list if get_nsrrid(filename) in age_filtered_df['nsrrid'].values]
+    print(f"Number of files after filtering by age > 13: {len(file_list)}")
+
+    filenames_edf = [filename for filename in file_list if filename.endswith('.edf')]
+    filenames_ann = [filename.replace('.edf', '-profusion.xml') for filename in filenames_edf]
+    filenames_ah = [filename.replace('.edf', '-nsrr.xml') for filename in filenames_edf]
+
+
+    edf_fnames = [os.path.join(BASE_PATH, cfg.data_dir, filename_edf) for filename_edf in filenames_edf]
+    ann_fnames = [os.path.join(BASE_PATH, cfg.annotation_dir,  filename_ann) for filename_ann in filenames_ann]
+    ah_fnames = [os.path.join(BASE_PATH, cfg.ah_dir, filename_ah) for filename_ah in filenames_ah]
+
+    edf_fnames.sort()   # sort edf file names list
+    ann_fnames.sort()   # sort ann file names list
+    ah_fnames.sort()    # sort ah file names list
+
+    input_list = zip(ann_fnames, ah_fnames, edf_fnames, range(len(ann_fnames)))    
+
+
+    # redo_list = ['shhs1-203353.npy', 'shhs1-203171.npy', 'shhs1-203301.npy', 'shhs1-203011.npy', 'shhs1-203300.npy', 'shhs1-203419.npy', 'shhs1-203286.npy', 'shhs1-203375.npy','shhs1-201820.npy']
+    # input_list = [args for args in input_list if os.path.basename(args[1]).replace('.edf', '.npy') in redo_list]
+    # print(len(input_list))
+    # exit(0)
+
+    # for args in input_list:
+    #     ann_fname, edf_fname, file_id = args
+    #     if 'shhs1-201820' in edf_fname:
+    #         process_data(args)
+        # break
+
+    with multiprocessing.Pool(8) as pool:
+        list(tqdm(pool.imap(process_data, input_list), total=len(ann_fnames)))
+
+

data/data_augments.py

@@ -0,0 +1,318 @@
+import torch
+import random
+import numpy as np
+import torch.nn.functional as F
+""" 
+Functions for Video Data augmentation. 
+Transforms are applied to the tensor frames of shape (frames_per_slide, height, width).
+"""
+class TensorRandomFlip:
+    """Randomly flips the tensor horizontally or vertically."""
+    def __init__(self, horizontal=True, vertical=False, p=0.5):
+        self.horizontal = horizontal
+        self.vertical = vertical
+        self.p = p
+
+    def __call__(self, frames):
+        # Get 2 random numbers for horizontal and vertical flips
+        rands = np.random.rand(2)
+        if rands[0] < self.p and self.horizontal:
+            frames = torch.flip(frames, dims=[2])
+        if rands[1] < self.p and self.vertical:
+            frames = torch.flip(frames, dims=[1])
+        return frames
+        
+
+class TensorNormalize01:
+    """Normalizes the tensor into 0 and 1"""
+    def __init__(self):
+        pass
+
+    def __call__(self, frames):
+        return frames / 255.0
+    
+
+class TensorNormalize:
+    """Normalizes the tensor with the given mean and std.
+        Mean and std not calculated yet.
+    """
+    def __init__(self, mean=0.5, std=0.5):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, frames):
+        # return (frames/255.0 - self.mean) / self.std        
+        for i in range(len(frames.shape) - 1):
+            m = torch.tensor(self.mean, dtype=torch.float32).unsqueeze(0)
+            s = torch.tensor(self.std, dtype=torch.float32).unsqueeze(0)
+
+        return (frames - m) / s
+
+class TensorRandomScaling:
+    """Randomly scales the tensor (image) within a specified range."""
+    def __init__(self, scale_range=(0.8, 1.2)):
+        """
+        Args:
+            scale_range (tuple): A tuple (min_scale, max_scale) specifying the range for scaling.
+        """
+        assert len(scale_range) == 2, "scale_range must be a tuple with (min_scale, max_scale)."
+        assert scale_range[0] > 0 and scale_range[1] > 0, "Scaling factors must be positive."
+        self.scale_range = scale_range
+
+    def __call__(self, frames):
+        """
+        Args:
+            frames (torch.Tensor): Tensor of shape (batch_size, height, width).
+        
+        Returns:
+            torch.Tensor: Scaled tensor of the same shape.
+        """
+        frames_per_slide, height, width = frames.shape
+        
+        # Randomly choose a scaling factor
+        scale = random.uniform(*self.scale_range)
+
+        # New height and width after scaling
+        new_height = int(height * scale)
+        new_width = int(width * scale)
+
+        # Add channel dimension
+        frames = frames.unsqueeze(1).float()
+
+        # Resize the tensor
+        resized_frames = F.interpolate(
+            frames,  # Add channel dimension: (frames_per_slide, 1, height, width)
+            size=(new_height, new_width),
+            mode='bilinear',
+            align_corners=False
+        ).squeeze(1)  # Remove channel dimension: (frames_per_slide, new_height, new_width)
+
+        # Center-crop or pad to return to original size
+        if scale > 1.0:  # Scale up: Center crop
+            start_y = (new_height - height) // 2
+            start_x = (new_width - width) // 2
+            scaled_frames = resized_frames[:, start_y:start_y + height, start_x:start_x + width]
+        else:  # Scale down: Pad
+            pad_y = (height - new_height) // 2
+            pad_x = (width - new_width) // 2
+            padding = (pad_x, width - new_width - pad_x, pad_y, height - new_height - pad_y)
+            scaled_frames = F.pad(resized_frames, padding, mode='constant', value=0)
+
+        return scaled_frames.squeeze(1).type_as(frames)
+
+
+class TensorCenterCrop:
+    """Center crop the tensor (image) within a specified range."""
+    def __init__(self, crop_size=(112,112)):
+        """
+        Args:
+            crop_size (tuple): A tuple (height, width) specifying image size after crop.
+        """
+        assert len(crop_size) == 2, "Target size must be a tuple with (height, width)."        
+        self.crop_size = crop_size
+
+    def __call__(self, frames):
+        """
+        Args:
+            frames (torch.Tensor): Tensor of shape (batch_size, height, width).
+        
+        Returns:
+            torch.Tensor: Scaled tensor of the same shape.
+        """
+        frames_per_slide, height, width = frames.shape
+        
+        crop_height = self.crop_size[0]
+        crop_width = self.crop_size[1]
+        # Randomly choose a scaling factor        
+        start_y = (height - crop_height) // 2
+        start_x = (width - crop_width) // 2
+        cropped_frames = frames[:, start_y:start_y + crop_height, start_x:start_x + crop_width]
+    
+
+        return cropped_frames.type_as(frames)
+
+class TensorRandomCrop:
+    """Center crop the tensor (image) within a specified range."""
+    def __init__(self, crop_size=(112,112)):
+        """
+        Args:
+            crop_size (tuple): A tuple (height, width) specifying image size after crop.
+        """
+        assert len(crop_size) == 2, "Target size must be a tuple with (height, width)."        
+        self.crop_size = crop_size
+
+    def __call__(self, frames):
+        """
+        Args:
+            frames (torch.Tensor): Tensor of shape (batch_size, height, width).
+        
+        Returns:
+            torch.Tensor: Scaled tensor of the same shape.
+        """
+        frames_per_slide, height, width = frames.shape
+        
+        crop_height = self.crop_size[0]
+        crop_width = self.crop_size[1]
+
+
+        start_y = random.randint(0, height - crop_height)
+        start_x = random.randint(0, width - crop_width)
+
+        cropped_frames = frames[:, start_y:start_y + crop_height, start_x:start_x + crop_width]
+    
+
+        return cropped_frames.type_as(frames)
+
+
+class TensorResize:
+    """Resizes the tensor (image)"""
+    def __init__(self, target_size=(112,112)):
+        """
+        Args:
+            target_size (tuple): A tuple (height, width) specifying image size after resize.
+        """
+        assert len(target_size) == 2, "scale_range must be a tuple with (min_scale, max_scale)."        
+        self.target_size = target_size
+
+    def __call__(self, frames):
+        """
+        Args:
+            frames (torch.Tensor): Tensor of shape (batch_size, height, width).
+        
+        Returns:
+            torch.Tensor: Resized tensor of the same shape.
+        """
+        frames_per_slide, height, width = frames.shape
+        
+
+        # New height and width after scaling
+        new_height = self.target_size[0]
+        new_width = self.target_size[1]
+
+        # Add channel dimension
+        frames = frames.unsqueeze(1).float()        
+
+        # Resize the tensor
+        resized_frames = F.interpolate(
+            frames,  # Add channel dimension: (frames_per_slide, 1, height, width)
+            size=(new_height, new_width),
+            mode='bilinear',
+            align_corners=False
+        ).squeeze(1)  # Remove channel dimension: (frames_per_slide, new_height, new_width)
+
+        return resized_frames.squeeze(1).type_as(frames)
+
+    
+
+# TransformRotation
+# [N, H, W] shape의 video tensor를 입력으로 받아서, angle 을 범위 내에서 랜덤하게 선택하고, angle로 모든 frame을 회전시킨다.
+# torchvision.transforms.RandomRotation을 활용하면 될 것 같다.
+import torchvision.transforms.functional as TF
+
+class TensorRandomRotation:
+    """Randomly rotates the tensor (image) within a specified angle range."""
+    def __init__(self, degrees):
+        """
+        Args:
+            degrees (sequence or float or int): Range of degrees to select from. If degrees is a number instead of sequence like (min, max), the range of degrees will be (-degrees, +degrees).
+        """
+        if isinstance(degrees, (int, float)):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            self.degrees = (-degrees, degrees)
+        else:
+            self.degrees = degrees
+
+    def __call__(self, frames):
+        """
+        Args:
+            frames (torch.Tensor): Tensor of shape (frames_per_slide, height, width).
+        
+        Returns:
+            torch.Tensor: Rotated tensor of the same shape.
+        """
+        # Randomly choose an angle
+        angle = random.uniform(self.degrees[0], self.degrees[1])
+        
+        # Add channel dimension for each frame
+        frames = frames.unsqueeze(1)  # Shape becomes [N, 1, H, W]
+        
+        # Rotate all frames at once
+        rotated_frames = TF.rotate(frames, angle)
+        
+        # Remove channel dimension
+        return rotated_frames.squeeze(1)
+
+class TensorRandomRotation90:
+    """Randomly rotates the tensor (image) within a specified angle range."""
+    def __init__(self):
+        pass
+
+    def __call__(self, frames):
+        """
+        Args:
+            frames (torch.Tensor): Tensor of shape (frames_per_slide, height, width).
+        
+        Returns:
+            torch.Tensor: Rotated tensor of the same shape.
+        """
+        # Randomly choose an angle
+        p = random.uniform(0, 1)
+
+        if p < 0.5:
+            return frames
+        else:
+            p2 = random.uniform(0, 1)
+            if p2 < 0.5:
+                angle = 90
+            else:
+                angle = -90
+        
+        # Add channel dimension for each frame
+        frames = frames.unsqueeze(1)  # Shape becomes [N, 1, H, W]
+        
+        # Rotate all frames at once
+        rotated_frames = TF.rotate(frames, angle)
+        
+        # Remove channel dimension
+        return rotated_frames.squeeze(1)
+
+class TensorRandomShift:
+    """Randomly shifts the tensor (image) within a specified range."""
+    def __init__(self, shift_range=(0.1, 0.1)):
+        """
+        Args:
+            shift_range (tuple): A tuple (shift_x, shift_y) specifying the range for shifting.
+        """
+        assert len(shift_range) == 2, "shift_range must be a tuple with (shift_x, shift_y)."
+        self.shift_range = shift_range
+
+    def __call__(self, frames):
+        """
+        Args:
+            frames (torch.Tensor): Tensor of shape (batch_size, height, width).
+        
+        Returns:
+            torch.Tensor: Shifted tensor of the same shape.
+        """
+        batch_size, height, width = frames.shape
+        
+        # Randomly choose shift values
+        shift_x = random.uniform(-self.shift_range[0], self.shift_range[0])
+        shift_y = random.uniform(-self.shift_range[1], self.shift_range[1])
+
+        # Add channel dimension
+        frames = frames.unsqueeze(1).float()
+
+        # Create the affine transformation matrix
+        theta = torch.tensor([[1.0, 0.0, shift_x], [0.0, 1.0, shift_y]], dtype=torch.float)
+        theta = theta.unsqueeze(0).repeat(batch_size, 1, 1)
+
+        # Create the grid
+        grid = F.affine_grid(theta, frames.size(), align_corners=False)
+
+        # Apply the grid to the frames
+        shifted_frames = F.grid_sample(frames, grid, align_corners=False)
+
+        # Remove channel dimension
+        return shifted_frames.squeeze(1)