球谐函数的简单介绍和PlenOctree中的代码实现

• 关于球谐函数的具体内容有很多很棒的介绍：

  • https://zhuanlan.zhihu.com/p/34343836?utm_id=0

  • https://zhuanlan.zhihu.com/p/351289217

    这里主要介绍实际的计算过程以及介绍代码为啥这样写，Plenoxel中的球谐函数是用CUDA写的，不过PlenOctree: https://github.com/sxyu/plenoctree/ 用的Python，还是比较容易理解的。

• 球谐函数在PlenOctree和Plenoxel以及大火的3D Gaussian中都用了，用来表示和观察方向相关的颜色。代码实现部分是将球谐系数转换为具体的RGB值。
  球谐函数就是用一组球谐系数对于球谐基加权求和，球谐基就是关于球坐标系下$(\theta ,\varphi )$的函数，下面的$y_l^m$代表第$l$阶的的SH Basis的第$m$个分量：
  （下面这个公式摘抄自PlenOctree：https://arxiv.org/pdf/2103.14024.pdf 的补充材料中)

$$c(d;k)=Sigmoid(\sum _{l=0}^{l_{max}}\sum _{m=-l}^l{K}_l^m{y}_l^m(d))$$

---

• STEP 1 $d\to (\theta ,\phi )$:

  $d$是$(x,y,z)$的单位向量，在单位球坐标系下：

• STEP 2 球谐函数的公式：( https://zhuanlan.zhihu.com/p/351289217 )

$$y_l^m(\theta ,\phi )=\{\begin{array}{cc}\sqrt{2}{K}_l^{m}cos(m\phi )P_l^m(cos\theta )& m>0\\ \sqrt{2}{K}_l^{m}sin(-m\phi )P_l^{-m}(cos\theta )& m<0\\ K_l^0{P}_l^0(cos\theta )& m=0\end{array}$$

$$\begin{gathered} P_n(x)=\frac{1}{2^n\cdot n!}\frac{d^n}{d{x}^n}[(x^2-1{)}^n] \\ {P}_l^m=(-1{)}^m(1-x^2{)}^{\frac{m}{2}}\frac{d^m}{d{x}^m}(P_l(x)) \\ {K}_l^m=\sqrt{\frac{(2l+1)(l-|m|)!}{4\pi (l+|m|)!}} \end{gathered}$$

• STEP 3 开始代入不同的$l,m$：

$$\begin{gathered} P_0(x)=1 \\ \{\begin{array}{c}{P}_0^0(x)=1\end{array} \end{gathered}$$

$$\begin{gathered} P_1(x)=x \\ \{\begin{array}{c}{P}_1^0(x)=x\\ P_1^1(x)=-(1-x^2{)}^{\frac{1}{2}}\end{array} \end{gathered}$$

$$\begin{gathered} P_2(x)=\frac{3x^2-1}{2} \\ \{\begin{array}{c}{P}_2^0(x)=\frac{3x^2-1}{2}\\ P_2^1(x)=-3x(1-x^2{)}^{\frac{1}{2}}\\ P_2^2(x)=3(1-x^2)\end{array} \end{gathered}$$

• STEP 4 求SH basis，即$y_l^m$:

  $$\{\begin{array}{c}-y_0^0=K_0^0\end{array}$$

  $$\{\begin{array}{c}{y}_1^{-1}=-\sqrt{2}{K}_1^{-1}sin(\phi )|sin\theta |=-\sqrt{2}{K}_1^{-1}y\\ y_1^0=K_1^{0}cos\theta =K_1^{0}z\\ y_1^1=-\sqrt{2}{K}_1^{-1}cos(\phi )|sin\theta |=-\sqrt{2}{K}_1^{-1}x\end{array}$$

  $$\{\begin{array}{c}{y}_2^{-2}=3\sqrt{2}{K}_2^{-2}sin(2\phi )si{n}^2\theta =6\sqrt{2}{K}_2^{-2}xy\\ y_2^{-1}=-3\sqrt{2}{K}_2^{-1}sin(\phi )|sin\theta |cos\theta =-3\sqrt{2}yz\\ y_2^0=\frac{3co{s}^2\theta -1}{2}=K_2^0\frac{2z^2-x^2-y^2}{2}\\ y_2^1=-3\sqrt{2}{K}_2^{1}cos(\phi )|sin\theta |cos\theta =-3\sqrt{2}{K}_2^{1}xz\\ y_2^2=3\sqrt{2}{K}_2^{2}cos(2\phi )si{n}^2\theta =3\sqrt{2}{K}_2^2(x^2-y^2)\end{array}$$

• STEP 5 求$K_l^m$:

  $$K_l^m=\sqrt{\frac{(2l+1)(l-|m|)!}{4\pi (l+|m|)!}}$$

  $$\begin{gathered} K_0^0=\frac{1}{\sqrt{4\pi }}=0.282094 \\ \sqrt{2}{K}_1^{-1}=\sqrt{2}{K}_1^1=K_1^0=\sqrt{\frac{3}{4\pi }}=0.4886 \\ \{\begin{array}{c}6\sqrt{2}{K}_2^{-2}=1.0925\\ -3\sqrt{2}{K}_2^{-1}=-1.0925\\ \frac{K_2^0}{2}=0.31539\\ -3\sqrt{2}{K}_2^1=-1.0925\\ 3\sqrt{2}{K}_2^2=0.54627\end{array} \end{gathered}$$

• 因此PlenOctree的代码

  ( https://github.com/sxyu/plenoctree/blob/master/nerf_sh/nerf/sh.py )

  就解释的通啦：

C0 = 0.28209479177387814
C1 = 0.4886025119029199
C2 = [
    1.0925484305920792,
    -1.0925484305920792,
    0.31539156525252005,
    -1.0925484305920792,
    0.5462742152960396
]
C3 = [
    -0.5900435899266435,
    2.890611442640554,
    -0.4570457994644658,
    0.3731763325901154,
    -0.4570457994644658,
    1.445305721320277,
    -0.5900435899266435
]
C4 = [
    2.5033429417967046,
    -1.7701307697799304,
    0.9461746957575601,
    -0.6690465435572892,
    0.10578554691520431,
    -0.6690465435572892,
    0.47308734787878004,
    -1.7701307697799304,
    0.6258357354491761,
]

def eval_sh(deg, sh, dirs):
    """
    Evaluate spherical harmonics at unit directions
    using hardcoded SH polynomials.
    Works with torch/np/jnp.
    ... Can be 0 or more batch dimensions.

    Args:
        deg: int SH deg. Currently, 0-3 supported
        sh: jnp.ndarray SH coeffs [..., C, (deg + 1) ** 2]
        dirs: jnp.ndarray unit directions [..., 3]

    Returns:
        [..., C]
    """
    assert deg <= 4 and deg >= 0
    assert (deg + 1) ** 2 == sh.shape[-1]
    C = sh.shape[-2]

    result = C0 * sh[..., 0]
    if deg > 0:
        x, y, z = dirs[..., 0:1], dirs[..., 1:2], dirs[..., 2:3]
        result = (result -
                C1 * y * sh[..., 1] +
                C1 * z * sh[..., 2] -
                C1 * x * sh[..., 3])
        if deg > 1:
            xx, yy, zz = x * x, y * y, z * z
            xy, yz, xz = x * y, y * z, x * z
            result = (result +
                    C2[0] * xy * sh[..., 4] +
                    C2[1] * yz * sh[..., 5] +
                    C2[2] * (2.0 * zz - xx - yy) * sh[..., 6] +
                    C2[3] * xz * sh[..., 7] +
                    C2[4] * (xx - yy) * sh[..., 8])

            if deg > 2:
                result = (result +
                        C3[0] * y * (3 * xx - yy) * sh[..., 9] +
                        C3[1] * xy * z * sh[..., 10] +
                        C3[2] * y * (4 * zz - xx - yy)* sh[..., 11] +
                        C3[3] * z * (2 * zz - 3 * xx - 3 * yy) * sh[..., 12] +
                        C3[4] * x * (4 * zz - xx - yy) * sh[..., 13] +
                        C3[5] * z * (xx - yy) * sh[..., 14] +
                        C3[6] * x * (xx - 3 * yy) * sh[..., 15])
                if deg > 3:
                    result = (result + C4[0] * xy * (xx - yy) * sh[..., 16] +
                            C4[1] * yz * (3 * xx - yy) * sh[..., 17] +
                            C4[2] * xy * (7 * zz - 1) * sh[..., 18] +
                            C4[3] * yz * (7 * zz - 3) * sh[..., 19] +
                            C4[4] * (zz * (35 * zz - 30) + 3) * sh[..., 20] +
                            C4[5] * xz * (7 * zz - 3) * sh[..., 21] +
                            C4[6] * (xx - yy) * (7 * zz - 1) * sh[..., 22] +
                            C4[7] * xz * (xx - 3 * yy) * sh[..., 23] +
                            C4[8] * (xx * (xx - 3 * yy) - yy * (3 * xx - yy)) * sh[..., 24])
    return result