MetricUNet: Synergistic Image- and Voxel-Level Learning for Precise CT Prostate Segmentation via Online Sampling

Fully convolutional networks (FCNs), including UNet and VNet, are widely-used network architectures for semantic segmentation in recent studies. However, conventional FCN is typically trained by the cross-entropy or Dice loss, which only calculates the error between predictions and ground-truth labels for pixels individually. This often results in non-smooth neighborhoods in the predicted segmentation. To address this problem, we propose a two-stage framework, with the first stage to quickly localize the prostate region, and the second stage to precisely segment the prostate by a multi-task UNet architecture. We introduce a novel online metric learning module through voxel-wise sampling in the multi-task network. Therefore, the proposed network has a dual-branch architecture that tackles two tasks: 1) a segmentation sub-network aiming to generate the prostate segmentation, and 2) a voxel-metric learning sub-network aiming to improve the quality of the learned feature space supervised by a combination of triplet and contrast loss function. Specifically, the voxel-metric learning sub-network sampled tuples (including triplets and pairs) in voxel-level through the intermediate feature maps. Unlike conventional deep metric learning methods that generate triplets or pairs in image-level before the training phase, our proposed voxel-wise tuples are sampled in an online manner and operated in an end-to-end fashion via multi-task learning. To evaluate the proposed method, we implement extensive experiments on a real CT image dataset consisting 339 patients. The ablation studies show that our method can effectively learn more representative voxel-level features compared with the conventional learning methods with cross-entropy or Dice loss. And the comparisons show that the proposed method outperforms the state-of-the-art methods by a reasonable margin.

[1]  Yaozong Gao,et al.  Locally-constrained boundary regression for segmentation of prostate and rectum in the planning CT images , 2015, Medical Image Anal..

[2]  Rongrong Guo,et al.  A combined learning algorithm for prostate segmentation on 3D CT images , 2017, Medical physics.

[3]  Xudong Lin,et al.  Deep Adversarial Metric Learning , 2018, 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.

[4]  David Zhang,et al.  Joint Learning of Single-Image and Cross-Image Representations for Person Re-identification , 2016, 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[5]  Li Fei-Fei,et al.  DenseCap: Fully Convolutional Localization Networks for Dense Captioning , 2015, 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[6]  James Philbin,et al.  FaceNet: A unified embedding for face recognition and clustering , 2015, 2015 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[7]  Weidong Cai,et al.  HD-Net: Hybrid Discriminative Network for Prostate Segmentation in MR Images , 2019, MICCAI.

[8]  Seyed-Ahmad Ahmadi,et al.  V-Net: Fully Convolutional Neural Networks for Volumetric Medical Image Segmentation , 2016, 2016 Fourth International Conference on 3D Vision (3DV).

[9]  Yinghuan Shi,et al.  Pelvic Organ Segmentation Using Distinctive Curve Guided Fully Convolutional Networks. , 2019, IEEE transactions on medical imaging.

[10]  Nima Tajbakhsh,et al.  UNet++: A Nested U-Net Architecture for Medical Image Segmentation , 2018, DLMIA/ML-CDS@MICCAI.

[11]  Dinggang Shen,et al.  Segmentation of Organs at Risk in thoracic CT images using a SharpMask architecture and Conditional Random Fields , 2017, 2017 IEEE 14th International Symposium on Biomedical Imaging (ISBI 2017).

[12]  Hervé Delingette,et al.  Automatic Segmentation of Bladder and Prostate Using Coupled 3D Deformable Models , 2007, MICCAI.

[13]  Ross B. Girshick,et al.  Focal Loss for Dense Object Detection , 2017, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[14]  Dinggang Shen,et al.  Learning-based deformable registration of MR brain images , 2006, IEEE Transactions on Medical Imaging.

[15]  Pheng-Ann Heng,et al.  CIA-Net: Robust Nuclei Instance Segmentation with Contour-aware Information Aggregation , 2019, IPMI.

[16]  Iasonas Kokkinos,et al.  DeepLab: Semantic Image Segmentation with Deep Convolutional Nets, Atrous Convolution, and Fully Connected CRFs , 2016, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[17]  Yinghuan Shi,et al.  Does Manual Delineation only Provide the Side Information in CT Prostate Segmentation? , 2017, MICCAI.

[18]  Liang Lin,et al.  Deep feature learning with relative distance comparison for person re-identification , 2015, Pattern Recognit..

[19]  Yaozong Gao,et al.  Hierarchical Representation For Ct Prostate Segmentation , 2019, 2019 IEEE 16th International Symposium on Biomedical Imaging (ISBI 2019).

[20]  Long Ang Lim,et al.  Foreground segmentation using convolutional neural networks for multiscale feature encoding , 2018, Pattern Recognit. Lett..

[21]  Trevor Darrell,et al.  Fully Convolutional Networks for Semantic Segmentation , 2017, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[22]  Tian Liu,et al.  Deeply supervised 3D fully convolutional networks with group dilated convolution for automatic MRI prostate segmentation , 2019, Medical physics.

[23]  Yaozong Gao,et al.  Accurate Segmentation of CT Male Pelvic Organs via Regression-Based Deformable Models and Multi-Task Random Forests , 2016, IEEE Transactions on Medical Imaging.

[24]  Luca Bertinetto,et al.  Fully-Convolutional Siamese Networks for Object Tracking , 2016, ECCV Workshops.

[25]  Dinggang Shen,et al.  Multi‐channel multi‐scale fully convolutional network for 3D perivascular spaces segmentation in 7T MR images , 2018, Medical Image Anal..

[26]  Hao Chen,et al.  DCAN: Deep contour‐aware networks for object instance segmentation from histology images , 2017, Medical Image Anal..

[27]  Thomas Brox,et al.  U-Net: Convolutional Networks for Biomedical Image Segmentation , 2015, MICCAI.

[28]  Yang Lei,et al.  Automated prostate segmentation of volumetric CT images using 3D deeply supervised dilated FCN , 2019, Medical Imaging: Image Processing.

[29]  Yang Lei,et al.  Synthetic MRI-aided multi-organ segmentation on male pelvic CT using cycle consistent deep attention network. , 2019, Radiotherapy and oncology : journal of the European Society for Therapeutic Radiology and Oncology.

[30]  Dinggang Shen,et al.  Hierarchical Fully Convolutional Network for Joint Atrophy Localization and Alzheimer's Disease Diagnosis Using Structural MRI , 2020, IEEE Transactions on Pattern Analysis and Machine Intelligence.

[31]  Jens Petersen,et al.  nnU-Net: a self-configuring method for deep learning-based biomedical image segmentation , 2020, Nature Methods.

[32]  Bhiksha Raj,et al.  SphereFace: Deep Hypersphere Embedding for Face Recognition , 2017, 2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).

[33]  Thomas Brox,et al.  3D U-Net: Learning Dense Volumetric Segmentation from Sparse Annotation , 2016, MICCAI.

[34]  Dinggang Shen,et al.  Deformable registration of brain tumor images via a statistical model of tumor-induced deformation , 2006, Medical Image Anal..

[35]  Andrew Zisserman,et al.  Deep Face Recognition , 2015, BMVC.

[36]  Luca Antiga,et al.  Automatic differentiation in PyTorch , 2017 .

[37]  Geoffrey E. Hinton,et al.  A Simple Framework for Contrastive Learning of Visual Representations , 2020, ICML.

[38]  Yi Li,et al.  R-FCN: Object Detection via Region-based Fully Convolutional Networks , 2016, NIPS.

[39]  Agne Paulauskaite-Taraseviciene,et al.  Enhancing Multi-tissue and Multi-scale Cell Nuclei Segmentation with Deep Metric Learning , 2020 .

[40]  Nanning Zheng,et al.  Person Re-identification by Multi-Channel Parts-Based CNN with Improved Triplet Loss Function , 2016, 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR).