{"slug": "hierarchical-nerf-with-jax3d-for-volumetric-rendering-novel-view-synthesis-and", "title": "Hierarchical NeRF with JAX3D for Volumetric Rendering, Novel-View Synthesis, and 3D Reconstruction", "summary": "A new tutorial demonstrates building an end-to-end hierarchical Neural Radiance Field (NeRF) using JAX, Flax, Optax, and the volume-rendering primitives in Google Research's jax3d library. The tutorial constructs a synthetic multi-view dataset, implements coarse and fine NeRF networks with positional encoding and view-direction conditioning, applies hierarchical importance sampling via sample_piecewise_constant_pdf, and trains with JAX JIT compilation, Adam optimization, exponential learning-rate decay, and gradient clipping. Evaluation covers novel-view synthesis with PSNR, depth and opacity visualization, 360-degree rendering, and marching-cubes geometry extraction.", "body_md": "In this **[tutorial](https://github.com/MARKTECHPOST-AI-MEDIA-INC/AI-Agents-Projects-Tutorials/blob/main/Computer%20Vision/jax3d_hierarchical_nerf_tutorial_Marktechpost.ipynb)**, we build an end-to-end hierarchical Neural Radiance Field (NeRF) using [** JAX**](https://github.com/google-research/jax3d), Flax, Optax, and the volume-rendering primitives provided by jax3d. We first construct a synthetic multi-view dataset from an analytic scene containing volumetric geometry and view-dependent radiance, using sample_along_rays and volume_rendering to establish the forward rendering process. We then implement a NeRF with positional encoding, skip connections, separate coarse and fine networks, and view-direction conditioning, followed by hierarchical importance sampling through sample_piecewise_constant_pdf. We train the model with JAX JIT compilation, Adam optimization, exponential learning-rate decay, and gradient clipping, and finally evaluate novel-view synthesis using PSNR, depth and opacity visualization, sampling diagnostics, 360-degree rendering, and marching-cubes geometry extraction.\n\n``` python\nimport os, sys, subprocess, importlib.util, functools, dataclasses, time, math\ndef _sh(cmd):\n   subprocess.run(cmd, shell=True, check=False,\n                  stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)\nprint(\"Installing dependencies ...\")\n_sh(f'{sys.executable} -m pip install -q \"etils[array-types,epy,etree,enp]\" '\n   f'chex flax optax scikit-image')\nREPO_DIR = \"/content/jax3d\" if os.path.isdir(\"/content\") else os.path.abspath(\"./jax3d\")\nif not os.path.isdir(REPO_DIR):\n   print(\"Cloning google-research/jax3d ...\")\n   _sh(f\"git clone -q --depth 1 https://github.com/google-research/jax3d.git {REPO_DIR}\")\ndef _load_module_by_path(name, path):\n   \"\"\"Load a single .py file without triggering the parent package __init__.\n   `from jax3d.math import volume_rendering` also works if you run\n   `pip install .` inside the clone, but that pulls in gin/tfds/etc.\n   \"\"\"\n   spec = importlib.util.spec_from_file_location(name, path)\n   mod = importlib.util.module_from_spec(spec)\n   sys.modules[name] = mod\n   spec.loader.exec_module(mod)\n   return mod\n_VR_PATH = os.path.join(REPO_DIR, \"jax3d\", \"jax3d\", \"math\", \"volume_rendering.py\")\nif not os.path.exists(_VR_PATH):\n   _VR_PATH = os.path.join(REPO_DIR, \"jax3d\", \"math\", \"volume_rendering.py\")\ntry:\n   j3vr = _load_module_by_path(\"j3d_volume_rendering\", _VR_PATH)\nexcept Exception as e:\n   raise SystemExit(\n       f\"Could not load {_VR_PATH}: {e}\\n\"\n       \"Try: pip install -U 'etils[array-types,epy,etree,enp]==1.9.4' and re-run.\"\n   )\nimport numpy as np\nimport jax\nimport jax.numpy as jnp\nimport flax.linen as nn\nimport optax\nfrom flax.training import train_state\nimport matplotlib.pyplot as plt\nfrom PIL import Image\nprint(\"jax\", jax.__version__, \"| device:\", jax.devices()[0].device_kind,\n     f\"({jax.devices()[0].platform})\")\nprint(\"jax3d volume_rendering API:\",\n     [n for n in (\"sample_along_rays\", \"volume_rendering\",\n                  \"sample_piecewise_constant_pdf\", \"sample_1d\")\n      if hasattr(j3vr, n)])\n@dataclasses.dataclass\nclass Config:\n   H: int = 64;            W: int = 64\n   n_train_views: int = 24; n_test_views: int = 3\n   cam_radius: float = 3.2; fov_deg: float = 40.0\n   near: float = 1.9;       far: float = 4.7\n   gt_samples: int = 256\n   n_coarse: int = 64;      n_fine: int = 64\n   deg_pos: int = 10;       deg_dir: int = 4\n   width: int = 128;        depth: int = 6;   skip: int = 3\n   batch_rays: int = 2048;  steps: int = 2500\n   lr_init: float = 5e-4;   lr_final: float = 5e-6\n   chunk: int = 4096\n   grid_res: int = 96\ncfg = Config()\nif jax.devices()[0].platform == \"cpu\":\n   print(\"\\n!! No GPU detected -- switching to a small CPU-friendly config.\")\n   print(\"   (Runtime > Change runtime type > T4 GPU for the full version.)\\n\")\n   cfg = dataclasses.replace(cfg, H=40, W=40, n_train_views=14, steps=400,\n                             gt_samples=128, n_coarse=32, n_fine=32,\n                             width=64, depth=4, skip=2, batch_rays=1024,\n                             chunk=1600, grid_res=64)\ndef _normalize(v, axis=-1):\n   return v / (np.linalg.norm(v, axis=axis, keepdims=True) + 1e-9)\ndef look_at(eye, target=(0., 0., 0.), up=(0., 0., 1.)):\n   \"\"\"OpenGL/NeRF convention camera-to-world: +x right, +y up, camera looks at -z.\"\"\"\n   eye, target, up = map(lambda a: np.asarray(a, np.float32), (eye, target, up))\n   fwd   = _normalize(target - eye)\n   right = _normalize(np.cross(fwd, up))\n   trueup = np.cross(right, fwd)\n   c2w = np.eye(4, dtype=np.float32)\n   c2w[:3, :3] = np.stack([right, trueup, -fwd], axis=1)\n   c2w[:3, 3] = eye\n   return c2w\ndef orbit_poses(n, radius, elev_lo=18., elev_hi=58., phase=0.0):\n   \"\"\"Golden-angle azimuths + monotone elevations => well-spread views on a dome.\"\"\"\n   i = np.arange(n, dtype=np.float64) + 0.5\n   az = 2 * np.pi * ((i * 0.6180339887) + phase)\n   elev = np.arcsin(np.linspace(np.sin(np.deg2rad(elev_lo)),\n                                np.sin(np.deg2rad(elev_hi)), n))\n   eyes = np.stack([radius * np.cos(elev) * np.cos(az),\n                    radius * np.cos(elev) * np.sin(az),\n                    radius * np.sin(elev)], axis=-1).astype(np.float32)\n   return np.stack([look_at(e) for e in eyes], axis=0)\ndef rays_from_pose(c2w, H, W, focal):\n   \"\"\"Returns (origins, dirs) of shape [H, W, 3]; dirs are unit-length, so the\n   depths returned by jax3d's sampler are true world-space distances.\"\"\"\n   i, j = np.meshgrid(np.arange(W, dtype=np.float32),\n                      np.arange(H, dtype=np.float32), indexing=\"xy\")\n   cam_dirs = np.stack([(i - W * .5 + .5) / focal,\n                        -(j - H * .5 + .5) / focal,\n                        -np.ones_like(i)], axis=-1)\n   dirs = _normalize(cam_dirs @ c2w[:3, :3].T)\n   origins = np.broadcast_to(c2w[:3, 3], dirs.shape)\n   return origins.astype(np.float32).copy(), dirs.astype(np.float32)\nFOCAL = 0.5 * cfg.W / math.tan(0.5 * math.radians(cfg.fov_deg))\n```\n\nWe set up the JAX3D environment, install the required dependencies, and load the volume_rendering module directly from the cloned repository. We configure GPU/CPU-adaptive training parameters and establish the camera model using pinhole intrinsics, look-at poses, and orbit-based camera placement. We then generate normalized world-space rays from each camera pose, providing the geometric foundation for the rendering pipeline.\n\n```\nLIGHT = jnp.asarray(_normalize(np.array([0.55, 0.75, 0.85], np.float32)))\n_SPHERES = [\n   (jnp.array([0.34, 0.02, -0.22]), 0.36, jnp.array([0.90, 0.24, 0.22])),\n   (jnp.array([-0.32, 0.28, 0.05]), 0.26, jnp.array([0.25, 0.78, 0.36])),\n   (jnp.array([-0.05, -0.36, 0.24]), 0.22, jnp.array([0.28, 0.40, 0.95])),\n]\ndef _sphere_field(pos, vdir, center, radius, albedo):\n   d = pos - center\n   dist = jnp.linalg.norm(d, axis=-1)\n   n = d / (dist[..., None] + 1e-8)\n   sigma = 80.0 * jax.nn.sigmoid((radius - dist) / 0.015)\n   v = -vdir\n   refl = 2.0 * jnp.sum(n * v, -1, keepdims=True) * n - v\n   spec = 0.65 * jnp.clip(jnp.sum(refl * LIGHT, -1), 0., 1.) ** 24\n   lamb = 0.35 + 0.65 * jnp.clip(jnp.sum(n * LIGHT, -1), 0., 1.)\n   rgb = jnp.clip(albedo * lamb[..., None] + spec[..., None], 0., 1.)\n   return sigma, rgb\ndef _floor_field(pos):\n   x, y, z = pos[..., 0], pos[..., 1], pos[..., 2]\n   m = (jax.nn.sigmoid((0.06 - jnp.abs(z + 0.62)) / 0.008)\n        * jax.nn.sigmoid((0.85 - jnp.abs(x)) / 0.01)\n        * jax.nn.sigmoid((0.85 - jnp.abs(y)) / 0.01))\n   checker = (jnp.floor(x * 3.0) + jnp.floor(y * 3.0)) % 2.0\n   rgb = jnp.where(checker[..., None] > 0.5,\n                   jnp.array([0.86, 0.86, 0.89]), jnp.array([0.22, 0.25, 0.30]))\n   return 80.0 * m, rgb\ndef gt_field(pos, vdir):\n   \"\"\"pos, vdir: [..., 3] -> (sigma [...], rgb [..., 3]). Density-weighted blend.\"\"\"\n   sig_sum = 0.0\n   col_sum = 0.0\n   for c, r, a in _SPHERES:\n       s, rgb = _sphere_field(pos, vdir, c, r, a)\n       sig_sum = sig_sum + s\n       col_sum = col_sum + s[..., None] * rgb\n   s, rgb = _floor_field(pos)\n   sig_sum = sig_sum + s\n   col_sum = col_sum + s[..., None] * rgb\n   return sig_sum, col_sum / (sig_sum[..., None] + 1e-8)\nWHITE_BG = jnp.ones((3,), jnp.float32)\n@jax.jit\ndef render_ground_truth(origins, dirs):\n   \"\"\"Fine-grained volumetric render of the analytic scene -> RGB + depth.\"\"\"\n   depths, positions = j3vr.sample_along_rays(\n       ray_origins=origins, ray_directions=dirs,\n       near=cfg.near, far=cfg.far,\n       sample_count=cfg.gt_samples, deterministic=True)\n   vdir = jnp.broadcast_to(dirs[..., None, :], positions.shape)\n   sigma, rgb = gt_field(positions, vdir)\n   out = j3vr.volume_rendering(\n       sample_values={\"rgb\": rgb}, sample_density=sigma, depths=depths,\n       background_values={\"rgb\": WHITE_BG})\n   return out.ray_values[\"rgb\"], out.ray_depth, out.ray_alpha\ndef build_dataset(poses):\n   O, D, C = [], [], []\n   for c2w in poses:\n       o, d = rays_from_pose(c2w, cfg.H, cfg.W, FOCAL)\n       rgb, _, _ = render_ground_truth(jnp.asarray(o), jnp.asarray(d))\n       O.append(o); D.append(d); C.append(np.asarray(rgb))\n   return (np.stack(O), np.stack(D), np.stack(C))\nprint(\"\\nRendering the synthetic multi-view dataset ...\")\nt0 = time.time()\ntrain_poses = orbit_poses(cfg.n_train_views, cfg.cam_radius, phase=0.00)\ntest_poses  = orbit_poses(cfg.n_test_views,  cfg.cam_radius, 26., 50., phase=0.41)\ntr_o, tr_d, tr_c = build_dataset(train_poses)\nte_o, te_d, te_c = build_dataset(test_poses)\nprint(f\"  {cfg.n_train_views} train + {cfg.n_test_views} test views \"\n     f\"at {cfg.H}x{cfg.W}  ({time.time()-t0:.1f}s)\")\nk = min(8, cfg.n_train_views)\nfig, axes = plt.subplots(1, k, figsize=(2 * k, 2.3))\nfor a, im, p in zip(axes, tr_c[:k], train_poses[:k]):\n   a.imshow(np.clip(im, 0, 1)); a.axis(\"off\")\n   a.set_title(f\"({p[0,3]:+.1f},{p[1,3]:+.1f},{p[2,3]:+.1f})\", fontsize=7)\nfig.suptitle(\"Training views (ground truth, rendered with jax3d.math.volume_rendering)\",\n            fontsize=11); plt.tight_layout(); plt.show()\nrays_o = jnp.asarray(tr_o.reshape(-1, 3))\nrays_d = jnp.asarray(tr_d.reshape(-1, 3))\nrays_c = jnp.asarray(tr_c.reshape(-1, 3))\nN_RAYS = rays_o.shape[0]\nprint(f\"  ray pool: {N_RAYS:,} rays\")\n```\n\nWe construct an analytic ground-truth scene containing soft-edged spheres, a patterned floor, and view-dependent specular radiance. We render this scene with JAX3D’s volume-rendering implementation to generate consistent RGB observations, depths, and opacity values across multiple camera views. We organize the resulting images into a flattened ray pool so that we can efficiently sample random rays during NeRF training.\n\n``` python\ndef posenc(x, deg):\n   \"\"\"NeRF sinusoidal encoding, with the raw input concatenated.\"\"\"\n   if deg == 0:\n       return x\n   scales = 2.0 ** jnp.arange(deg, dtype=x.dtype)\n   xb = (x[..., None, :] * scales[:, None]).reshape(*x.shape[:-1], -1)\n   return jnp.concatenate([x, jnp.sin(xb), jnp.cos(xb)], axis=-1)\nclass NeRFMLP(nn.Module):\n   width: int; depth: int; skip: int; deg_pos: int; deg_dir: int\n   @nn.compact\n   def __call__(self, pos, dirs):\n       inp = posenc(pos, self.deg_pos)\n       x = inp\n       for i in range(self.depth):\n           x = nn.relu(nn.Dense(self.width)(x))\n           if i == self.skip:\n               x = jnp.concatenate([x, inp], axis=-1)\n       sigma = nn.softplus(nn.Dense(1)(x)[..., 0] - 1.0)\n       h = jnp.concatenate([nn.Dense(self.width)(x), posenc(dirs, self.deg_dir)], -1)\n       rgb = nn.sigmoid(nn.Dense(3)(nn.relu(nn.Dense(self.width // 2)(h))))\n       return sigma, rgb\nmodel = NeRFMLP(cfg.width, cfg.depth, cfg.skip, cfg.deg_pos, cfg.deg_dir)\n```\n\nWe implement the NeRF representation using sinusoidal positional encoding for both spatial coordinates and viewing directions. We use a deep Flax MLP with a skip connection to predict non-negative volumetric density from position while conditioning RGB on the viewing direction. We therefore separate view-independent geometry from view-dependent appearance, allowing the model to represent both scene structure and specular effects.\n\n``` php\ndef render_rays(params, origins, dirs, rng, deterministic):\n   \"\"\"Coarse pass -> importance-resample -> fine pass. All sampling and\n   compositing comes from jax3d.math.volume_rendering.\"\"\"\n   rng_c, rng_f = jax.random.split(rng)\n   depths_c, pos_c = j3vr.sample_along_rays(\n       ray_origins=origins, ray_directions=dirs,\n       near=cfg.near, far=cfg.far, sample_count=cfg.n_coarse,\n       deterministic=deterministic, rng=rng_c)\n   dirs_c = jnp.broadcast_to(dirs[:, None, :], pos_c.shape)\n   sigma_c, rgb_c = model.apply(params[\"coarse\"], pos_c, dirs_c)\n   out_c = j3vr.volume_rendering(\n       sample_values={\"rgb\": rgb_c}, sample_density=sigma_c, depths=depths_c,\n       background_values={\"rgb\": WHITE_BG})\n   mid = 0.5 * (depths_c[..., 1:] + depths_c[..., :-1])\n   bin_edges = jnp.concatenate([depths_c[..., :1], mid, depths_c[..., -1:]], -1)\n   t_fine = j3vr.sample_piecewise_constant_pdf(\n       bin_edges=bin_edges, weights=out_c.sample_weights,\n       sample_count=cfg.n_fine, deterministic=deterministic, rng=rng_f)\n   t_fine = jax.lax.stop_gradient(t_fine)\n   depths_f = jnp.sort(jnp.concatenate([depths_c, t_fine], -1), axis=-1)\n   pos_f = origins[:, None, :] + depths_f[..., None] * dirs[:, None, :]\n   dirs_f = jnp.broadcast_to(dirs[:, None, :], pos_f.shape)\n   sigma_f, rgb_f = model.apply(params[\"fine\"], pos_f, dirs_f)\n   out_f = j3vr.volume_rendering(\n       sample_values={\"rgb\": rgb_f}, sample_density=sigma_f, depths=depths_f,\n       background_values={\"rgb\": WHITE_BG})\n   aux = {\"depths_c\": depths_c, \"weights_c\": out_c.sample_weights, \"t_fine\": t_fine}\n   return out_c, out_f, aux\ndef mse_to_psnr(x):\n   return -10.0 * jnp.log10(jnp.maximum(x, 1e-10))\n```\n\nWe implement the core hierarchical renderer by first sampling coarse points along each ray and compositing their densities and colors through JAX3D’s volume-rendering operator. We convert the resulting coarse rendering weights into a piecewise-constant probability distribution and importance-sample additional fine points around high-contribution regions. We combine and sort the coarse and fine samples before performing the final fine-network rendering, while stopping gradients through the sampling operation.\n\n```\nkey = jax.random.PRNGKey(0)\nkey, k1, k2 = jax.random.split(key, 3)\ndummy_p = jnp.zeros((1, 1, 3)); dummy_d = jnp.zeros((1, 1, 3))\nparams = {\"coarse\": model.init(k1, dummy_p, dummy_d),\n         \"fine\":   model.init(k2, dummy_p, dummy_d)}\nn_params = sum(x.size for x in jax.tree.leaves(params))\nprint(f\"\\nModel: {n_params/1e6:.2f}M parameters (coarse + fine networks)\")\nschedule = optax.exponential_decay(cfg.lr_init, cfg.steps,\n                                  cfg.lr_final / cfg.lr_init)\ntx = optax.chain(optax.clip_by_global_norm(1.0), optax.adam(schedule))\nstate = train_state.TrainState.create(apply_fn=model.apply, params=params, tx=tx)\n@jax.jit\ndef train_step(state, o, d, target, rng):\n   def loss_fn(p):\n       out_c, out_f, _ = render_rays(p, o, d, rng, deterministic=False)\n       l_c = jnp.mean((out_c.ray_values[\"rgb\"] - target) ** 2)\n       l_f = jnp.mean((out_f.ray_values[\"rgb\"] - target) ** 2)\n       return l_c + l_f, l_f\n   (loss, l_fine), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)\n   return state.apply_gradients(grads=grads), loss, l_fine\nprint(f\"Training {cfg.steps} steps x {cfg.batch_rays} rays \"\n     f\"({cfg.n_coarse} coarse + {cfg.n_coarse + cfg.n_fine} fine samples/ray) ...\")\nhistory = []\nt0 = time.time()\nfor step in range(1, cfg.steps + 1):\n   key, k_idx, k_render = jax.random.split(key, 3)\n   idx = jax.random.randint(k_idx, (cfg.batch_rays,), 0, N_RAYS)\n   state, loss, l_fine = train_step(state, rays_o[idx], rays_d[idx],\n                                    rays_c[idx], k_render)\n   if step % 25 == 0 or step == 1:\n       history.append((step, float(mse_to_psnr(l_fine))))\n   if step % max(1, cfg.steps // 10) == 0 or step == 1:\n       print(f\"  step {step:5d}/{cfg.steps} | loss {float(loss):.5f} \"\n             f\"| train PSNR {float(mse_to_psnr(l_fine)):5.2f} dB \"\n             f\"| {time.time()-t0:6.1f}s\")\nprint(f\"Done in {time.time()-t0:.1f}s\")\n```\n\nWe initialize independent coarse and fine NeRF networks and optimize them jointly with Adam, using exponential learning-rate decay and global gradient clipping. We supervise both rendering stages against ground-truth ray colors, encouraging the coarse network to learn useful sampling distributions while improving the final fine reconstruction. We run the training step with JAX JIT compilation and monitor the fine-network PSNR throughout optimization.\n\n``` python\n@jax.jit\ndef render_chunk(params, o, d, rng):\n   _, out_f, aux = render_rays(params, o, d, rng, deterministic=True)\n   return out_f.ray_values[\"rgb\"], out_f.ray_depth, out_f.ray_alpha, aux\ndef render_image(params, origins, dirs, rng):\n   \"\"\"Chunked full-image render with padding, so only one shape gets compiled.\"\"\"\n   o = jnp.asarray(origins.reshape(-1, 3)); d = jnp.asarray(dirs.reshape(-1, 3))\n   R = o.shape[0]; rgb, dep, alp = [], [], []\n   for i in range(0, R, cfg.chunk):\n       oc, dc = o[i:i + cfg.chunk], d[i:i + cfg.chunk]\n       pad = cfg.chunk - oc.shape[0]\n       if pad:\n           oc = jnp.concatenate([oc, jnp.tile(oc[-1:], (pad, 1))], 0)\n           dc = jnp.concatenate([dc, jnp.tile(dc[-1:], (pad, 1))], 0)\n       c, dp, a, _ = render_chunk(params, oc, dc, rng)\n       n = cfg.chunk - pad\n       rgb.append(c[:n]); dep.append(dp[:n]); alp.append(a[:n])\n   s = (cfg.H, cfg.W)\n   return (np.asarray(jnp.concatenate(rgb)).reshape(*s, 3),\n           np.asarray(jnp.concatenate(dep)).reshape(*s),\n           np.asarray(jnp.concatenate(alp)).reshape(*s))\nh = np.array(history)\nplt.figure(figsize=(6, 3))\nplt.plot(h[:, 0], h[:, 1], lw=1.6)\nplt.xlabel(\"step\"); plt.ylabel(\"train PSNR (dB)\")\nplt.title(\"Fine-network training PSNR\"); plt.grid(alpha=.3)\nplt.tight_layout(); plt.show()\nprint(\"\\nRendering held-out test views ...\")\nkey, k_eval = jax.random.split(key)\npsnrs = []\nfig, axes = plt.subplots(cfg.n_test_views, 4,\n                        figsize=(11, 2.7 * cfg.n_test_views), squeeze=False)\nfor v in range(cfg.n_test_views):\n   pred, depth, alpha = render_image(state.params, te_o[v], te_d[v], k_eval)\n   p = float(mse_to_psnr(np.mean((pred - te_c[v]) ** 2))); psnrs.append(p)\n   depth_vis = depth + (1.0 - alpha) * cfg.far\n   for a, (im, ttl, kw) in zip(axes[v], [\n           (np.clip(te_c[v], 0, 1), \"ground truth\", {}),\n           (np.clip(pred, 0, 1), f\"NeRF  ({p:.2f} dB)\", {}),\n           (depth_vis, \"depth (ray_depth)\", dict(cmap=\"turbo\",\n                                                 vmin=cfg.near, vmax=cfg.far)),\n           (alpha, \"opacity (ray_alpha)\", dict(cmap=\"gray\", vmin=0, vmax=1))]):\n       a.imshow(im, **kw); a.set_title(ttl, fontsize=9); a.axis(\"off\")\nplt.suptitle(f\"Novel-view synthesis   |   mean PSNR = {np.mean(psnrs):.2f} dB\",\n            fontsize=12)\nplt.tight_layout(); plt.show()\nprint(f\"  mean held-out PSNR: {np.mean(psnrs):.2f} dB\")\ncy, cx = cfg.H // 2, cfg.W // 2\no1 = jnp.asarray(te_o[0][cy, cx])[None]; d1 = jnp.asarray(te_d[0][cy, cx])[None]\no1 = jnp.tile(o1, (cfg.chunk, 1)); d1 = jnp.tile(d1, (cfg.chunk, 1))\n_, _, _, aux = render_chunk(state.params, o1, d1, k_eval)\ndc = np.asarray(aux[\"depths_c\"][0]); wc = np.asarray(aux[\"weights_c\"][0])\ntf = np.asarray(aux[\"t_fine\"][0])\nfig, ax = plt.subplots(figsize=(8, 3))\nax.bar(dc, wc, width=(cfg.far - cfg.near) / cfg.n_coarse * .9,\n      alpha=.55, label=\"coarse weights (the PDF)\")\nax.plot(tf, np.full_like(tf, wc.max() * .06), \"|\", ms=16, color=\"crimson\",\n       label=\"fine samples (sample_piecewise_constant_pdf)\")\nax.set_xlabel(\"depth along ray\"); ax.set_ylabel(\"weight\")\nax.set_title(\"Importance resampling concentrates samples on the surface\")\nax.legend(fontsize=8); plt.tight_layout(); plt.show()\nprint(\"\\nRendering 360-degree orbit ...\")\nn_frames = 24 if jax.devices()[0].platform != \"cpu\" else 8\nframes = []\nfor t in range(n_frames):\n   az = 2 * np.pi * t / n_frames; el = np.deg2rad(32.0)\n   eye = cfg.cam_radius * np.array([np.cos(el) * np.cos(az),\n                                    np.cos(el) * np.sin(az), np.sin(el)])\n   o, d = rays_from_pose(look_at(eye), cfg.H, cfg.W, FOCAL)\n   rgb, _, _ = render_image(state.params, o, d, k_eval)\n   frames.append((np.clip(rgb, 0, 1) * 255).astype(np.uint8))\ngif_path = os.path.join(os.getcwd(), \"nerf_orbit.gif\")\npil = [Image.fromarray(f).resize((cfg.W * 3, cfg.H * 3), Image.NEAREST) for f in frames]\npil[0].save(gif_path, save_all=True, append_images=pil[1:], duration=90, loop=0)\ntry:\n   from IPython.display import Image as IPImage, display\n   display(IPImage(filename=gif_path))\nexcept Exception:\n   pass\nprint(\"  saved\", gif_path)\nprint(\"\\nExtracting isosurface from the learned density field ...\")\ntry:\n   from skimage import measure\n   g = np.linspace(-1.0, 1.0, cfg.grid_res, dtype=np.float32)\n   X, Y, Z = np.meshgrid(g, g, g, indexing=\"ij\")\n   pts = np.stack([X, Y, Z], -1).reshape(-1, 3)\n   @jax.jit\n   def density_at(p):\n       s, _ = model.apply(state.params[\"fine\"], p, jnp.zeros_like(p))\n       return s\n   vol = np.concatenate([np.asarray(density_at(jnp.asarray(pts[i:i + 65536])))\n                         for i in range(0, pts.shape[0], 65536)])\n   vol = vol.reshape(cfg.grid_res, cfg.grid_res, cfg.grid_res)\n   step = (cfg.far - cfg.near) / (cfg.n_coarse + cfg.n_fine)\n   level = float(-np.log(0.5) / step)\n   if not (vol.min() < level < vol.max()):\n       level = float(np.percentile(vol, 99.0))\n   verts, faces, _, _ = measure.marching_cubes(vol, level=level)\n   verts = -1.0 + verts * (2.0 / (cfg.grid_res - 1))\n   fig = plt.figure(figsize=(6, 6)); ax = fig.add_subplot(111, projection=\"3d\")\n   ax.plot_trisurf(verts[:, 0], verts[:, 1], verts[:, 2], triangles=faces,\n                   cmap=\"viridis\", lw=0.0, antialiased=False, alpha=.95)\n   ax.set_box_aspect((1, 1, 1))\n   ax.set_xlim(-1, 1); ax.set_ylim(-1, 1); ax.set_zlim(-1, 1)\n   ax.view_init(elev=24, azim=-58)\n   ax.set_title(f\"Marching cubes on learned density  (sigma = {level:.1f}, \"\n                f\"{len(faces):,} faces)\", fontsize=10)\n   plt.tight_layout(); plt.show()\nexcept Exception as e:\n   print(\"  isosurface step skipped:\", e)\nprint(\"\\n\" + \"=\" * 70)\nprint(f\"FINAL held-out PSNR: {np.mean(psnrs):.2f} dB   ({n_params/1e6:.2f}M params, \"\n     f\"{cfg.steps} steps)\")\nprint(\"jax3d functions exercised: sample_along_rays, volume_rendering, \"\n     \"sample_piecewise_constant_pdf\")\nprint(\"=\" * 70)\n```\n\nWe evaluate the trained representation through chunked novel-view rendering and measure reconstruction quality with held-out PSNR, along with depth and opacity maps. We visualize how hierarchical sampling concentrates fine samples around important surfaces, then generate a 360-degree orbit GIF to inspect the learned radiance field from multiple viewpoints. We finally query the learned density on a 3D grid and apply marching cubes to extract an approximate geometric isosurface.\n\nIn conclusion, we demonstrated the complete inverse-rendering pipeline by learning a continuous density and radiance field from synthetic multi-view observations and reconstructing it through hierarchical volume rendering. We used the coarse network to identify informative regions along each ray and the fine network to concentrate additional samples around high-contribution surfaces. At the same time, view-direction encoding allows us to model view-dependent appearance. In the final evaluation stages, we measured novel-view reconstruction quality with PSNR, inspected learned depth and opacity, visualized importance-sampling behavior, generated a 360-degree orbit, and extracted an approximate learned geometry with marching cubes. Overall, we showed how the mathematical components of jax3d integrate with modern JAX-based neural-network training to form a compact yet technically complete NeRF reconstruction system.\n\nCheck out **[the FULL CODES here](https://github.com/MARKTECHPOST-AI-MEDIA-INC/AI-Agents-Projects-Tutorials/blob/main/Computer%20Vision/jax3d_hierarchical_nerf_tutorial_Marktechpost.ipynb)**. All credit goes to the researcher of this project. Also, feel free to follow us on **[Twitter](https://x.com/intent/follow?screen_name=marktechpost)** and don’t forget to join our **[150k+ML SubReddit](https://www.reddit.com/r/machinelearningnews/)** and Subscribe to **[our Newsletter](https://magic.beehiiv.com/v1/f5e63dd4-5653-4f09-83e2-321a8b1ba526?email={{email}})**. Wait! are you on telegram? [now you can join us on telegram as well.](https://t.me/machinelearningresearchnews)\n\nNeed to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? [Connect with us](https://forms.gle/wbash1wF6efRj8G58)", "url": "https://wpnews.pro/news/hierarchical-nerf-with-jax3d-for-volumetric-rendering-novel-view-synthesis-and", "canonical_source": "https://www.marktechpost.com/2026/09/13/hierarchical-nerf-with-jax3d-for-volumetric-rendering-novel-view-synthesis-and-3d-reconstruction/", "published_at": "2026-09-13 19:46:38+00:00", "updated_at": "2026-09-13 19:50:38.022618+00:00", "lang": "en", "topics": ["artificial-intelligence", "computer-vision", "neural-networks", "ai-research", "developer-tools"], "entities": ["JAX", "jax3d", "Flax", "Optax", "Google Research", "NeRF", "Adam", "PSNR"], "alternates": {"html": "https://wpnews.pro/news/hierarchical-nerf-with-jax3d-for-volumetric-rendering-novel-view-synthesis-and", "markdown": "https://wpnews.pro/news/hierarchical-nerf-with-jax3d-for-volumetric-rendering-novel-view-synthesis-and.md", "text": "https://wpnews.pro/news/hierarchical-nerf-with-jax3d-for-volumetric-rendering-novel-view-synthesis-and.txt", "jsonld": "https://wpnews.pro/news/hierarchical-nerf-with-jax3d-for-volumetric-rendering-novel-view-synthesis-and.jsonld"}}