Physics-Informed Graph Neural Networks for Frequency-Aware Optical Aberration Correction

TL;DR

This paper introduces ZRNet, a physics-informed graph neural network that predicts Zernike coefficients and restores microscopy images, achieving state-of-the-art results in correcting optical aberrations by leveraging optical physics and Fourier domain alignment.

Contribution

The work presents a novel Zernike Graph module and Frequency-Aware Alignment loss, explicitly modeling physical relationships and improving correction accuracy in microscopy imaging.

Findings

Achieves state-of-the-art performance in microscopy aberration correction.

Demonstrates robustness to sensor noise and generalizes to real experimental data.

Effectively models physical wavefront distortions using graph neural networks.

Abstract

Optical aberrations significantly degrade image quality in microscopy, particularly when imaging deeper into samples. These aberrations arise from distortions in the optical wavefront and can be mathematically represented using Zernike polynomials. Existing methods often address only mild aberrations on limited sample types and modalities, typically treating the problem as a black-box mapping without leveraging the underlying optical physics of wavefront distortions. We propose ZRNet, a physics-informed framework that jointly performs Zernike coefficient prediction and optical image Restoration. We contribute a Zernike Graph module that explicitly models physical relationships between Zernike polynomials based on their azimuthal degrees-ensuring that learned corrections align with fundamental optical principles. To further enforce physical consistency between image restoration and Zernike prediction, we introduce a Frequency-Aware Alignment (FAA) loss, which better aligns Zernike coefficient prediction and image features in the Fourier domain. Extensive experiments on CytoImageNet demonstrates that our approach achieves state-of-the-art performance in both image restoration and Zernike coefficient prediction across diverse microscopy modalities and biological samples with complex, large-amplitude aberrations. We further validate on experimental PSF data from a physical microscope and demonstrate robustness to realistic sensor noise, confirming generalisation beyond simulated conditions. Code is available at https://github.com/janetkok/ZRNet.