The AIRI plug-and-play algorithm for image reconstruction in radio-interferometry: variations and robustness (2312.07137v3)
Abstract: Plug-and-Play (PnP) algorithms are appealing alternatives to proximal algorithms when solving inverse imaging problems. By learning a Deep Neural Network (DNN) denoiser behaving as a proximal operator, one waives the computational complexity of optimisation algorithms induced by sophisticated image priors, and the sub-optimality of handcrafted priors compared to DNNs. Such features are highly desirable in radio-interferometric (RI) imaging, where precision and scalability of the image reconstruction process are key. In previous work, we introduced AIRI, PnP counterpart to the unconstrained variant of the SARA optimisation algorithm, relying on a forward-backward algorithmic backbone. Here, we introduce variations of AIRI towards a more general and robust PnP paradigm in RI imaging. Firstly, we show that the AIRI denoisers can be used without any alteration to instantiate a PnP counterpart to the constrained SARA optimisation algorithm itself, relying on a primal-dual forward-backward algorithmic backbone, thus extending the remit of the AIRI paradigm. Secondly, we show that AIRI algorithms are robust to strong variations in the nature of the training dataset, with denoisers trained on medical images yielding similar reconstruction quality to those trained on astronomical images. Thirdly, we develop a functionality to quantify the model uncertainty introduced by the randomness in the training process. We validate the image reconstruction and uncertainty quantification functionality of AIRI algorithms against the SARA family and CLEAN, both in simulation and on real data of the ESO 137-006 galaxy acquired with the MeerKAT telescope. AIRI code is available in the BASPLib code library on GitHub.
- K. Bredies, K. Kunisch, and T. Pock, “Total generalized variation,” SIAM Journal on Imaging Sciences, vol. 3, no. 3, pp. 492–526, 2010.
- J.-F. Cai, B. Dong, S. Osher, and Z. Shen, “Image restoration: total variation, wavelet frames, and beyond,” Journal of the American Mathematical Society, vol. 25, no. 4, pp. 1033–1089, 2012.
- J. Zbontar, F. Knoll, A. Sriram, T. Murrell, Z. Huang, M. J. Muckley, A. Defazio, R. Stern, P. Johnson, M. Bruno et al., “fastmri: An open dataset and benchmarks for accelerated mri,” arXiv preprint arXiv:1811.08839, 2018.
- T. M. Quan, T. Nguyen-Duc, and W.-K. Jeong, “Compressed sensing mri reconstruction using a generative adversarial network with a cyclic loss,” IEEE transactions on medical imaging, vol. 37, no. 6, pp. 1488–1497, 2018.
- J. Liang, J. Cao, G. Sun, K. Zhang, L. Van Gool, and R. Timofte, “Swinir: Image restoration using swin transformer,” in Proceedings of the IEEE/CVF international conference on computer vision, 2021, pp. 1833–1844.
- S. V. Venkatakrishnan, C. A. Bouman, and B. Wohlberg, “Plug-and-play priors for model based reconstruction,” in IEEE Global Conference on Signal and Information Processing, 2013, pp. 945–948.
- T. Meinhardt, M. Moller, C. Hazirbas, and D. Cremers, “Learning proximal operators: Using denoising networks for regularizing inverse imaging problems,” in Proceedings of the IEEE International Conference on Computer Vision, 2017, pp. 1781–1790.
- K. Zhang, Y. Li, W. Zuo, L. Zhang, L. Van Gool, and R. Timofte, “Plug-and-play image restoration with deep denoiser prior,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 44, no. 10, pp. 6360–6376, 2022.
- R. Liu, Y. Sun, J. Zhu, L. Tian, and U. S. Kamilov, “Recovery of continuous 3d refractive index maps from discrete intensity-only measurements using neural fields,” Nature Machine Intelligence, vol. 4, no. 9, pp. 781–791, 2022.
- E. Ryu, J. Liu, S. Wang, X. Chen, Z. Wang, and W. Yin, “Plug-and-play methods provably converge with properly trained denoisers,” in International Conference on Machine Learning. PMLR, 2019, pp. 5546–5557.
- J.-C. Pesquet, A. Repetti, M. Terris, and Y. Wiaux, “Learning maximally monotone operators for image recovery,” SIAM Journal on Imaging Sciences, vol. 14, no. 3, pp. 1206–1237, 2021.
- Y. Wiaux, L. Jacques, G. Puy, A. M. Scaife, and P. Vandergheynst, “Compressed sensing imaging techniques for radio interferometry,” Monthly Notices of the Royal Astronomical Society, vol. 395, no. 3, pp. 1733–1742, 2009.
- E. J. Candes, J. K. Romberg, and T. Tao, “Stable signal recovery from incomplete and inaccurate measurements,” Communications on Pure and Applied Mathematics: A Journal Issued by the Courant Institute of Mathematical Sciences, vol. 59, no. 8, pp. 1207–1223, 2006.
- R. E. Carrillo, J. D. McEwen, and Y. Wiaux, “Sparsity averaging reweighted analysis (sara): a novel algorithm for radio-interferometric imaging,” Monthly Notices of the Royal Astronomical Society, vol. 426, no. 2, pp. 1223–1234, 2012.
- T. J. Cornwell, “Multiscale clean deconvolution of radio synthesis images,” IEEE Journal of selected topics in signal processing, vol. 2, no. 5, pp. 793–801, 2008.
- A. Offringa, B. McKinley, N. Hurley-Walker, F. Briggs, R. Wayth, D. Kaplan, M. Bell, L. Feng, A. Neben, J. Hughes et al., “Wsclean: an implementation of a fast, generic wide-field imager for radio astronomy,” Monthly Notices of the Royal Astronomical Society, vol. 444, no. 1, pp. 606–619, 2014.
- A. Offringa and O. Smirnov, “An optimized algorithm for multiscale wideband deconvolution of radio astronomical images,” Monthly Notices of the Royal Astronomical Society, vol. 471, no. 1, pp. 301–316, 2017.
- P. Arras, H. L. Bester, R. A. Perley, R. Leike, O. Smirnov, R. Westermann, and T. A. Enßlin, “Comparison of classical and bayesian imaging in radio interferometry-cygnus a with clean and resolve,” Astronomy & Astrophysics, vol. 646, p. A84, 2021.
- M. Terris, A. Dabbech, C. Tang, and Y. Wiaux, “Image reconstruction algorithms in radio interferometry: From handcrafted to learned regularization denoisers,” Monthly Notices of the Royal Astronomical Society, vol. 518, no. 1, pp. 604–622, 2023.
- A. Dabbech, M. Terris, A. Jackson, M. Ramatsoku, O. M. Smirnov, and Y. Wiaux, “First AI for deep super-resolution wide-field imaging in radio astronomy: unveiling structure in eso 137-006,” The Astrophysical Journal Letters, vol. 939, no. 1, p. L4, 2022.
- A. Chambolle and T. Pock, “A first-order primal-dual algorithm for convex problems with applications to imaging,” Journal of mathematical imaging and vision, vol. 40, pp. 120–145, 2011.
- B. C. Vũ, “A splitting algorithm for dual monotone inclusions involving cocoercive operators,” Advances in Computational Mathematics, vol. 38, no. 3, pp. 667–681, 2013.
- S. Hurault, A. Leclaire, and N. Papadakis, “Gradient step denoiser for convergent plug-and-play,” in International Conference on Learning Representations, 2022. [Online]. Available: https://openreview.net/forum?id=fPhKeld3Okz
- ——, “Proximal denoiser for convergent plug-and-play optimization with nonconvex regularization,” in International Conference on Machine Learning. PMLR, 2022, pp. 9483–9505.
- J. Hertrich, S. Neumayer, and G. Steidl, “Convolutional proximal neural networks and plug-and-play algorithms,” Linear Algebra and its Applications, vol. 631, pp. 203–234, 2021.
- A. Repetti, M. Terris, Y. Wiaux, and J.-C. Pesquet, “Dual forward-backward unfolded network for flexible plug-and-play,” in 2022 30th European Signal Processing Conference (EUSIPCO). IEEE, 2022, pp. 957–961.
- H. Y. Tan, S. Mukherjee, J. Tang, and C.-B. Schönlieb, “Provably convergent plug-and-play quasi-newton methods,” arXiv preprint arXiv:2303.07271, 2023.
- S. Boyd, N. Parikh, E. Chu, B. Peleato, J. Eckstein et al., “Distributed optimization and statistical learning via the alternating direction method of multipliers,” Foundations and Trends® in Machine learning, vol. 3, no. 1, pp. 1–122, 2011.
- F. Heide, M. Steinberger, Y.-T. Tsai, M. Rouf, D. Pająk, D. Reddy, O. Gallo, J. Liu, W. Heidrich, K. Egiazarian et al., “Flexisp: A flexible camera image processing framework,” ACM Transactions on Graphics (ToG), vol. 33, no. 6, pp. 1–13, 2014.
- S. Ono, “Primal-dual plug-and-play image restoration,” IEEE Signal Processing Letters, vol. 24, no. 8, pp. 1108–1112, 2017.
- S. K. Shastri, R. Ahmad, and P. Schniter, “Autotuning plug-and-play algorithms for mri,” in 2020 54th Asilomar Conference on Signals, Systems, and Computers. IEEE, 2020, pp. 1400–1404.
- V. Missaglia, F. Massaro, A. Capetti, M. Paolillo, R. Kraft, R. D. Baldi, and A. Paggi, “Watcat: a tale of wide-angle tailed radio galaxies,” Astronomy & Astrophysics, vol. 626, p. A8, 2019.
- D. Wittor, S. Ettori, F. Vazza, K. Rajpurohit, M. Hoeft, and P. Domínguez-Fernández, “Exploring the spectral properties of radio relics–i: integrated spectral index and mach number,” Monthly Notices of the Royal Astronomical Society, vol. 506, no. 1, pp. 396–414, 2021.
- K. Zhang, Y. Li, J. Liang, J. Cao, Y. Zhang, H. Tang, R. Timofte, and L. Van Gool, “Practical blind denoising via swin-conv-unet and data synthesis,” arXiv preprint, 2022.
- K. Wei, A. I. Avilés-Rivero, J. Liang, Y. Fu, H. Huang, and C.-B. Schönlieb, “Tfpnp: Tuning-free plug-and-play proximal algorithms with applications to inverse imaging problems.” J. Mach. Learn. Res., vol. 23, no. 16, pp. 1–48, 2022.
- A. G. Wilber, A. Dabbech, M. Terris, A. Jackson, and Y. Wiaux, “Scalable precision wide-field imaging in radio interferometry: Ii. airi validated on askap data,” arXiv preprint arXiv:2302.14149, 2023.
- A. Repetti, M. Pereyra, and Y. Wiaux, “Scalable bayesian uncertainty quantification in imaging inverse problems via convex optimization,” SIAM Journal on Imaging Sciences, vol. 12, no. 1, pp. 87–118, 2019.
- R. Laumont, V. D. Bortoli, A. Almansa, J. Delon, A. Durmus, and M. Pereyra, “Bayesian imaging using plug & play priors: when langevin meets tweedie,” SIAM Journal on Imaging Sciences, vol. 15, no. 2, pp. 701–737, 2022.
- S. Mukherjee, A. Hauptmann, O. Öktem, M. Pereyra, and C.-B. Schönlieb, “Learned reconstruction methods with convergence guarantees: a survey of concepts and applications,” IEEE Signal Processing Magazine, vol. 40, no. 1, pp. 164–182, 2023.
- C. Zhang and B. Jin, “Probabilistic residual learning for aleatoric uncertainty in image restoration,” arXiv preprint arXiv:1908.01010, 2019.
- T. I. Liaudat, M. Mars, M. A. Price, M. Pereyra, M. M. Betcke, and J. D. McEwen, “Scalable bayesian uncertainty quantification with data-driven priors for radio interferometric imaging,” 2023.
- D. Narnhofer, A. Effland, E. Kobler, K. Hammernik, F. Knoll, and T. Pock, “Bayesian uncertainty estimation of learned variational mri reconstruction,” IEEE Transactions on Medical Imaging, vol. 41, no. 2, pp. 279–291, 2021.
- S. Lahlou, M. Jain, H. Nekoei, V. I. Butoi, P. Bertin, J. Rector-Brooks, M. Korablyov, and Y. Bengio, “Deup: Direct epistemic uncertainty prediction,” arXiv preprint arXiv:2102.08501, 2021.
- C. Ekmekci and M. Cetin, “What does your computational imaging algorithm not know?: A plug-and-play model quantifying model uncertainty,” in Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021, pp. 4018–4027.
- A. Kendall and Y. Gal, “What uncertainties do we need in bayesian deep learning for computer vision?” Advances in neural information processing systems, vol. 30, 2017.
- C. F. Baumgartner, K. C. Tezcan, K. Chaitanya, A. M. Hötker, U. J. Muehlematter, K. Schawkat, A. S. Becker, O. Donati, and E. Konukoglu, “Phiseg: Capturing uncertainty in medical image segmentation,” in Medical Image Computing and Computer Assisted Intervention–MICCAI 2019: 22nd International Conference, Shenzhen, China, October 13–17, 2019, Proceedings, Part II 22. Springer, 2019, pp. 119–127.
- Y. Xue, S. Cheng, Y. Li, and L. Tian, “Reliable deep-learning-based phase imaging with uncertainty quantification,” Optica, vol. 6, no. 5, pp. 618–629, 2019.
- V. Edupuganti, M. Mardani, S. Vasanawala, and J. Pauly, “Uncertainty quantification in deep mri reconstruction,” IEEE Transactions on Medical Imaging, vol. 40, no. 1, pp. 239–250, 2020.
- Y. Kwon, J.-H. Won, B. J. Kim, and M. C. Paik, “Uncertainty quantification using bayesian neural networks in classification: Application to ischemic stroke lesion segmentation,” in Medical Imaging with Deep Learning, 2022.
- S. Gao and X. Zhuang, “Bayesian image super-resolution with deep modeling of image statistics,” IEEE Transactions on Pattern Analysis and Machine Intelligence, vol. 45, no. 2, pp. 1405–1423, 2022.
- J. Jonas and M. Team, “The meerkat radio telescope,” MeerKAT Science: On the Pathway to the SKA, p. 1, 2016.
- K. Dabov, A. Foi, V. Katkovnik, and K. Egiazarian, “Image denoising by sparse 3-d transform-domain collaborative filtering,” IEEE Transactions on image processing, vol. 16, no. 8, pp. 2080–2095, 2007.
- K. Zhang, W. Zuo, Y. Chen, D. Meng, and L. Zhang, “Beyond a gaussian denoiser: Residual learning of deep cnn for image denoising,” IEEE transactions on image processing, vol. 26, no. 7, pp. 3142–3155, 2017.
- D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” in International Conference on Learning Representations, 2015.
- M. Ramatsoku, M. Murgia, V. Vacca, P. Serra, S. Makhathini, F. Govoni, O. Smirnov, L. Andati, E. De Blok, G. Józsa et al., “Collimated synchrotron threads linking the radio lobes of eso 137-006,” Astronomy & Astrophysics, vol. 636, p. L1, 2020.
- A. Repetti and Y. Wiaux, “A non-convex perspective on calibration and imaging in radio interferometry,” in Wavelets and Sparsity XVII, vol. 10394. SPIE, 2017, pp. 392–404.
- P.-A. Thouvenin, A. Abdulaziz, A. Dabbech, A. Repetti, and Y. Wiaux, “Parallel faceted imaging in radio interferometry via proximal splitting (faceted hypersara): I. algorithm and simulations,” Monthly Notices of the Royal Astronomical Society, vol. 521, no. 1, pp. 1–19, 2023.