Skip to content

3D U-net, Attention U-net, Res U-net, Attention Res U-net, and MSRes U-net are implemented and compared for emulation of current density induced during transcranial direct current stimulation (tDCS).

License

Notifications You must be signed in to change notification settings

jxf0623/DeeptDCS

Repository files navigation

DeeptDCS: Deep learning-based estimation of currents induced during transcranial direct current stimulation

This is an implementation of a deep learning-based transcranial direct current stimulation (tDCS) emulator named DeeptDCS (https://arxiv.org/abs/2205.01858) on Python 3, Keras, and TensorFlow. The emulator leverages Attention U-net taking the volume conductor models (VCMs) of human head tissues as inputs and outputting the three-dimensional current density distribution across the entire head.

3D U-net and its four variants (Attention U-net, Res U-net, Attention Res U-net, and MSRes U-net) are implemented and their performance are compared.

The code is designed to be easy to extend. If you use it in your research, please consider citing our paper (bibtex below).

Fig 1. DeeptDCS workflow

Fig 2. DeeptDCS results

Requirements

  • Python 3.6
  • TensorFlow 2.2.0
  • Keras 2.3.0-tf
  • Other common libraries

Data structure

Data is structured as follows:

├── dir(subject ID 1)
│  ├── dir(Montage Position 1)
|  |  ├── dir(1)
|  |  |  ├── file(field_cond.mat)
|  |  ├── dir(2)
|  |  |  ├── file(field_cond.mat)
|  |  ├── ...
│  |  └──dir(X)
│  ├── dir(Montage Position 2)
|  ├── ...
│  └── dir(Montage Position Y)
├── dir(subject ID 2)
├── ...
└── dir(subject ID Z)

An example dataset can be downloaded from DeeptDCS_data_samples*.

Model Zoo

This model_zoo* contains pre-trained models presented in our DeeptDCS paper.

To use the pre-trained model, please rename the model as 'DeeptDCS.hdf5' and copy it to folder './saved_model/'.

1. Models pre-trained on 59,000 samples constructed from 59 MRIs and five 5*5cm2 square montage positions.

2. Attention U-net models fine-tuned on

Training and Testing

Neural network selection

U-net and its four variations are implemented. The default network is Attention U-net. To change the network, please search self.model = self.attn_unet_3d(self.num_gpu) in model.py and replace the command by one of the following:

self.model = self.unet_3d(self.num_gpu)
self.model = self.MSResUnet_3d(self.num_gpu)
self.model = self.ResUnet_3d(self.num_gpu)
self.model = self.attn_unet_3d(self.num_gpu)
self.model = self.attn_ResUnet_3d(self.num_gpu)

1. Test pre-trained DeeptDCS Attention U-net model on the example dataset

python train DeeptDCS.py --varification_test=True

2. Train and test a new model from scratch

  • Clear or delete directory './saved_model/'
  • In dataloader.py, change self.subjectIDs, self.electrode_positions, self.data_size_1position_1subject according to the new dataset.
python train DeeptDCS.py --data_path=/path/to/dataset/

3. Test a well-trained model

  • Copy the model to './saved_model/' and rename it as 'DeeptDCS.hdf5'
  • In DeeptDCS.py, set args.epochs = 0
python train DeeptDCS.py --data_path=/path/to/dataset/

4. Transfer learning for non-trained montage configurations

  • Copy the pre-trained model to './saved_model/' and rename it as 'DeeptDCS.hdf5'
  • In dataloader.py, change self.subjectIDs, self.electrode_positions, self.data_size_1position_1subject according to the new dataset.
python train DeeptDCS.py --data_path=/path/to/transfer_learning_dataset/

5. Train and test on custom dataset

  • Change dataloader.py to load data according to the costume dataset structure.

* The aforementioned data samples and pre-trained models are also available at 百度云 (提取码tDCS).

Citation

Please use this bibtex to cite our paper:

@ARTICLE{9915436,  
  author={Jia, Xiaofan and Sayed, Sadeed Bin and Hasan, Nahian Ibn and Gomez, Luis J. and Huang, Guang-Bin and Yucel, Abdulkadir C.},  
  journal={IEEE Transactions on Biomedical Engineering},   
  title={DeeptDCS: Deep Learning-Based Estimation of Currents Induced During Transcranial Direct Current Stimulation},   
  year={2022},  
  volume={},  
  number={},  
  pages={1-11},  
  doi={10.1109/TBME.2022.3213266}
}

Acknowledgment

DeeptDCS emulator is developed based on IntelAI/unet.

TO DO

change citation

About

3D U-net, Attention U-net, Res U-net, Attention Res U-net, and MSRes U-net are implemented and compared for emulation of current density induced during transcranial direct current stimulation (tDCS).

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages