Bayesian uncertainty analysis for underwater 3D reconstruction with neural radiance fields

NeRFs represents a novel approach leveraging deep learning techniques to reconstruct 3D scenes. However, because of the effects of backscatter and attenuation lead to the complex light propagation in underwater, makes accurate 3D scene reconstruction by traditional NeRFs difficult. To address this issue, researchers have proposed various improvements methods. WaterNeRF [15] uses NeRFs to additionally learn water column parameters, and combines them with light transport model [16] for color correction. By matching the corrected color distribution to the reference image distribution and using the Sinkhorn loss function [17] for training, it enables to produce the consistent color-corrected results from varying viewpoints during the inference stage. U2NeRF [18] extends UPIFM [19]. It modifies the generalizable NeRF transformer (GNT) [20] to predict scene radiance, direct and back scatter transmission maps, and employs variational autoencoder (VAE) to predict the global background light component. By integrating these four components with the image formation model [19], it reconstructs the original underwater image. WaterHE-NeRF [21] designs a novel water-ray tracing field based on the Retinex model [22] to learn color, density, and illuminance attenuation. By controlling the intensity of illuminance attenuation, it can generate both degraded and restored multi-view images simultaneously. SeaThru-NeRF [5] based on the SeaThru [23], replaces the traditional single opaque object density with the sum of the object density and the medium density. The final pixel color consists of the object radiance attenuated by the medium and the accumulated medium radiance along the light. It can not only generate realistic images from new viewpoints but also utilize the learned object and medium parameters to remove the effects of the scattering medium, reconstructing the true appearance and color of the scene as if the image were taken in the air.

In deep learning, uncertainty refers to the confidence level in the model’s prediction results. Effective estimation and management of uncertainty can not only improve the model’s predictive performance but also enhance its robustness to anomalies [24, 25]. Variational inference [26, 27, 28] approximates the intractable posterior distribution by introducing a prespecified family of variational distribution and minimizes the evidence lower bound (ELBO) to reduce the discrepancy between them. While this approach provides comprehensive uncertainty estimates, it has high computational complexity. Deep ensemble strategies [29, 30] use multiple independently trained models combined with diverse data subsets, model structures, or initialisation methods to reduce the variance of a single model, thus improving generalisation and robustness. However, training multiple models requires more computational resources and storage space, and may introduce additional prediction delays. MC-dropout-based methods [31, 32, 33, 34] simulate the effect of Bayesian neural networks by performing multiple forward propagation in the inference stage, and statistical properties of different predictions (such as mean and variance) are used to estimate the uncertainty of the model. However, multiple forward propagation leads to high computational cost. Chen et al. [35, 36, 37, 38] and Qu et al. [39] combine model order reduction techniques with deep learning to expedite the sampling process for uncertainty quantification. Compared with aforementioned uncertainty quantification methods, the Laplace approximation [40, 41] does not require retraining the model or modifying the training process. Instead, it approximates the posterior distribution with a Gaussian distribution around the mode of the posterior, thereby effectively reducing computational costs. Specifically, by performing a Taylor expansion of the log-posterior and retaining terms up to the second order at the mode, a Gaussian distribution is obtained to approximate the posterior. The mean of this Gaussian distribution is the mode, and the covariance is the negative inverse of the Hessian matrix at the mode.

In practical applications, NeRFs needs to address uncertainty issues to improve model reliability and robustness. Therefore, researchers have explored various approaches. ActiveNeRF [42] models the radiance values at each location as a Gaussian distribution instead of a single value, enabling NeRFs to provide reasonable high-variance predictions in unobserved regions. S-NeRF [43] utilizes Bayesian variational inference to achieve uncertainty quantification related to outputs in tasks such as novel view synthesis or depth estimation. CF-NeRF [44] employs conditional normalizing flows and latent variable modelling to flexibly learn the radiance field distribution without any prior assumptions. It estimates uncertainty by evaluating the predicted mean and variance during the inference process. Sünderhauf et al. [45] proposes to quantify prediction uncertainty by independently training multiple NeRFs with different initial parameters on the same dataset [46, 47]. In addition to considering the variance in RGB colors, it introduces an epistemic uncertainty term based on termination probability to capture uncertainty in unobserved regions. FG-NeRF [48] decouples NeRFs into deterministic and probabilistic branches, using Flow-GAN [49] to model the probabilistic branch to avoid independence assumptions. Additionally, it employs a patch adversarial training strategy. These designs enable it to achieve more accurate uncertainty estimation in complex scenes. ProbNeRF [50] employs the hamiltonian monte carlo (HMC) [51] method during testing to perform posterior inference on NeRFs’ parameters for given views. It accurately infers the 3D geometry and appearance of objects from a single or several views while quantifying the associated uncertainty. Recursive-NeRF [52] draws inspiration from level of detail (LOD). Each network layer predicts the uncertainty of query points: for points with low uncertainty, results are directly output without passing to deeper layers; for points with high uncertainty, they are passed to the next, more powerful layer for further processing. This strategy significantly improves rendering efficiency while ensuring the quality of view synthesis. Unlike the above methods, we introduce an additional perturbation field to avoid retraining the model or modifying the training process, significantly reducing the computational cost.

Traditional underwater image generation models based on atmospheric models result in significant errors. To address this issue of underwater image color restoration, SeaThru [23] employs a revised optical model that accounts for the attenuation and scattering characteristics of light propagation in water, thereby generating images that are closer to true colors.

The underwater image generation model is described as follows:

where $\mathbf{c}{\in}\{RGB\}$ is color channel, $I_{\mathbf{c}}$ is the image captured by the camera, $D_{\mathbf{c}}$ is the direct signal, and $B_{\mathbf{c}}$ is the backscatter.

Traditional methods assume a uniform attenuation coefficient for light throughout the entire scene, which is a coarse approximation. Actually, the direct signal and backscatter are controlled by different attenuation coefficients, which also depend on factors such as object distance and reflectance. Therefore, SeaThru extends Eq. (1) as follows:

where $J_{\mathbf{c}}$ is the clear scene that should be captured ideally, $B_{\mathbf{c}}^{\infty}$ represents veiling light, $z$ is the distance from the object to the camera, $\beta^{D}$ and $\beta^{B}$ are the attenuation coefficients and backscatter coefficients, respectively. The vectors $\mathbf{v}_{D}$ and $\mathbf{v}_{B}$ represent the dependence of $\beta^{D}$ and $\beta^{B}$ on factors such as distance, reflectance.

SeaThru analyzes multi-view images to estimate the distance of each pixel, thereby inferring the effects of scattering and attenuation. It then corrects color distortions in the image by applying these estimated values.

NeRFs is a novel 3D reconstruction technique that reconstructs high-quality 3D scenes from collections of 2D images and camera poses. It achieves high-resolution rendering and image synthesis by representing scenes as a dense volumetric radiance field function and leveraging volumetric rendering techniques.

Specifically, as shown in the Fig. 2, NeRFs utilizes MLP to learn a mapping function $F$ that maps any given 3D spatial point $\mathbf{x}=(x,y,z)$ and viewing direction $\mathbf{d}=(\theta,\phi)$ to a color value $\mathbf{c}=(r,g,b)$ and a density value $\sigma$:

where $\boldsymbol{\varphi}$ are the learnable parameters of the MLP.

To render the continuous 5D scene structure represented by color $\mathbf{c}$ and density $\sigma$ into a 2D image, NeRFs uses the volumetric rendering equation to integrate the color and density values along a camera ray $\mathbf{r}(t)=\mathbf{o}+t\mathbf{d}$, where $\mathbf{o}$ is the camera center and $t\in\mathbb{R}_{+}$.

The expected color obtained along the ray $\mathbf{r}$ is given by:

Here, the integration is bounded between the near bound $t_{n}$ and the far bound $t_{f}$, and the accumulated transmittance $T(t)$ between them is defined as:

In practice, NeRFs utilizes the quadrature rule [53] to discretize the integration range $[t_{n},t_{f}]$ into $N$ intervals $\{[t_{i},t_{i+1}]\}_{i=1}^{N}$ (where $t_{n}=t_{1}<\cdots<t_{N}=t_{f}$), and assumes that the density $\sigma$ and color $\mathbf{c}$ are constant within each interval. Thus:

and

where $\delta_{i}=t_{i+1}-t_{i}$ is the distance between adjacent sample points.

NeRFs adjusts the parameters of the MLP during training by optimizing the squared distance between the expected color $\hat{C}(\mathbf{r})$ and the ground truth $C^{\mathrm{gt}}(\mathbf{r})$.

Inspired by SeaThru [23], SeaThru-NeRF not only considers the opaque objects in traditional NeRFs but also treats the medium as a semi-transparent entity, introducing separate color and density parameters for both the medium and the objects. This method extends the capability of NeRFs to handle scattering media such as underwater (Fig. 3) through mapping function $F$.

Specifically, it replaces the traditional single opaque object density $\sigma(t)$ with the sum of the object density $\sigma^{\mathrm{obj}}(t)$ and the medium density $\sigma^{\mathrm{med}}(t)$.

where $\mathbf{c}^{\mathrm{obj}}(t)$ and $\mathbf{c}^{\mathrm{med}}(t)$ are the color of the object and medium at $t$, respectively. When $\sigma^{\mathrm{med}}=0$, it simplifies to the traditional NeRFs case.

Adopting the same discretization strategy as NeRFs:

According to the rendering equation Eq. (10), the color contributions can be divided into object and medium components, reflecting their different contributions to the final color.

where

and

To simplify the model, SeaThru-NeRF assumes the color and density of the medium are constant along the ray $\mathbf{r}$. Additionally, because of $\sigma^{\mathrm{med}}\gg\sigma^{\mathrm{obj}}$ before the object and $\mathbf{\sigma}^{\mathrm{med}}\ll\sigma^{\mathrm{obj}}$ at the object [5], thus

In the aforementioned discussion, the object component and the backscatter component used the same attenuation coefficient. According to SeaThru, the effective $\sigma^{\mathrm{med}}$ for the object component $\hat{C}^{\mathrm{obj}}(\mathbf{r})$ and the backscatter component $\hat{C}^{\mathrm{med}}(\mathbf{r})$ are different. Therefore, in the final model, different parameters are used for each component—$\sigma^{\mathrm{attn}}$ for the object component $\hat{C}^{\mathrm{obj}}(\mathbf{r})$ and $\sigma^{\mathrm{bs}}$ for the backscatter component $\hat{C}^{\mathrm{med}}(\mathbf{r})$.

Like NeRFs, SeaThru-NeRF through minimize the squared distance between expected color and ground truth to optimize the network parameters for each ray $\mathbf{r}$ sampled from image $\mathbf{I}_{n}$ of training set images $\mathcal{I}=\{\mathbf{I}\}_{n=0}^{\mathrm{N}}$. In a Bayesian perspective, this is equivalent to assuming a Gaussian likelihood $p(C_{\boldsymbol{\varphi}}|\boldsymbol{\varphi})\sim\mathcal{N}(C_{n}^{\mathrm{gt}},\frac{1}{2})$ and inferring $\boldsymbol{\varphi}^{*}$, the mode of the posterior distribution

which, by Bayes’ rule, is the same as minimizing the negative log-likelihood

Laplace approximation obtains the optimal network weights $\boldsymbol{\omega}^{*}$ through pre-trained model, then approximates the posterior distribution of the network parameters as a multivariate Gaussian distribution centered at $\boldsymbol{\omega}^{*}$, i.e.,

where $\mathrm{\Sigma}$ is the covariance matrix.

According to the Laplace approximation, we consider the second-order Taylor expansion of the objective function $h(\boldsymbol{\omega})=-\log p(\boldsymbol{\omega}|\mathcal{I})$ at $\boldsymbol{\omega}^{*}$:

where $H(\boldsymbol{\omega}^{*})$ is the Hessian matrix of $h(\boldsymbol{\omega})$ at $\boldsymbol{\omega}^{*}$, and since $\boldsymbol{\omega}^{*}$ is an extremum point of $h(\boldsymbol{\omega})$ , the first-order derivative term is zero.

By comparing the Eq. (24) with the usual log squared exponential Gaussian likelihood of the multivariate Gaussian distribution, we can derive the expression for the covariance matrix $\Sigma$:

However, directly equating $\boldsymbol{\omega}$ with $\boldsymbol{\varphi}$ is impracticable. The high correlation between model parameters makes accurately estimating the covariance matrix $\Sigma$ of the parameter distribution challenging. Additionally, even with an accurate $\Sigma$, converting it into geometrically meaningful distribution necessitates an expensive sampling process, resulting in significant computational overhead.

To address these issues, we introduce a reparameterization method based on a perturbation field in Section 4.2, which is equivalent to adding a differentiable spatial deformation module before the MLP. This reparameterization method makes the parameters more amenable to Laplace approximation.

As illustrated in Fig. 4, a space (green area) exists where the green line segment in Fig. 4(a) can be perturbed to any shape within this region without impacting the reconstruction loss (Fig. 4(b)). During model training, different random seeds may cause convergence to various configurations within this space. Consequently, for a pre-trained reconstruction model, due to the limited training data, certain regions of the scene have perturbable spaces that do not affect the reconstruction loss. The degree of allowable perturbation indicates the model’s uncertainty level in that region. By systematically applying perturbations throughout the entire 3D scene and measuring the maximum tolerated perturbation at each spatial position, one can obtain the model’s uncertainty distribution across the entire space (Fig. 4(c)).

Inspired by the above, we introduce a parametrized perturbation field $\mathcal{D}{:}\mathbb{R}^{D}\to\mathbb{R}^{D}$, which can be interpreted as a perturbation transformation applied to the coordinates before inputting them into the MLP (Fig. 5). The parameters of the perturbation field $\mathcal{D}$ are represented as $\boldsymbol{\omega}\in\mathbb{R}^{M^{D}\times D}$ , where $M$ denotes the size of the grid used to store displacement vectors and $D$ represents vector dimension. For any spatial coordinate $\mathbf{x}$, its perturbation is computed by trilinear interpolation of displacement vectors from neighboring grid points:

Next, we reparameterize the optimized MLP neural network by perturbing each coordinate:

where $\sigma_{{\boldsymbol{\varphi}}^{*}}^{\mathrm{obj}}$ and $\mathbf{c}_{{\boldsymbol{\varphi}}^{*}}^{\mathrm{obj}}$ are the optimized density and radiance of object, respectively. And as mentioned in Section 3.3, the density and color associated with the medium are constants at each spatial coordinate, so they are not affected by perturbations.

Based on the reparameterized $\tilde{\sigma}_{\boldsymbol{\omega}}^{\mathrm{obj}}$ and $\tilde{\mathbf{c}}_{\boldsymbol{\omega}}^{\mathrm{obj}}$ , the perturbed predicted pixel color can be obtained as:

where

and accumulated transmittance of object can be obtained as:

Since we still want the predicted color ${\tilde{C}}_{\boldsymbol{\omega}}(\mathbf{r})$ to be as close to $\mathrm{C}_{n}^{\mathrm{gt}}$ as possible after reparameterization, we assume a likelihood function of the same form as the original model, i.e., $\tilde{C}_{\boldsymbol{\omega}}|{\boldsymbol{\omega}}\sim\mathcal{N}(C_{n}^{\mathrm{gt}},\frac{1}{2})$. Moreover, since ${\boldsymbol{\varphi}}^{*}$ are the optimal parameters obtained during the training of the MLP network, the model has achieved optimal performance at these parameters. Therefore, when perturbations are applied to new parameters, we expect these perturbations to neither significantly improve nor degrade the model’s performance. Hence, we impose a regularization Gaussian prior $\boldsymbol{\omega}\sim\mathcal{N}(0,\lambda^{-1})$ on the new parameters $\boldsymbol{\omega}$.

Under these assumptions, the negative log-likelihood function of the posterior distribution $p(\boldsymbol{\omega}|\mathcal{I})$ is given by:

Here, $\mathbb{E}_{n}\mathbb{E}_{\mathbf{r}\sim\mathbf{I}_{n}}\|\tilde{C}_{\boldsymbol{\omega}}(\mathbf{r})-C_{n}^{\mathrm{gt}}(\mathbf{r})\|_{2}^{2}$ represents the reconstruction error of the pixel colors, and $\lambda\|\boldsymbol{\omega}\|^{2}$ is the regularization term.

When $\boldsymbol{\omega}=0$, we have $\tilde{\sigma}_{0}^{\mathrm{obj}}(\mathbf{x})=\sigma_{{\boldsymbol{\varphi}}^{*}}^{\mathrm{obj}}(\mathbf{x})$ and $\tilde{\mathbf{c}}_{0}^{\mathrm{obj}}(\mathbf{x})=\mathbf{c}_{{\boldsymbol{\varphi}}^{*}}^{\mathrm{obj}}(\mathbf{x},\mathbf{d})$, thus $\tilde{C}_{0}(\mathbf{r})=C_{{\boldsymbol{\varphi}}^{*}}(\mathbf{r})$. Therefore, $\boldsymbol{\omega}=0$ minimizes the reconstruction error and is the mode of the posterior distribution $p(\boldsymbol{\omega}|\mathcal{I})$.

According to the Laplace approximation, we perform a second-order Taylor expansion around the mode, which yields

where $\mathbf{H}(0)$ is the Hessian matrix of $h(\boldsymbol{\omega})$ evaluated at zero.

Because directly computing the second derivatives in the Hessian matrix is computationally expensive, we approximate it using the Fisher information.

For any parameterized family of probability distributions $p_{\boldsymbol{\omega}}$ , the Hessian matrix of its log-likelihood function with respect to the parameters ${\boldsymbol{\omega}}$ is related to the Fisher information as follows:

where $h(\mathbf{X};\boldsymbol{\omega})$ is defined as the log-likelihood function.

Furthermore, under reasonable regularity conditions, the Fisher information can also be defined as:

We use the random variable $(\mathbf{r},\mathbf{y})$ to correspond to a camera ray $\mathbf{r}$ and its corresponding ground truth $\mathbf{y}=C_{n}^{\mathrm{gt}}(\mathbf{r})$. Thus,

where $\epsilon_{\boldsymbol{\omega}}(\mathbf{r})=\|\tilde{C}_{\boldsymbol{\omega}}(\mathbf{r})-C_{n}^{\mathrm{gt}}(\mathbf{r})\|^{2}$. And

represents the Jacobian matrix of the predicted color with respect to the parameter $\boldsymbol{\omega}$, which can be computed through backpropagation.

Furthermore, based on the properties of conditional expectations:

According to the likelihood ${\mathcal{N}}(C_{n}^{\mathrm{gt}},{\frac{1}{2}})$, it holds that $\mathbb{E}_{\mathbf{y}|\mathbf{r}}\left[\epsilon_{\boldsymbol{\omega}}(\mathbf{r})\right]$ is nothing more than $\frac{1}{2}$, so:

Here, we approximate the expectation by sampling $R$ rays.

We can obtain:

From Eq. (35), we can finally obtain:

Since each entry of the parameter vector $\boldsymbol{\omega}$ corresponds to a vertex in the grid, its influence is limited to the cell containing that vertex, making $\mathbf{H}(\boldsymbol{\omega})$ inherently sparse and thereby minimizing the number of related parameters. Similar to the approach by Ritter et al. [54], we approximate $\Sigma$ by considering only the diagonal elements of $\mathbf{H}$.

where $\Sigma$ encodes the spatial uncertainty of the radiance field. Intuitively, it represents the extent to which the geometry of NeRF can be altered without compromising the quality of reconstruction.

By computing the diagonal entries of $\Sigma$, we obtain the marginal variance vector $\boldsymbol{\sigma}=(\sigma_{x},\sigma_{y},\sigma_{z})$. At each grid vertex, $\boldsymbol{\sigma}$ defines a spatial ellipsoid representing the region within which deformations can occur with minimal reconstruction cost. The norm of the vector $\sigma=\|\boldsymbol{\sigma}\|_{2}$ is a positive scalar that measures the local spatial uncertainty of the radiance field at each grid vertex.

Through this method, we can define the spatial uncertainty field $\mathcal{U}:\mathbb{R}^{3}\to\mathbb{R}^{+}$, expressed as (Fig. 6):

Strictly speaking, as described above, $\mathcal{U}$ measures the uncertainty at $(1+\mathcal{D}_{\boldsymbol{\omega}})^{-1}(\mathbf{x})$, not at $\mathbf{x}$; however, for a trained SeaThru-NeRF where $\mathcal{D}_{{\boldsymbol{\omega}}^{*}}=0$, these points are effectively the same.

Fig. 7 illustrates the uncertainty quantification results across five datasets on SeaThru-NeRF and SeaThru-NeRF-lite. The uncertainty estimates are visually conveyed using color, where the color intensity represents the level of uncertainty: bluer colors indicate higher uncertainty, and redder colors indicate lower uncertainty. In addition, as shown in Table 1, our uncertainty quantification results exhibit excellent performance across AUSE metrics. At the same time, as observed in Table 1, the PSNR, SSIM, and LPIPS metrics of the model show negligible differences compared to the original model (base). It is in line with our emphasis in Section 4.2 on performing uncertainty quantification without significantly impacting reconstruction loss. This indicates that we can provide additional confidence information while maintaining the original reconstruction performance, which provides important technical support and guarantee for practical application. For example, in AUVs navigation, high-quality environmental reconstruction is crucial for path planning and obstacle detection. Additional uncertainty information can assist in making more cautious decisions when faced with uncertain or hazardous situations.

In this work, we utilize Bayesian Laplace approximation to compute the covariance matrix $\Sigma$ for the perturbed field parameters $\boldsymbol{\omega}$. The diagonal elements of $\Sigma$ represent the uncertainty in each dimension. We render these diagonal elements as a new volumetric data $\mathcal{U}(\mathbf{x})$, where regions with larger $\mathcal{U}(\mathbf{x})$ values indicate higher uncertainty. These high uncertainty areas often correspond to artifacts in the rendering results. By applying a threshold to $\mathcal{U}(\mathbf{x})$, we can identify and remove these artifacts, resulting in clearer reconstruction outcomes (Fig. 14). Consequently, this method not only effectively quantifies the uncertainty in each dimension but also post-processes the reconstruction results.