WaveMo: Learning Wavefront Modulations to See Through Scattering

TL;DR

The paper tackles imaging through scattering media by learning wavefront modulations within a differentiable, end-to-end framework that jointly optimizes a proxy reconstruction network and a modulation generator. Learned modulations preserve higher-frequency information, improving reconstruction quality for both the proxy network and, when paired with unsupervised iterative methods, unseen scattering conditions. The approach decouples modulations from the reconstruction network at deployment, enabling use with other algorithms while maintaining performance gains. Results on simulated data and real tissue imaging demonstrate substantial PSNR/SSIM improvements and strong generalization, with a project page detailing supplementary materials.

Abstract

Imaging through scattering media is a fundamental and pervasive challenge in fields ranging from medical diagnostics to astronomy. A promising strategy to overcome this challenge is wavefront modulation, which induces measurement diversity during image acquisition. Despite its importance, designing optimal wavefront modulations to image through scattering remains under-explored. This paper introduces a novel learning-based framework to address the gap. Our approach jointly optimizes wavefront modulations and a computationally lightweight feedforward "proxy" reconstruction network. This network is trained to recover scenes obscured by scattering, using measurements that are modified by these modulations. The learned modulations produced by our framework generalize effectively to unseen scattering scenarios and exhibit remarkable versatility. During deployment, the learned modulations can be decoupled from the proxy network to augment other more computationally expensive restoration algorithms. Through extensive experiments, we demonstrate our approach significantly advances the state of the art in imaging through scattering media. Our project webpage is at https://wavemo-2024.github.io/.

Figures (12)

• Figure 1: Learned Wavefront Modulations. Top: During acquisition, we can modulate the wavefront of scattered light by using a spatial light modulator (SLM), and capture a set of image measurements useful for scene reconstruction. Bottom left: We propose learning modulations that enhance our ability to recover the scattered scene. Bottom right: Our learned modulations drastically improve the reconstruction quality of a state-of-the-art method feng2023neuws that previously applies randomly chosen modulations.
• Figure 2: Overview of the Proposed End-to-end Learning Framework. During training, we jointly optimize the implicit neural representation for the wavefront modulations and the proxy reconstruction network in an end-to-end fashion. During inference, one can apply the learned wavefront modulations to any reconstruction algorithms for imaging through scattering, either a trained feed-forward reconstruction network or an unsupervised iterative optimization algorithm feng2023neuws. The former approach is far faster and performs better when test data and training data fall into similar distributions, while the latter generalizes better to unseen distributions of target scenes.
• Figure 3: MTF Comparison. We compare the MTF of learned wavefront modulations vs. the MTF of modulations randomly drawn from the same distribution of the scattering media. The X-axis represents spatial frequency, and the Y-axis represents the modulation transfer (the higher the better). Our learned wavefront modulations exhibit a higher MTF compared to unoptimized ones, especially over higher frequency bands, suggesting the learned modulations preserve more high-frequency information.
• Figure 4: Proxy Network Reconstruction on Simulated Scattering. Simulated imaging results through different aberrations. Learned modulations lead to a better quality.
• Figure 5: Proxy Network Reconstruction on Physical Scattering. Experimental results of imaging objects through scattering media by using our proxy reconstruction network. The left columns are in-distribution adipose tissue slides (zoomed-in region labeled with red boxes); the right columns are out-of-distribution targets. Learned modulations yield superior imaging quality for both in-distribution and out-of-distribution scenes, with the former out-performing the latter.