GeoNeRF: Generalizing NeRF with Geometry Priors

In plain words

A neural radiance field (NeRF) is a learned 3D model of one scene that can draw that scene from new camera positions. A standard NeRF has to be trained again for every new scene, which is slow and needs many photos. GeoNeRF is trained once on many scenes, so it can render a new scene without that per-scene training.

Its main idea is to work out the geometry first, the way classic multi-view stereo does. For each nearby photo, the model builds a stack of depth guesses and checks which guess agrees with the other photos. From this it learns which photos can see a given point in space, and it blends their colours accordingly. A final step turns the result into density and colour along each camera ray, and the ray is then drawn.

The authors also describe a version that accepts a sparse depth map, of the kind depth sensors produce, as extra input. A short fine-tuning run of 1k or 10k steps adapts the general model to one scene, and the authors describe the result as competitive with a per-scene NeRF.

TL;DR

GeoNeRF renders novel views of unseen scenes without per-scene optimisation §1. For each source view it builds three-level cascaded cost volumes with group-wise correlation §3.1 eq. 2, fuses them with four attention layers and an along-ray auto-encoder §3.3 eq. 4 eq. 5, then volume-renders colour and depth eq. 7 eq. 8. Author-reported on LLFF: 25.44 PSNR and 0.839 SSIM with no per-scene optimisation, against 25.13 and 0.817 for IBRNet, which uses 10 source views tab. 1 tab. 2. A 10k-step per-scene fine-tune reaches 26.58 PSNR on LLFF tab. 1. The official code is GPL-3.0 and differs from the paper in several loss weights and constants; see Paper vs code.

Problem

Per-scene NeRF training is lengthy and needs densely captured images of each scene §1. Earlier generalizing NeRF models condition the renderer on features from source views, but their understanding of scene geometry and occlusion is limited, which causes artifacts §1. MVSNeRF builds a low-resolution plane-swept cost volume, which the authors report makes detailed images difficult and does not handle occlusion §1. The method assumes source views with known camera poses and intrinsics, because its cost volumes are built by homography warping between views §3.1.

Method

The model has two stages, a geometry reasoner and a renderer §1 fig. 2. For a novel view, it takes V nearby source views, each with a neighbour set of similar views §3.1.

Geometry reasoner §3.1. Each source image goes through a feature pyramid network that produces features at three scales, with C channels at the finest scale eq. 1. The reasoner follows CasMVSNet: depth planes are swept at three levels, starting with a coarse sweep over the full depth range and then narrowing around the depths estimated at the previous level §3.1 app. §1. Two changes are made §3.1. An extra output head produces multi-level 3D features Φ alongside the depth maps. Group-wise correlation §3.1 replaces variance as the matching measure, which lowers memory use. A 3D hourglass network regularises each cost volume and outputs a depth map and 3D features per level eq. 2.

Sampling §3.2. For each ray, N_c points are placed uniformly between near and far. A probability p0 over depth is estimated from the full-resolution volume: it is the chance that a point is covered by that volume. The paper draws N_f more points from p0, so only one network is used, not a coarse and a fine network as in NeRF §3.2.

Renderer §3.3. For each sample point, the model interpolates the full-resolution 2D feature and the three 3D features from every source view. An occlusion mask M is set to zero for a view when the point lies behind that view's full-resolution depth map or outside its image §3.3. Each point then gets one view-dependent token per source view, built from the interpolated features, and one view-independent token, built from the mean and variance of the 2D features across views eq. 3. Four stacked multi-head attention layers run over these tokens, shared across all points on all rays. Occluded views are excluded from the attention scores eq. 4 §3.3. The view-independent tokens are then taken along each ray and passed through a convolutional auto-encoder along the ray dimension; a two-layer perceptron turns the result into density σ eq. 5. The authors argue that convolution along the ray is a better inductive bias than attention between samples, and cheaper §3.3.

For colour, each view-dependent token is concatenated with an encoding of the angle between the novel ray and the line from source camera to point. A perceptron gives one weight per view, a softmax over the visible views normalises them, and the point colour is the weighted sum of the source colours eq. 6. The ray colour and a ray depth come from the standard volume rendering sums eq. 7 eq. 8 §3.3.

Losses §3.4. Training uses mean squared error on the rendered ray colours eq. 9. Where ground-truth depth exists, the model adds a smooth-L1 ray-depth loss with weight 0.1 eq. 10 and a multi-level depth-map loss with weight 2^-l at level l, with total weight 1 eq. 11. Where it does not, the model uses self-supervised depth §3.4. The rendered depth of each target ray is re-projected into each source view. A source depth is pulled toward the re-projected depth only where the target colour agrees with the source colour and the source patch is textured §3.4. This loss has weight 0.1. The overall loss is written as L = L_c + 0.1 L_d + λ Σ_l L_D^(l), with λ = 1.0 for ground-truth depth and 0.1 for self-supervised depth eq. 13. Fine-tuning on one scene keeps only the colour loss §3.4.

RGBD variant GeoNeRF+D §3.5. Motivated by embedded depth sensors, a sparse, noisy, low-resolution depth map (H/4 by W/4) is converted into a binary volume B with one slice per coarse depth plane. This volume is concatenated with the coarsest cost volume before its 3D network eq. 14. The authors say the model is robust to the quality and sparsity of the depth input §3.5.

# GeoNeRF: colour and depth of one novel-view ray  [§3] [fig. 2]
# images: V source images, each with its camera; ray: origin and direction

def render_ray(images, cameras, ray):
    # 2D feature pyramid per source view, three scales  [§3.1] [eq. 1]
    pyramids = [feature_pyramid(image) for image in images]
    # Cascaded cost volumes and depth maps per source view  [§3.1] [eq. 2]
    volumes, depth_maps = cascaded_volumes(pyramids, cameras)
    # N_c uniform depths, then N_f more drawn from the coverage  [§3.2]
    depths = sample_depths(ray, volumes, cameras)
    points = ray.origin + depths * ray.direction
    # Which source views see each point, for occlusion  [§3.3]
    visible = occlusion_masks(points, depth_maps, cameras)
    # One global token and one token per view, for every point  [eq. 3]
    global_token, view_tokens = make_tokens(points, volumes, pyramids)
    tokens = concatenate(view_tokens, global_token)  # (N, V+1, d)
    # Four attention layers shared across points, masked by visibility  [eq. 4]
    for layer in attention_layers:
        tokens = layer(tokens, visible)
    # Along-ray auto-encoder on the global token, then an MLP  [eq. 5]
    density = relu(mlp_sigma(auto_encoder(tokens[:, -1])))  # (N,)
    # Angle between the ray and each source camera's view of the point  [eq. 6]
    angles = sine_of_angles(points, ray, cameras)  # (N, V)
    # Softmax over the visible views gives the colour weights  [eq. 6]
    weights = masked_softmax(mlp_color(tokens[:, :-1], angles), visible)
    colours = (weights * source_colours(points, images)).sum(dim=1)  # (N, 3)
    # Classical volume rendering along the ray  [eq. 7] [eq. 8]
    return volume_render(colours, density, depths)

# Depths along one ray: N_c uniform, N_f from the coverage  [§3.2]
# coverage[i]: how many source views contain coarse depth i in their volume

def sample_depths(near, far, coverage, n_coarse, n_fine):
    # Uniform coarse depths between near and far  [§3.2]
    coarse = linspace(near, far, n_coarse)
    # The code uses coverage as logits: weight = exp(coverage)
    # [code: utils/utils.py#L362-L363]
    index = categorical(logits=coverage).sample(n_fine)
    # Jitter in each cell; t can dip below 0  [code: utils/utils.py#L365-L368]
    t = (index - uniform(0, 1)) / (n_coarse - 1)
    fine = near * (1 - t) + far * t
    return sort(concatenate(coarse, fine))

# Training loss for one batch of rays  [§3.4] [eq. 9] [eq. 10] [eq. 11] [eq. 13]
# Fine-tuning on one scene uses only the first term

def training_loss(batch, fine_tuning):
    # Pixel loss on the sampled rays  [eq. 9]
    loss = mean((rendered_colour - gt_colour) ** 2)
    if fine_tuning:
        return loss
    # Ray-depth loss, only for rays that have ground-truth depth  [eq. 10]
    loss = loss + 0.1 * smooth_l1(rendered_depth, gt_ray_depth)
    if batch.has_gt_depth:
        # Depth maps at three levels, weight 2^-level  [eq. 11]
        for level in range(3):
            weight = 2 ** -level
            depth_error = smooth_l1(depth_map[level], gt_depth[level])
            loss = loss + weight * depth_error
    else:
        # Self-supervised depth, weight 0.1  [§3.4] [eq. 13]
        loss = loss + 0.1 * self_supervised_depth_loss(batch, rendered_depth)
    return loss

Hyperparameters

Results

Author-reported. Setup: CVPR 2022 README.md; training on DTU, LLFF and IBRNet scenes, evaluation on 16 DTU, 8 LLFF and 8 NeRF-synthetic test scenes §4. Top block: no per-scene optimisation. Bottom block: fine-tuned on each scene (caption) tab. 1.

The authors describe GeoNeRF-10k as competitive with per-scene NeRF, and say GeoNeRF-1k reaches 98.15% of GeoNeRF-10k's average performance §4.1 tab. 1.

tab. 1 Notes from the paper: DTU baseline numbers are borrowed from the MVSNeRF paper, and DTU metrics are computed on foreground pixels with ground-truth depth inside the scene bounds. The MVSNeRF fine-tuned row overfits on the standard test set: the authors say its NeRF-dataset PSNR would drop from 27.07 to 20.02 there tab. 1.

Source-view sensitivity (LLFF) tab. 2. GeoNeRF with V = 3, 4, 5, 6 and 9 source views reaches 24.33, 25.05, 25.25, 25.37 and 25.44 PSNR. IBRNet with 10 views reaches 25.13.

Discarding the nearest views (LLFF) tab. 3. Removing the K nearest source views degrades quality. PSNR is 25.44 (K = 0), 24.18 (K = 2), 23.35 (K = 4), 22.74 (K = 6) and 22.06 (K = 8).

RGBD variant (DTU) tab. 4. GeoNeRF reaches 31.34 PSNR, 0.959 SSIM and 0.060 LPIPS. GeoNeRF+D, with the depth input, reaches 31.58, 0.961 and 0.057.

Ablation (NeRF synthetic and LLFF) app. §4. Removing a component lowers the score. Without self-supervision, PSNR is 28.10 on NeRF synthetic and 25.37 on LLFF. Without the attention renderer, 27.69 and 24.95. Without the AE, 23.53 and 24.92. A single cost volume gives 26.60 and 24.60. The full model gives 28.33 and 25.44.

Per-scene breakdown app. §3. Per-scene tables for LLFF and NeRF synthetic (PSNR, SSIM, LPIPS) are in the supplement, at 10k and 1k fine-tuning steps.

Paper vs code

Code: idiap/GeoNeRF at e6249fd (GPL-3.0, a copyleft licence), cited by path and line only. Its header also incorporates CasMVSNet_pl (GPL-3.0) and MIT and Apache-2.0 code from other projects. The arXiv metadata used here is v2 (2022-03-21). The repository does not ship the pretrained checkpoints; the README links them on Google Drive. Differences found:

• Depth-map loss weight. The paper weights level l by 2^-l eq. 11. The code uses 2^(1-l), which doubles every level's weight utils.py utils.py.
• Ray-depth normaliser. The paper averages over the rays with ground truth depth, 1/|R_d| eq. 10. The code sums over those rays and divides by the whole batch size utils.py.
• Renderer MLPs. The paper describes two-layer perceptrons for density and colour §3.3. The code uses three linear layers each, with ELU between them and a ReLU on density self_attn_renderer.py self_attn_renderer.py.
• Global token. The paper uses one linear transform of mean and variance §3.3. The code uses two linear layers with ELU self_attn_renderer.py.
• View tokens. The code also feeds the visibility bit into each view token self_attn_renderer.py. The paper does not mention it.
• Attention layers. The paper describes multi-head attention alone §3.3. The code wraps each layer with a feed-forward block, residual connections and LayerNorm self_attn_renderer.py.
• Angle feature. The paper says "angle" §3.3. The code uses the sine of the angle, computed as sqrt(1 - cos²) rendering.py.
• Fine samples. The paper draws from p0, the probability that a point is covered by the full-resolution volume §3.2. The code treats the coverage count as logits, so the weight is exp(count), and the draw is random even at evaluation utils.py.
• Self-supervision. The code uses a colour difference below 0.02 and a patch standard deviation above 0.05, where the paper gives a variance threshold and no values utils.py utils.py. It also checks the target colour against the warped source colour utils.py.
• Binary RGBD volume. The paper uses a binary indicator volume eq. 14. The code places each depth value across the two nearest planes with linear weights geo_reasoner.py.
• Occlusion. The code marks a view occluded when the point is more than 0.2 (normalised depth) in front of the depth map rendering.py rendering.py.
• No visible view. When no source view sees a point, the code re-enables all views self_attn_renderer.py self_attn_renderer.py.
• Volume rendering. The code adds 1e-10 inside the transmittance product rendering.py. The paper writes the exponential of a sum eq. 7.
• Training data. The paper samples one scene per iteration §4. The code samples with fixed dataset weights get_datasets.py. DTU training also applies a random rectangular cutout to the depth input dtu.py.
• Warm-up and schedule. The code uses a 500-step linear warm-up and a minimum learning rate of 1e-5 options.py run_geo_nerf.py.
• Constant bias. The density head's bias is initialised to a constant rather than a random draw self_attn_renderer.py self_attn_renderer.py.
• Near and far. The code sets NeRF synthetic to [2, 6] nerf.py and derives LLFF bounds from the file, scaled by 0.75 of the nearest value llff.py. The paper states neither.
• Positional encoding. The code encodes the angle with 4 frequencies rendering.py. The paper does not state the number.

Pitfalls

• Dense GPU code. The renderer and sampling code call .cuda() directly self_attn_renderer.py utils.py. Porting to CPU or another accelerator needs changes beyond the model.
• Self-supervised depth is fiddly. It needs ground-truth colours, source views warped with camera matrices, 5x5 patch statistics and four masks. None of the thresholds are in the paper §3.4 utils.py.
• Fine samples can leave the near-far range. The jitter formula can produce depths slightly below near utils.py. Clamp or follow the code.
• Evaluation is not deterministic. The fine samples are random at test time and the code sets no seed utils.py.
• Occlusion uses detached depth. Training detaches the depth used for occlusion, so occlusion masks do not push gradients into the geometry run_geo_nerf.py.
• Training resolution changes the cost. Feature maps at width 800 or more use checkpointing geo_reasoner.py geo_reasoner.py.
• Fine-tuning is short but sensitive. The paper reports gains at 1k and 10k steps on one scene tab. 1. The MVSNeRF caveat tab. 1 shows how easily fine-tuning results can be overstated.
• Sparse source views lower quality. Quality drops as the number of source views falls and as the nearest views are removed tab. 2 tab. 3.
• Failures in textureless areas. The authors report that stereo fails in textureless regions, and the geometry priors then mislead the renderer app. §5 app. fig. 4.

Lineage

• Builds on: CasMVSNet for cascaded cost volumes §3.1; group-wise correlation §3.1; feature pyramid networks §3.1; NeRF's volume rendering §2; Transformer multi-head attention §3.3. Each is cited in the paper and not linked here.
• Compared against: pixelNeRF, IBRNet, MVSNeRF and per-scene NeRF tab. 1.
• Concurrent work named by the authors: RGBD-Net, NeuralMVS and NeuRay §1.