LeFusion: Lesion-Focused Diffusion Model. The top illustrates the training process of LeFusion, while the bottom shows the inference. During training, LeFusion avoids learning unnecessary background generation using a lesion-focused loss. In inference, by combining forward-diffused real backgrounds with reverse-diffused generated foregrounds, LeFusion ensures high-quality background generation. Additionally, we introduce histogram-based texture control to handle multi-peak lesions and multi-channel decomposition for multi-class lesions. (arXiv)
We utilized the LIDC dataset (PubMed), which includes 1,010 chest CT scans. From these, we extracted 2,624 pathology regions of interest (ROIs) related to lung nodules to train the LeFusion Model. The dataset is divided into 808 cases for training, containing 2,104 lung nodule ROIs, and 202 cases for testing, containing 520 lung nodule ROIs. This portion of the dataset is located in LIDC-IDRI\Pathological
, with the test.txt
listing the data used for testing.
Additionally, we provide 20 normal ROIs from healthy patients, representing areas where lung nodules typically appear. This data is located in LIDC-IDRI\Normal
, where Image
contains the healthy images, and Mask
includes the corresponding masks generated by matching lung and ground truth masks, which can be used to generate lesions. You can simulate lesion generation on the Normal dataset.
Furthermore, we provide pre-generated images with lesions based on the LIDC-IDRI\Normal
dataset. These images are stored in LIDC-IDRI\Demo
, where Image_i
represents the images generated under the control information hist_i. The pre-trained weights used to generate these images are available in the pre-trained weights mentioned below (HuggingFace🤗).
├── LIDC-IDRI
├── Pathological
│ ├── Image
│ ├── Mask
│ └── test.txt
├── Normal
│ ├── Image
│ └── Mask
└── Demo
├── Image
│ ├── Image_1
│ ├── Image_2
│ └── Image_3
└── Mask
├── Mask_1
├── Mask_2
└── Mask_3
- Create a virtual environment
conda create -n lefusion python=3.10
and activate itconda activate lefusion
- Download the code
git clone https://github.com/M3DV/LeFusion.git
- Check if your pip version is 22.3.1. If it is not, install pip version 22.3.1
pip install pip==22.3.1
- Enter the LeFusion folder
cd LeFusion/LeFusion_LIDC
and runpip install -r requirements.txt
-
Download the LIDC_IDRI dataset (HuggingFace🤗)
In our study, the LeFusion Model focuses on the generation of lung nodule regions.If you want to train a Diffusion Model to synthesize lung nodules, you can use the LIDC-IDRI dataset that has already been processed by us to train the LeFusion Model. Just put the LIDC-IDRI dataset to
LeFusion/LeFusion_LIDC/data
.✨Note: Before running the following command, make sure you are inside the
LeFusion/LeFusion_LIDC
folder.mkdir data cd data mkdir LIDC-IDRI cd LIDC-IDRI wget https://huggingface.co/datasets/YuheLiuu/LIDC-IDRI/resolve/main/Pathological.tar -O Pathological.tar tar -xvf Pathological.tar wget https://huggingface.co/datasets/YuheLiuu/LIDC-IDRI/resolve/main/Normal.tar -O Normal.tar tar -xvf Normal.tar wget https://huggingface.co/datasets/YuheLiuu/LIDC-IDRI/resolve/main/Demo.tar -O Demo.tar tar -xvf Demo.tar
-
Download the pre-trained LeFusion Model (HuggingFace🤗)
We offer the pre-trained LeFusion Model, which has been trained for 50,001 steps on the LIDC-IDRI dataset. This pre-trained model can be directly used for Inference if you do not want to re-train the LeFusion Model. Simply download it to
LeFusion/LeFusion_model
.cd ../.. mkdir LeFusion_model cd LeFusion_model wget https://huggingface.co/YuheLiuu/LeFusion/resolve/main/LIDC_LeFusion_Model/model-50.pt -O model-50.pt
If you have downloaded the pre-trained model, you can skip the training step and proceed directly to inference!
Start training:
✨Note: Before running the following command, make sure you are inside the
LeFusion/LeFusion_LIDC
folder.
test_txt_dir=data/LIDC-IDRI/Pathological/test.txt
dataset_root_dir=data/LIDC-IDRI/Pathological/Image
train_num_steps=50001
python train/train.py dataset.test_txt_dir=$test_txt_dir dataset.root_dir=$dataset_root_dir model.train_num_steps=$train_num_steps
Notably, data_path
actually refers to the directory location of the corresponding images. Additionally, the corresponding label directory should be placed in the same folder as the image directory and should be named Mask
.
Our model was trained for 50,000 steps using five 40GB A100 GPUs, taking two and a half days. However, we found that the model performs very well after 20,000 steps. Therefore, when training a model on your own, anywhere between 20,000 to 50,000 steps would yield good results. Additionally, by default, we save the weights every 1,000 steps, and you can modify the relevant parameters in LeFusion/train/config
.
Start inference:
✨Note: Before running the following command, make sure you are inside the
LeFusion/LeFusion_LIDC
folder.
test_txt_dir=data/LIDC-IDRI/Pathological/test.txt
dataset_root_dir=data/LIDC-IDRI/Normal/Image/
target_img_path=data/LIDC-IDRI/gen/Image/
target_label_path=data/LIDC-IDRI/gen/Mask/
jump_length=5
jump_n_sample=5
batch_size=2
python test/inference.py batch_size=$batch_size test_txt_dir=$test_txt_dir dataset_root_dir=$dataset_root_dir target_img_path=$target_img_path target_label_path=$target_label_path schedule_jump_params.jump_length=$jump_length schedule_jump_params.jump_n_sample=$jump_n_sample
Three folders, Image_1, Image_2, and Image_3, will be generated under the target_img_path
directory, each representing images generated under the control of hist_1, hist_2, and hist_3 respectively. Similarly, three folders will be generated under the Mask directory, but unlike the Image folders, files with the same name in each of the three Mask folders contain the same mask.
For jump_length and jump_n_sample, larger values generally result in longer image generation times. We found that when these two parameters are between 2 and 10, the generated images maintain good quality. When both parameters are set to 2, it takes about 40 seconds to generate an image using a 40G A100 GPU.
The first image is a healthy image from LIDC-IDRI/Normal
. The second image is the corresponding generated mask, where lesions will be generated in the areas marked by the mask. Image_1, Image_2, and Image_3 are the lesions generated when the control information is set to Hist_1, Hist_2, and Hist_3, respectively.
@misc{zhang2024lefusioncontrollablepathologysynthesis,
title={LeFusion: Controllable Pathology Synthesis via Lesion-Focused Diffusion Models},
author={Hantao Zhang and Yuhe Liu and Jiancheng Yang and Shouhong Wan and Xinyuan Wang and Wei Peng and Pascal Fua},
year={2024},
eprint={2403.14066},
archivePrefix={arXiv},
primaryClass={eess.IV},
url={https://arxiv.org/abs/2403.14066},
}
Some of our code is modified based on medicaldiffusion and RePaint, and we greatly appreciate the efforts of the respective authors for providing open-source code. We also thank DiffTumor for providing the segmentation model code.
Community Contribution: 3D Slicer Extension for LeFusion
For those who work with medical imaging and seek to bring LeFusion's inpainting model closer to real-world clinical practice, we are excited to introduce a community contribution: a 3D Slicer extension! This extension leverages our inpainting model as the backend, offering practical applications for radiologists and other medical professionals.
Special thanks to @pedr0sorio for developing this valuable tool.
✅ The preprocessed LIDC-IDRI dataset 🚀
✅ The LeFusion model applied to LIDC-IDRI 🚀
🔲 The DiffMask model used for generating mask
🔲 The preprocessed EMIDC dataset
🔲 The LeFusion model applied to EMIDC