AlbUNet-1D-2D-Tensorflow-Keras

This repository contains 1D and 2D Signal Segmentation Model Builder for AlbUNet and several of its variants developed in Tensorflow-Keras. The code supports Deep Supervision, AutoEncoder mode, Guided Attention and other options. The segmentation models can be used for binary or multiclass segmentation, or for regression tasks.

Models supported [1]

1. AlbUNet18
2. AlbUNet34
3. AlbUNet50
4. AlbUNet101
5. AlbUNet152

AlbUNet

AlbUNet has a ResNet based Encoder and traditional UNet based Decoder, as shown in the Figure below for ALbUNet34, which uses ResNet34 as the Encoder.

AlbUNet Architecture

Supported Features

The speciality about this model is its flexibility, such as:

1. The user can choose any of the 5 available AlbUNet variants for either 1D or 2D Segmentation tasks.
2. The models can be used for Binary or Multi-Class Classification, or Regression type Segmentation tasks.
3. The models allow Deep Supervision [2] with flexibility during Segmentation.
4. The segmentation models can also be used as Autoencoders [3] for Feature Extraction.
5. The Segmentation Models can be Attention Guided [4].
6. Number of input kernel/filter, commonly known as the Width of the model can be varied.
7. Number of classes for Classification tasks and number of extracted features for Regression tasks can be varied.
8. Number of Channels in the Input Dataset can be varied.

Mentionable that the 2D version of AlbUNet can also be used in Transfer Learning from previously trained weights (e.g., ImageNet), just the encoder blocks should be replaced with the trained model layers.

References

[1] A. Shvets, V. Iglovikov, A. Rakhlin, and A. A. Kalinin, “Angiodysplasia detection and localization using deep convolutional neural networks,” arXiv.org, 21-Apr-2018. [Online]. Available: https://arxiv.org/abs/1804.08024. [2] Zhou, Z., Siddiquee, M., Tajbakhsh, N., & Liang, J. (2021). UNet++: Redesigning Skip Connections to Exploit Multiscale Features in Image Segmentation. Arxiv-vanity.com. Retrieved 30 August 2021, from https://www.arxiv-vanity.com/papers/1912.05074/.
[3] Zhou, Z., Siddiquee, M., Tajbakhsh, N., & Liang, J. (2021). UNet++: A Nested U-Net Architecture for Medical Image Segmentation. arXiv.org. Retrieved 30 August 2021, from https://arxiv.org/abs/1807.10165.
[4] M. Noori, A. Bahri, and K. Mohammadi, “Attention-guided version of 2D UNET for automatic brain tumor segmentation,” arXiv.org, 04-Apr-2020. [Online]. Available: https://arxiv.org/abs/2004.02009.