Bi-Directional ConvLSTM U-Net with Densely Connected Convolutions

Deep auto-encoder-decoder network for medical image segmentation with state of the art results on skin lesion segmentation, lung segmentation, and retinal blood vessel segmentation. This method applies bidirectional convolutional LSTM layers in U-net structure to non-linearly encode both semantic and high-resolution information with non-linearly technique. Furthermore, it applies densely connected convolution layers to include collective knowledge in representation and boost convergence rate with batch normalization layers. If this code helps with your research please consider citing the following papers:

R. Azad, M. Asadi, Mahmood Fathy and Sergio Escalera "Bi-Directional ConvLSTM U-Net with Densely Connected Convolutions ", ICCV, 2019, download link.

M. Asadi, R. Azad, Mahmood Fathy and Sergio Escalera "Multi-level Context Gating of Embedded Collective Knowledge for Medical Image Segmentation", The first two authors contributed equally. arXiv:2003.05056, 2020, download link.

Please consider starring us, if you found it useful. Thanks

Updates

• July 20, 2020: SEDU model added to the Skin Lesion segmentation code (inside models.py), now you can use this model for higher performance on skin lesion segmentation, inside the train file call the SEDU_Net_D3 model).
• March 5, 2020: An extended version of the network has been released(Complete implemenation for SKin Lesion Segmentation on ISIC 217, Skin Lesion Segmentation PH2 Dataset and cell nuclei along with the network implementation will be update soon).
• December 4, 2019: Document Image Binarization using BCDU-Net on DIBCO Challenges has been implemented, best performance on DIBCO series link
• Augest 28, 2019: First release (Complete implemenation for SKin Lesion Segmentation on ISIC 218, Retina Blood Vessel Segmentation and Lung segmentation dataset added.)
• Augest 27, 2019: Paper Accepted in the ICCV workshop 2019 (Oral presentation).

Prerequisties and Run

This code has been implemented in python language using Keras libarary with tensorflow backend and tested in ubuntu OS, though should be compatible with related environment. following Environement and Library needed to run the code:

• Python 3
• Keras - tensorflow backend

Run Demo

For training deep model for each task, go to the related folder and follow the bellow steps:

Skin Lesion Segmentation

1- Download the ISIC 2018 train dataset from this link and extract both training dataset and ground truth folders inside the dataset_isic18. 2- Run Prepare_ISIC2018.py for data preperation and dividing data to train,validation and test sets. 3- Run train_isic18.py for training BCDU-Net model using trainng and validation sets. The model will be train for 100 epochs and it will save the best weights for the valiation set. You can also train U-net model for this dataset by changing model to unet, however, the performance will be low comparing to BCDU-Net. 4- For performance calculation and producing segmentation result, run evaluate.py. It will represent performance measures and will saves related figures and results in output folder.

Retina Blood Vessel Segmentation

1- Download Drive dataset from this link and extract both training and test folders in a folder name DRIVE (make a new folder with name DRIVE) 2- Run prepare_datasets_DRIVE.py for reading whole data. This code will read all the train and test samples and will saves them as a hdf5 file in the DRIVE_datasets_training_testing folder. 3- The next step is to extract random patches from the training set to train the model, to do so, Run save_patch.py, it will extract random patches with size 64*64 and will save them as numpy file. This code will use help_functions.py, spre_processing.py and extract_patches.py functions for data normalization and patch extraction.
4- For model training, run train_retina.py, it will load the training data and will use 20% of training samples as a validation set. The model will be train for 50 epochs and it will save the best weights for the valiation set. 4- For performance calculation and producing segmentation result, run evaluate.py. It will represent performance measures and will saves related figures and results in test folder. Note: For image pre-processing and patch extraction we used this github's code.

Lung Segmentation

1- Download the Lung Segmentation dataset from Kaggle link and extract it. 2- Run Prepare_data.py for data preperation, train/test seperation and generating new masks around the lung tissues. 3- Run train_lung.py for training BCDU-Net model using trainng and validation sets (20 percent of the training set). The model will be train for 50 epochs and it will save the best weights for the valiation set. You can train either BCDU-net model with 1 or 3 densly connected convolutions. 4- For performance calculation and producing segmentation result, run evaluate_performance.py. It will represent performance measures and will saves related figures and results.

Quick Overview

Structure of the Bidirection Convolutional LSTM that used in our network

Structure of the BConvLSTM+SE that used in our network (MCGU-Net)

Results

For evaluating the performance of the proposed method, Two challenging task in medical image segmentaion has been considered. In bellow, results of the proposed approach illustrated.

Task 1: Retinal Blood Vessel Segmentation

Performance Comparision on Retina Blood Vessel Segmentation

In order to compare the proposed method with state of the art appraoches on retinal blood vessel segmentation, we considered Drive dataset.

Retinal blood vessel segmentation result on test data

Skin Lesion Segmentation

Performance Evalution on the Skin Lesion Segmentation task

Skin Lesion Segmentation results

Lung Segmentation

Performance Evalution on the Lung Segmentation task

Lung Segmentation results

Cell Nuclei Segmentation

Cell Nuclei Segmentation results

Model weights

You can download the learned weights for each task in the following table.

Query

All implementation done by Reza Azad. For any query please contact us for more information.

[email protected]