[Paper Review] Medical Image Segmentation Using a U-Net type of Architecture

Eshal Zahra, B Mohammed Ali+1 more|arXiv (Cornell University)|May 11, 2020

Advanced Neural Network Applications19 references4 citations

TL;DR

This paper proposes a modified U-Net architecture that enhances feature representation in the bottleneck layer by introducing a fully supervised, two-layer fully connected (FC) sub-network trained with pixel-wise cross-entropy loss. By injecting explicit supervision at the encoder's bottleneck, the model improves semantic feature learning, achieving high segmentation accuracy on MRI and CT-scan data—specifically 97.6% accuracy on CT scans and 91.3% on MRI scans—demonstrating that supervised bottleneck training can match or exceed standard U-Net performance.

ABSTRACT

Deep convolutional neural networks have been proven to be very effective in image related analysis and tasks, such as image segmentation, image classification, image generation, etc. Recently many sophisticated CNN based architectures have been proposed for the purpose of image segmentation. Some of these newly designed networks are used for the specific purpose of medical image segmentation, models like V-Net, U-Net and their variants. It has been shown that U-Net produces very promising results in the domain of medical image segmentation.However, in this paper, we argue that the architecture of U-Net, when combined with a supervised training strategy at the bottleneck layer, can produce comparable results with the original U-Net architecture. More specifically, we introduce a fully supervised FC layers based pixel-wise loss at the bottleneck of the encoder branch of U-Net. The two layer based FC sub-net will train the bottleneck representation to contain more semantic information, which will be used by the decoder layers to predict the final segmentation map. The FC layer based sub-net is trained by employing the pixel-wise cross entropy loss, while the U-Net architectures trained by using L1 loss.

Motivation & Objective

• To improve feature representation in U-Net's bottleneck layer for better medical image segmentation.
• To investigate whether supervised training at the bottleneck can outperform or match self-supervised U-Net training.
• To reduce reliance on extensive data augmentation by enhancing feature learning through explicit supervision.
• To evaluate the method on diverse medical imaging modalities, specifically MRI and CT-scan.
• To demonstrate that a fully connected sub-network at the bottleneck can effectively guide decoder performance.

Proposed method

• The network uses a standard encoder-decoder U-Net architecture with 3×3 convolutional filters, ReLU activation, and 2×2 max pooling for downsampling.
• A two-layer fully connected (FC) sub-network is inserted at the bottleneck to predict the ground-truth segmentation map using linear transformation.
• The FC sub-network is trained with pixel-wise cross-entropy loss to provide explicit supervision to the encoder's learned features.
• The main U-Net branch is trained using L1 loss to optimize the final segmentation output.
• Skip connections are preserved between corresponding encoder and decoder layers to preserve spatial details.
• The entire network is trained end-to-end, combining supervised loss at the bottleneck with standard segmentation loss in the decoder.

Experimental results

Research questions

• RQ1 Can explicit supervision at the bottleneck layer improve feature representation in a U-Net-based architecture for medical image segmentation?
• RQ2 Does training the bottleneck with a fully connected sub-network yield better segmentation performance than standard self-supervised U-Net training?
• RQ3 How does the proposed method compare to baseline U-Net in terms of sensitivity, specificity, and accuracy on MRI and CT-scan data?
• RQ4 Can the addition of a supervised bottleneck reduce dependency on data augmentation in low-data medical imaging scenarios?
• RQ5 Does the FC-based bottleneck enhance semantic feature learning without degrading spatial resolution or increasing model complexity?

Key findings

• The proposed model achieved 97.6% accuracy on CT-scan images, with 96.1% specificity and 97.2% sensitivity.
• On MRI images, the model reached 91.3% accuracy, with 92.6% specificity and 93.9% sensitivity.
• The supervised bottleneck significantly improved feature learning, enabling the decoder to produce more accurate segmentation maps.
• The method demonstrated strong generalization across different imaging modalities, including MRI and CT-scan.
• The integration of a fully connected sub-network at the bottleneck layer provided a performance boost comparable to or better than standard U-Net training.