实现球谐函数

在上一篇文章中,我们已经硬编码了球谐函数的系数。在这篇文章中,我们继续完成球谐函数的编写。

在动手写代码之前,我们要明确目标:我们要预测的是 RGB 三个通道的颜色。

如果我们选择使用 3 阶(Degree 3) 的球谐函数,这意味着每个通道需要 16 个系数来共同决定颜色。


视角方向

球谐函数的核心逻辑是,如果视角不变化,返回相同的颜色;如果视角稍微改变,颜色也会随之变化。

因此,我们需要计算当前相机观察高斯点的方向向量。

  1. # 1. 提取相机在世界坐标系中的位置 (最后一列前三个元素)
  2. camera_center = camera_to_world[:3, 3]
  3.  
  4. # 2. 计算并归一化视角方向 (点到相机的向量)
  5. view_dir = points - camera_center
  6. view_dir = view_dir / (torch.norm(view_dir, dim=-1, keepdim=True) + 1e-12)
  7.  
  8. # 3. 提取 x, y, z 分量
  9. x, y, z = view_dir[:, 0], view_dir[:, 1], view_dir[:, 2]

编写球谐多项式 (SH Polynomials)

这是整个过程中最枯燥但也最体现内功的一步。我们需要将空间坐标 x, y, z 转化为 16 个球谐基函数 Y_0Y_{15}

为了避免重复计算并提升性能,我们会预先声明一些中间变量(如 xy, yz, xz 等)。

  1. # 4. 预计算高阶项乘积以提升效率
  2. xx, yy, zz = x * x, y * y, z * z
  3. xy, yz, xz = x * y, y * z, x * z
  4.  
  5. # 5. 计算 16 个球谐基函数 (Y0 至 Y15)
  6. Y0 = torch.full_like(x, SH_C0)
  7.  
  8. # 1
  9. Y1 = -SH_C1_y * y
  10. Y2 = SH_C1_z * z
  11. Y3 = -SH_C1_x * x
  12.  
  13. # 2
  14. Y4 = SH_C2_xy * xy
  15. Y5 = -SH_C2_yz * yz
  16. Y6 = SH_C2_zz * (3.0 * zz - 1.0)
  17. Y7 = -SH_C2_xz * xz
  18. Y8 = SH_C2_xx_yy * (xx - yy)
  19.  
  20. # 3
  21. Y9 = -SH_C3_y_3x2_y2 * y * (3.0 * xx - yy)
  22. Y10 = SH_C3_xyz * xy * z
  23. Y11 = -SH_C3_y_zz_x2_y2 * y * (4.0 * zz - xx - yy)
  24. Y12 = SH_C3_zz_x2_y2 * z * (2.0 * zz - 3.0 * xx - 3.0 * yy)
  25. Y13 = -SH_C3_x_zz_x2_y2 * x * (4.0 * zz - xx - yy)
  26. Y14 = SH_C3_z_x2_y2 * z * (xx - yy)
  27. Y15 = -SH_C3_x_x2_3y2 * x * (xx - 3.0 * yy)
  28.  
  29. # 6. 将 Y 堆叠为 (N, 16)
  30. Y = torch.stack([Y0, Y1, Y2, Y3, Y4, Y5, Y6, Y7, Y8, Y9, Y10, Y11, Y12, Y13, Y14, Y15], dim=-1)

重组系数矩阵

高斯模型训练后,颜色系数被分为了两部分:代表基础颜色(0 阶)的 FDC,以及代表视角依赖颜色(1-3 阶)的 FREST。它们的数据结构不同,我们需要将它们“打包”到一个统一的矩阵中,形状为 (N, 16, 3)

  1. # 7. 重组系数矩阵为 (N, 16, 3)
  2. N = points.shape[0]
  3. sh = torch.empty((N, 16, 3), dtype=points.dtype, device=points.device)
  4.  
  5. # 填充 0 阶系数 (f_dc)
  6. sh[:, 0, :] = f_dc
  7.  
  8. # 填充 1-3 阶系数 (f_rest)
  9. sh[:, 1:, 0] = f_rest[:, :15]
  10. sh[:, 1:, 1] = f_rest[:, 15:30]
  11. sh[:, 1:, 2] = f_rest[:, 30:]

颜色融合

现在,我们手里有了“预训练好的颜色系数矩阵 sh (形状 N, 16, 3)” 以及 “根据当前视角算出的基底矩阵 Y (形状 N, 16)”。

如何得出最终的 RGB 颜色?只需要将它们相乘并求和。

  1. # 8. 相乘求和并进行 Sigmoid 激活得到最终 RGB 颜色
  2. raw_rgb = torch.sum(sh * Y.unsqueeze(-1), dim=1)

可视化验证

至此我们的向前渲染流程就都完成了。

我们加入新的球谐函数逻辑进行可视化验证。可视化的完整代码可查看 Commit ccbb63d:visualize_bonsai_gaussian.py。显示效果还是有问题,我们之后再看。