Deep Learning segmentation suite dessigned for 2D microscopy image segmentation

This repository provides researchers with a code to try different encoder-decoder configurations for the binary segmentation of 2D images in a video. It offers regular 2D U-Net variants and recursive approaches by combining ConvLSTM on top of the encoder-decoder.

Citation

If you found this code useful for your research, please, cite the corresponding preprint:

Estibaliz Gómez-de-Mariscal, Hasini Jayatilaka, Özgün Çiçek, Thomas Brox, Denis Wirtz, Arrate Muñoz-Barrutia, Search for temporal cell segmentation robustness in phase-contrast microscopy videos, arXiv 2021 (arXiv:2112.08817).

@misc{gómezdemariscal2021search,
      title={Search for temporal cell segmentation robustness in phase-contrast microscopy videos}, 
      author={Estibaliz Gómez-de-Mariscal and Hasini Jayatilaka and Özgün Çiçek and Thomas Brox and Denis Wirtz and Arrate Muñoz-Barrutia},
      year={2021},
      eprint={2112.08817},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}

Quick guide

Installation

Clone this repository and create all the required libraries

git clone https://github.com/esgomezm/microscopy-dl-suite-tf
pip3 install -r microscopy-dl-suite-tf/dl-suite/requirements.txt

Download or place your data in an accessible directory

Download the example data from Zenodo. Place the training, validation and test data in three independent folders. Each of them should contain an inputand labels folder. For 2D images, the name of the images should be raw_000.tif and instance_ids_000.tif for the input and ground truth images respectively. If the ground truth is given as videos, then the inputs and labels should have the same name.

Create a configuration .json file with all the information for the model architecture and training.

Check out some examples of configuration files. You will need to update the paths to the training, validation and test datasets. All the details for this file is given here.

Run model training

Run the file train.py indicating the path to the configuration JSON that contains all the information. This script will also test the model with the images provided in the "TESPATH" field of the configuration file.

python microscopy-dl-suite-tf/dl-suite/train.py 'microscopy-dl-suite-tf/examples/config/config_mobilenet_lstm_5.json' 

Run model testing

If you only want to run the test step, it is also possible with the test.py:

python microscopy-dl-suite-tf/dl-suite/test.py 'microscopy-dl-suite-tf/examples/config/config_mobilenet_lstm_5.json' 

Cell tracking from the instance segmentations

Videos with instance segmentations can be easily tracked with TrackMate. TrackMate is compatible with cell splitting, merging, and gap filling, making it suitable for the task.

The cells in our 2D videos exit and enter the focus plane, so we fill part of the gaps caused by these irregularities. We apply a Gaussian filter along the time axis on the segmented output images. The filtered result is merged with the output masks of the model as follows: all binary masks are preserved, and the positive values of the filtered image are included as additional masks. Those objects smaller than 100 pixels are discarded.

This processing is contained in the file tracking.py, in the section called Process the binary images and create instance segmentations. TrackMate outputs a new video with the information of the tracks given as uniquely labelled cells. Then, such information can be merged witht he original segmentation (without the temporal interpolation), using the code section Merge tracks and segmentations.

Technicalities

Available model architectures

• 'mobilenet_mobileunet_lstm': A pretrained mobilenet in the encoder with skip connections to the decoder of a mobileunet and a ConvLSTM layer at the end that will make the entire architecture recursive.
• 'mobilenet_mobileunet': A pretrained mobilenet in the encoder with skip connections to the decoder of a mobileunet (2D).
• 'unet_lstm': 2D U-Net with ConvLSTM units in the contracting path.
• 'categorical_unet_transpose': 2D U-Net for different labels ({0}, {1}, ...) with transpose convolutions instead of upsampling.
• 'categorical_unet_fc_dil': 2D U-Net for different labels ({0}, {1}, ...) with fully connected dilated convolutions.
• 'categorical_unet_fc': 2D U-Net for different labels ({0}, {1}, ...) with fully connected convolutions.
• 'categorical_unet': 2D U-Net for different labels ({0}, {1}, ...).
• 'unet' or "None": 2D U-Net with a single output.

Programmed loss-functions

When the output of the network has just one channel: foreground prediction

• Weighted binary cross-entropy: define a map of weights for the segmentation and feed the network with three images: input, weights and output.

  • Inspired by:
    • https://stackoverflow.com/questions/48555820/keras-binary-segmentation-add-weight-to-loss-function/48577360#48577360
    • https://stackoverflow.com/questions/55213599/u-net-with-pixel-wise-weighted-cross-entropy-input-dimension-errors
  • Open post: https://stackoverflow.com/questions/61225857/pixel-wise-weighted-loss-function-in-keras-tensorflow-2-0

• Weighted binary cross-entropy + DICE coefficient: Fix a weight value for foreground and combine it with DICE coefficient to create a loss function focused on false negatives (i.e. the tendency is to avoid them).

  • Taken from https://www.kaggle.com/lyakaap/weighing-boundary-pixels-loss-script-by-keras2
  • Check description here: https://lars76.github.io/neural-networks/object-detection/losses-for-segmentation/
  • Alternative loss functions: https://github.com/kohrah/DSBowl2018/blob/master/src/zoo_losses_K.py

• Binary cross-entropy: keras classical binary cross-entropy

When the output of the network has two channels: background and foreground prediction

• (Weighted) categorical cross-entropy: keras classical categorical cross-entropy
• Sparse categorical cross-entropy: same as the categorical cross-entropy but it allows the user to enter labelled grounf truth with a single channel and as many labels as classes, rather tha in a one-hote encoding fashion.

Prepare the data

• If you want to create a set of ground truth data with the format specified in the Cell Tracking Challenge, you can use the script prepare_videos_ctc.py.
• If you want to create 2D images from the videos, you can use the script prepare_data.py.
• In the folder additional_scripts you will find ImageJ macros or python code to keep processing the data to generate borders around the segmented cells for example.

Parameter configuration in the configuration.json

Notes about the code reused from different sources or the CNN architecture definitions

U-Net for binary segmentation U-Net architecture for TensorFlow 2 based on the example given in https://www.kaggle.com/advaitsave/tensorflow-2-nuclei-segmentation-unet