Como construir um NeRF hierárquico do zero com JAX3D, Flax e Optax

O que mudou: reconstrução 3D saiu dos laboratórios

Há poucos anos, reconstruir uma cena em 3D a partir de imagens exigia hardware especializado, câmeras calibradas e semanas de processamento. Hoje, um Neural Radiance Field (NeRF) — uma rede neural que aprende a representar densidade e cor de um volume — faz o trabalho em uma GPU comum, e as bibliotecas JAX, Flax e Optax tornaram o código compacto e rápido via compilação JIT.

Este tutorial constrói, do zero, um NeRF hierárquico usando o jax3d (repositório da Google Research) para as primitivas de volume rendering. Você vai entender cada componente — do modelo de câmera ao marching cubes — e sair com um pipeline completo de síntese de novas visões.

O que você ganha (e o que não ganha)

✅ Prós

• Pipeline completo de inverse rendering: da cena analítica à extração de geometria.
• Uso das primitivas matemáticas prontas do jax3d (sample_along_rays, volume_rendering, sample_piecewise_constant_pdf).
• Treinamento acelerado por JIT com Adam, decaimento exponencial de taxa e gradient clipping.
• Adaptação automática entre GPU e CPU (config reduzida para execução sem GPU).
• Avaliação completa: PSNR, mapas de profundidade/opacidade, órbita 360° e isosuperfície.

⚠️ Contras

• É um cenário sintético: a cena é gerada analiticamente, não reconstruída de fotos reais.
• NeRF puro é mais lento que abordagens mais novas como 3D Gaussian Splatting para cenas grandes.
• O código assume familiaridade razoável com Python e álgebra linear (raios, poses de câmera).

Requisitos

Passo 1 — Ambiente e dependências

Instale as dependências e clone o repositório jax3d. Em vez de instalar o pacote completo (que puxa gin e tfds desnecessariamente), carregamos apenas o módulo de volume rendering pelo caminho:

import os, sys, subprocess, importlib.util

def _sh(cmd):
    subprocess.run(cmd, shell=True, check=False, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)

_sh(f'{sys.executable} -m pip install -q "etils[array-types,epy,etree,enp]" chex flax optax scikit-image')

REPO_DIR = "./jax3d"
if not os.path.isdir(REPO_DIR):
    _sh("git clone -q --depth 1 https://github.com/google-research/jax3d.git ./jax3d")

A escolha de carregar o módulo por caminho (spec_from_file_location) evita disparar o __init__ do pacote, que importaria gin/tfds — economiza memória e instalação.

Passo 2 — Configuração e modelo de câmera

O Config centraliza resolução, número de visões, distância near/far e hiperparâmetros das redes. O modelo de câmera usa a convenção OpenGL/NeRF (+x direita, +y cima, câmera olha para −z) e posiciona as câmeras em órbita com ângulo áureo para espalhar bem as visões:

def look_at(eye, target=(0.,0.,0.), up=(0.,0.,1.)):
    fwd   = target - eye; fwd /= np.linalg.norm(fwd)
    right = np.cross(fwd, up); right /= np.linalg.norm(right)
    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 orbit_poses(n, radius, elev_lo=18., elev_hi=58., phase=0.0):
    i = np.arange(n) + 0.5
    az = 2 * np.pi * ((i * 0.6180339887) + phase)      # ângulo áureo
    elev = np.arcsin(np.linspace(np.sin(np.deg2rad(elev_lo)),
                                 np.sin(np.deg2rad(elev_hi)), 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([look_at(e) for e in eyes], axis=0)

Cada pose é convertida em raios normalizados (origens e direções) via matriz de intrínsecos pinhole — a base geométrica de todo o pipeline.

Passo 3 — Cena analítica e dataset sintético

Em vez de fotos reais, o tutorial define uma cena com três esferas de bordas suaves, um piso quadriculado e reflexos especulares dependentes da visão. A função gt_field soma densidade e cor de cada primitiva, e o render usa sample_along_rays + volume_rendering do jax3d para produzir RGB, profundidade e opacidade:

@jax.jit
def render_ground_truth(origins, dirs):
    depths, positions = j3vr.sample_along_rays(
        ray_origins=origins, ray_directions=dirs,
        near=cfg.near, far=cfg.far,
        sample_count=cfg.gt_samples, deterministic=True)
    sigma, rgb = gt_field(positions, jnp.broadcast_to(dirs[..., None, :], positions.shape))
    out = j3vr.volume_rendering(
        sample_values={"rgb": rgb}, sample_density=sigma, depths=depths,
        background_values={"rgb": WHITE_BG})
    return out.ray_values["rgb"], out.ray_depth, out.ray_alpha

O dataset é um “pool de raios” achatado — cada raio vira uma amostra de treino, o que permite sortear lotes aleatórios de forma eficiente.

Passo 4 — O MLP do NeRF

A rede usa codificação posicional senoidal para coordenadas e direções, uma MLP Flax com skip connection, e separa geometria (densidade, independente da visão) de aparência (cor, dependente da direção):

def posenc(x, deg):
    if deg == 0:
        return x
    scales = 2.0 ** jnp.arange(deg, dtype=x.dtype)
    xb = (x[..., None, :] * scales[:, None]).reshape(*x.shape[:-1], -1)
    return jnp.concatenate([x, jnp.sin(xb), jnp.cos(xb)], axis=-1)

class NeRFMLP(nn.Module):
    width: int; depth: int; skip: int; deg_pos: int; deg_dir: int
    @nn.compact
    def __call__(self, pos, dirs):
        inp = posenc(pos, self.deg_pos)
        x = inp
        for i in range(self.depth):
            x = nn.relu(nn.Dense(self.width)(x))
            if i == self.skip:
                x = jnp.concatenate([x, inp], axis=-1)
        sigma = nn.softplus(nn.Dense(1)(x)[..., 0] - 1.0)
        h = jnp.concatenate([nn.Dense(self.width)(x), posenc(dirs, self.deg_dir)], -1)
        rgb = nn.sigmoid(nn.Dense(3)(nn.relu(nn.Dense(self.width // 2)(h))))
        return sigma, rgb

A densidade sai de softplus (sempre não negativa) e a cor de sigmoid (entre 0 e 1). A separação visão-independente/visão-dependente é o que permite modelar brilhos especulares sem contaminar a geometria.

Passo 5 — Rendering hierárquico (coarse → fine)

O coração do NeRF: primeiro a rede “coarse” amostra pontos uniformes ao longo do raio; os pesos resultantes viram uma PDF por partes, da qual a rede “fine” faz importance sampling ao redor das superfícies relevantes:

def render_rays(params, origins, dirs, rng, deterministic):
    rng_c, rng_f = jax.random.split(rng)
    depths_c, pos_c = j3vr.sample_along_rays(
        ray_origins=origins, ray_directions=dirs, near=cfg.near, far=cfg.far,
        sample_count=cfg.n_coarse, deterministic=deterministic, rng=rng_c)
    sigma_c, rgb_c = model.apply(params["coarse"], pos_c, jnp.broadcast_to(dirs[:, None, :], pos_c.shape))
    out_c = j3vr.volume_rendering(
        sample_values={"rgb": rgb_c}, sample_density=sigma_c, depths=depths_c,
        background_values={"rgb": WHITE_BG})
    mid = 0.5 * (depths_c[..., 1:] + depths_c[..., :-1])
    bin_edges = jnp.concatenate([depths_c[..., :1], mid, depths_c[..., -1:]], -1)
    t_fine = j3vr.sample_piecewise_constant_pdf(
        bin_edges=bin_edges, weights=out_c.sample_weights,
        sample_count=cfg.n_fine, deterministic=deterministic, rng=rng_f)
    t_fine = jax.lax.stop_gradient(t_fine)
    ...  # junta coarse+fine, renderiza com a rede fine e retorna ambos os estágios

O stop_gradient impede que o gradiente flua pela operação de amostragem — ela não é diferenciável e não deve interferir no treino.

Passo 6 — Treinamento com JIT, Adam e clipping

As redes coarse e fine são treinadas em conjunto, com exponential decay da taxa de aprendizado, gradient clipping global e supervisão nos dois estágios:

schedule = optax.exponential_decay(cfg.lr_init, cfg.steps, cfg.lr_final / cfg.lr_init)
tx = optax.chain(optax.clip_by_global_norm(1.0), optax.adam(schedule))

@jax.jit
def train_step(state, o, d, target, rng):
    def loss_fn(p):
        out_c, out_f, _ = render_rays(p, o, d, rng, deterministic=False)
        l_c = jnp.mean((out_c.ray_values["rgb"] - target) ** 2)
        l_f = jnp.mean((out_f.ray_values["rgb"] - target) ** 2)
        return l_c + l_f, l_f
    (loss, l_fine), grads = jax.value_and_grad(loss_fn, has_aux=True)(state.params)
    return state.apply_gradients(grads=grads), loss, l_fine

O @jax.jit compila o passo de treino uma única vez, e o PSNR da rede fine é monitorado a cada 25 passos para acompanhar a convergência.

Passo 7 — Avaliação: PSNR, órbita 360° e marching cubes

A avaliação renderiza visões de teste não vistas, mede PSNR, inspeciona profundidade/opacidade e gera um GIF de órbita 360°. Por fim, consulta a densidade aprendida numa grade 3D e aplica marching cubes para extrair a geometria aproximada — provando que a rede realmente aprendeu a forma da cena, não apenas a cor.

Casos de uso reais

• Inspeção industrial: reconstruir peças e equipamentos em 3D a partir de poucas fotos para medição e controle de qualidade.
• E-commerce e imobiliário: gerar visualizações 360° de produtos e imóveis sem sessão fotográfica completa.
• Arqueologia e patrimônio: digitalizar objetos e sítios frágeis preservando a forma geométrica.
• Games e efeitos visuais: criar assets 3D a partir de vídeo real, reduzindo modelagem manual.

Troubleshooting

• ❌ “No GPU detected” → o código reduz resolução e passos para CPU. Troque o runtime para GPU T4 para a versão completa.
• ❌ Erro ao importar volume_rendering → verifique o caminho do módulo dentro do clone; a estrutura pode variar entre jax3d/jax3d/math/ e jax3d/math/. O código testa ambos.
• ❌ Memória insuficiente durante o render → reduza chunk e batch_rays; o render é dividido em blocos para caber em GPUs menores.
• ❌ PSNR estagnado → confirme que n_fine está sendo usado e que o gradient clipping não está mascarando gradientes grandes demais.
• ❌ Marching cubes vazio → o nível de isosuperfície cai para o percentil 99 do volume se o valor padrão não interceptar a densidade.

FAQ

• Preciso de GPU para rodar? Não. O código detecta CPU e aplica uma configuração reduzida — mas espere treinos bem mais longos.
• Por que usar jax3d em vez de implementar o rendering na mão? As primitivas de amostragem e composição são testadas e diferenciáveis, o que remove uma classe inteira de bugs.
• Isso funciona com fotos reais? O exemplo é sintético, mas o mesmo pipeline aceita fotos reais se você calibrar as poses de câmera (ex.: COLMAP).
• Qual a diferença para 3D Gaussian Splatting? NeRF é mais simples e compacto; Gaussian Splatting costuma renderizar mais rápido e escala melhor para cenas grandes.
• O que é “hierárquico” aqui? Duas redes: a coarse distribui amostras uniformes; a fine concentra amostras extras nas regiões de alta densidade.

Para onde isso vai

O campo está migrando de NeRFs densos para representações híbridas — Gaussian Splatting, hash grids e campos implícitos esparsos — que renderizam em tempo real. Ainda assim, o NeRF hierárquico continua sendo o modelo mental canônico para entender volume rendering diferenciável, e o ecossistema JAX/Flax está se consolidando como alternativa de alta performance ao PyTorch nessa área. Em 2027, espere que a reconstrução 3D a partir de vídeo esteja acessível a qualquer desenvolvedor com uma GPU de entrada.