An MRF-UNet Product of Experts for Image Segmentation

While convolutional neural networks (CNNs) trained by back-propagation have seen unprecedented success at semantic segmentation tasks, they are known to struggle on out-of-distribution data. Markov random fields (MRFs) on the other hand, encode simpler distributions over labels that, although less flexible than UNets, are less prone to over-fitting. In this paper, we propose to fuse both strategies by computing the product of distributions of a UNet and an MRF. As this product is intractable, we solve for an approximate distribution using an iterative mean-field approach. The resulting MRF-UNet is trained jointly by back-propagation. Compared to other works using conditional random fields (CRFs), the MRF has no dependency on the imaging data, which should allow for less over-fitting. We show on 3D neuroimaging data that this novel network improves generalisation to out-of-distribution samples. Furthermore, it allows the overall number of parameters to be reduced while preserving high accuracy. These results suggest that a classic MRF smoothness prior can allow for less over-fitting when principally integrated into a CNN model. Our implementation is available at this https URL.

[1]  Jayaram K. Udupa,et al.  Intensity standardization simplifies brain MR image segmentation , 2009, Comput. Vis. Image Underst..

[2]  Jose Dolz,et al.  3D fully convolutional networks for subcortical segmentation in MRI: A large-scale study , 2016, NeuroImage.

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

[4]  Jaime S. Cardoso,et al.  Elastic deformations for data augmentation in breast cancer mass detection , 2018, 2018 IEEE EMBS International Conference on Biomedical & Health Informatics (BHI).

[5]  Ender Konukoglu,et al.  A Lifelong Learning Approach to Brain MR Segmentation Across Scanners and Protocols , 2018, MICCAI.

[6]  Alan L. Yuille,et al.  Learning Deep Structured Models , 2014, ICML.

[7]  Konstantinos Kamnitsas,et al.  Unsupervised domain adaptation in brain lesion segmentation with adversarial networks , 2016, IPMI.

[8]  Daniel K Sodickson,et al.  Assessment of the generalization of learned image reconstruction and the potential for transfer learning , 2019, Magnetic resonance in medicine.

[9]  Mário A. T. Figueiredo,et al.  Conditional Random Fields as Recurrent Neural Networks for 3D Medical Imaging Segmentation , 2018, ArXiv.

[10]  Amod Jog,et al.  PSACNN: Pulse sequence adaptive fast whole brain segmentation , 2019, NeuroImage.

[11]  Quoc V. Le,et al.  Unsupervised Data Augmentation for Consistency Training , 2019, NeurIPS.

[12]  Geoffrey E. Hinton.  Training Products of Experts by Minimizing Contrastive Divergence , 2002, Neural Computation.

[13]  Vibhav Vineet,et al.  Conditional Random Fields as Recurrent Neural Networks , 2015, 2015 IEEE International Conference on Computer Vision (ICCV).

[14]  Xiao Han,et al.  Atlas Renormalization for Improved Brain MR Image Segmentation Across Scanner Platforms , 2007, IEEE Transactions on Medical Imaging.

[15]  Radford M. Neal.  Pattern Recognition and Machine Learning , 2007, Technometrics.

[16]  Adrian V. Dalca,et al.  Learning image registration without images , 2020 .

[17]  Adrian V. Dalca,et al.  A Learning Strategy for Contrast-agnostic MRI Segmentation , 2020, MIDL.

[18]  Konstantinos Kamnitsas,et al.  Efficient multi‐scale 3D CNN with fully connected CRF for accurate brain lesion segmentation , 2016, Medical Image Anal..

[19]  J. W. Modestino,et al.  Use of the mean-field approximation in an EM-based approach to unsupervised stochastic model-based image segmentation , 1992, [Proceedings] ICASSP-92: 1992 IEEE International Conference on Acoustics, Speech, and Signal Processing.

[20]  Stephen M. Smith,et al.  Segmentation of brain MR images through a hidden Markov random field model and the expectation-maximization algorithm , 2001, IEEE Transactions on Medical Imaging.

[21]  Jun Liu,et al.  A Regularized Convolutional Neural Network for Semantic Image Segmentation , 2019, Analysis and Applications.

[22]  Simon K. Warfield,et al.  Normalization of joint image-intensity statistics in MRI using the Kullback-Leibler divergence , 2004, 2004 2nd IEEE International Symposium on Biomedical Imaging: Nano to Macro (IEEE Cat No. 04EX821).

[23]  Adrian V. Dalca,et al.  Partial Volume Segmentation of Brain MRI Scans of any Resolution and Contrast , 2020, MICCAI.

[24]  Victor Alves,et al.  Brain Tumor Segmentation Using Convolutional Neural Networks in MRI Images , 2016, IEEE Transactions on Medical Imaging.

[25]  Karl J. Friston,et al.  Unified segmentation , 2005, NeuroImage.

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

[27]  John Ashburner,et al.  Nonlinear Markov Random Fields Learned via Backpropagation , 2019, IPMI.

[28]  Anders M. Dale,et al.  Sequence-independent segmentation of magnetic resonance images , 2004, NeuroImage.

[29]  M. Tribus,et al.  Probability theory: the logic of science , 2003 .

[30]  Koenraad Van Leemput,et al.  Automated model-based tissue classification of MR images of the brain , 1999, IEEE Transactions on Medical Imaging.