The project is a tool to build Bone Suppression model, written in tensorflow
What is Bone Suppression?
Bone suppression is an autoencoder-like model for eliminating bone shadow from Chest X-ray images. The model require two types of dataset: normal and bone-suppression X-ray images. The target model can suppress bone shadow from Chest X-ray images, help Radiologists diagnose better lung related diseases. Although there are some softwares supporting bone suppression (ClearRead, CareStream), this project is a practical open source in computer vision and deep learning.
- Preprocessing data, including registration and augmentation.
- Train/test by following the quickstart. You can get a model with performance close to the paper.
- Visualize your training result with tensorboard
The project requires Python>=3.5
.
I have trained on an instance with 1 NVIDIA GTX 1080Ti (11GB VRAM)
and it takes approximately 14 hours.
- You can download the dataset here. This dataset includes 3 parts:
JSRT
dataset inpng
format,BSE_JSRT
dataset inpng
format, andaugmented
dataset which can be trained directly. - To register the dataset, make sure you set
data_registration
totrue
, and the input images are read fromsource_dir
(JSRT) andtarget_dir
(BSE_JSRT). The registered images will be saved toregistered_output_dir
intosource
andtarget
subdirectories. - To augment the dataset, make sure you set
data_augmentation
totrue
, thesource_dir
andtarget_dir
will be used to augment.The total data after augmentation for source/target
=augmentation_seed
X total number of images insource_dir
ortarget_dir
. The augmented images will be saved tosource
andtarget
subdirectories ofaugmented_output_dir
with.png
extension.
source_folder
andtarget_folder
are folders to load training images.- If you want to continue training from your last model, set
use_trained_model
to true andtrained_model
to your model path. output_model
is where you save your model during training andoutput_log
is where you save the tensorboard checkpoints.- The other parameters is set following the published paper
If you want to start testing without training from scratch, you can use the model I have trained. The model has loss value: 0.01409, MSE: 7.1687e-4, MS-SSIM: 0.01517
Note that currently this project can only be executed in Linux and macOS. You might run into some issues in Windows.
- Create & activate a new python3 virtualenv. (optional)
- Install dependencies by running
pip install -r requirements.txt
. - Run
python preprocessing.py
to preprocess dataset. If you want to change your config path:
python preprocessing.py --config <config path>
- Run
python train.py
to train a new model. If you want to change your config path:
python train.py --config <config path>
During training, you can use Tensorboard to visualize the results:
tensorboard --logdir=<output_log in train.cfg>
- Run
python test.py
to evaluate your model on specific image. To change default parameters, you can use:
python test.py --model <model_path> --config <model config path> --input <image path> --output <output image path>
I would like to thank LoudeNOUGH for scratch training script and Hussam Habbreeh (حسام هب الريح) for sharing his experiences on this task.
Chuong M. Huynh ([email protected])
MIT