Review — Shape-Aware Organ Segmentation by Predicting Signed Distance Maps
Shape-Aware Organ Segmentation by Predicting Signed Distance Maps,
SDM, by Pennsylvania State University, Tencent Hippocrates Research Lab, and Shandong Cancer Hospital and Institute
2020 AAAI, Over 60 Citations (Sik-Ho Tsang @ Medium)
Medical Imaging, Medical Image Analysis, Image Segmentation, Multi-Task Learning
- Existing approaches often produce results that do not capture the overall shape of the target organ and often lack smoothness.
- By converting the segmentation task into predicting an Signed Distance Map (SDM), better smoothness and continuity in shape are obtained. Specifically, an approximated Heaviside function to train the model by predicting SDMs and segmentation maps simultaneously.
- SDM Learning
1. SDM Learning
1.1. Model Architecture
- 3D U-Net variant with 6 downsampling operations, with the largest receptive field has size 64³.
- The model contains 6 skip connections between different scales of feature maps.
- Leaky ReLU is used instead of ReLU.
- The deconvolution is replaced by trilinear upsampling followed by convolution.
- More importantly, Group Normalization is used instead of Batch Normalization as in previous works. The Group Normalization is designed for training with smaller batch sizes.
1.2. Standard Loss Function
- Dice loss is used as Lseg, to measure the overlapping between ground-truth and predicted segmentation maps and is defined as:
- where N is the number of classes, t denotes the t-th organ class. yt and pt represent the ground-truth annotation and model prediction, respectively.
1.3. SDM Learning
- Given a target organ and a point x in the 3D medical image, the Signed Distance Map (SDM) which maps R³ to R is defined as:
- where S represents the surface of the target organ, Ωin and Ωout denote the region inside and outside the target organ, respectively. In other words, the absolute value of SDM indicates the distance from the point to the closest point on organ surface, while the sign indicates either inside or outside the organ. Note that the zero distance or zero level set means that the point is on the surface of the organ.
- The ground-truth SDM is approximated using Danielsson’s algorithm (Danielsson 1980) based on the ground-truth segmentation map.
- (Please read details of Danielsson’s algorithm form lecture notes by Stanford University.)
- Φ(x) is further normalized to be in the range [-1, 1] for each input image and use the tanh activation in the output layer. The normalization is done by dividing the SDM by the maximum positive value for points outside the organ, or by the minimum negative value for points inside the organ.
- The Heaviside function is used for generating segmentation map, but it is not differentiable.
- A smooth approximation to the Heaviside function is proposed, which is defined as:
- where k controls the steepness of the curve and closeness to the original Heaviside function, larger k means closer approximation.
- k=1500 due to the fact that such value guarantees around 99.9% overlapping between converted segmentation map and the original segmentation map.
- Sole L1 loss will have unstable training. Thus, L1 loss is combined with the proposed regression loss based on a product that is defined as:
- where yt represents the ground-truth SDM and pt denotes the predicted SDM. The intuition behind taking the product of prediction and ground-truth is that we want to penalize the output SDM for having the wrong sign.
- The total loss is:
- where λ=10.
2.1. Authors’ Collected Hippocampus Segmentation Dataset
The proposed jointing training with SDM and segmentation map achieves best performances in all evaluation metrics.
- The learned SDM with only SDM training obtains smoothest contours, while the joint training of SDM and segmentation map predicts more accurate organ boundary. Overall, they both preserve the shape of hippocampus and align well with the ground-truth SDM.
Segmentation by jointly training with SDM and binary map supervision get best performances.
2.2. MICCAI Head and Neck Auto Segmentation Challenge 2015 Dataset
- For both Dice and HD95, the proposed methods improve upon previous state-of-the-arts significantly, especially in small organs such as Chiasm, left and right Optic Nerve.
- Compared with the proposed backbone network trained with Dice loss only, the joint training model has a slightly lower Dice score, while still outperforms other state-of-the-art methods by a large margin.
2.3. Ablation Study
- The proposed backbone network trained with only segmentation prediction generally performs well, however, it produces some isolated false positives far away from the actual organ.
- The backbone network trained with only SDM prediction has smooth outputs, but does not converge on the small organs.
- The joint training with the proposed SDM loss converges well on all organs and preserves continuous shape.
- The backbone network trained with the proposed loss achieves best scores in all evaluation metrics.
While the network needs to predict the segmentation and SDM, it is a kind of multi-task learning.