Plug-and-Play Priors for Reconstruction-based Placental Image Registration
Jiarui Xing, Ulugbek Kamilov, Wenjie Wu, Yong Wang, Miaomiao Zhang
PPlug-and-Play Priors for Reconstruction-basedPlacental Image Registration
Jiarui Xing , Ulugbek Kamilov , , Wenjie Wu , , Yong Wang , and MiaomiaoZhang , Electrical and Computer Engineering, University of Virginia, Charlottesville, USA Computer Science and Engineering, Washington University in St. Louis, USA Electrical and Systems Engineering, Washington University in St. Louis, USA Biomedical Engineering, Washington University in St. Louis, USA. Obstetrics and Gynecology, Washington University in St. Louis, USA. Computer Science, University of Virginia, Charlottesville, USA
Abstract.
This paper presents a novel deformable registration frame-work, leveraging an image prior specified through a denoising function,for severely noise-corrupted placental images. Recent work on plug-and-play (PnP) priors has shown the state-of-the-art performance of recon-struction algorithms under such priors in a range of imaging applications.Integration of powerful image denoisers into advanced registration meth-ods provides our model with a flexibility to accommodate datasets thathave low signal-to-noise ratios (SNRs). We demonstrate the performanceof our method under a wide variety of denoising models in the contextof diffeomorphic image registration. Experimental results show that ourmodel substantially improves the accuracy of spatial alignment in ap-plications of 3D in-utero diffusion-weighted MR images (DW-MRI) thatsuffer from low SNR and large spatial transformations.
Placental pathology, such as immune cell infiltration and inflammation [4], is acommon reason for preterm labor. It occurs in around 11% percent of worldpregnancies.
Diffusion-weighted magnetic resonance imaging (DW-MRI) is anon-invasive technique that is extensively used to monitor placental health andto assess its function throughout the entire pregnancy. However, this methodis quite susceptible to motion artifacts caused by maternal breathing and fetalmovements [13]. Additionally, DW-MRI scans often suffer from noise and severeartifacts induced by low signal-to-noise ratios (SNRs) at high b-values [22, 32].To address these issues, a noise-robust registration algorithm is needed.Many attempts have been made to develop registration methods that arerobust to image noise [12, 14, 17, 25]. A traditional approach integrates an imagereconstruction algorithm for removing noise and artifacts as a pre-processingstep to the registration task [29]. Further improvements can be achieved bydeveloping a joint framework that alternates between image reconstruction andregistration [12,16,28,32]. The most widely used image reconstruction algorithms a r X i v : . [ ee ss . I V ] S e p J. Xing et al. are based on optimization of an objective function that includes a regularizationterm for mitigating noise. Recently, however, the interest in the area has shiftedtowards a more flexible approach, known as plug-and-play priors (PnP) [30],that regularizes the problem using off-the-shelf image denoising algorithms. Ithas been shown that the combination of reconstruction algorithms with ad-vanced denoisers, such as non-local means [7] or block matching and 3D filtering(BM3D) [11], leads to the state-of-the-art performance for various imaging prob-lems [8, 9, 18, 26].In this paper, we extend the current family of joint reconstruction-registrationalgorithms by introducing a new method for deformable image registration called
PnP-RR (where RR stands for registration - reconstruction ). Our algorithm lever-ages PnP image priors, which makes it robust for registering severely noise-corrupted images. PnP-RR is very easy to implement by using a wide variety ofexisting algorithms with minimal effort to modify the infrastructure. We demon-strate how PnP priors can be used to mix and match a wide variety of existingreconstruction models with the state-of-the-art registration algorithm on both2D synthetic data and real 3D images. To show the effectiveness of the algorithmin improving the performance of spatial alignment for severely noise-corruptedimages, we test on 3D in-utero DW-MRI scans, affected by a low signal-to-noise(SNR) ratio and large motions. In this section, we briefly review the mathematical foundation of image reg-istration. Consider a d -dimensional image I defined as a continuous mapping I : Ω → R d , where Ω is the image domain. The transformation φ : Ω → Ω deforms a source image S by function composition S ◦ φ − , where ◦ denotesresampling. The goal of image registration is to find an optimal transformation φ , such that the deformed image S ◦ φ − is similar to a target image T .The desired transformation φ is typically computed by minimizing an energyfunction E ( φ ) = dist( S ◦ φ − , T ) + reg( φ ) . Here, the distance function dist( · , · )measures the dissimilarity between two images, such as sum-of-squared differ-ences of image intensities [3], mutual information [15], and normalized crosscorrelation [1]. The regularization term reg( · ) guarantees the smoothness ofthe transformation. A very original function φ is defined as a linear function φ ( x ) = x + u ( x ), where x ∈ Ω and u is a displacement vector field. With theregularity being set to (cid:107) Lu (cid:107) L ( L is a differential operator), the optimization ofthe energy function E over u arrives at a solution for elastic registration [6].However, such algorithm is not able to avoid geometric artifacts (e.g., fold-ing, tearing, or flipping) of the transformations, especially when large deforma-tion occurs, and may destroy the topology of local structures [10]. Instead, anelegant algorithm called large deformation diffeomorphic metric mapping (LD-DMM) was developed to ensure a smooth and invertible smooth mapping of φ between images [3]. The regularization term is defined as an integration overtime-dependent velocity fields derived from the transformations. We have the lug-and-Play Priors for Reconstruction-based Placental Image Registration 3 objective function of LDDMM asarg min v t σ (cid:13)(cid:13) S ◦ φ − − T (cid:13)(cid:13) L + (cid:90) ( Lv t , v t ) d t, s.t. d φ t d t = v t ( φ t ) , (1)where σ is a weighting parameter, and ( · , · ) acts similar to an inner product.The optimization of the original LDDMM is solved by gradient-based methodover the entire time sequence of v t , which is computationally expensive on high-dimensional images (e.g, a 3D placental MRI with the size of 128 ). Later, ageodesic shooting algorithm [19, 31] shows that once given an initial velocity v ,the shortest path of φ can be uniquely determined by integrating the geodesicevolution equation (also known as Euler-Poincare differential equation (EPDiff))defined by d v t d t = − K (cid:2) ( Dv t ) T · m t + Dm t · v t + m t · div v t (cid:3) , (2)where K is an inverse operator of the differential operator L , m t = Lv t is amomentum vector living in the dual space of v t , D denotes a Jacobian matrix,and div is a divergence operator.The optimization of Eq. (1) can be equivalently reformulated asarg min v σ (cid:107) S ◦ φ − − T (cid:107) L + ( Lv , v ) , s.t. d φ t d t = v t ( φ t ) & Eq. (2) . (3)This effectively shrinks the searching space from a time collection of { v t } to asingle initial point v , thus significantly reducing the computational complexityof the entire optimization.It has been recently demonstrated that the initial velocity v can be effi-ciently captured via a discrete low-dimensional bandlimited representation inthe Fourier space [33]. We develop our model by employing this fast registrationalgorithm named FLASH, which is the start-of-the-art variant of LDDMM withgeodesic shooting algorithm [34, 35]. In this section, we introduce a novel noise-robust registration model that incor-porates a PnP prior as an additional image regularizer. We show that the ourmodel can be implemented using two independent software modules – one forimage reconstruction and the other for image registration. Therefore, changingthe prior model only involves the implementation of image reconstruction. Thatis to say, our framework can be used to match a wide variety of priors with asuitable registration model.
We first consider the following joint objective function that builids on Eq. (3)to combine image regularization with deformable registration F ( v , ˜ T ) = 1 σ (cid:107) S ◦ φ − − ˜ T (cid:107) L + ( Lv , v ) + λ R ( ˜ T ) + λ (cid:107) T − ˜ T (cid:107) L , (4) J. Xing et al. where T is the target image, ˜ T is the reconstructed image, R ( · ) is the regulariza-tion term characterizing the prior on the image, λ is the parameter controllingthe strength of regularization, and λ controls the fidelity of the reconstructedand noisy images.In order to solve the problem (4) efficiently, we adopt an alternating min-imization approach [20], where v is first minimized for a fixed ˜ T under theconstraints in Eq. (3) and vice versa, as follows v k = arg min v F ( v , ˜ T k − ) , s.t. d φ t d t = v t ( φ t ) and Eq. (2) , (5a)˜ T k = arg min ˜ T F ( v k , ˜ T ) , (5b)where k denotes the k -th iteration.By ignoring the terms independent of v , the step (5a) can be expressed as v k = register σ ( S, ˜ T k − ) = arg min v σ (cid:107) S ◦ φ − − ˜ T k − (cid:107) L + ( Lv , v ) , where we didn’t explicitly write the constraints for better readability. Note thatthis step precisely matches the deformable image registration problem in Eq. (3).Similarly, the step (5b) can be simplified to the following form˜ T k = prox τ R ( Z k ) = arg min ˜ T (cid:107) ˜ T − Z k (cid:107) L + τ R ( ˜ T ) , (6)where we define Z k = λ T + (1 /σ )( S ◦ φ − ) λ + (1 /σ ) and τ = λ λ + (1 /σ )) . The minimization problem (6) is widely known as the proximal operator [21]and corresponds to an image denoising formulates as R ( · ) regularized optimiza-tion. For many popular regularizers, such as (cid:96) -norm or total variation penalty,the proximal operator either has a closed form solution or can be efficientlyimplemented [2], without differentiating R ( · ). Our alternating minimization algorithm in Eq. (5) iteratively refines a denoisedimage ˜ T k by applying the proximal operator defined in Eq. (6). Recently, themathematical equivalence of the proximal operator to image denoising has in-spired Venkatakrishnan et al. [30] to introduce a powerful PnP framework forimage reconstruction. The key idea of PnP is to replace the proximal operator inan iterative algorithm with a state-of-the-art image denoiser (e.g., BM3D), whichdoes not necessarily have a corresponding regularization function R ( · ). This im-plies that PnP methods generally lose interpretability as optimization problems.Nonetheless, the framework has gained in popularity due to its effectiveness in a lug-and-Play Priors for Reconstruction-based Placental Image Registration 5 range of applications. Additionally, several recent publications have theoreticallycharacterized the convergence and fixed points of PnP algorithms [8,9,24,26,27].Algorithm 1 summarizes our PnP-RR algorithm for joint image reconstruc-tion and registration. The fixed point ( v ∗ , ˜ T ∗ ) of PnP-RR is defined by a balancebetween denoising and registration operators, rather than the minimum of a costfunction. This makes the algorithm easy to adapt to specific datasets by simplyswapping denoisers or registration operators. We corroborate the performance ofPnP-RR in the next section by applying it to the challenging problem of imageregistration under severe amounts of noise. Algorithm 1
PnP-RR input: Source image S , target image T , parameters λ , λ , and σ set: τ = λ / (2( λ + (1 /σ )))3: for k = 1 , , . . . do v k ← register σ ( S, ˜ T k − ) (cid:46) registration step5: Z k ← ( λ T + (1 /σ )( S ◦ φ − )) / ( λ + (1 /σ ))6: ˜ T k ← denoise τ ( Z k ) (cid:46) denoising step7: end for To evaluate our proposed method, we test its performance with three existingreconstruction algorithms - total variation (TV) [23], total generalized variation(TGV) [5], and BM3D [11] on both synthetic 2D images and real 3D placentalDW-MRI scans with different b-values.We compare our method with the state-of-the-art fast registration methodFLASH [34] (downloaded from: https://bitbucket.org/FlashC/flashc). In all ex-periments, we set L as a Laplacian operator, e.g., L = − ( α (cid:52) + I ) c with apositive weight parameter α = 1 . c = 3 .
0. We set σ = 0 .
015 and the number of time integration steps n = 10 across all algorithms.We also perform registration-based segmentation and examine the resulting seg-mentation accuracy of the algorithm. To evaluate volume overlap between thepropagated segmentation A and the manual segmentation B for placenta, wecompute the Dice Similarity Coefficient DSC ( A, B ) = 2( | A | ∩ | B | ) / ( | A | + | B | ),where ∩ denotes an intersection of two regions. Data.
For 2D synthetic images, we generate a collection of binary images withresolution 100 . We then add white Gaussian noise with standard deviation σ = 0 . ± J. Xing et al.
30 channel phase-array torso coil (FOV = 386 × × − mm , 3 mm isotropic voxels, interleaved slice acquisition, TR = 14600 ms , TE = 62 ms , FlipAngle = 90 ◦ ). Multiple scans with different b values ( b = 0 , , , s/mm )were tested and the placenta were manually delineated for images with b = 0by radiologists. All DW-MRIs are of dimension 128 × ×
50 and underwentbias field correction, co-registration with affine transformations and intensitynormalization.
Experiments.
We first run an experiment on 2D synthetic data registering froma clean source image to a noisy target image, and compare the performance ofour method with the baseline algorithm FLASH. For the denoisers, we cross-validate a variety of different parameters and set λ = 0 . , λ = 0 .
067 for TV.Similarly, we have λ = 0 . , λ = 0 .
015 for TGV, and λ = 0 . , λ = 0 . b = 0) are considered as source images, while others with highb-values (typically noisy images) are target images. After testing a set of differentparameters, we set λ = 0 . , λ = 0 . λ = 0 . , λ = 0 . λ = 0 . , λ = 0 . Results.
Fig. 1 displays the registration results of the baseline algorithm andour model with different denoisers. It shows that our method achieves bettertransformations that nicely deform the source image fairly close to the targetimage, without being affected by the noises. (a) Source (b) Target(without noise) (c) Target(d) FLASH (e) Ours(TV) (f) Ours(TGV) (g) Ours(BM3D)
Fig. 1: Top: source image, clean target image, and noisy target image; Bottom:registration results from the baseline method FLASH and our model with TV,TGV, and BM3D denoisers. lug-and-Play Priors for Reconstruction-based Placental Image Registration 7
Fig. 2 demonstrates an example of the transformed segmentation of placenta(outlined in magenta) estimated by all algorithms. It clearly shows that thesegmentations produced by our algorithm align better with the manual segmen-tation (outlined in blue) than the baseline algorithm. Our model provides muchreliable segmentation than the baseline algorithm, especially on the left part ofthe placenta where relatively large deformation occurs. (d) FLASH (e) Ours(TV) (f) Ours(TGV) (g) Ours(BM3D)(a) Source (c) Target
Fig. 2: Top: source and target images; Bottom: comparison of estimate segmen-tations of all algorithms overlapped with manually labeled delineation.Fig. 3 shows another advantage of our model compared to two-step ap-proaches where image reconstruction is preformed before registration. We com-pute average dice scores with different parameter settings on both methods. Ourhigher average dice scores with smaller variations indicate that the proposedalgorithm is more robust to parameter-tuning.
In this paper, we presented a novel reconstruction-based registration algorithm,named PnP-RR, for severely noise-corrupted images. Our method is the first tointroduce PnP priors, represented through denoising functions, into the state-of-the-art registration framework. In contrast to previous approaches, our modelhas the flexibility to allow any reconstruction algorithm integrated with the reg-istration task. This provides a much more robust way to register images withlow SNRs and large motions. The theoretical tools developed in our work arebroadly applicable to a wide variety of joint reconstruction-registration algo-rithms. In addition, our method can be easily implemented through the currentimplementation of registration and reconstruction algorithms. Future researchwill involve collecting more dataset on placental images and exploring othercutting-edge denoisers, such as deep learning based approaches.
J. Xing et al.
Fig. 3: Comparison of averaged Dice score estimated from two-step approachesand ours.
Acknowledgement.
This work was supported by NIH grant R01HD094381,NIH grant R01AG053548, and BrightFocus Foundation A2017330S.
References
1. Avants, B.B., Epstein, C.L., Grossman, M., Gee, J.C.: Symmetric diffeomorphicimage registration with cross-correlation: evaluating automated labeling of elderlyand neurodegenerative brain. Medical image analysis (1), 26–41 (2008)2. Beck, A., Teboulle, M.: Fast gradient-based algorithms for constrained total vari-ation image denoising and deblurring problems. IEEE transactions on image pro-cessing (11), 2419–2434 (2009)3. Beg, M.F., Miller, M.I., Trouv´e, A., Younes, L.: Computing large deformationmetric mappings via geodesic flows of diffeomorphisms. International journal ofcomputer vision (2), 139–157 (2005)4. Blencowe, H., Cousens, S., Chou, D., Oestergaard, M., Say, L., Moller, A., Kinney,M., Lawn, J.: Born too soon: the global epidemiology of 15 million preterm births.Reproductive health (1), S2 (2013)5. Bredies, K., Kunisch, K., Pock, T.: Total generalized variation. SIAM Journal onImaging Sciences (3), 492–526 (2010)6. Broit, C.: Optimal registration of deformed images (1981)7. Buades, A., Coll, B., Morel, J.M.: A non-local algorithm for image denoising. In:2005 IEEE Computer Society Conference on Computer Vision and Pattern Recog-nition (CVPR’05). vol. 2, pp. 60–65. IEEE (2005)8. Buzzard, G.T., Chan, S.H., Sreehari, S., Bouman, C.A.: Plug-and-play unplugged:Optimization free reconstruction using consensus equilibrium. SIAM J. ImagingSci. (3), 2001–2020 (2018)9. Chan, S.H., Wang, X., Elgendy, O.A.: Plug-and-play ADMM for image restoration:Fixed-point convergence and applications. IEEE Trans. Comp. Imag. (1), 84–98(March 2017)lug-and-Play Priors for Reconstruction-based Placental Image Registration 910. Christensen, G.E.: Deformable shape models for anatomy (1994)11. Dabov, K., Foi, A., Egiazarian, K.: Video denoising by sparse 3d transform-domaincollaborative filtering. In: 2007 15th European Signal Processing Conference. pp.145–149. IEEE (2007)12. Han, J., Berkels, B., Rumpf, M., Hornegger, J., Droske, M., Fried, M., Scorzin, J.,Schaller, C.: A variational framework for joint image registration, denoising andedge detection. In: Bildverarbeitung f¨ur die Medizin 2006, pp. 246–250. Springer(2006)13. Le Bihan, D., Poupon, C., Amadon, A., Lethimonnier, F.: Artifacts and pitfalls indiffusion mri. Journal of Magnetic Resonance Imaging: An Official Journal of theInternational Society for Magnetic Resonance in Medicine (3), 478–488 (2006)14. Lempitsky, V., Rother, C., Blake, A.: Logcut-efficient graph cut optimization formarkov random fields. In: 2007 IEEE 11th International Conference on ComputerVision. pp. 1–8. IEEE (2007)15. Leventon, M., Wells III, W.M., Grimson, W.E.L.: Multiple view 2d-3d mutual in-formation registration. In: Image Understanding Workshop. vol. 20, p. 21. Citeseer(1997)16. Lombaert, H., Cheriet, F.: Simultaneous image de-noising and registration us-ing graph cuts: Application to corrupted medical images. In: 2012 11th Interna-tional Conference on Information Science, Signal Processing and their Applications(ISSPA). pp. 264–268. IEEE (2012)17. Lombaert, H., Cheriet, F.: Simultaneous image denoising and registration usinggraph cuts (07 2012)18. Meinhardt, T., Moeller, M., Hazirbas, C., Cremers, D.: Learning proximal opera-tors: Using denoising networks for regularizing inverse imaging problems. In: Proc.IEEE Int. Conf. Comp. Vis. (ICCV). pp. 1799–1808. Venice, Italy (October 22-29,2017)19. Miller, M.I., Trouv´e, A., Younes, L.: Geodesic shooting for computational anatomy.Journal of mathematical imaging and vision (2), 209–228 (2006)20. Nocedal, J., Wright, S.J.: Numerical Optimization. Springer, 2 edn. (2006)21. Parikh, N., Boyd, S.: Proximal algorithms. Foundations and Trends in Optimiza-tion (3), 123–231 (2014)22. Partridge, S.C., McDonald, E.S.: Diffusion weighted mri of the breast: Protocol op-timization, guidelines for interpretation, and potential clinical applications. Mag-netic resonance imaging clinics of North America (3), 601 (2013)23. Rudin, L.I., Osher, S., Fatemi, E.: Nonlinear total variation based noise removalalgorithms. Physica D: nonlinear phenomena (1-4), 259–268 (1992)24. Ryu, E.K., Liu, J., Wnag, S., Chen, X., Wang, Z., Yin, W.: Plug-and-play meth-ods provably converge with properly trained denoisers. In: Proc. 36th Int. Conf.Machine Learning (ICML). Long Beach, CA, USA (Jun 2019)25. Sanches, J.M., Marques, J.S.: Joint image registration and volume reconstructionfor 3d ultrasound. Pattern Recognition Letters (4-5), 791–800 (2003)26. Sreehari, S., Venkatakrishnan, S.V., Wohlberg, B., Buzzard, G.T., Drummy, L.F.,Simmons, J.P., Bouman, C.A.: Plug-and-play priors for bright field electron to-mography and sparse interpolation. IEEE Trans. Comp. Imag. (4), 408–423 (De-cember 2016)27. Sun, Y., Wohlberg, B., Kamilov, U.S.: An online plug-and-play algorithm for reg-ularized image reconstruction. IEEE Trans. Comput. Imaging (2019)28. Telea, A., Preusser, T., Garbe, C., Droske, M., Rumpf, M.: A variational approachto joint denoising, edge detection and motion estimation. In: Joint Pattern Recog-nition Symposium. pp. 525–535. Springer (2006)0 J. Xing et al.29. Tomaˇzeviˇc, D., Likar, B., Pernuˇs, F.: Reconstruction-based 3d/2d image registra-tion. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 231–238. Springer (2005)30. Venkatakrishnan, S.V., Bouman, C.A., Wohlberg, B.: Plug-and-play priors formodel based reconstruction. In: 2013 IEEE Global Conference on Signal and In-formation Processing. pp. 945–948. IEEE (2013)31. Vialard, F.X., Risser, L., Rueckert, D., Cotter, C.J.: Diffeomorphic 3d image reg-istration via geodesic shooting using an efficient adjoint calculation. InternationalJournal of Computer Vision (2), 229–241 (2012)32. Vishnevskiy, V., Stoeck, C., Sz´ekely, G., Tanner, C., Kozerke, S.: Simultaneousdenoising and registration for accurate cardiac diffusion tensor reconstruction frommri. In: International Conference on Medical Image Computing and Computer-Assisted Intervention. pp. 215–222. Springer (2015)33. Zhang, M., Fletcher, P.T.: Finite-dimensional lie algebras for fast diffeomorphic im-age registration. In: International Conference on Information Processing in MedicalImaging. pp. 249–260. Springer (2015)34. Zhang, M., Fletcher, P.T.: Fast diffeomorphic image registration via fourier-approximated lie algebras. International Journal of Computer Vision127