Source-linked AI summary
KiU-Net: Towards Accurate Segmentation of Biomedical Images using Over-complete Representations
Jeya Maria Jose, Vishwanath Sindagi, Ilker Hacihaliloglu, Vishal M. Patel
TL;DR
The paper addresses U-Net’s reduced performance on small anatomical landmarks with blurred or noisy boundaries. It introduces over-complete Ki-Net and combines it with U-Net in KiU-Net, achieving better segmentation performance alongside faster convergence and fewer parameters.
Problem
U-Net-based biomedical segmentation can perform poorly on small anatomical structures and blurred, noisy boundaries, where fine details are difficult to capture.
Method
KiU-Net augments U-Net with spatially over-complete Ki-Net and uses cross-scale fusion to combine information from both architectures.
Results
4% improvement in DICE accuracy over U-Net and 2% improvement over the state-of-the-art method were reported.
Takeaways & Limitations
KiU-Net captures fine edges and small anatomical structures while using fewer parameters and achieving faster convergence.
Abstract
from arXiv · showhide
Due to its excellent performance, U-Net is the most widely used backbone architecture for biomedical image segmentation in the recent years. However, in our studies, we observe that there is a considerable performance drop in the case of detecting smaller anatomical landmarks with blurred noisy boundaries. We analyze this issue in detail, and address it by proposing an over-complete architecture (Ki-Net) which involves projecting the data onto higher dimensions (in the spatial sense). This network, when augmented with U-Net, results in significant improvements in the case of segmenting small anatomical landmarks and blurred noisy boundaries while obtaining better overall performance. Furthermore, the proposed network has additional benefits like faster convergence and fewer number of parameters. We evaluate the proposed method on the task of brain anatomy segmentation from 2D Ultrasound (US) of preterm neonates, and achieve an improvement of around 4% in terms of the DICE accuracy and Jaccard index as compared to the standard-U-Net, while outperforming the recent best methods by 2%. Code: https://github.com/jeya-maria-jose/KiU-Net-pytorch .
1 Introduction
Biomedical ultrasound segmentation is difficult because small anatomical structures have unclear, noisy, or incomplete boundaries, while standard U-Net representations can miss fine details. The paper addresses this limitation with over-complete Ki-Net and a cross-scale fusion architecture, KiU-Net.
- Clinical motivation: Preterm-neonate brain ultrasound supports diagnosis but contains unclear ventricular boundaries, shading artifacts, variable anatomy, and measurement variability.These conditions affect both qualitative assessment and manual landmark measurements.
- Existing approaches: 70.8%, 80%, and 76.5% DICE values were reported for prior traditional segmentation methods, with one method requiring 54 minutes.The cited traditional approaches were described as time consuming or insufficiently robust under challenging scan conditions.
- U-Net limitation: U-Net can miss fine details because its under-complete encoder reduces spatial dimensionality and its increasing receptive fields emphasize higher-level features.The paper notes that tiny structures require smaller receptive fields.
- Proposed direction: Over-complete representations project intermediate data into higher spatial dimensions and are explored here to address the limited focus on small structures.The paper identifies over-complete representations as relatively unexplored for segmentation.
- Proposed direction: KiU-Net combines Ki-Net with standard U-Net through cross-scale fusion to capture finer shapes and edges while targeting state-of-the-art segmentation performance.The paper summarizes the architecture as combining under-complete and over-complete deep networks.
2 Proposed Method
The proposed method uses an over-complete Ki-Net alongside U-Net to preserve fine spatial details while retaining high-level shape features. Cross Residual Fusion combines complementary features from both branches throughout the network.
- Motivation: U-Net’s increasing receptive field reduces deeper filters’ ability to learn fine edges and textures, producing less-sharp segmentation boundaries.Convolution and max-pooling enlarge receptive fields as depth increases, reducing filter focus on fine details.
- Ki-Net: Ki-Net projects intermediate representations into higher spatial dimensions using encoder upsampling and decoder max-pooling.Upsampling follows each encoder convolution, while max-pooling in the decoder reduces dimensionality toward the input size.
- Ki-Net: Restricting Ki-Net’s receptive fields makes deeper filters focus on smaller regions and fine edges across encoder layers.Compared with U-Net, Ki-Net filters become more localized as depth increases and capture fine edges at increased resolution.
- KiU-Net: KiU-Net combines Ki-Net’s fine-edge feature maps with U-Net’s high-level shape feature maps through a parallel architecture.The combined network is designed to improve overall segmentation accuracy by exploiting complementary representations from both branches.
- Cross Residual Fusion: Cross Residual Fusion merges features at each block, extracting complementary information and forwarding it to both network branches.Block-level fusion supports gradient flow across both branches, while the final decoder features are added and passed through a 1×1 convolution to produce the segmentation mask.
- Training: Training uses pixel-wise binary cross entropy between predicted and ground-truth outputs.The loss is defined over image locations using prediction and ground-truth values at each pixel.
3 Experiments and results
Experiments evaluate KiU-Net on preterm-neonate ultrasound segmentation, compare it with existing approaches, and examine ablations, convergence, and qualitative predictions.
- Dataset: The dataset comprises 1,629 annotated ultrasound images from 20 premature neonates, including subjects with intraventricular hemorrhage and healthy subjects.Ventricles and septum pellucidum were manually segmented by an expert ultrasonographer.
- Quantitative results: 4% improvement in DICE accuracy over U-Net and 2% over the state-of-the-art method were reported for the proposed method.The comparison is reported in Table 1.
- Qualitative results: KiU-Net predicts very small masks precisely and detects edges better than the other methods in qualitative comparisons.The method also performs well across different anatomy-structure sizes.
- Efficiency: KiU-Net converges faster than standard U-Net and has an inference time of 8 ms for one test image.The proposed network also uses fewer parameters than the compared methods.
- Ablation study: The over-complete architecture alone performs worse than the under-complete architecture because it captures edges but not most high-level features.Combining the branches improves performance over either architecture alone.
- Ablation study: Adding cross residual fusion blocks at each block level produces further improvements, demonstrating the effectiveness of the cross-fusion strategy.Figure 8 presents qualitative improvements after adding each major block.
4 Conclusion
KiU-Net augments a U-Net-based under-complete architecture with an over-complete Ki-Net and cross-scale residual fusion to improve segmentation of fine edges and small structures.
- Conclusion: KiU-Net combines an under-complete U-Net with an over-complete Ki-Net to capture fine edges and small anatomical structures.The architecture is designed for segmentation cases where these details are typically missed.
- Conclusion: Cross-scale residual blocks provide a more effective fusion of information from the two network branches.The conclusion also reports fewer parameters and faster convergence.
- Conclusion: The proposed method achieves better performance than recent methods on a complex dataset containing both small and large segmentation masks.
Experiments on other modalities
The authors test KiU-Net beyond ultrasound on GLAS microscopic images and RITE retinal fundus images, reporting improved segmentation performance over U-Net.
- Experimental scope: The authors evaluate KiU-Net across other modalities using two public datasets.
- GLAS Dataset: On GLAS, KiU-Net achieves a 4% improvement in dice accuracy over U-Net without pre-processing or post-processing.The dataset contains 165 H&E-stained microscopic images, split into 85 training and 80 testing images.
- GLAS Dataset: Qualitative GLAS comparisons indicate that KiU-Net captures edges better and produces better segmentation predictions than Seg-Net and U-Net.
- RITE Dataset: On RITE, KiU-Net achieves significant improvement in dice accuracy over U-Net.The authors report no specific pre-processing steps or specific loss functions for this experiment.
- RITE Dataset: The RITE experiment predicts vessel segmentation from resized retinal fundus images and compares KiU-Net with Seg-Net and U-Net.The dataset contains 40 image sets, split into 20 training and 20 testing sets.
Network Architecture
The implementation uses separate Ki-Net and U-Net branches within KiU-Net, with configurations reported for both branches and ultrasound images sized at 128 × 128.
- Network Architecture: KiU-Net uses distinct Ki-Net and U-Net branches whose configurations are reported separately.
- Network Architecture: The ultrasound train and test images use H = 128 and W = 128.