Cell Nuclei Segmentation from Histology images using cGAN

Overview

Data

The database consists of 50 512*512 image pairs and was split into 43 and 7 train and test sets (folders Slide_10 and Slide_11), respectively. In order to use the code in this repository, the histology images and corresponding segmentation maps need to be merged as shown below. .

Baseline

The baseline architecture was U-Net: Convolutional Networks for Biomedical Image Segmentation. The model was trained for 200 epochs using binary crossentropy loss. The dice similarity coefficient was 0.637. Since the database is relatively small for training a deep neural network model, I tried on-the-fly data augmentation by randomly rotating and flipping both input image and the segmentation map during training. The performance improved by 9% as a result of this augmentation.

cGAN-based Segmentation

In order to improve the performance, I designed a conditional Generative Adversarial Network (cGAN) shown below. The generator design follows the encoder-decoder structure with skip-connections, as in U-Net. The discriminator on the other hand uses patch information to resolve details in the prediction map. With cGAN, the dice similarity coefficient increases to 0.723.

Results

Qualitative

Quantitative

Model	Dice Measure
Unet with no data augmentation (Baseline)	0.637
Unet with data augmentation	0.693
cGAN with data augmentation	0.723

Dependencies

• Pytorch >= 0.4.0
• Python >= 3.5
• Numpy

To train and reproduce the results:

• Set path to the training and validation data and the code.
• Run main.py to train model using cGAN
• Run unet_train.py to train model using Unet

For testing, run:

• Set path to the trained model
• Run test.py
• Predicted maps from test image will be in /test_results