Self-Supervised Deep Learning for Improved Image-Based Wave-Front Sensing

Abstract

Phase retrieval from supervised learning neural networks is restricted due to the problem of obtaining labels. To address this situation, in the present paper, we propose a phase retrieval model of self-supervised physical deep learning combined with a complete physical model to represent the image-formation process. The model includes two parts: one is MobileNet V1, which is used to map the input samples to the Zernike coefficients, the other one is an optical imaging system and it is used to obtain the point spread function for training the model. In addition, the loss function is calculated based on the similarity between the input and the output to realize self-supervised learning. The root-mean-square (RMS) of the wave-front error (WFE) between the input and reconstruction is 0.1274 waves in the situation of D/r0 = 20 in the simulation. By comparison, The RMS of WFE is 0.1069 waves when using the label to train the model. This method retrieves numerous wave-front errors in real time in the presence of simulated detector noise without relying on label values. Moreover, this method is more suitable for practical applications and is more robust than supervised learning. We believe that this technology has great applications in free-space optical communication.

1. Introduction

As shown in Figure 1, space light needs to be coupled into single-mode fiber at the receiving end in free-space optical communication (FSOC). However, wave-front aberrations generated by atmospheric turbulence affect the beam quality, then deteriorate the fiber coupling efficiency and the quality of communication. Therefore, it is necessary to detect and correct the wave-front aberration. The measurement of the wave-front aberration in FSOC differs from other scenarios in that the wave-front aberration changes continuously. Therefore, an accurate real-time measurement is required in FSOC. There are two main methods for measuring wave-front aberrations. The first method uses additional wave-front sensors, such as a Hartmann sensor or interferometer [1,2,3], to monitor the wave-front slope. The other method uses the spot quality as the objective function and optimizes the objective function with continuous iteration, such as an image-based sensor [4,5,6,7]. The real-time performance of this second method is poor, so its application range is very limited. Deep learning is applied to the image-based sensor in order to improve the real-time performance of this method.

Image-based wave-front sensing uses parametric physical models and a measured point spread function (PSF) with nonlinear optimization for calculation [8,9,10]. Artificial neural networks are nonlinear, autonomous, and regulated information-processing systems for learning and generalized nonlinear mapping between inputs and outputs [11,12,13]. At the end of the twentieth century, artificial neural networks were used to measure the optical phase distortion caused by atmospheric disturbance [14,15,16], although the structure of the neural networks in the system was too simple at that time, resulting in poor generalization. Since then, significant research has optimized the neural network structure.

Convolutional neural networks (CNNs) [17,18,19] use convolutional kernels and down-sampling to perform machine-learning tasks on images, which reduces the dimensionality of large data volumes and retains image features. In another approach, Paine et al. used Inception V3 for phase detection [20], but the PSFs used in the training set were set at the focal point with less effective information, which led to fitting error. Nishizaki proposed the generalized wave-front sensing framework [21] to estimate a single intensity distribution and showed that the proper preprocessing of a single intensity image improves the calculation and fitting accuracy.

Compared to the output of CNNs is the class label of the entire image, UNet is a pixel-level classification that outputs the pixel category and is often used on biomedical images [22,23]. Given that image data in this task are often small, Ciresan et al. trained a CNN using a sliding window to provide the pixel and its surroundings as input to predict the class label of each pixel. UNets are also widely used in the wave-front detection field [24,25,26], but the network structure is complex and too slow for real-time tasks.

To summarize, conventional supervised learning neural networks have good detection accuracy, but the many labeled data samples for training greatly increase application costs and difficulty. The essence of this approach is to map PSF image pixels to Zernike coefficients. UNet uses data enhancement to solve the problem of few training images, but it cannot extract the control parameters required for phase correction, such as Zernike coefficients.

Unsupervised learning uses pretexts to mine its own supervised information from large-scale unsupervised data and learns valuable representations for downstream tasks. In another approach, Wang proposed phase imaging with an untrained neural network [27] to iterate over a single damaged image, learn the prior information on the image, and restore the image. However, when the important features of an image are damaged, the image is not easily repaired. Because training in advance is not possible, every time the image is processed, it must be iterated many times to achieve an ideal result. Emrah Bostan proposed using an untrained neural network to restore the phase without a label [28]. Poor generalization is the same problem as above, because offline processing does not require high real-time performance. However, high real-time performance is important in the measurement of wave-front aberration in free-space optical communication. Ramos [29] proposed the use of blind convolution to restore multiple targets, but the network structure is complex, the calculation time is long, and the accuracy is poor. Tobias Liaudat proposed a paradigm shift in the data-driven modeling of the instrumental response field of telescopes by adding a differentiable optical forward model into the modeling framework [30]. This approach is also an unsupervised training based on optical modeling, but the aim is to simplify the building of the instrumental response model.

To solve these problems, we propose a phase retrieval method based on self-supervised physical deep learning (PDL), which maps the input data to hidden vectors through a deep neural network and combines the optical imaging physical model to fit the output. In addition, a retrieval model is established according to the mapping relationship between input and output. Compared with traditional CNNs, this method does not rely on labels, but only establishes the relationship between the collected light spot samples through the intrinsic characteristics of the PSF. Compared with UNet, the Zernike coefficient is more quickly extracted, giving this method better real-time performance. This method also avoids the technical bottleneck of traditional supervised learning that requires many label samples, which facilitates practical applications. Compared with the existing self-supervised learning network, the reasoning for new information requires only one step and the accuracy is higher, but the method we proposed can only be applied to the detection of point objects.

2. Method

As shown in Figure 2, PDL includes three parts: network encoding, a hidden layer vector, and network decoding.

2.1. Encoding Part

The encoder maps the input samples to the hidden layer vector. We tested the ImageNet classification CNN and chose MobileNet [31], which is more suitable for mobile terminals. MobileNet consists of a 3 × 3 standard convolution, followed by a stacked depth-wise separable convolution, some of which involves down-sampling through strides = 2. Average pooling is used to change the feature to 1 × 1, and a fully connected layer is added according to the predicted category size. The batch size is 32. Since we are not dealing with a classification problem, we removed the last SoftMax layer. The core consists of a decomposable depth-wise separable convolution, which not only reduces the computational complexity of the model, but also reduces the size of the model.

For real mobile application scenarios that require low latency, networks, such as MobileNet, are the focus of continuous research. Depth-wise separable convolution is a factorized convolution operation that can be decomposed into two smaller operations: depth-wise convolution and pointwise convolution. Depth-wise convolution uses a different convolution kernel for each input channel and pointwise convolution uses a 1 × 1 convolution kernel. The former uses depth-wise convolution to separately convolve different input channels and then uses pointwise convolution to combine the above outputs.

2.2. Decoding Part

The model of optical imaging system is shown in Figure 3 and it is the decoding part in the network. The object plane is regarded as the weighted superposition of primitive objects according to the imaging principle of convolution. For the image plane, this is the coherent superposition of the diffraction image produced by the corresponding object point [32]. The relationship between the image plane and the object plane can be expressed as the following convolution:

$$g_i(x,y)=g_o(x_0,y_0)\ast h(x,y),$$

(1)

where $g_i(x,y)$ is the complex optical amplitude distribution of the image plane; $g_o(x_0,y_0)$ is the complex optical amplitude of the object plane under ideal conditions; x and y are in the image plane and $x_0$ and $y_0$ are in the object plane; ∗ is the convolution operator; and h is the impulse response function of the system, also called the PSF, which is used to describe the two-dimensional distribution of light on the focal plane from a point-source imaging system. For a point target, this is equivalent to a linear space-invariant incoherent imaging system. The object and image plane functions satisfy the linear transformation of intensity:

$$I_i(x,y)={\displaystyle \iint I_o(u,v)}h(x-u,v-y)dudv,$$

(2)

where $I_i$ is the intensity distribution system at the image plane; $I_o$ is the ideal intensity distribution system at the object plane; and u and v are the coordinates of the pupil plane. In the case of point targets, the object–image relation of the system and the PSF in the case of the diffraction limitation is expressed as

$$g_i(x,y)=h(x,y)=\left|\mathrm{FFT}\right\{P(u,v)\}{|}^2,$$

(3)

where FFT is the fast Fourier transform operator and P is the generalized pupil function. In the case of aberration, the pupil function on the image plane at the exit-pupil position is expressed as

$$P(u,v)=O(u,v)e^{2\pi i\varphi (u,v)},$$

(4)

where $O(u,v)$ is the aperture function; the value inside the aperture is unity; the value outside the aperture is zero; and $\varphi (u,v)$ is the phase-distribution function at the pupil plane. The Zernike mode method makes use of a set of polynomials that reflect the orthogonal properties of a circular region. Therefore, the Zernike mode method is usually used as an orthogonal basis for wave-front reconstruction. Wave-front phase functions (x, y) in the circular domain can be expanded into Zernike polynomials as follows:

$$\phi (x,y)=\sum _{i=1}^{i=\infty }{a}_i{z}_i(x,y),$$

(5)

where $a_i$ is the coefficient of the pattern of term i and $z_i$ is the Zernike pattern corresponding to term i. The matrix distribution of the Zernike coefficient is composed of the all-order polynomials and the corresponding coefficients:

$$Z_i(r,\theta )=\{\begin{array}{l}{N}_n^m{R}_n^{\left|m\right|}(r)\cos (m\theta ), m\ge 0\\ -N_n^m{R}_n^{\left|m\right|}(r)\sin (m\theta ), m\le 0,\end{array}$$

(6)

where $0\le r\le 1$, $0\le \theta \le 2\pi$; n is a non-negative integer; the step length m = 2; the outputs are between −n and n; and $R_n^{\left|m\right|}$ is the radial function for the polynomials. The hidden layer vector we used in the network is a Zernike coefficient. The aberration can be calculated with the Zernike coefficients and polynomials. Then, the PSF can be obtained by the optical imaging system for training the model.

2.3. Loss Function

We strived to make the inputs and outputs of the network consistent by using the correlation coefficient between the input and output PSF as a loss function. The formula for calculating the correlation coefficients is

$$r=\frac{{\displaystyle \sum _m{\displaystyle \sum _n(A_{mn}-\overline{A})}}(B_{mn}-\overline{B})}{\sqrt{\left({{\displaystyle \sum _m{\displaystyle \sum _n(A_{mn}-\overline{A})}}}^2\right)\left({{\displaystyle \sum _m{\displaystyle \sum _n(B_{mn}-\overline{B})}}}^2\right)}},$$

(7)

where A and B are the pixel values of the two PSFs, and $\overline{A}$ and $\overline{B}$ are the mean values of A and B, respectively. When the model training is completed, the model interferes without the decoding part. The output of the network is a Zernike coefficient and can be used to correct aberrations. Therefore, there is no difference in the resource consumption and time consumption, compared to the supervised networks.

3. Simulation Demonstration

3.1. Simulation Demonstration of 3–20 Orders of Zernike Polynomials

The simulation data set is generated by using the Zernike pattern method. The piston error (The Zernike polynomial $Z_0$) is ignored because it is the translation of the beam phase and does not affect the beam quality. The tip and tilt terms ($Z_1$ and $Z_2$), which can be quickly estimated by using the centroid algorithm or other registration methods, are also ignored. The PSF simulation model in this article sets the pupil size to 8.6 mm, the focal length of the lens to 150 mm, and the defocus to 8 mm. The wavelength of the optical source is 532 nm, and to simulate the actual environment, Gaussian white noise with a mean value of 0 and a variance of 0.01 is added to the PSF. In the case of different mode coefficients, 22,000 distorted spot images are generated with D/r0 = 20, 20,000 out-of-focus PSFs are input into the self-supervised learning neural network, and PDL is applied according to the correlation coefficient between the input and output PSFs. A total of 2000 PSFs were tested, and the Zernike coefficients were inferred on the basis of the mapping relationship.

Figure 4 and Figure 5 shows the test results, and the PSF images are normalized. The image size is 256 × 256 and the pixel size is 5.5 µm. Moreover,

Table 1. Comparison of accuracies (RMSE) of equivalent wave-fronts from the Zernike coefficients estimated in the experiments between PDL and MobileNet (3–20 orders).

Zernike Order	RMSE of Testing Set	RMSE of MobileNet	RMSE of PDL
3–20	0.8284λ	0.0447λ	0.0648λ

lists the root mean square errors (RMSEs) obtained from testing and inference. To reflect the retrieval effect of PDL more clearly, we compared the results with those of the MobileNet supervised learning method. For coefficients 3–20, the RMSE of the original distortion wave-front is 0.8284λ. The correlation coefficient of the test set was 97%, and the RMSE of the one-time calibration is 0.0648λ. In contrast, the supervised corrected wave-plane RMSE is 0.0447λ. Although the detection accuracy is slightly lower compared with supervised learning, no a priori knowledge of the Zernike coefficients is required, which facilitates practical applications.

To test the generalization ability of the network, 3000 distorted spot images are generated with D/r0 interval 5 from 5 to 30. The network trained before is used to inference the 3000 distorted spot images.

Table 2. Comparison of the accuracies (RMSE) of the equivalent wave-fronts from the Zernike coefficients estimated in the experiments between PDL and MobileNet with different D/r0 (3–20 order).

D/r0	RMSE of Testing Set	RMSE of MobileNet	RMSE of PDL
30	1.1991λ	0.1176λ	0.1469λ
25	0.9962λ	0.0621λ	0.0762λ
20	0.8606λ	0.0462λ	0.0634λ
15	0.6651λ	0.0396λ	0.0519λ
10	0.4641λ	0.0276λ	0.0352λ
5	0.2615λ	0.0232λ	0.0256λ

shows the results. The results show that the generalization ability of the network is good.

3.2. Simulation Demonstration of 3–64 Orders of Zernike Polynomials

The influence of higher-order wave-front aberrations on the network performance is considered as follows: the 3–64 Zernike order is used to train the network; the RMSE of the original distortion wave-front is 0.8852λ; and the corrected wave front RMSE of PDL is 0.1274λ and the RMSE of MobileNet is 0.1069λ. Figure 6,

Table 3. Comparison of accuracies (RMSE) of equivalent wave-fronts from the Zernike coefficients estimated in the experiments between PDL and MobileNet (3–64 orders).

Zernike Order	RMSE of Testing Set	RMSE of MobileNet	RMSE of PDL
3–64	0.8852λ	0.1069λ	0.1274λ

and

Table 4. Comparison of the accuracies (RMSE) of the equivalent wave-fronts from the Zernike coefficients estimated in the experiments between PDL and MobileNet with different D/r0 (3–64 orders).

D/r0	RMSE of Testing Set	RMSE of MobileNet	RMSE of PDL
30	1.2021λ	0.3052λ	0.3487λ
25	1.0438λ	0.1887λ	0.2136λ
20	0.8784λ	0.1084λ	0.1255λ
15	0.676λ	0.0857λ	0.0984λ
10	0.4864λ	0.0671λ	0.0696λ
5	0.2715λ	0.0610λ	0.0611λ

shows the test results, and the PSF images are normalized. The image size is 256 × 256 and the pixel size is 5.5 µm. The results show that the performance of the self-supervised network is also close to the supervised network in the high-order case.

4. Experimental Demonstration

Figure 7 shows the optical platform of the phase retrieval experiment, which used a liquid crystal spatial light modulator (LCSLM) to simulate the PSF. A 532 nm laser served as a light source. A collimator expanded the beam, which was normal to the incident aperture (A), polarizer (P), and the LCSLM. The size of the aperture was limited to 8.6 mm to ensure that the incident light was evenly distributed over the center of the LCSLM. The BS changed the direction of the light beam modulated by the LCSLM. Phase distortion was simulated by loading different grayscale phase screens onto the LCSLM. Due to a small gap between the pixels, higher diffraction orders and a high-energy zero-order spot exist at the output. An X tilt was added to displace the spot from the central zero-order spot. The modulated light was imaged through the lens (L) before impinging on the CCD. The effective resolution of the CCD acquisition window was 256 × 256, which was compressed to 128 × 128 for training and testing. To ensure the detection accuracy of the neural network and the range of the CCD field of view, we set the CCD defocus to 8 mm. To ensure clear imaging with the CCD, a suitable attenuator was positioned in front of the CCD. The model pluto-2-NIR-011LCSLM used in this experiment was produced by HoloEye and has a pixel resolution of 1920 × 1080, with 8 µm pixels. As with the simulations, we trained with 20,000 experimental images and tested with 2000 PSFs. The overall tilt of the experiment used the Zernike coefficients of order 1–20 to produce a test set correlation coefficient of 95%. Figure 8 shows the calculation results obtained by inputting the collected light spot into the experimental training model. The image size is 256 × 256 and the pixel size is 5.5 µm. The light spots are similarly shaped. A Jetson TX2 from NVIDIA was used for inference calculations. After using tensor optimization on the TX2, the single-inference time is 3.72 ms, which is adequate for real-time performance.

5. Discussion

The experimental results are the same as we expected. The detection accuracy of the self-supervised deep learning model proposed in this paper is slightly inferior to the supervised deep learning model. However, the self-supervised model reduces the difficulty of sample acquisition, so it is easier to apply in practice. The supervised model is related to the method proposed in reference [22]. We made some adjustments to the method, according to the new application scenario. Due to the change of atmosphere turbulence intensity, the result of the supervised model is different from those in reference [22].

Compared to other self-supervised models based on the optical imaging system mentioned in the Introduction, we proposed a new network model. Additionally, a new loss function, the correlation coefficient, was used to train the network with batch-size in practice. Training with batch-size presents our network with a high detection accuracy and good generalization ability.

However, the training of the method needs to know the optical parameters, such as optical aperture, the focal length of lens, etc. In most application scenarios, these parameters are known.

6. Conclusions

Herein, we proposed a self-supervised model that retrieves the phase without labels. The model is generalizable, can be used with different networks, and can solve different problems. In addition, for reasoning, it does not affect the time consumption of the network. The model is suitable for practical applications of deep learning for phase detection and other optical problems. Although the detection accuracy is slightly less than that provided by supervised learning, no a priori knowledge of the Zernike coefficients is required, which facilitates practical applications.

In future work, we will validate the effectiveness of PDL in satellite-to-ground laser communication. Additionally, we will use different loss functions and train them to estimate other terms. PDL will be adapted to different wavelengths, apertures, and focal lengths, in order to directly measure spots of different wavelengths, apertures, and focal lengths without retraining or changing the models.