Source-linked AI summary
KiU-Net: Overcomplete Convolutional Architectures for Biomedical Image and Volumetric Segmentation
Jeya Maria Jose Valanarasu, Vishwanath A. Sindagi, Ilker Hacihaliloglu, Vishal M. Patel
TL;DR
U-Net-based segmentation can miss small structures and precise boundaries because deeper layers emphasize high-level features as receptive fields grow. The paper proposes KiU-Net, combining overcomplete and undercomplete branches, and reports stronger performance across five datasets with fewer parameters and faster convergence.
Problem
U-Net-based architectures often fail to segment small structures and blurred or noisy boundaries accurately.
Method
KiU-Net combines an overcomplete Kite-Net branch for fine details with an undercomplete U-Net branch for high-level features, while KiU-Net 3D extends the design to volumetric segmentation.
Results
KiU-Net achieves a significant performance boost over compared methods under the same training pipeline and is evaluated across five datasets and modalities.
Takeaways & Limitations
The paper reports fewer parameters, faster convergence, and lower memory requirements while maintaining better performance than U-Net in the examined settings.
Abstract
from arXiv · showhide
Most methods for medical image segmentation use U-Net or its variants as they have been successful in most of the applications. After a detailed analysis of these "traditional" encoder-decoder based approaches, we observed that they perform poorly in detecting smaller structures and are unable to segment boundary regions precisely. This issue can be attributed to the increase in receptive field size as we go deeper into the encoder. The extra focus on learning high level features causes the U-Net based approaches to learn less information about low-level features which are crucial for detecting small structures. To overcome this issue, we propose using an overcomplete convolutional architecture where we project our input image into a higher dimension such that we constrain the receptive field from increasing in the deep layers of the network. We design a new architecture for image segmentation- KiU-Net which has two branches: (1) an overcomplete convolutional network Kite-Net which learns to capture fine details and accurate edges of the input, and (2) U-Net which learns high level features. Furthermore, we also propose KiU-Net 3D which is a 3D convolutional architecture for volumetric segmentation. We perform a detailed study of KiU-Net by performing experiments on five different datasets covering various image modalities like ultrasound (US), magnetic resonance imaging (MRI), computed tomography (CT), microscopic and fundus images. The proposed method achieves a better performance as compared to all the recent methods with an additional benefit of fewer parameters and faster convergence. Additionally, we also demonstrate that the extensions of KiU-Net based on residual blocks and dense blocks result in further performance improvements. The implementation of KiU-Net can be found here: https://github.com/jeya-maria-jose/KiU-Net-pytorch
I. INTRODUCTION
Traditional encoder-decoder networks often miss small anatomical structures and precise boundaries, motivating KiU-Net’s combination of overcomplete and undercomplete branches. The paper extends this design to volumetric segmentation and evaluates it across five modalities.
- Motivation: U-Net and related architectures often segment large structures well but miss small masks and blurred or noisy boundaries.The authors attribute this weakness to insufficient focus on low-level features in deeper networks.
- Proposed architecture: KiU-Net combines an overcomplete Kite-Net branch for fine details with an undercomplete U-Net branch for high-level features.Kite-Net projects inputs into higher spatial dimensions and is paired with U-Net in a multi-branch architecture.
- Volumetric extension: KiU-Net 3D extends the approach to volumetric segmentation using a 3D convolution-based overcomplete network for low-level information.The extension is introduced specifically for volumetric segmentation.
- Architectural variants: Residual-block and dense-block variants, Res-KiUNet and Dense-KiUNet, are proposed to improve network learning.These variants replace or augment the learning architecture with residual connections and dense blocks.
- Evaluation: The method is evaluated on five datasets spanning ultrasound, MRI, CT, microscopic, and fundus images to assess performance across modalities.The datasets include BraTS, LiTS, GlaS, RITE, and Brain Anatomy Segmentation.
II. RELATED WORK
The related work centers on encoder-decoder segmentation networks, methods designed to preserve local detail, and modality-specific biomedical segmentation systems. Prior work spans full-resolution, geodesic, topology-aware, and U-Net-derived approaches.
- Scope: The paper reviews medical image segmentation methods with emphasis on datasets used in its experiments and methods for capturing fine details.The review focuses on relevant deep-learning approaches for the evaluated tasks.
- Compared methods: Compared architectures include Seg-Net, U-Net, U-Net++, sSE-UNet, CPFNet, FullNet, and DeepGIoS.The comparison includes both standard encoder-decoder models and methods targeting local details.
- Local-detail methods: FullNet maintains full-resolution feature maps without max-pooling or downsampling to focus more on boundaries.Its design avoids increasing receptive-field size through downsampling.
- Overcomplete representations: Signal-processing literature introduced overcomplete representations using more basis functions than input samples to flexibly capture data structure.The paper connects this earlier representation-learning idea to its architectural motivation.
- Biomedical applications: Prior biomedical applications include U-Net and PSP-Net for ultrasound brain anatomy, contour- and topology-aware models for glands, and 2D and 3D U-Net methods for brain tumors.The reviewed tasks span ultrasound, microscopic, and MRI segmentation.
III. METHOD
The method addresses U-Net’s difficulty with tiny structures and high-curvature boundaries by constraining receptive-field growth through an overcomplete branch. KiU-Net 3D applies the same design principle to volumetric scans.
- Issues with traditional encoder-decoder networks: U-Net-based models can miss tiny structures and inaccurately segment boundaries, even when their overall predictions for large structures are good.The authors report this behavior for ultrasound anatomy and volumetric lesion segmentation.
- Issues with traditional encoder-decoder networks: Downsampling in U-Net’s encoder increases deep-layer receptive fields, shifting emphasis toward high-level features instead of small masks and fine edges.The first few convolutional layers retain the strongest low-level focus.
- Overcomplete Networks: Kite-Net uses bilinear upsampling in its encoder to project inputs into higher spatial dimensions and constrain receptive-field growth.Its decoder uses convolution followed by max-pooling, reversing the usual encoder-decoder ordering.
- KiU-Net 3D: KiU-Net 3D implements the overcomplete branch with 3D convolutions followed by trilinear upsampling in the encoder and 3D max-pooling in the decoder.Both operations use coefficient 2 and ReLU activation in the described blocks.
3) Architecture Details of KiU-Net 3D:
KiU-Net 3D combines Kite-Net 3D and U-Net 3D through cross-residual feature blocks at multiple scales. Their final feature maps are fused to produce the volumetric segmentation prediction.
- Cross Residual Feature Block: CRFBs combine features from Kite-Net 3D and U-Net 3D at each encoder and decoder level.The two branches learn different features, which are combined to exploit their complementary representations.
- Cross Residual Feature Block: Cross-residual features are extracted from each branch and added to the other branch’s original features.The resulting complementary features are forwarded between the two networks.
- Output Fusion: The final feature maps from both branches are added and passed through a 1 × 1 3D convolution to obtain the prediction.This fusion occurs after processing the input 3D voxel through Kite-Net 3D and U-Net 3D.
IV. EXPERIMENTS
The experiments evaluate KiU-Net and KiU-Net 3D for 2D image and 3D volumetric medical image segmentation. The brain-anatomy ultrasound dataset supports automatic delineation relevant to diagnosis and prognosis.
- Experimental Scope: The study evaluates KiU-Net and KiU-Net 3D through experiments on 2D medical images and 3D medical volumes.The evaluation is designed to compare the proposed networks in both segmentation settings.
- Brain Anatomy Segmentation (US): Brain-anatomy ultrasound segmentation targets brain ventricles and the septum pellucidum in scans from premature neonates.The dataset was collected from 20 neonates younger than one year after institutional review board approval.
- Brain Anatomy Segmentation (US): Automatic segmentation of these structures is intended to support accurate diagnosis and prognosis of preterm brain disorders.The passage connects the target anatomy to intraventricular hemorrhage and septo-optic dysplasia.
1) Datasets:
The 2D experiments span microscopic, fundus, and ultrasound imagery, using standardized preprocessing and training settings. Evaluation includes a comparison table and binary cross-entropy optimization.
- Brain Anatomy Segmentation (US): The ultrasound brain-anatomy dataset contains 1629 annotated images, with 1300 for training and 329 for testing.Images were resized to 128 × 128 before processing.
- Gland Segmentation (Microscopic): The GLAS microscopic gland dataset contains 165 images, divided into 85 training and 80 testing images.The images were resized to 128 × 128.
- Retinal Nerve Segmentation (Fundus): The RITE fundus dataset contains 40 images, split evenly into 20 training and 20 testing images.The images were resized to 128 × 128, and vessel extraction supports analysis of retinal morphology.
- Training and Implementation: The 2D networks use binary cross-entropy loss between predictions and ground truth without application-specific or metric loss functions.Training uses batch size 1, learning rate 0.001, Adam, and up to 300 epochs or convergence.
- Evaluation: Table I presents performance comparisons for 2D image segmentation against existing methods.The supplied passage identifies the table’s comparison scope but does not provide its numerical entries.
1) Datasets:
The volumetric experiments use BraTS MRI and LiTS CT datasets, train KiU-Net 3D with voxel-wise cross-entropy, and report quantitative segmentation outcomes using established evaluation procedures.
- Volumetric Datasets: LiTS provides contrast-enhanced abdominal CT scans annotated for liver and liver lesions.Training uses 109 scans, while testing uses 21 scans from the specified LiTS range.
- Training and Implementation: KiU-Net 3D is trained with cross-entropy loss over the scan’s 3D voxel coordinates and class dimension.The formulation supports the four-class BraTS segmentation setting.
- Evaluation: The results section reports quantitative evaluations against recent approaches using widely used medical-segmentation metrics.It also describes qualitative visualization of sample predictions to analyze comparative performance.
1) KiU-Net:
KiU-Net is evaluated against established medical image and volumetric segmentation architectures using quantitative and qualitative comparisons across multiple datasets. The experiments report stronger performance for KiU-Net 3D and statistically significant results under a common training pipeline.
- KiU-Net is compared with Seg-Net, U-Net, U-Net++, sSE-UNet, CPFNet, FullNet, and DeepGIoS using the same training pipeline.All methods are trained from scratch for fair comparison.
- Qualitative comparisons cover brain anatomy, gland, and retinal nerve segmentation across the brain US, GLAS, and RITE datasets.Figure 4 compares SegNet, UNet, UNet++, and KiU-Net.
- KiU-Net 3D outperforms the other methods on quantitative BraTS and LiTS volumetric segmentation metrics.The reported metrics include Dice accuracy and Hausdorff distance for BraTS, alongside liver segmentation metrics for LiTS.
- Qualitative volumetric comparisons cover brain tumor and liver segmentation using 2D slices and 3D predictions from BraTS and LiTS.Figure 5 compares 2D and 3D variants of SegNet, UNet, and KiU-Net.
- All reported paired t-test p-values are below 0.05 for LiTS tumor segmentation and BraTS whole-tumor accuracy.The tests compare the proposed method with baseline methods using Dice accuracy.
B. Qualitative Results
Qualitative results show that KiU-Net captures small structures and sharp boundaries more precisely than traditional networks in image and volumetric segmentation. The study also examines architectural variants and implementation properties for further improvement.
- Qualitative Results: KiU-Net accurately segments small ventricles and produces sharper edges than Seg-Net, U-Net, and U-Net++ on image segmentation examples.The RITE results also perform reasonably well despite the dataset's low number of training images.
- Qualitative Results: KiU-Net 3D predicts brain-tumor surfaces and edges more precisely than competing methods, which smooth the sharp tumor boundaries.The BraTS visualizations compare 2D slices and 3D scan predictions with the ground truth.
- Qualitative Results: On LiTS, KiU-Net is reported as effective for segmenting larger regions as well as small lesions in BraTS.The qualitative comparison uses 3D scan predictions comprising 48 2D images for LiTS and 155 for BraTS.
- Further Improvements: The study examines parameter count, convergence rate, and memory requirements to characterize KiU-Net's network properties.These properties are studied because the architecture is intended for both image and volumetric segmentation.
- Further Improvements: Residual and other KiU-Net variants are evaluated because the authors report further improvements over the base architecture.Res-KiU-Net uses residual connections in both branches at every convolutional block and level.
1) Res-KiUNet:
Res-KiUNet and Dense-KiUNet extend KiU-Net with residual connections or dense blocks, improving performance on GlaS while KiU-Net uses substantially fewer parameters and converges faster than compared networks.
- Res-KiUNet: Residual connections in Res-KiUNet are intended to improve network learning by propagating gradients to initial layers faster.The residual formulation is y = F(x) + x.
- Variants: Res-KiUNet and Dense-KiUNet provide further performance improvements over KiU-Net on the GlaS dataset.The comparison reports Dice and Jaccard metrics in Table IV and visualizes predictions in Fig. 7.
- Model complexity: KiU-Net uses approximately 10× fewer parameters than U-Net and approximately 40× fewer than SegNet.KiU-Net uses a three-layer network with 32, 64, and 128 filters, compared with five-layer encoder-decoder networks.
- Dense-KiUNet: Dense-KiUNet places dense blocks after every convolutional layer in both branches and fuses the branch outputs before prediction.Each dense block uses four convolutional layers with concatenated feature maps.
- Convergence: KiU-Net converges faster than SegNet, U-Net, and U-Net++ in loss comparisons on GlaS, with similar trends observed on other datasets.Faster convergence is associated with lower training complexity.
D. Ablation Study
The ablation study compares under-complete and over-complete branches, with and without skip connections, and examines how depth and filter count affect memory and performance.
- D. Ablation Study: The ablation study evaluates under-complete and over-complete convolutional architectures both without and with skip connections on the brain anatomy ultrasound dataset.The compared configurations include UC, OC, UC + SK, UC + OC with SK, and KiU-Net.
- Memory and capacity: Reducing network depth and filter count lowers memory requirements while maintaining decent performance.Memory depends on network depth and the number of filters.
- Experimental settings: The reduced-depth experiments keep the number of filters fixed, whereas the filter-count experiments keep the depth fixed at three layers.These settings are evaluated on the GlaS dataset in Tables VII and VIII.
F. Dependence on batch-size
Experiments examine batch size and image resolution as practical settings for U-Net, U-Net++, and KiU-Net. KiU-Net benefits substantially from a larger batch size in the reported GlaS comparison, while resolution changes produce only slight performance differences.
- F. Dependence on batch-size: With GPU memory limited to 11 GB, U-Net and U-Net++ use batch size 4 while KiU-Net uses batch size 1.The comparison is performed on the GlaS dataset.
- F. Dependence on batch-size: KiU-Net achieves a significant performance boost compared with U-Net and U-Net++ when trained with the higher batch-size setting.The corresponding performance change for U-Net and U-Net++ is described as negligible.
- G. Dependence on resolution of image: Images in the main experiments are resized to 128 × 128 for Brain US, GlaS, and RITE, while higher-resolution experiments use 512 × 512 images.The higher-resolution setting is also evaluated for RITE and compared with 128 × 128 on GlaS.
- G. Dependence on resolution of image: Changing image resolution produces only a slight performance change for both U-Net and KiU-Net on the reported GlaS comparison.The passage also describes experiments using 512 × 512 images for RITE.