Physics-enhanced machine learning for virtual fluorescence microscopy

In this work, we collected and processed data across two different experimental setups. Both setups used an off-the-shelf LED array (Adafruit product ID 607), which provided the capability to turn on specific LED array patterns at different colors (e.g., red, green and blue) for multi-angle and multi-spectral specimen illumination (Figure 1A). In each setup, the LED array was placed sufficiently far from the sample plane such that the illumination from each LED could be modelled as a plane wave at the sample, propagating at a unique angle corresponding to its relative position. This type of illumination has been previously used for FPM [32, 34], as well as phase contrast [26, 27] and super-resolved 3D imaging [35, 36, 37], among other enhancement techniques. As noted above, illumination from a wide variety of angles and colors provides diverse information about biological specimens, which are primarily transparent, that is not available under normal illumination (see Figure 2). However, it is not directly obvious which type of illumination is best-suited for mapping non-fluorescent image data into fluorescence stained image data. The optimal patterns will likely depend upon which sample features are fluorescently labeled, as well as properties of the specimen’s complex index of refraction. In this work, we consider two separate fluorescent labeled sample categories - one in which the cell nucleus is labeled, and a second in which the cell membrane is labeled - to verify this hypothesis. In addition, the specific inference task may also impact the optimal illumination distributions (i.e., segmentation versus image formation) in a non-obvious way. To solve for such task specific illumination patterns, we propose to use a modified DNN, which includes a ”physical layer” [24] for joint hardware optimization.

In this type of approach, the supervised machine learning network itself jointly determines the optimal distribution of LEDs’ brightness and colors to use for sample illumination, and to perform the subsequent image-to-image inference task. As we describe below, the process of LED illumination can be modeled by a single physical layer, which we prepend to the front of a DNN U-Net architecture for improved image-to-image inference. After the network training is complete, the optimized weights within the physical layer (in this case, representing the angular and spectral distribution of the sample illumination) inform us of a better optical design for each specific inference task at hand. Figure 3 shows how this process is realized for both training and inference.

To mathematically model image formation, we express the $j$th specimen of interest as the complex function $o_{j}(r,\lambda)$. Given $N$ LEDs within the array, we can denote the amplitude of the plane wave emitted by each colored LED emitting at center wavelength $\lambda$ as a weight $\sqrt{w_{n}(\lambda)}$ for $n=1,\dots,N$. Each LED in the array illuminates the sample with a coherent plane wave from a unique angle $\theta_{n}$. Since, the LEDs are mutually incoherent with respect to each other, the image formed by an illumination composed of multiple LEDs is equivalent to the incoherent sum of images obtained by illuminating the member LEDs individually. Assuming a thin specimen, for example, we may write the $j$th detected quasi-monochromatic image $I_{j}^{\prime}$ as the incoherent sum,

where, $\sqrt{w_{n}(\lambda)}e^{i\mathbf{k}_{n}\cdot\mathbf{r}}$ describes the plane wave generated by the $n$th LED with intensity $w_{n}(\lambda)$ at wavelength $\lambda$ across the sample plane with coordinate $\mathbf{r}$, $\mathbf{k}_{n}=\frac{2\pi}{\lambda}\sin{\theta_{n}}$ denotes the plane wave transverse wavevector with respect to the optical axis, and $h(r)$ denotes the microscope’s coherent point-spread function and describes the imaging system blur.

As the LED brightness $w_{n}(\lambda)$ is a scalar quantity, we can factor it out of the summation in Eq. 1. Furthermore, if we denote the image of the $j$th sample formed when it is illuminated by the $n$th LED at a fixed brightness and wavelength as $I_{j,n}(\lambda)=|o_{j}(\lambda)\cdot e^{i\mathbf{k}_{n}\cdot\mathbf{r}}\star h|^{2}$, then we arrive at a simple linear model for image formation for the $j$th sample:

The detected image under illumination from a particular LED pattern, where each LED has a brightness $w_{n}(\lambda)$, is equal to the weighted sum of images captured by turning on each LED individually. The detected image $I_{j}^{\prime}$ is then entered into the neural network for processing. Equation 2 represents our physical layer for microscope illumination, and the fully differentiable optimization pipeline is shown in Figure 3.

To find the set of weights $w_{n}(\lambda)$ which best parameterize the transform in Eq. 2 for the subsequent inference task, we must have access to all $N$ uniquely illuminated images $I_{j,n}$ during network training. However, once training is complete, the set of optimized weights $\{w_{n}\}$ are then mapped onto the physical LED matrix to allow for acquisition of a single optimally illuminated image, which is then processed by the remainder of the DNN for the fluorescent image inference task.

We used a consistent U-Net architecture that we prepended with a physical layer to model Equation. 2. The exact architecture and layer configurations, used in all the reported experiments, is detailed in Appendix B. The weights within the physical layer were unconstrained, meaning they could take on both positive and negative values. To experimentally realize an unconstrained set of illumination weights $\mathbf{w}$, we captured two images (instead of one) - a first with the positive set of weights, and a second with the negative set, before subtracting the second image from the first for the final result.

Although the experimentally captured measurements inherently contain Poisson noise, the digital simulation of multi-LED images used during training will have an artifactually inflated SNR (e.g., averaging $N$ images produces a $\sqrt{N}$ improvement in SNR). To compensate for this we introduced a Noise Layer which adds dynamically generated Gaussian random noise to the data after the first physical layer:

Additive noise is modelled on a per-pixel level with Equation 3, using a hyperparameter $k$ to control the scale of the noise in proportion to the pixel intensity. Note that the variance of the random noise is proportional to the image pixel intensity itself, and thus is consistent with a Poisson noise model.

Finally, an $L1$ penalty within the network cost function (absolute deviation from zero) was applied to the physical layer weights $w_{n}$. This $L1$ penalty term is given a small weight proportional to the magnitude of the gradients to drive weights to zero, if and only if they are not significantly contributing to the task performance (see details in Appendix B). The addition of this type of penalty reduces variance across random seeds and aids in interpretation of the resulting LED patterns.

To understand the impact of the optimized illumination pattern $\mathbf{w}$, we varied the task difficulty by modifying the fidelity of the target label. The 7 bit ground truth fluorescence image was converted to 7 different images with different bit depths, ranging from 1 bit to 7 bits of dynamic range. Here, we hypothesize that reconstruction of a 1-bit fluorescence image, which is effectively a segmentation mask, is easier than accurately recovering a 7 bit fluorescence image. It should be noted that the input bright-field images (the training data) remains same for all tested cases. This allowed the same data to be used to train a neural network to perform both binary image segmentation task (1 bit precision), and to separately train another network to perform a full fluorescent image inference task (7 bit precision), and to also consider all tasks that span these two common efforts. Figure 4 shows how by varying the precision of the desired network output, it is possible to vary the difficulty of the subsequent image inference task, from producing a prediction that contains $1$ bits to $7$ bits per pixel.

To achieve a balance of pixel values across each discretized histogram within the label images, we employed a global k-means quantization strategy, where $k=2^{bits}$. To find the mean values a naive k-means algorithm was applied on a flattened version of the entire dataset, treating each pixel value as an independent value. The initial mean values were uniformly distributed across the space between zero and one: $m_{i}=\{\frac{i}{k}|1\leq i\leq k\}$. Mean values were iteratively adjusted until either all values had converged (with a threshold of: $1\mathrm{e}{-05}$) or $20$ iterations passed.

The inference results were only rounded to meet the appropriate precision of the target bit level, with no further post processing (k-means or otherwise).

The data used in the first experiment was originally acquired by a FPM microscope [39] and contains 90% confluent HeLa cells, stained with a fluorescent nuclear label (DAPI). The employed microscope included a two-lens arrangement with an f = 200 mm tube lens (ITL200, Thorlabs) and an f = 50 mm Nikon lens (f/1.8D AF Nikkor). This two-lens setup has a collection numerical aperture (NA) of 0.085 with 3.87x magnification. The sample was placed at the front focal plane of the f = 50 mm lens and a CCD detector (pixel size 5.5 $\mu$m, Prosilica GX6600) captured images of the sample under LED matrix illumination (one LED at a time, as noted above). The LED array was placed 80 mm behind the sample with 32x32 individually addressable elements (pitch size 4 mm), of which only the inner 15x15 were utilized. Full-FOV images from all illumination angles and colors were divided into 442 unique $256\times 256\times 675$ datacubes. The set of datacubes was randomly split into training, testing, and validation sets containing 356, 48, and 48 images respectively.

For the second set of experiments, we captured images of Pan 16 pancreatic cancer cells stained using CellMask Green plasma membrane stain (C37608). Similarly to the dataset described in Section IV-B, an identical LED array was placed 60mm beneath the biological sample and of which the inner 15x15 grid was used. An Olympus PlanN 0.25NA 10x objective lens was used in conjunction with a Basler Ace (acA4024-29um) CMOS sensor. To capture the matching fluorescence data a set of Thorlabs filters (MDF-GFP - GFP Excitation, Emission, and Dichroic Filters) were inserted into the optical setup. Figure 1 shows a photo of the setup used. The data was split into test/train/validation sets such that data from each sample only existed in one of the sets, providing a robust means of gauging generalization and overfitting. When split into $256\times 256$ images there were 820/108/108 samples in the train, validation and test sets respectively.

Across both tasks and all bit depth configurations, we found that the optimized illumination patterns determined by our physical layer outperformed all other illumination arrangements (coherent and incoherent bright-field, dark-field, random and phase contrast). Figure 5 plots the mean squared error (MSE) of each inference task as a function of inference bit depth (i.e., inference difficulty), cell type, and physical layer parameterization. This main result shows that jointly optimized microscope hardware can lead to superior image inference performance across a wide range of tasks - ranging from image segmentation to virtual fluorescence imaging.

Although MSE provides an accurate measure of the task performance, it is less suited to understand the perceptual quality of the results. We thus also computed the structural similarity index (SSIM) [40] across the test set for both tasks (results shown in Figure 5). Here we can see that the reduction in mean squared error brought by the optimized illumination carries onto a higher perceptual quality.

By varying the bit depth of the output label, we were able to test our approach with seven different configurations. We first observe that the gap in performance (both MSE and SSIM – Figure 5) between the inference results created by the optimized LED pattern (”learned”), and the alternative standard LED illumination patterns, is quite consistent. Given that the pixel values within all of the output labels are normalized to a $[0,1]$ range, this result is somewhat consistent with our expectations. The increase in difficulty of the task is partially offset through the decrease in the ”cost” of being wrong.

To get a better understanding of the interaction of task difficulty and performance we examined the relative performance of the ”learned” method compared to the best of the ”standard” methods. To do this we constructed a relative performance metric, $R_{MSE}$, which is the quotient of the minimum MSE among the standard patterns’ MSE and the learned configurations MSE, for each bit depth $K$. See equations 4 and 5.

Within Figure 6 we show that the relative performance of the learned LED configuration improves with difficulty when compared to the best standard configuration. This illustrates how the joint optimization process provides a consistent performance improvement that becomes more important as the difficulty of the task increases.

In addition to improved performance, the DNN-optimized LED illumination patterns also offer several interesting physical insights. To aid in interpretation of the patterns, each LED was normalized by it’s exposure (as shown in Figure 7), as using different exposures per LED enabled us to fully utilize the dynamic range of our image sensor. Figure 8 displays the learned patterns as a function of inference task (1-bit to 7-bit) on a per-color-channel basis for both tested specimens.

These illumination patterns highlight the importance of two primary features for transmitting visible light information that is most relevant to a particular fluorescent specimen and stain. First, producing phase contrast by illuminating differently within the bright-field (the inner 3x3 LEDs for the HeLa task, and the inner 5x5 LEDs for the PAN task) and dark-field appears important across all tasks. Second, providing color contrast (e.g., green for the positive image and red for the negative image) also appears important, most notably for the HeLa nuclei prediction task. The pattern differences on a per-task (HeLa vs. PAN) can be attributed both to the differences between setups (notably that the PAN task setup included more bright-field LEDs) as well as the differences between samples. It should be noted that for the PAN task the LED patterns appear shifted, this is due to an indexing error which occurred during the data-preprocessing steps and does not hold any physical significance.

We are also able to compare the distribution of optimized LED patterns as a function of task difficulty (i.e., output label bit depth). Although the importance of establishing a learned pattern to improve network performance is clear (section V-B), the trend among the spatial patterns themselves is less so. We found that, in general, the lower bit-depth patterns generated images which put more emphasis on higher frequency information, and vice versa. Figure 9 shows how the average resolved spatial frequency power decreases with increased bit depth, although this trend was noisy and several bit-depths were exceptions. The values were determined by finding the first moment of each illuminated images’ spatial frequency power representation along it’s radius, more details are provided in Appendix A. We hope to examine more possible trends in future work.

Although there is only a small amount of variation present in the overall system performance, the LED patterns themselves do vary across random seeds. Figure 12 shows examples and statistics of this variance across both tasks. We hypothesize that there are two factors driving this variance. First, images acquired from high-angle illumination (particularly dark-field) had lower intensity values on average and therefore contributed less to the synthesized image. This translates to a higher degree of allowed variance in the optimized result, since as the brightness value of an LED illuminating from a high angle doesn’t impact the image as much as the more centrally located LEDs, and hence the gradients driving its value will be fairly weak. The use of the L1 norm penalty reduced this variance across both tasks.

Second, there is clearly a coupling between the solution of the neural network and the parameterization of the LED array. The joint optimization procedure followed in this work uses this coupling to get higher performance results, and shows that higher performance is consistently obtained. But, as an adverse effect, it also decreases our ability to draw conclusions from the end parameterization of the LED array. This is because when a deep neural network is being trained, it can arrive at one of many local minima (finding the global minima is highly unlikely). Due to the network arriving and terminating at a local minima, the LED parameterization is such that the lighting is optimized for that specific local minima, instead of what might be globally optimal. Within section V-D, we comment on one method of alleviating the impact of arriving at local optima through the use of targeted regularization.

Many of the trials presented here were run both with and without the addition of regularization on the weights representing the illumination pattern. Upon comparison, we found that the addition of this regularization did not impact our target performance statistic (MSE on the test set), however it did have a large effect on the LED patterns themselves. Most notably, when the LED weights are left unconstrained, the network tends towards an illumination pattern configuration using many LEDs at high dark-field angles. When constrained to use fewer LEDs, the optimized pattern instead utilizes bright-field and dark-field LEDs at lower angles. Although information contributed from higher angles may still be useful overall, the L1 norm prioritizes illumination angles that consistently contribute large amounts of information. In this respect a standard L1 norm may not be ideal for this kind of optimization, investigation into a more suitable regularization strategy will be investigated in future work.

In Figure 10, we show side-by-side comparisons for both cases across a representative set of bit depths. This comparison clearly shows that not only do the regularized patterns contain less within-trial spatial variance, but also they are potentially more spatially interpretable. We postulate that this is because the added regularization is having the desired effect of forcing the optimization to find a parameterization that is robust to perturbations (caused by the addition of noise) while simultaneously having the minimum overall energy (caused by the L1 penalty). While performance in the simulated environment remains the same (with and without regularization), we hypothesize that these spatially smoother LED patterns are more robust to changes in the eventual physical setup. However, we leave testing this hypothesis for future work.