HazeMatching: Dehazing Light Microscopy Images with Guided Conditional Flow Matching

TL;DR

HazeMatching is proposed, a novel iterative method for dehazing light microscopy images, which effectively balances the fidelity of the dehazing results and the realism of individual predictions by adapting the conditional flow matching framework by guiding the generative process with a hazy observation in the conditional velocity field.

Abstract

Fluorescence microscopy is a major driver of scientific progress in the life sciences. Although high-end confocal microscopes are capable of filtering out-of-focus light, cheaper and more accessible microscopy modalities, such as widefield microscopy, can not, which consequently leads to hazy image data. Computational dehazing is trying to combine the best of both worlds, leading to cheap microscopy but crisp-looking images. The perception-distortion trade-off tells us that we can optimize either for data fidelity, e.g. low MSE or high PSNR, or for data realism, measured by perceptual metrics such as LPIPS or FID. Existing methods either prioritize fidelity at the expense of realism, or produce perceptually convincing results that lack quantitative accuracy. In this work, we propose HazeMatching, a novel iterative method for dehazing light microscopy images, which effectively balances these objectives. Our goal was to find a balanced trade-off between the fidelity of the dehazing results and the realism of individual predictions (samples). We achieve this by adapting the conditional flow matching framework by guiding the generative process with a hazy observation in the conditional velocity field. We evaluate HazeMatching on 5 datasets, covering both synthetic and real data, assessing both distortion and perceptual quality. Our method is compared against 11 baselines, achieving a consistent balance between fidelity and realism on average. Additionally, with calibration analysis, we show that HazeMatching produces well-calibrated predictions. Note that our method does not need an explicit degradation operator to exist, making it easily applicable on real microscopy data. All data used for training and evaluation and our code will be publicly available under a permissive license.

Figures (32)

• Figure 1: HazeMatching proposes a way to utilize Conditional Flow Matching to remove out-of-focus light (haze) from fluorescence microscopy data.
• Figure 2: Data fidelity vs. realism. Each row corresponds to one dataset, i.e.Zebrafish (top), Organoids1 (middle), and Organoids2 (bottom). For each dataset, we show PSNR vs. LPIPS (left) and PSNR vs. FID (center), capturing the trade-off between pixel-level fidelity and perceptual quality. Our goal is to find a method that leads to high fidelity (high PSNR) while also leading to realistic looking predictions (low LPIPS/FID). HazeMatching results are highlighted with an additional red circle around its blue marker. Results tables (right) further quantify all results. Methods displayed in red are point-predicting baselines, while blue methods are generative posterior models (see main text). Note that HazeMatching consistently places itself close to the bottom left corner of all plots, i.e. achieving the desired balanced performance across metrics and datasets. Additional results can be found in Supplementary Section \ref{['sup:results']}.
• Figure 3: HazeMatching models are well calibrated. RMSE vs. predicted RMV is shown for Zebrafish, Organoids1, and Organoids2 (left to right). The dashed line is $y = x$. Blue and purple circles show native calibration and calibrated results obtained by learning an additional calibration factor and offset, respectively. Shaded areas denoting standard error. Learned calibration parameters (scaling and offset) are shown below each plot. Additional results for Neuron and Microtubule are shown in Section \ref{['sup:results']}.
• Figure 4: Qualitative results. We show representative results for three datasets. For each dataset we show (a) the full input ($1024 \times 1024$) and a selected crop ($128 \times 128$ yellow box), (b) the selected hazy crop, (c) the non-hazy ground truth crop, (d–o) predictions by all baseline methods (see Section \ref{['subsec:baselines']}), and (p) results obtained with HazeMatching. Results with red borders are predictions by point-predictors, while methods with blue borders are results by generative posterior models (see also Figure \ref{['fig:plot_main']} and main text). We present additional results in the Supplementary Section \ref{['sup:results']}.
• Figure S1: Overview of our training and inference. The figure illustrates the different distributions involved in our approach. The top panel shows how samples are drawn from the base noisy distribution $p_0(\mathbf{x})$, the target non-hazy image distribution $p_{M_1}(\mathbf{x}|s^i)$, and the hazy source image distribution $p_{M_0}(\mathbf{x}|s^i)$, and the computing of $\mathbf{x}_t$ with sampling time as $t \sim U[0,1]$. The lower left panel presents the training scheme, where the model is trained to learn a mapping between these distributions using conditional flow matching with inputs $concat[\mathbf{x}_t,\mathbf{x}_{M_0}]$ and $t$. The lower right panel depicts the inference process, where given a degraded observation $\mathbf{x}_{M_0}$, the trained model takes as inputs $concat[\mathbf{x}_0,\mathbf{x}_{M_0}]$ ($\mathbf{x}_0 \sim p_0(\mathbf{x})$) and $t$ an predicts the time dependent velocity field per pixel ($d\mathbf{x}_t/dt$) which is integrated using an (Euler) ODE solver iteratively using our proposed dehazing function $D$ to produce one prediction.