@jax.jit
def render_chunk(params, o, d, rng):
_, out_f, aux = render_rays(params, o, d, rng, deterministic=True)
return out_f.ray_values["rgb"], out_f.ray_depth, out_f.ray_alpha, aux
def render_image(params, origins, dirs, rng):
"""Chunked full-image render with padding, so only one shape gets compiled."""
o = jnp.asarray(origins.reshape(-1, 3)); d = jnp.asarray(dirs.reshape(-1, 3))
R = o.shape[0]; rgb, dep, alp = [], [], []
for i in range(0, R, cfg.chunk):
oc, dc = o[i:i + cfg.chunk], d[i:i + cfg.chunk]
pad = cfg.chunk - oc.shape[0]
if pad:
oc = jnp.concatenate([oc, jnp.tile(oc[-1:], (pad, 1))], 0)
dc = jnp.concatenate([dc, jnp.tile(dc[-1:], (pad, 1))], 0)
c, dp, a, _ = render_chunk(params, oc, dc, rng)
n = cfg.chunk - pad
rgb.append(c[:n]); dep.append(dp[:n]); alp.append(a[:n])
s = (cfg.H, cfg.W)
return (np.asarray(jnp.concatenate(rgb)).reshape(*s, 3),
np.asarray(jnp.concatenate(dep)).reshape(*s),
np.asarray(jnp.concatenate(alp)).reshape(*s))
h = np.array(history)
plt.figure(figsize=(6, 3))
plt.plot(h[:, 0], h[:, 1], lw=1.6)
plt.xlabel("step"); plt.ylabel("train PSNR (dB)")
plt.title("Fine-network training PSNR"); plt.grid(alpha=.3)
plt.tight_layout(); plt.show()
print("\nRendering held-out test views ...")
key, k_eval = jax.random.split(key)
psnrs = []
fig, axes = plt.subplots(cfg.n_test_views, 4,
figsize=(11, 2.7 * cfg.n_test_views), squeeze=False)
for v in range(cfg.n_test_views):
pred, depth, alpha = render_image(state.params, te_o[v], te_d[v], k_eval)
p = float(mse_to_psnr(np.mean((pred - te_c[v]) ** 2))); psnrs.append(p)
depth_vis = depth + (1.0 - alpha) * cfg.far
for a, (im, ttl, kw) in zip(axes[v], [
(np.clip(te_c[v], 0, 1), "ground truth", {}),
(np.clip(pred, 0, 1), f"NeRF ({p:.2f} dB)", {}),
(depth_vis, "depth (ray_depth)", dict(cmap="turbo",
vmin=cfg.near, vmax=cfg.far)),
(alpha, "opacity (ray_alpha)", dict(cmap="gray", vmin=0, vmax=1))]):
a.imshow(im, **kw); a.set_title(ttl, fontsize=9); a.axis("off")
plt.suptitle(f"Novel-view synthesis | mean PSNR = {np.mean(psnrs):.2f} dB",
fontsize=12)
plt.tight_layout(); plt.show()
print(f" mean held-out PSNR: {np.mean(psnrs):.2f} dB")
cy, cx = cfg.H // 2, cfg.W // 2
o1 = jnp.asarray(te_o[0][cy, cx])[None]; d1 = jnp.asarray(te_d[0][cy, cx])[None]
o1 = jnp.tile(o1, (cfg.chunk, 1)); d1 = jnp.tile(d1, (cfg.chunk, 1))
_, _, _, aux = render_chunk(state.params, o1, d1, k_eval)
dc = np.asarray(aux["depths_c"][0]); wc = np.asarray(aux["weights_c"][0])
tf = np.asarray(aux["t_fine"][0])
fig, ax = plt.subplots(figsize=(8, 3))
ax.bar(dc, wc, width=(cfg.far - cfg.near) / cfg.n_coarse * .9,
alpha=.55, label="coarse weights (the PDF)")
ax.plot(tf, np.full_like(tf, wc.max() * .06), "|", ms=16, color="crimson",
label="fine samples (sample_piecewise_constant_pdf)")
ax.set_xlabel("depth along ray"); ax.set_ylabel("weight")
ax.set_title("Importance resampling concentrates samples on the surface")
ax.legend(fontsize=8); plt.tight_layout(); plt.show()
print("\nRendering 360-degree orbit ...")
n_frames = 24 if jax.devices()[0].platform != "cpu" else 8
frames = []
for t in range(n_frames):
az = 2 * np.pi * t / n_frames; el = np.deg2rad(32.0)
eye = cfg.cam_radius * np.array([np.cos(el) * np.cos(az),
np.cos(el) * np.sin(az), np.sin(el)])
o, d = rays_from_pose(look_at(eye), cfg.H, cfg.W, FOCAL)
rgb, _, _ = render_image(state.params, o, d, k_eval)
frames.append((np.clip(rgb, 0, 1) * 255).astype(np.uint8))
gif_path = os.path.join(os.getcwd(), "nerf_orbit.gif")
pil = [Image.fromarray(f).resize((cfg.W * 3, cfg.H * 3), Image.NEAREST) for f in frames]
pil[0].save(gif_path, save_all=True, append_images=pil[1:], duration=90, loop=0)
try:
from IPython.display import Image as IPImage, display
display(IPImage(filename=gif_path))
except Exception:
pass
print(" saved", gif_path)
print("\nExtracting isosurface from the learned density field ...")
try:
from skimage import measure
g = np.linspace(-1.0, 1.0, cfg.grid_res, dtype=np.float32)
X, Y, Z = np.meshgrid(g, g, g, indexing="ij")
pts = np.stack([X, Y, Z], -1).reshape(-1, 3)
@jax.jit
def density_at(p):
s, _ = model.apply(state.params["fine"], p, jnp.zeros_like(p))
return s
vol = np.concatenate([np.asarray(density_at(jnp.asarray(pts[i:i + 65536])))
for i in range(0, pts.shape[0], 65536)])
vol = vol.reshape(cfg.grid_res, cfg.grid_res, cfg.grid_res)
step = (cfg.far - cfg.near) / (cfg.n_coarse + cfg.n_fine)
level = float(-np.log(0.5) / step)
if not (vol.min() < level < vol.max()):
level = float(np.percentile(vol, 99.0))
verts, faces, _, _ = measure.marching_cubes(vol, level=level)
verts = -1.0 + verts * (2.0 / (cfg.grid_res - 1))
fig = plt.figure(figsize=(6, 6)); ax = fig.add_subplot(111, projection="3d")
ax.plot_trisurf(verts[:, 0], verts[:, 1], verts[:, 2], triangles=faces,
cmap="viridis", lw=0.0, antialiased=False, alpha=.95)
ax.set_box_aspect((1, 1, 1))
ax.set_xlim(-1, 1); ax.set_ylim(-1, 1); ax.set_zlim(-1, 1)
ax.view_init(elev=24, azim=-58)
ax.set_title(f"Marching cubes on learned density (sigma = {level:.1f}, "
f"{len(faces):,} faces)", fontsize=10)
plt.tight_layout(); plt.show()
except Exception as e:
print(" isosurface step skipped:", e)
print("\n" + "=" * 70)
print(f"FINAL held-out PSNR: {np.mean(psnrs):.2f} dB ({n_params/1e6:.2f}M params, "
f"{cfg.steps} steps)")
print("jax3d functions exercised: sample_along_rays, volume_rendering, "
"sample_piecewise_constant_pdf")
print("=" * 70)
