PD_SSL_ZOO

This is the codebase for the paper "Enhancing 3D Dopamine Transporter Imaging as a Biomarker for Parkinson's Disease via Self-Supervised Learning with Diffusion Models".

Publication

Enhancing 3D Dopamine Transporter Imaging as a Biomarker for Parkinson's Disease via Self-Supervised Learning with Diffusion Models

Jongjun Won¹, Grace Yoojin Lee¹, Sungyang Jo¹, Jihyun Lee¹, Sangjin Lee¹, Jae Seung Kim¹, Changhwan Sung¹, Jungsu S. Oh¹, Kyum-Yil Kwon², Soo Bin Park², Joonsang Lee¹, Jieun Yum¹, Sun Ju Chung¹, and Namkug Kim¹

¹ Asan Medical Center, ² Soonchunhyang University Seoul Hospital
Cell Reports Medicine (Acceptance, to appear in 2025)

Contents

This repository is composed of

1_UPSTREAM
2_DOWNSTREAM
3_RECONSTRUCTION
4_LATENT_MANIPULATION

Our overall workflow code parts are mainly in "1_UPSTREAM" and "2_DOWNSTREAM," illustrated below:

This repository is based on other repositories of MONAI, lucidrains, and eladrich.

Monai Generative Models : HWDAE, WDDAE, DDAE, HDAE

Monai/research-contribution (DisAE) : DisAE, SimMIM

lucidrains/StyleGAN2-pytorch : StyleGAN2

lucidrains/denoising_diffusion_pytorch : WDDAE, DDAE

eladrich/Pixel2Style2Pixel : P2S2P

requirements

pip install -r requirements.txt

Pretrained upstream model weights & synthetic samlpe scans

One Drive/Weight

One Drive/Synthetic FP-CIT-PET Samples

Train & Test

There are directories for each upstream model and downstream task.

Models : 1_HWDAE, 2_WDDAE, 3_DDAE, 4_P2S2P, 5_DisAE, 6_HDAE, 7_SimMIM

Tasks : 1_EP, 2_PMP, 3_SOY

Upstream:

For the pre-training stage of SSL models.

-> Please refer the "main.py" in the each folders of "/1_UPSTREAM/Models/."

python main.py --batch_size 2 --log_dir <log_dir>

Downstream:

For the linear probing, training from scratch, or fine-tuning stages of downstream tasks from the upstream models.

-> Please refer the "main.py" in the each folders of "/2_DOWNSTREAM/Tasks/."

python main.py --batch_size <batch_size> --name <model_name> --log_dir <log_dir> --data_per <data_percentage> --linear_mode <linear | scratch | fine_tuning> 

Generation, Reconstruction and Latent Manipulation:

Unconditional image generation of "Models" in [WDDAE, DDAE, StyleGAN2].

-> Please refer the "generate.py" in the each folders of "/3_RECONSTRUCTION/0_GENERATION/Models/."

python generate.py

Image Reconstruction of "Models" in [HWDAE, HDAE, P2S2P, DisAE, SimMIM]

-> Please refer the "RECONSTRUCTION.ipynb" in the each folders of "/3_RECONSTRUCTION/Models/."

Latent Manipulation of HWDAE.

-> Please refer the "HWDAE_MANIPULATION.ipynb" in the folder of "/4_HWDAE_MANIPULATION/."