使用 JAX3D 的分層神經輻射場：體積渲染、新視角合成與 3D 重建

In this tutorial, we build an end-to-end hierarchical Neural Radiance Field (NeRF) using JAX, 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 samplealongrays and volumerendering 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 samplepiecewiseconstantpdf.

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.

Copy CodeCopiedUse a different Browserimport os, sys, subprocess, importlib.util, functools, dataclasses, time, math def sh(cmd): subprocess.run(cmd, shell=True, check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL) print("Installing dependencies ...") sh(f'{sys.

executable} -m pip install -q "etils[array-types,epy,etree,enp]" ' f'chex flax optax scikit-image') REPODIR = "/content/jax3d" if os.path.isdir("/content") else os.path.abspath("./jax3d") if not os.path.isdir(REPODIR): print("Cloning google-research/jax3d ...

") sh(f"git clone -q --depth 1 https://github.com/google-research/jax3d.git {REPODIR}") def loadmodulebypath(name, path): """Load a single .py file without triggering the parent package init.from jax3d.math import volumerendering also works if you run pip install .

inside the clone, but that pulls in gin/tfds/etc.""" spec = importlib.util.specfromfilelocation(name, path) mod = importlib.util.modulefromspec(spec) sys.modules[name] = mod spec.loader.execmodule(mod) return mod VRPATH = os.path.join(REPODIR, "jax3d", "jax3d", "math", "volumerendering.

py") if not os.path.exists(VRPATH): VRPATH = os.path.join(REPODIR, "jax3d", "math", "volumerendering.py") try: j3vr = loadmodulebypath("j3dvolumerendering", VRPATH) except Exception as e: raise SystemExit( f"Could not load {VRPATH}: {e}\n" "Try: pip install -U 'etils[array-types,epy,etree,enp]==1.9.

4' and re-run." ) import numpy as np import jax import jax.numpy as jnp import flax.linen as nn import optax from flax.training import trainstate import matplotlib.pyplot as plt from PIL import Image print("jax", jax.version, "| device:", jax.devices()[0].devicekind, f"({jax.devices()[0].

platform})") print("jax3d volumerendering API:", [n for n in ("samplealongrays", "volumerendering", "samplepiecewiseconstantpdf", "sample1d") if hasattr(j3vr, n)]) @dataclasses.dataclass class Config: H: int = 64; W: int = 64 ntrainviews: int = 24; ntestviews: int = 3 camradius: float = 3.

2; fovdeg: float = 40.0 near: float = 1.9; far: float = 4.

7 gtsamples: int = 256 ncoarse: int = 64; nfine: int = 64 degpos: int = 10; degdir: int = 4 width: int = 128; depth: int = 6; skip: int = 3 batchrays: int = 2048; steps: int = 2500 lrinit: float = 5e-4; lrfinal: float = 5e-6 chunk: int = 4096 gridres: int = 96 cfg = Config() if jax.devices()[0].

platform == "cpu": print("\n!!No GPU detected -- switching to a small CPU-friendly config.") print(" (Runtime > Change runtime type > T4 GPU for the full version.)\n") cfg = dataclasses.

replace(cfg, H=40, W=40, ntrainviews=14, steps=400, gtsamples=128, ncoarse=32, nfine=32, width=64, depth=4, skip=2, batchrays=1024, chunk=1600, gridres=64) def normalize(v, axis=-1): return v / (np.linalg.norm(v, axis=axis, keepdims=True) + 1e-9) def lookat(eye, target=(0., 0., 0.), up=(0., 0., 1.

)): """OpenGL/NeRF convention camera-to-world: +x right, +y up, camera looks at -z.""" eye, target, up = map(lambda a: np.asarray(a, np.float32), (eye, target, up)) fwd = normalize(target - eye) right = normalize(np.cross(fwd, up)) trueup = np.cross(right, fwd) c2w = np.eye(4, dtype=np.

float32) c2w[:3, :3] = np.stack([right, trueup, -fwd], axis=1) c2w[:3, 3] = eye return c2w def orbitposes(n, radius, elevlo=18., elevhi=58., phase=0.0): """Golden-angle azimuths + monotone elevations => well-spread views on a dome.""" i = np.arange(n, dtype=np.float64) + 0.5 az = 2 np.pi ((i 0.

6180339887) + phase) elev = np.arcsin(np.linspace(np.sin(np.deg2rad(elevlo)), np.sin(np.deg2rad(elevhi)), n)) eyes = np.stack([radius np.cos(elev) np.cos(az), radius np.cos(elev) np.sin(az), radius np.sin(elev)], axis=-1).astype(np.float32) return np.

stack([lookat(e) for e in eyes], axis=0) def raysfrompose(c2w, H, W, focal): """Returns (origins, dirs) of shape [H, W, 3]; dirs are unit-length, so the depths returned by jax3d's sampler are true world-space distances.""" i, j = np.meshgrid(np.arange(W, dtype=np.float32), np.arange(H, dtype=np.

float32), indexing="xy") camdirs = np.stack([(i - W .5 + .5) / focal, -(j - H .5 + .5) / focal, -np.oneslike(i)], axis=-1) dirs = normalize(camdirs @ c2w[:3, :3].T) origins = np.broadcastto(c2w[:3, 3], dirs.shape) return origins.astype(np.float32).copy(), dirs.astype(np.float32) FOCAL = 0.5 cfg.

W / math.tan(0.5 math.radians(cfg.fovdeg)) We set up the JAX3D environment, install the required dependencies, and load the volumerendering 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.

Copy CodeCopiedUse a different BrowserLIGHT = jnp.asarray(normalize(np.array([0.55, 0.75, 0.85], np.float32))) SPHERES = [ (jnp.array([0.34, 0.02, -0.22]), 0.36, jnp.array([0.90, 0.24, 0.22])), (jnp.array([-0.32, 0.28, 0.05]), 0.26, jnp.array([0.25, 0.78, 0.36])), (jnp.array([-0.05, -0.36, 0.

24]), 0.22, jnp.array([0.28, 0.40, 0.95])), ] def spherefield(pos, vdir, center, radius, albedo): d = pos - center dist = jnp.linalg.norm(d, axis=-1) n = d / (dist[..., None] + 1e-8) sigma = 80.0 jax.nn.sigmoid((radius - dist) / 0.015) v = -vdir refl = 2.0 jnp.

sum(n v, -1, keepdims=True) n - v spec = 0.65 jnp.clip(jnp.sum(refl LIGHT, -1), 0., 1.) ** 24 lamb = 0.35 + 0.65 jnp.clip(jnp.sum(n LIGHT, -1), 0., 1.) rgb = jnp.clip(albedo lamb[..., None] + spec[..., None], 0., 1.) return sigma, rgb def floorfield(pos): x, y, z = pos[..., 0], pos[..., 1], pos[...

, 2] m = (jax.nn.sigmoid((0.06 - jnp.abs(z + 0.62)) / 0.008) jax.nn.sigmoid((0.85 - jnp.abs(x)) / 0.01) jax.nn.sigmoid((0.85 - jnp.abs(y)) / 0.01)) checker = (jnp.floor(x 3.0) + jnp.floor(y 3.0)) % 2.0 rgb = jnp.where(checker[..., None] > 0.5, jnp.array([0.86, 0.86, 0.89]), jnp.array([0.22, 0.25, 0.

30])) return 80.0 m, rgb def gtfield(pos, vdir): """pos, vdir: [..., 3] -> (sigma [...], rgb [..., 3]).Density-weighted blend.""" sigsum = 0.0 colsum = 0.0 for c, r, a in SPHERES: s, rgb = spherefield(pos, vdir, c, r, a) sigsum = sigsum + s colsum = colsum + s[...

, None] rgb s, rgb = floorfield(pos) sigsum = sigsum + s colsum = colsum + s[..., None] rgb return sigsum, colsum / (sigsum[..., None] + 1e-8) WHITEBG = jnp.ones((3,), jnp.float32) @jax.jit def rendergroundtruth(origins, dirs): """Fine-grained volumetric render of the analytic scene -> RGB + depth.

""" depths, positions = j3vr.samplealongrays( rayorigins=origins, raydirections=dirs, near=cfg.near, far=cfg.far, samplecount=cfg.gtsamples, deterministic=True) vdir = jnp.broadcastto(dirs[..., None, :], positions.shape) sigma, rgb = gtfield(positions, vdir) out = j3vr.

volumerendering( samplevalues={"rgb": rgb}, sampledensity=sigma, depths=depths, backgroundvalues={"rgb": WHITEBG}) return out.rayvalues["rgb"], out.raydepth, out.rayalpha def build_dataset(poses): O, D, C = [], []