Learning to Sense for Coded Diffraction Imaging

Abstract

In this paper, we present a framework to learn illumination patterns to improve the quality of signal recovery for coded diffraction imaging. We use an alternating minimization-based phase retrieval method with a fixed number of iterations as the iterative method. We represent the iterative phase retrieval method as an unrolled network with a fixed number of layers where each layer of the network corresponds to a single step of iteration, and we minimize the recovery error by optimizing over the illumination patterns. Since the number of iterations/layers is fixed, the recovery has a fixed computational cost. Extensive experimental results on a variety of datasets demonstrate that our proposed method significantly improves the quality of image reconstruction at a fixed computational cost with illumination patterns learned only using a small number of training images.

Keywords: coded diffraction imaging; learned sensors; phase retrieval.

Figure 1

Pipeline of our proposed framework at inference time. Our framework mainly contains two components: (1) a learnable sensing system that updates the illumination patterns during training time, but at inference time the learned illumination patterns are fixed; (2) a fixed unrolled network that runs phase retrieval process to recover the original signal x form measurements Y. The number of layers in the network is fixed to K. Steps at every iteration are fixed and depicted as an unrolled network (details can be found in Algorithm 1).

Figure 2

Reconstructed images using random and learned illumination patterns ($T=4$), along with ground truth (GT) in (a,b) and corresponding learned illumination patterns (c,d). PSNR is shown on top of every reconstruction. Below each dataset, we show the histograms of the PSNRs of all images with random patterns (shown in blue) and learned patterns (shown in orange). The dashed vertical line indicates the mean of all PSNRs.

Figure 3

Comparison of the reconstruction quality with random (in blue) and learned (in red) illumination patterns for different values of $K=1,\dots ,200$. We plot the average PSNR in a bright color and the PSNR of randomly selected 100 samples in light shadows.

Figure 4

Reconstruction quality vs. number of iterations (layers) at test time (i.e., K is different for training and testing with $T=4$). We show an error bar of $\pm 0.25\sigma$ for each dataset. In (a,b), we fixed K (K = 10, 20) and tested using different K. In (c), we trained and tested using the same number of layers.

Figure 5

Reconstruction quality vs. noise level of the measurements for different datasets ($T=4$). Here, we show a shaded error bar of $\pm 0.25\sigma$ for each dataset.

Figure 6

Test results on images shifted to bottom right by 5 pixels; from left to right: MNIST, F. MNIST, and CIFAR10.

Figure 7

Test results on images rotated by 90${}^{\circ }$; from left to right: MNIST, F. MNIST, and CIFAR10.