Skip to content

hmchuong/ML-BoneSuppression

Repository files navigation

Bone Suppression from Chest Radiographs

The project is a tool to build Bone Suppression model, written in tensorflow

CAM example image

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.

In this project you can

  1. Preprocessing data, including registration and augmentation.
  2. Train/test by following the quickstart. You can get a model with performance close to the paper.
  3. Visualize your training result with tensorboard

Requirements

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.

Configuration

  1. You can download the dataset here. This dataset includes 3 parts: JSRT dataset in png format, BSE_JSRT dataset in png format, and augmented dataset which can be trained directly.
  2. To register the dataset, make sure you set data_registration to true, and the input images are read from source_dir (JSRT) and target_dir (BSE_JSRT). The registered images will be saved to registered_output_dir into source and target subdirectories.
  3. To augment the dataset, make sure you set data_augmentation to true, the source_dir and target_dir will be used to augment. The total data after augmentation for source/target = augmentation_seed X total number of images in source_dir or target_dir. The augmented images will be saved to source and target subdirectories of augmented_output_dir with .png extension.
  1. source_folder and target_folder are folders to load training images.
  2. If you want to continue training from your last model, set use_trained_model to true and trained_model to your model path.
  3. output_model is where you save your model during training and output_log is where you save the tensorboard checkpoints.
  4. The other parameters is set following the published paper

Pretrained model

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

Quickstart

Note that currently this project can only be executed in Linux and macOS. You might run into some issues in Windows.

  1. Create & activate a new python3 virtualenv. (optional)
  2. Install dependencies by running pip install -r requirements.txt.
  3. Run python preprocessing.py to preprocess dataset. If you want to change your config path:
python preprocessing.py --config <config path>
  1. 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>
  1. 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>

Acknowledgement

I would like to thank LoudeNOUGH for scratch training script and Hussam Habbreeh (حسام هب الريح) for sharing his experiences on this task.

Author

Chuong M. Huynh ([email protected])

License

MIT