Spaces:
Running
on
A10G
Running
on
A10G
| # Copyright 2021 The PlenOctree Authors. | |
| # Redistribution and use in source and binary forms, with or without | |
| # modification, are permitted provided that the following conditions are met: | |
| # | |
| # 1. Redistributions of source code must retain the above copyright notice, | |
| # this list of conditions and the following disclaimer. | |
| # | |
| # 2. Redistributions in binary form must reproduce the above copyright notice, | |
| # this list of conditions and the following disclaimer in the documentation | |
| # and/or other materials provided with the distribution. | |
| # | |
| # THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" | |
| # AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE | |
| # IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE | |
| # ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE | |
| # LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR | |
| # CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF | |
| # SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS | |
| # INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN | |
| # CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) | |
| # ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE | |
| # POSSIBILITY OF SUCH DAMAGE. | |
| import torch | |
| 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 | |
| coeff = (deg + 1) ** 2 | |
| assert sh.shape[-1] >= coeff | |
| 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 | |
| def RGB2SH(rgb): | |
| return (rgb - 0.5) / C0 | |
| def SH2RGB(sh): | |
| return sh * C0 + 0.5 |