Complex Wavelet-Based Transformer for Neurodevelopmental Disorder Diagnosis via Direct Modeling of Real and Imaginary Components

This repository contains the official implementation of the paper "Complex Wavelet-Based Transformer for Neurodevelopmental Disorder Diagnosis via Direct Modeling of Real and Imaginary Components." (Medical Image Analysis, 2026)

Ah-Yeong Jeong*, Da-Woon Heo*, and Heung-Il Suk^†
*Equally contributed, ^†Corresponding author

Getting Started

Installation

# Clone the repository
git clone https://github.com/ayjxxng/CWBrainNet.git
cd CWBrainNet

# Create and activate the environment
conda create -n cwbrainnet python=3.10
conda activate cwbrainnet

# Install dependencies
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
pip install -r requirements.txt

Data Preparation

This project expects preprocessed dataset files in .npz format and corresponding index files in .pkl format.

1. Place your dataset and index files in a directory (e.g., dataset/).
2. Update the configuration files in conf/dataset/ to point to your data, or use the default paths:
   • conf/dataset/ABIDE.yaml: Default path ./dataset/ABIDE.npz
   • conf/dataset/ADHD.yaml: Default path ./dataset/ADHD.npz

Training

To train the model, use the provided scripts in the scripts/ directory:

# Train on ABIDE dataset
bash scripts/train_abide.sh

# Train on ADHD dataset
bash scripts/train_adhd.sh

Citation

If you find this work useful, please cite our paper:

@article{JEONG2026103914,
title = {Complex wavelet-based Transformer for neurodevelopmental disorder diagnosis via direct modeling of real and imaginary components},
journal = {Medical Image Analysis},
volume = {109},
pages = {103914},
year = {2026},
issn = {1361-8415},
doi = {https://doi.org/10.1016/j.media.2025.103914},
url = {https://www.sciencedirect.com/science/article/pii/S1361841525004608},
author = {Ah-Yeong Jeong and Da-Woon Heo and Heung-Il Suk},
}

Acknowledgments

This work was supported by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) No. RS-2022-II220959 ((Part 2) FewShot Learning of Causal Inference in Vision and Language for Decision Making), Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government(MSIT) (No. RS-2019-II190079, Artificial Intelligence Graduate School Program(Korea University)), and Institute of Information & communications Technology Planning & Evaluation (IITP) under the artificial intelligence star fellowship support program to nurture the best talents (IITP-2025-RS-2025-02304828) grant funded by the Korea government(MSIT).