Source-linked AI summary
Psi-Net: Shape and boundary aware joint multi-task deep network for medical image segmentation
Balamurali Murugesan, Kaushik Sarveswaran, Sharath M Shankaranarayana, Keerthi Ram, Mohanasankar Sivaprakasam
TL;DR
U-Net-like medical segmentation networks can produce coarse, discontinuous boundaries and fail to incorporate shape information. Psi-Net addresses this with three parallel decoders and a joint loss, and experiments report improved segmentation, shape, and boundary metrics.
Problem
U-Net-like networks can produce coarse, non-smooth segmentations, lose spatial information, and omit shape information from pixel-wise classification losses.
Method
Psi-Net uses one encoder and three parallel decoders for mask prediction, contour detection, and distance map estimation, trained with a joint loss.
Results
Experiments report better segmentation, boundary, and shape metrics than other evaluated networks on optic cup and disc and polyp segmentation tasks.
Takeaways & Limitations
The paper concludes that jointly learning mask, contour, and distance-map predictions preserves shape and improves boundary outputs and segmentation performance.
Abstract
from arXiv · showhide
Image segmentation is a primary task in many medical applications. Recently, many deep networks derived from U-Net have been extensively used in various medical image segmentation tasks. However, in most of the cases, networks similar to U-net produce coarse and non-smooth segmentations with lots of discontinuities. To improve and refine the performance of U-Net like networks, we propose the use of parallel decoders which along with performing the mask predictions also perform contour prediction and distance map estimation. The contour and distance map aid in ensuring smoothness in the segmentation predictions. To facilitate joint training of three tasks, we propose a novel architecture called Psi-Net with a single encoder and three parallel decoders (thus having a shape of $Ψ$), one decoder to learns the segmentation mask prediction and other two decoders to learn the auxiliary tasks of contour detection and distance map estimation. The learning of these auxiliary tasks helps in capturing the shape and the boundary information. We also propose a new joint loss function for the proposed architecture. The loss function consists of a weighted combination of Negative Log likelihood and Mean Square Error loss. We have used two publicly available datasets: 1) Origa dataset for the task of optic cup and disc segmentation and 2) Endovis segment dataset for the task of polyp segmentation to evaluate our model. We have conducted extensive experiments using our network to show our model gives better results in terms of segmentation, boundary and shape metrics.
I. INTRODUCTION
Psi-Net addresses limitations of U-Net-like segmentation by jointly predicting masks, contours, and distance maps with parallel decoders. Its auxiliary tasks and joint loss are designed to improve shape, boundary smoothness, and segmentation performance.
- Motivation: U-Net-like networks can produce uneven, discontinuous boundaries, lose spatial information through undersampling, and omit shape information from pixel-wise losses.Cross-entropy can also create class-imbalance problems when the background dominates medical images.
- Proposed approach: Psi-Net uses one encoder and three decoders for parallel mask prediction, contour detection, and distance map estimation.The architecture is described as having the shape of Ψ.
- Proposed approach: Contour detection and distance map estimation serve as auxiliary tasks that regularize mask prediction and produce refined masks with smooth boundaries.These tasks provide shape and boundary information during joint learning.
- Proposed approach: The joint loss combines Negative Log Likelihood for mask and contour prediction with Mean Square Error for distance estimation.This loss is intended to train the three tasks together.
- Evaluation: The evaluation uses Jaccard, Dice, Hausdorff distance, and trimap-based boundary metrics across optic cup and disc and polyp segmentation tasks.The datasets are ORIGA and a GIANA polyp segmentation dataset.
- Results: The comparative study reports better segmentation, shape, and boundary metrics than competing networks, producing masks with smoother boundaries.The reported comparison includes segmentation, shape similarity, and boundary evaluation.
A. Architecture
Psi-Net combines one contracting encoder with three structurally similar decoder paths. The decoders predict masks, contours, and distance maps, while a joint loss uses classification and regression objectives to regularize mask prediction.
- Architecture: Psi-Net has one contracting encoder and three expansive decoder paths whose arrangement resembles Ψ.The encoder repeatedly downsamples feature maps, while decoder layers concatenate corresponding encoder features to retain multi-scale information.
- Architecture: The three decoders perform mask segmentation, contour extraction, and distance map estimation as two classification tasks and one regression task.The distance decoder outputs one channel, while the mask and contour decoders output one channel per input class.
- Loss function: The loss uses Negative Log Likelihood for mask and contour decoders and Mean Square Error for the distance decoder.Mask prediction is regularized by both contour and distance-map predictions.
1) Mask:
The mask and contour losses measure pixel-wise classification error using predicted probabilities for the true labels after softmax activation.
- 1) Mask:: Lmask measures pixel-wise classification error over image-space pixel positions.The predicted probability corresponds to the true mask label after softmax activation.
- 2) Contour:: Lcontour likewise measures pixel-wise classification error for contour labels.It uses the predicted probability for the true contour label after softmax activation.
3) Distance:
The distance loss is a pixel-wise mean square error comparing the estimated distance map with its ground-truth counterpart.
- 3) Distance:: Ldistance computes pixel-wise mean square error between estimated and ground-truth distance maps.The estimated map follows sigmoid activation, while the ground-truth map is denoted D(x).
A. Dataset and Pre-processing
The approach was validated on optic cup and disc segmentation using ORIGA and polyp segmentation using the GIANA dataset, with contour and distance targets derived from masks.
- A. Dataset and Pre-processing: The study validates the segmentation approach on optic cup and disc segmentation and polyp segmentation.These are the two stated application areas.
- A. Dataset and Pre-processing: ORIGA provides 650 color fundus images with ground-truth optic disc and cup segmentations at 256 × 256 resolution.Ellipse fitting is applied to the output segmentation mask.
- A. Dataset and Pre-processing: The GIANA polyp dataset contains 912 images with ground-truth masks.
- A. Dataset and Pre-processing: Contour targets are estimated from connected-component boundaries and dilated with a disk filter of radius 5.Distance targets are obtained by applying a Euclidean distance transform to the mask.
- A. Dataset and Pre-processing: Models are trained for 150 epochs with Adam at learning rate 1e-4 and batch size 4.Experiments use an NVIDIA GeForce GTX 1060 GPU with 6GB RAM.
C. Evaluation metrics
Evaluation compares predicted outputs with ground truth using segmentation overlap, shape similarity, and boundary-focused measures.
- C. Evaluation metrics: The evaluation defines A as the method output and B as the actual ground truth.
- 1) Segmentation evaluation: Jaccard index measures segmentation overlap as intersection size divided by union size.It is also known as intersection over union, or IoU.
- 2) Shape Similarity: Hausdorff distance measures shape similarity between the segmented object and the ground-truth object.
D. Results and Discussion
The study compares Psi-Net’s three-decoder design with single- and two-decoder networks for mask, contour, and distance-map learning.
- Psi-Net uses one encoder and three parallel decoders for mask prediction, contour detection, and distance map estimation.The compared alternatives use one decoder for mask prediction, or two decoders for mask plus contour or distance estimation.
1) Standard Evaluation:
Psi-Net is evaluated using segmentation, shape, and boundary metrics against three related network designs. The reported results show stronger segmentation and shape performance, with lower boundary error across trimap widths.
- Segmentation metrics: Psi-Net outperformed the 1Enc 1Dec M, 1Enc 2Dec MC, and 1Enc 2Dec MD networks on Dice and Jaccard.The authors attribute the improvement to combining contour detection and distance map estimation as two auxiliary regularizers.
- Cross-task comparison: The authors report substantial improvement for polyp segmentation compared with optic cup and disc segmentation.They describe polyp segmentation as relatively difficult because polyp size and shape vary substantially.
- Shape metrics: Psi-Net captured shape information better than the other networks according to Hausdorff distance.The authors report that auxiliary tasks preserve shape, with distance map estimation capturing shape better than contour extraction.
- Boundary metrics: Psi-Net produced fewer misclassified pixels within trimaps of different widths, indicating improved boundary accuracy.Boundary error was measured within narrow bands surrounding ground-truth object boundaries.
4) Qualitative comparison:
Qualitative comparisons on polyp images show that Psi-Net produces smooth masks without outliers and handles multi-instance objects effectively.
- Qualitative comparison: Psi-Net and 1Enc 2Dec MD produced smooth masks without outliers, unlike 1Enc 1Dec M and 1Enc 2Dec MC.The comparison is shown for outputs from the polyp dataset.
- Qualitative comparison: Psi-Net handled multi-instance object segmentation well, whereas 1Enc 1Dec MD failed in that case.The authors state that Psi-Net performed as well as 1Enc 2Dec MC for this example.
- Qualitative comparison: Psi-Net outperformed the other networks in an additional qualitative example shown in the fourth row.
- Conclusion: The proposed architecture combines mask prediction, contour extraction, and distance map estimation to preserve shape and improve boundary outputs and segmentation performance.