Compact JAX NeRF: Hierarchical Coarse→Fine Rendering, Mesh Extraction and 360° Preview

A compact JAX NeRF you can run and learn from

A compact JAX/Flax NeRF that renders novel views, produces a 360° GIF, and extracts a mesh, all from a tiny synthetic dataset, is a practical template for teams prototyping view synthesis and geometry extraction. The notebook implements a full hierarchical NeRF pipeline in JAX + Flax + Optax and uses google-research’s jax3d primitives for sampling and volume rendering; it’s a useful engineering reference if you want a runnable, end-to-end example of how all the pieces fit together.

Why product and engineering leaders should care

  • Shows how a JAX-based ML stack (JAX + Flax + Optax + jax3d) composes into a differentiable renderer and reconstruction pipeline, helpful if your team already targets JAX for production or research.
  • Provides a repeatable, synthetic testbed for debugging architecture and sampling choices without the mess of real‑world captures.
  • Includes a simple isosurface extraction flow (marching cubes + a practical isovalue heuristic) so teams can move quickly from learned density → mesh for early validation.

What the notebook actually does

It synthesizes a controlled dataset (three analytic soft spheres + a patterned floor), uses jax3d.math.volume_rendering to render ground-truth RGB/depth/alpha, trains a canonical coarse→fine NeRF (positional encoding, skip connections, view-direction conditioning), and evaluates with PSNR, depth/opacity visualizations, importance-sampling diagnostics, a 360° orbit GIF, and marching‑cubes mesh extraction. The implementation is available on GitHub: https://github.com/MARKTECHPOST-AI-MEDIA-INC/AI-Agents-Projects-Tutorials/blob/main/Computer%20Vision/jax3d_hierarchical_nerf_tutorial_Marktechpost.ipynb

Key primitives and the hierarchical flow

  • sample_along_rays (stratified or deterministic sampling between near and far)
  • volume_rendering (composites sampled sigma+rgb into ray RGB, expected depth, and alpha)
  • sample_piecewise_constant_pdf (importance resampling: convert coarse weights → PDF → fine samples)

Flow (high level):

  • Coarse: sample n_coarse depth locations along each ray and evaluate the coarse MLP.
  • Composite: run volume_rendering on coarse outputs to get coarse colors and per-bin weights.
  • Resample: build a PDF from coarse weights and draw n_fine importance samples using sample_piecewise_constant_pdf (these are additional samples, not replacements).
  • Fine: combine coarse+fine depths, evaluate the fine MLP at those positions, and composite for final color/depth.
  • MSE on both coarse and fine RGB outputs, optimize both networks jointly with Optax Adam plus gradient clipping and an exponential LR decay.

“Coarse pass -> importance-resample -> fine pass. All sampling and compositing comes from jax3d.math.volume_rendering.”

Config and the knobs that matter (notebook defaults)

The notebook sets a comprehensive config. Key defaults (quoted exactly from the notebook) include:

  • Image size: H = 64, W = 64
  • Views: n_train_views = 24, n_test_views = 3
  • Camera: cam_radius = 3.2, fov_deg = 40.0
  • Sampling: gt_samples = 256, n_coarse = 64, n_fine = 64
  • MLP: deg_pos = 10, deg_dir = 4, width = 128, depth = 6, skip = 3
  • Training: batch_rays = 2048, steps = 2500, lr_init = 5e-4 → lr_final = 5e-6
  • Chunking: chunk = 4096 (used to limit compilation shapes / memory)
  • Geometry grid: grid_res = 96

There is an automatic CPU fallback with smaller defaults (for example H/W = 40, fewer steps, grid_res = 64) so the notebook remains runnable without a GPU. The focal length used for the pinhole camera is computed as:

focal_px = 0.5 * W_px / tan(0.5 * fov_rad)

i.e., converting horizontal field-of-view in radians to focal length in pixels.

Scene model, density, and the marching-cubes heuristic

The synthetic scene uses three soft spheres with view‑dependent specular lighting. The notebook’s density formulation is explicit and tuned for numeric stability; for example it uses:

sigma = 80.0 * sigmoid((radius – dist) / 0.015)

For mesh extraction the notebook samples a 3D grid over the scene domain (linspace(-1, 1, grid_res)) and applies skimage.measure.marching_cubes. To pick an isovalue it uses a simple, practical derivation:

  • step = (far – near) / (n_coarse + n_fine)
  • Assuming a uniform slab of density sigma over one sampling step gives transmittance ≈ exp(-sigma * step). Setting transmittance=0.5 ⇒ sigma ≈ -ln(0.5)/step.
  • So the notebook uses level = -ln(0.5) / step; if that level is outside the sampled volume range it falls back to the 99th percentile of the density field as a safety net.

This heuristic places the isosurface where a single sample step would absorb half the light. It’s sensible for synthetic scenes but sensitive to density scaling, the chosen sample counts, and grid_res, so sweep the isovalue and grid resolution when you move to different data.

Practical training and evaluation notes

  • The NeRF MLP uses positional (sinusoidal) encoding for positions and directions, skip connections, softplus for density, and a sigmoid for RGB. Concretely, the notebook shows the pattern: Dense(1) → softplus → subtract 1.0 to form sigma. If you port this, ensure sigma is clamped or validated afterward because softplus(raw) – 1.0 can produce small negative values, and volume rendering expects non‑negative densities.
  • Chunking (cfg.chunk) is necessary for JAX because XLA-compiled functions prefer consistent shapes and the first compilation is slow. Use small chunks to sanity-check shapes and then increase chunk sizes for throughput.
  • PSNR is reported via mse_to_psnr = -10.0 * log10(max(mse, 1e-10)). This assumes images scaled to [0, 1]. If you use 8-bit images remember to include peak^2 scaling or convert to floats in [0, 1].
  • The notebook prints final summary lines such as: “FINAL held-out PSNR: {np.mean(psnrs):.2f} dB ({n_params/1e6:.2f}M params, {cfg.steps} steps)” and lists the jax3d functions exercised: sample_along_rays, volume_rendering, sample_piecewise_constant_pdf.

Limitations, gotchas, and debugging tips

  • PSNR ≠ geometric truth. PSNR measures pixel fidelity, not geometric correctness. Debug tip: overlay predicted depth edges on ground-truth alpha masks to spot “floating” density or smearing.
  • Softplus − 1.0 for sigma. That subtraction shifts densities toward zero but can produce negative values. Debug tip: clamp sigma to >=0 before rendering, or log the sigma min/max during training to detect problematic values.
  • Marching cubes sensitivity. The chosen level (−ln(0.5)/step or 99th percentile fallback) is a heuristic. Debug tip: sweep isovalues and grid_res and visualize. If you have analytic geometry (like the spheres), compute Chamfer distance to pick the best threshold.
  • JAX compilation and shapes. Expect a long first compilation. Debug tip: run a single small-step compile with tiny chunk to validate shapes before full training.
  • API/version fragility. The notebook imports jax3d from the google-research repo. Import paths and function signatures can change, so pin jax/jax3d/flax/optax versions or the notebook commit to avoid surprises.

Three experiments to try next

  • Quantitative geometry: compute Chamfer distance between the marching-cubes mesh and the analytic sphere ground truth while sweeping isovalue and grid_res.
  • Speed vs quality: replace positional encoding with a multiresolution hash encoding (instant-ngp style) to measure training time and final fidelity tradeoffs.
  • Real data robustness: swap synthetic images for a small real multi-view capture, add pose refinement or bundle adjustment, and compare visual fidelity and geometry recovery.

Quick implementation checklist

  • Open the notebook in Colab or your Jupyter environment; switch the runtime to GPU (for the full defaults the notebook suggests a T4).
  • If jax3d import fails, the notebook attempts to clone google-research/jax3d and prints a pip suggestion, pin versions if possible.
  • Start with the CPU fallback config to verify behavior, then move to GPU for full fidelity.
  • Monitor PSNR, visualize depth/alpha, and inspect coarse weights to confirm importance resampling is concentrating samples around surfaces.
  • When extracting geometry, sweep isovalues and grid_res and, if available, compute geometric metrics (Chamfer) against ground truth.

Questions leaders and engineers will ask, short, honest answers

  • Can I run this as-is on a laptop CPU?

    Yes. The notebook supplies CPU-friendly defaults (smaller H/W, fewer steps, reduced model depth/width and grid_res). Expect slower iterations and lower fidelity compared to a GPU run.

  • Does this work with real multi-view photos out of the box?

    No. The notebook uses synthetic, perfectly posed images. Real data requires pose estimation, photometric preprocessing, and additional robustness to specularities and noise.

  • Is the marching-cubes mesh a ground-truth surface?

    Not necessarily. The mesh is an isosurface of the learned density; topology and accuracy depend on density scaling, grid resolution, and the chosen isovalue. Treat it as an approximation and validate with geometric metrics when accuracy matters.

  • How much compute should I budget?

    The notebook suggests using a GPU (e.g., T4) for the full settings (H/W=64, steps=2500). Exact runtimes depend on hardware and chunking; expect long first-run compile times with JAX but much faster iterations thereafter on modern GPUs.

  • Where does this fit in the NeRF ecosystem?

    It’s a faithful JAX/Flax implementation of the canonical hierarchical NeRF pipeline (coarse→fine sampling + volume rendering). For production or speed you’ll want to explore instant-ngp (multiresolution hash encoding), mip‑NeRF (anti-aliasing), or TensoRF (compact factorized reps).

Bottom-line takeaways

  • This notebook is a clean, runnable JAX example that shows how sampling, compositing, and hierarchical importance resampling assemble into a NeRF training pipeline.
  • It’s a great prototype and debugging playground, but synthetic demos don’t eliminate the additional engineering needed for real-world captures (pose noise, higher resolution, speed optimizations, and geometric validation).
  • Next steps for teams: run the notebook, swap to a small real capture, add pose refinement, and measure geometry with Chamfer distance while experimenting with alternative encodings for speed.
0