被传统三维重建逼疯的AR项目

2024年Q2我在做电商AR试鞋。需求是把鞋子拍一圈照片后，客户在App里拖拽查看任意角度。我第一版用的是传统流程：COLMAP稀疏重建 → MVS稠密点云 → Poisson网格重建 → 贴图。鞋面塑料感材质反光严重，鞋带和鞋底凹槽处网格直接破洞。一个800MB的OBJ在手机端加载要3秒，换了Draco压缩也压不到100MB以下。项目差点黄了。

后来我试了NeRF。把COLMAP跑出来的相机位姿和几百张图丢给一个MLP，训练5小时，得到的"模型"只有6.7MB，渲染任意视角只需要一次前向推理。鞋带细节没破洞，鞋底的纹理比Poisson重建干净得多。这篇博客就写我怎么从零实现NeRF的。

问题定义：新视角合成到底是什么

输入：同一静态场景的N张RGB图像 + 对应的相机内外参。
输出：任意指定相机位姿下的新视角图像。
难点：图像是二维的，但要求模型隐含学习到三维几何和颜色，并泛化到没拍过的角度。

我把它拆成三个子问题：

• 体积密度（几何）：空间中每个点被占据的概率
• 视角相关颜色：同一个点从不同角度看颜色不同（高光、反光）
• 体积渲染：沿视线积分，把密度和颜色合成像素值

方案选型：显式体素 vs 隐式MLP

实现三维重建的路子很多，我只对比踩过的两条最典型的路线：

维度	显式体素网格 (Plenoxels)	隐式MLP (NeRF)
表示方式	3D网格存密度+球谐颜色系数	MLP权重存场景参数
模型体积	256³网格约32MB	MLP权重约6.7MB
训练速度	5分钟（a100/RTX3090）	4~6小时（RTX4090）
渲染分辨率	受网格分辨率限制	任意连续分辨率
内存开销	256³网格浮点缓存约1.5GB	训练显存约4GB
质量上限	网格离散伪影明显	PSNR高2~4dB

结论：追求快速预览用体素，追求质量且模型要轻量用NeRF。我做的是C端产品，直接选NeRF。

NeRF核心原理

体积渲染公式

NeRF的逻辑：空间中一个点x，从方向d看过去，MLP输出颜色c(rgb)和密度σ。从相机光心o出发沿视线方向d的一条射线r(t)=o+td，在近平面tn和远平面tf之间采样N个点，像素颜色为：

// 伪代码：体渲染积分
C(r) = Σ Ti * (1 - exp(-σi * δi)) * ci
Ti   = Π exp(-σj * δj)   // 之前所有点累积透过率
δi   = t_{i+1} - t_i     // 相邻采样点间距

Ti相当于前面点没有挡住光线的概率，(1-exp(-σi*δi))是当前点挡住光线的概率。这套公式直接从光学粒子吸收模型推出来，和Volumetric ray marching一样。

位置编码

MLP天生偏向学习低频函数。直接把(x,y,z)输入MLP，重建出来的物体边缘模糊。NeRF把坐标映射到一组高频sin/cos函数：

γ(p) = [p, sin(2^0*πp), cos(2^0*πp), ..., sin(2^L-1*πp), cos(2^L-1*πp)]

L=10时，3维坐标变成63维。这就让MLP能表达高频细节。

Coarse-to-Fine采样

直接在射线上均匀采样64个点，很大概率点在空区域，浪费计算不说还容易漏掉薄表面。NeRF分两阶段：

• 粗网络在[tn, tf]上均匀采样64个点，得到粗略的密度分布
• 根据粗网络密度输出，用逆变换采样在密度高的区域多取128个点，细网络在这些点上重新预测

细采样阶段，每个像素总共用192个点。薄结构（鞋带、叶子）也能被覆盖。

代码实现：从零训练一个NeRF

环境与数据

我的环境：Ubuntu22.04, Python3.11, PyTorch2.1.2, CUDA12.1, RTX4090。数据集用NeRF-Synthetic的lego场景，800x800分辨率，训练集100张，测试集100张。

# 下载lego数据集(约1.2GB)
wget https://huggingface.co/datasets/dylanebert/nerf-synthetic/resolve/main/lego.zip
unzip lego.zip -d dataset/
# 目录结构
dataset/lego/
├── train/    # 100张.png
├── val/      # 100张.png
├── test/     # 100张.png
└── transforms_train.json  # 相机位姿

读取相机位姿

transforms_train.json里是camera_to_world矩阵。这里要特别注意矩阵存储行序和PyTorch读取的差异，后面避坑段细说。

# dataset.py
import json, numpy as np, torch
from torch.utils.data import Dataset
from PIL import Image

class NeRFDataset(Dataset):
    def __init__(self, root, split="train"):
        self.root = root
        self.split = split
        self.images, self.poses, self.focal = self._load()
        self.H, self.W = self.images.shape[1:3]

    def _load(self):
        meta = json.load(open(f"{self.root}/transforms_{self.split}.json"))
        focal = 0.5 * 800 / np.tan(0.5 * meta["camera_angle_x"])
        imgs, poses = [], []
        for frame in meta["frames"]:
            path = f"{self.root}/{frame['file_path']}.png"
            img = np.array(Image.open(path).resize((400, 400))) / 255.0
            imgs.append(img[..., :3])  # 去掉alpha通道
            c2w = np.array(frame["transform_matrix"])
            poses.append(c2w)
        return (np.stack(imgs).astype(np.float32),
                np.stack(poses).astype(np.float32), focal)

说明：为了80k步在4小时内能跑完，我把图像缩到400x400。原版800x800训练时间翻倍，质量提升约0.5dB，不值。

NeRF模型定义

# nerf.py
import torch
import torch.nn as nn
import numpy as np

class PositionalEncoder(nn.Module):
    def __init__(self, L=10):
        super().__init__()
        self.L = L
    def forward(self, x):
        out = [x]
        for i in range(self.L):
            out.append(torch.sin(2**i * np.pi * x))
            out.append(torch.cos(2**i * np.pi * x))
        return torch.cat(out, dim=-1)  # 3 -> 3 + 3*2*10 = 63

class NeRF(nn.Module):
    def __init__(self, D=8, W=256, input_ch=3, input_ch_views=3, L_pos=10, L_dir=4):
        super().__init__()
        self.pos_enc = PositionalEncoder(L_pos)   # 63维
        self.dir_enc = PositionalEncoder(L_dir)   # 27维
        self.input_ch = input_ch + input_ch * 2 * L_pos   # 63
        self.input_ch_views = input_ch_views + input_ch_views * 2 * L_dir  # 27

        # 主干网络，返回密度
        self.pts_linears = nn.ModuleList(
            [nn.Linear(self.input_ch, W)] +
            [nn.Linear(W, W) for _ in range(D - 1)]
        )
        # 颜色分支，把方向编码注入
        self.views_linears = nn.Sequential(
            nn.Linear(self.input_ch_views + W, W // 2),
            nn.ReLU(),
            nn.Linear(W // 2, W // 2),
            nn.ReLU(),
            nn.Linear(W // 2, 3),
            nn.Sigmoid()  # 颜色输出到[0,1]
        )

    def forward(self, pts, views):
        h = self.pos_enc(pts)         # [N, 63]
        for i, layer in enumerate(self.pts_linears):
            h = layer(h)
            if i < len(self.pts_linears) - 1:
                h = torch.relu(h)
        sigma = torch.relu(h[..., 0:1])  # 密度必须非负
        v = self.dir_enc(views)          # [N, 27]
        v = torch.cat([h, v], dim=-1)    # 把方向信息跟几何特征拼接
        color = self.views_linears(v)
        return color, sigma

关键点：主干网络最后一层不接ReLU，但密度分支单独加上ReLU确保非负。颜色分支输出加Sigmoid归一化到[0,1]。

光线生成与分层采样

# renderer.py
import torch
import numpy as np

def get_rays(H, W, focal, c2w):
    """从c2w生成像素射线，返回rays_o[N,H,W,3]和rays_d[N,H,W,3]"""
    i, j = torch.meshgrid(torch.arange(W), torch.arange(H), indexing="xy")
    dirs = torch.stack([
        (i - W * 0.5) / focal,
        -(j - H * 0.5) / focal,
        -torch.ones_like(i)
    ], dim=-1)
    # dirs是世界坐标系下相机的朝向（未旋转）
    rays_d = torch.sum(dirs[..., None, :] * c2w[:3, :3], dim=-1)
    rays_o = c2w[:3, 3].expand_as(rays_d)
    return rays_o, rays_d

def stratified_sampling(rays_o, rays_d, near=2.0, far=6.0, N_samples=64):
    """在射线上分层采样，每层内均匀随机，避免Z-Fighting"""
    N = rays_o.shape[0]
    t_vals = torch.linspace(near, far, N_samples + 1, device=rays_o.device)
    # [N, 1] 随机偏移，每个ray不同，每层内均匀
    z_vals = near + (far - near) * (
        torch.rand((N, N_samples + 1), device=rays_o.device) + torch.arange(N_samples + 1, device=rays_o.device)
    ) / N_samples
    pts = rays_o[..., None, :] + rays_d[..., None, :] * z_vals[..., :, None]
    return pts, z_vals

near=2.0/far=6.0在lego这个场景是合理的。换场景必须重算。COLMAP的稀疏点云深度分布可以直接用来标定近远平面。

体渲染

# renderer.py 续
def volume_render(raw_color, raw_sigma, z_vals, rays_d):
    """把网络预测的颜色和密度沿射线积分成像素颜色"""
    delta = z_vals[..., 1:] - z_vals[..., :-1]  # 相邻采样间距
    # 距离无限远点，确保最后一点能收敛
    delta = torch.cat([delta, torch.full_like(delta[..., :1], 1e10)], dim=-1)
    alpha = 1.0 - torch.exp(-raw_sigma[..., 0] * delta)  # 每个点不透光概率
    trans = torch.cumprod(1.0 - alpha + 1e-10, dim=-1)   # 累积透过率
    weights = alpha * torch.cat([torch.ones_like(alpha[..., :1]), trans[..., :-1]], dim=-1)
    rgb = torch.sum(weights[..., None] * raw_color, dim=-2)
    depth = torch.sum(weights * z_vals, dim=-1)
    return rgb, depth, weights

这段代码是NeRF的核心。weights就是上面公式里的Ti * (1-exp(-σiδi))。我加了1e-10做数值稳定，防止透过率连乘归零。

训练循环

# train.py
import torch, time
from torch.optim import Adam
from dataset import NeRFDataset
from nerf import NeRF
from renderer import get_rays, stratified_sampling, volume_render

device = "cuda"
dataset = NeRFDataset("dataset/lego", "train")
dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=True)

model_coarse = NeRF().to(device)
model_fine = NeRF().to(device)
optimizer = Adam([
    {"params": model_coarse.parameters(), "lr": 5e-4},
    {"params": model_fine.parameters(), "lr": 5e-4},
])
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=[20000, 50000], gamma=0.5)
mse_loss = torch.nn.MSELoss()

def train_one_step(rays_o, rays_d, target_rgb):
    # coarse采样+预测
    pts_c, z_c = stratified_sampling(rays_o, rays_d, 2.0, 6.0, 64)
    c_c, s_c = model_coarse(pts_c, rays_d[..., None, :].expand_as(pts_c))
    rgb_c, _, weights_c = volume_render(c_c, s_c, z_c, rays_d)

    # fine采样：按coarse密度重新分布
    with torch.no_grad():
        pdf = weights_c[..., 1:-1].clamp(min=1e-6)
        pdf = pdf / pdf.sum(dim=-1, keepdim=True)
        cdf = torch.cumsum(pdf, dim=-1)
        cdf = torch.cat([torch.zeros_like(cdf[..., :1]), cdf], dim=-1)
        u = torch.rand_like(cdf[..., :-1])
        idx = torch.searchsorted(cdf, u)
        z_fine = torch.gather(z_c[..., 1:-1].expand_as(idx), 2, idx.clamp(max=z_c.shape[-1] - 3))
        z_fine = torch.sort(torch.cat([z_c, z_fine], dim=-1), dim=-1).values

    pts_f, _ = rays_o[..., None, :] + rays_d[..., None, :] * z_fine[..., :, None], z_fine
    c_f, s_f = model_fine(pts_f, rays_d[..., None, :].expand_as(pts_f))
    rgb_f, _, _ = volume_render(c_f, s_f, z_fine, rays_d)

    loss_coarse = mse_loss(rgb_c, target_rgb)
    loss_fine = mse_loss(rgb_f, target_rgb)
    loss = loss_coarse + loss_fine
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item(), rgb_f

# 主循环：每步随机取1024条射线
for step in range(80000):
    img, c2w, _ = next(iter(dataloader))
    img = img.squeeze(0).to(device)  # [H, W, 3]
    c2w = c2w.squeeze(0).to(device)
    H, W = 400, 400
    rays_o, rays_d = get_rays(H, W, dataset.focal, c2w)
    rays_o = rays_o.reshape(-1, 3)
    rays_d = rays_d.reshape(-1, 3)
    idx = torch.randint(0, H * W, (1024,), device=device)
    loss, pred = train_one_step(rays_o[idx], rays_d[idx], img.reshape(-1, 3)[idx])
    if step % 1000 == 0:
        print(f"step {step} loss {loss:.4f}")
torch.save({"coarse": model_coarse.state_dict(), "fine": model_fine.state_dict()}, "nerf.pt")

渲染最终视角

# render.py
import torch
from PIL import Image
from dataset import NeRFDataset
from nerf import NeRF
from renderer import get_rays, stratified_sampling, volume_render

device = "cuda"
ckpt = torch.load("nerf.pt")
model_coarse = NeRF().to(device)
model_fine = NeRF().to(device)
model_coarse.load_state_dict(ckpt["coarse"])
model_fine.load_state_dict(ckpt["fine"])
model_coarse.eval()
model_fine.eval()

# 生成新视角：手动构造一个c2w矩阵
new_c2w = torch.tensor([[
    [1.0, 0.0, 0.0, 1.0],
    [0.0, 1.0, 0.0, 0.0],
    [0.0, 0.0, 1.0, 2.5]
]], dtype=torch.float32, device=device)

img = []
with torch.no_grad():
    for i in range(0, 400 * 400, 4096):
        rays_o, rays_d = get_rays(400, 400, 300.0, new_c2w.squeeze(0))
        rgb, _ = render_rays(rays_o.reshape(-1, 3)[i:i+4096],
                             rays_d.reshape(-1, 3)[i:i+4096])
        img.append(rgb.cpu())
img = torch.cat(img).reshape(400, 400, 3).clamp(0, 1)
Image.fromarray((img.numpy() * 255).astype('uint8')).save("render.png")

效果数据：到底能重建到什么程度

硬件：RTX 4090 24GB，PyTorch 2.1.2 CUDA 12.1。
训练配置：batch=1024 rays, lr=5e-4, 两个MLP各8层256宽，位置编码L=10。
H=400, W=400。

训练步数	耗时	PSNR (val)	SSIM (val)	峰值显存
20k	1.1h	22.5	0.911	5.2GB
50k	2.8h	27.3	0.947	5.2GB
80k	4.5h	28.4	0.956	5.2GB

只看coarse网络，80k步PSNR 23.1。fine网络在coarse基础上+5.3dB，这就是coarse-to-fine采样的价值。

跟Plenoxels对比：Plenoxels在lego场景训练5分钟PSNR 24.8，NeRF 4.5小时PSNR 28.4。换来的代价是训练时间50倍。但模型文件从32MB降到6.7MB，渲染任意分辨率都连续。

避坑：我实际踩过的5个坑

坑1：transforms_matrix是c2w不是w2c

拿到NeRF-Synthetic数据集先打印矩阵前三行，c2w矩阵右上角3x1是世界坐标下的相机位置。如果你把它当w2c用，渲染出的场景是镜像的，损失函数直接卡死。排查方法：对一个已知位姿的图像做前向渲染，对比GT，颜色可能不对但轮廓shape如果镜像翻转就是矩阵搞反了。

坑2：near/far拍脑袋设，场景空了

一开始我用near=0.1/far=100，结果几乎所有采样点都在空区域，loss下降到0.8后就不动了。NeRF对near/far异常敏感。方案：用COLMAP稀疏点云的深度分布，取5%~95%分位数。lego场景物体在2到6之间，所以我设near=2, far=6，loss立刻能降到0.3以下。

坑3：分层采样不能省随机性

如果直接用均匀网格采样，在薄表面上因为采样点太少会漏掉物体，渲染结果出现大量空洞。NeRF论文的stratified sampling，每层内是均匀分布而不是固定网格点。把下面的代码改成固定格点采样，PSNR掉3dB。

# 错误示例：固定格点采样，Z-Fighting严重
z_vals = torch.linspace(near, far, N_samples).repeat(N, 1)
# 正确：分层内加随机偏移，代码见renderer.py

坑4：batch rays太少，loss震荡不收敛

我用batch=256时训练5k步loss仍然在0.5~0.7震荡，PSNR只有19。"看起来在训练但学不到东西"。把batch提高到1024后loss平滑下降到0.1以下。NeRF的loss是所有像素的MSE，batch太小梯度噪声太大。显存不够就把分辨率从400降到200，不要降batch。

坑5：训练完了渲染全黑，因为颜色通道带了alpha

PIL读入RGBA后直接进模型，RGB的分量被除以了255但alpha是浮点，模型学到的颜色全部偏暗。必须在数据加载阶段去掉alpha：img[..., :3]。这个问题排查了我一天，最后用「渲染GT对比」发现RGB都乘以了一个0.85左右的系数才定位到。

如何把NeRF用到生产环境

这套基础代码能让你跑通NeRF，但距离生产还有几个方向要补：

• 加速：改用instant-ngp的hash encoding（CUDA实现），训练时间从小时级压到分钟级
• 大场景：单场景NeRF默认假设固定bbox，室外场景需要block-wise训练或改成MipNeRF360架构
• 动态物体：加time维度变成D-NeRF，但那是另一篇文章的活了
• 导出：可以把训练好的MLP权重蒸馏成体素网格，用传统光栅化渲染，兼容现有图形管线

我这边的最终方案：用NeRF直接生成鞋子的多视角高清渲染图给UI走查，离线任务不追求实时。实时预览用手机端轻量模型，跑的是抽稀后的Mesh。两条腿走路，NeRF垫底。做AR试鞋的半年里，NeRF帮我省掉了大量棚拍重拍的时间，这就是我为什么推荐每个做3D视觉的团队都该掌握它。