High-quality reconstruction of real space structures from X-ray holography by a conditional Wavelet Flow

Digital holographic reconstruction has been a challenging task for many years due to its strict
requirements of prior knowledge of experimental setups as well as additional filtering algorithms
to suppress the zeroth-order light and twin image problems. Inspired by the data-driven
deep learning method, which is directly able to learn the non-linearity mapping between two
variables, a variety of research has been done in this area by pioneers. However, previous
methods are mainly constrained to discriminative models and they ignore the ill-posed nature
of holographic reconstruction. To address this problem, in this work, we define a novel variant
of normalizing flow, named conditional Wavelet Flow (cWavelet Flow), to reconstruct original
real space structures from digital X-ray holograms with a high degree of generative diversity
and quality. To pave the ground for the construction of cWavelet Flow, we first reproduce the
architecture and part of the experimental results of Wavelet Flow [2] based on PyTorch and
then extend it as cWavelet Flow by attaching an additional conditioning network on top of it.
The conditioning network, which consists of an optional pre-trained backbone network and a
head network, is constructed as an exceptionally lightweight yet performant structure. Such an
architectural design allows cWavelet Flow to directly model the conditional data distribution
of high resolution up to 1024 × 1024, which is almost impossible with the flow-based models
developed previously. Furthermore, another appealing point of cWavelet Flow is its highly
efficient training process. In comparison to other state-of-the-art baseline models like cINN
and U-Net, an improvement of up to 11.5× and 126.2× fewer FLOPs are achieved, respectively,
while maintaining reconstruction quality comparable to these baseline models.