U-Net##: A Powerful Novel Architecture for Medical Image Segmentation

Official PyTorch implementation of the MICAD 2022 conference paper: "U-Net##: A Powerful Novel Architecture for Medical Image Segmentation"

Information

In this paper, we propose a powerful novel architecture named U-Net##, which consists of multiple overlapping U-Net pathways and has the strategies of sharing feature maps between parallel neural networks, using auxiliary convolutional blocks for additional feature extractions and deep supervision, so that it performs as a boosted U-Net model for medical image segmentation.

UPDATE: New U-Net## Architecture

Changes:

• 3D model version is added.
• Number of filters is reduced by 50%.
• Feature map additions are directed upward.
• BatchNorm layers are replaced with InstanceNorm.
• Weighted summation is implemented for deep supervision.
• Skip connections on the ceiling network pathway are omitted.
• Feature map additions from distant network pathways are omitted.
• Deep supervision outputs are obtained from the ceiling network pathway.
• Final Sigmoid functions are commented out and passed to the loss functions.

Results:

• Better performance.
• ~4x faster training speed.
• Reduced parameters from 43M to 8.85M
• Parameter number of the 3D model is 26.5M

New Architecture

Block Diagram of the New U-Net## Architecture

Convolutional Blocks of the New U-Net## Architecture

Original Architecture

Block Diagram of the Original U-Net## Architecture

Convolutional Blocks of the Original U-Net## Architecture

Results from the Original Architecture

The U-Net## model is evaluated on the TCIA-LGG Segmentation Dataset from The Cancer Imaging Archive (TCIA) to segment the brain regions with FLAIR abnormalities on the related brain MRI images.

Some Output Images Predicted by the Trained Models

Score Results of the Trained Models

Comparison of the Dice Score Changes

How to Use

# Import available loss functions if needed
from loss.loss import DiceLoss, BCEDiceLoss

# Import model architectures
from model.UNetSharpSharp2D import UNetSharpSharp2D  # 2D model variant
from model.UNetSharpSharp3D import UNetSharpSharp3D  # 3D model variant

# Instantiate the 2D model
model_2d = UNetSharpSharp2D(in_ch=3, out_ch=1, deep_supervision=True)

# Instantiate the 3D model
model_3d = UNetSharpSharp3D(in_ch=3, out_ch=1, deep_supervision=True)

Citation

If you find this work useful for your research, please consider citing:

@InProceedings{10.1007/978-981-16-6775-6_19,
author={Korkmaz, Fırat},
editor={Su, Ruidan and Zhang, Yudong and Liu, Han and F Frangi, Alejandro},
title={U-Net##: A Powerful Novel Architecture for Medical Image Segmentation},
booktitle={Medical Imaging and Computer-Aided Diagnosis},
year={2023},
publisher={Springer Nature Singapore},
address={Singapore},
pages={231--241},
isbn={978-981-16-6775-6}
}