BS-Diff: Effective Bone Suppression in CXRs via Conditional Diffusion Models

🧨 Congratulations! Our paper has been accepted by ISBI 2024(Oral Presentation)!

Proposed method

We spend a lot of time collecting and summarizing relevant papers and datasets, where you can find them at https://github.com/diaoquesang/A-detailed-summarization-about-bone-suppression-in-Chest-X-rays

This code is a PyTorch implementation of our paper "BS-Diff: Effective Bone Suppression in CXRs via Conditional Diffusion Models".

Our proposed framework comprises two stages: a conditional diffusion model (CDM) equipped with a U-Net architecture and a simple enhancement module that incorporates an autoencoder. It can not only generate soft tissue images with a high bone suppression ratio but also possess the capability to capture fine image information and spatial features, while preserving overall structures. The figure below shows our proposed network.

The bone-suppressed images generated by our method

Comparison performance with previous works (visualization)

Clinical evaluation

The results below demonstrated that our soft-tissues can clearly preserve the visibility of pulmonary vessels and central airways and greatly suppress bones, significantly improving the clinician’s performance in finding lung lesions. Each criterion has a maximum score of 3.

Clinical Evaluation Criteria		Junior clinian	Intermediate clinian	Senior clinian
Pulmonary vessels visibility	Clearly displayed (3)	2	3	3
	Displayed (2)
	Not displayed (1)
Central airway visibility	Lobar and intermediate bronchi (3)	2	3	2
	Main bronchus and rump (2)
	Trachea (1)
Degree of bone sup- pression	Nearly perfect suppression (3)	2	3	2
	Unsuppressed bones less than 5 (2)
	5 or more bones unsuppressed (1)

Pre-requisties

• Linux

• Python>=3.7

• NVIDIA GPU (memory>=6G) + CUDA cuDNN

Download the dataset

Now, we only provide three paired images with CXRs and soft-tissues via pre-processing. Soon, we will make them available to the public after data usage permission. Three paired images are located at

├─ Data
│    ├─ BS_Aug
│    │    ├─ 0.png
│    │    ├─ 1.png
│    │    └─ 2.png
│    ├─ CXR_Aug
│    │    ├─ 0.png
│    │    ├─ 1.png
│    │    └─ 2.png

Getting started to evaluate

Install dependencies

pip install -r requirements.txt

Download the checkpoint

Due to the fact that our proposed model comprises two stages, you need to download both stages' checkpoints to successfully run the codes! These two files can be found in the following link :

https://drive.google.com/drive/folders/1cDlXJ7Sh4k05aM_tvzor9_F_TPCeIGMN?usp=sharing

Evaluation

To do the evaluation process, first run the following command in stage 1 (the conditional diffusion model):

python Test.py

Then, you will get a series of images generated by the conditional diffusion model. After that, run the following command in stage 2 with these images as inputs.

python Hybrid_autoencodereval.py

Train by yourself

If you want to train our model by yourself, you are primarily expected to split the whole dataset into training, validation, and testing. You can find the codes in Data Spliting directory and run the following commands one by one:

python txt.py
python split.py

Then, you can run the following command in stage 1:

python Train.py

Then after finishing stage 1, you can use the generated output of stage 1 to train our stage (enhancement module) by running the following command:

python Hybridloss_autoencoder.py

These two files are located at

├─ Stage1
│    └─ Train.py
├─ Stage2
│    ├─ Hybridloss_autoencoder.py
│    └─ pytorch_msssim.py

Evaluation metrics

You can also run the following commands about evaluation metrics in our experiment including PSNR, SSIM, MSE and BSR:

python metrics.py