Physics-Informed Graph Neural Network with Frequency-Aware Learning for Optical Aberration Correction

TL;DR

This work presents ZRNet, a physics-informed framework that jointly restores microscopy images degraded by depth-dependent aberrations and predicts Zernike coefficients. It introduces a Zernike Graph to model couplings among Zernike modes based on azimuthal degree and a Fourier-domain Frequency-Aware Alignment loss to enforce consistency between restoration and aberration prediction. On CytoImageNet, ZRNet delivers state-of-the-art restoration quality and accurate, high-order Zernike estimation while reducing computational overhead. The approach demonstrates the value of integrating physical wavefront priors with deep learning for robust, cross-modality optical aberration correction.

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. Code is available at https://github.com/janetkok/ZRNet.

Figures (5)

• Figure 1: Comparison of aberration severity and correction performance. The top row shows aberrated images: those generated using implementations from previous works and an aberrated image from our approach. Our method applies approximately twice as many Zernike coefficients with larger amplitudes, resulting in a more severely aberrated image. The bottom row displays the corresponding corrected images from the baseline method and our method, respectively, along with the ground truth for reference.
• Figure 2: Overall pipeline: (1) synthesis of aberrated psfs using Zernike polynomials, (2) generation of phase-diverse images by convolving ground truth samples with aberrated psfs across multiple focal planes, (3) ZRNet that simultaneously restores images and predicts Zernike coefficients in a single forward pass, (4) with supervision across three domains: image space, frequency space, and Zernike coefficients.
• Figure 3: Our Zernike graph architecture for processing Zernike modes based on their azimuthal degrees. The framework processes information in four stages: (1) Node initialisation and intra-group exchange between Zernike nodes sharing the same azimuthal degree, (2) Information aggregation within each azimuthal degree group, (3) Inter-group information exchange via a fully-connected graph between different azimuthal degree groups, and (4) Information feedback to individual Zernike nodes to generate final latent representations.
• Figure 4: Qualitative comparison of sota image restoration networks on CytoImageNet hua2021cytoimagenet with optical aberration applied. ZRNet exhibits a higher level of fine detail after reconstruction, closer to the ground truth images.
• Figure 5: Zernike coefficients prediction comparing ground truth to ZRNet predictions.