Advancing 3d reconstruction with hierarchical nerf and jax3d

Researchers have developed a hierarchical Neural Radiance Field (NeRF) using JAX, Flax, Optax, and jax3d to enable volumetric rendering, novel-view synthesis, and 3D reconstruction. The approach constructs a synthetic multi-view dataset from an analytic scene with volumetric geometry and view-dependent radiance, using sample_along_rays and volume_rendering for forward rendering. The model incorporates positional encoding, skip connections, and view-direction conditioning, with hierarchical importance sampling through sample_piecewise_constant_pdf. Training involves JAX JIT compilation, Adam optimization, exponential learning-rate decay, and gradient clipping, while evaluation focuses on PSNR, depth and opacity visualization, and marching-cubes geometry extraction. The system sets up a JAX3D environment, configures GPU/CPU-adaptive parameters, and uses pinhole intrinsics and orbit-based camera placement. It generates normalized world-space rays and constructs a ground-truth scene with soft-edged spheres and patterned floors. The NeRF representation uses sinusoidal positional encoding and a deep Flax MLP with skip connections to predict density and condition RGB on viewing direction. The hierarchical renderer samples coarse and fine points, combines them, and stops gradients during sampling. Training optimizes both networks jointly, with PSNR monitoring throughout optimization. The final evaluation measures novel-view reconstruction quality with PSNR, visualizes importance-sampling behavior, generates a 360-degree orbit GIF, and extracts geometry using marching cubes. The study demonstrates an end-to-end inverse-rendering pipeline, integrating jax3d’s mathematical components with modern JAX-based training for a complete NeRF reconstruction system.