diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index 4cf18c47a..63facf0fd 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -44,6 +44,91 @@ from diff_rast.rasterize import RasterizeGaussians from diff_rast.project_gaussians import ProjectGaussians from diff_rast.sh import SphericalHarmonics, num_sh_bases +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] + """ + 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, + ] + + 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 random_quat_tensor(N, **kwargs): u = torch.rand(N, **kwargs) v = torch.rand(N, **kwargs) @@ -103,9 +188,9 @@ class GaussianSplattingModelConfig(ModelConfig): """threshold of positional gradient norm for densifying gaussians""" densify_size_thresh: float = 0.01 """below this size, gaussians are *duplicated*, otherwise split""" - sh_degree_interval: int = 1000 + sh_degree_interval: int = 500 """every n intervals turn on another sh degree""" - max_screen_size: int = 300 + max_screen_size: int = 150 """maximum screen size of a gaussian, in pixels""" random_init: bool = False """whether to initialize the positions uniformly randomly (not SFM points)""" @@ -324,8 +409,8 @@ class GaussianSplattingModel(Model): # cull transparent ones culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze() if self.step > self.config.refine_every * self.config.reset_alpha_every: - # cull huge ones that are in the center of the scene (a 2-unit box is twice as large as all the cameras) - toobigs = (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze() & (self.means.amax(dim=-1) < 2) + # cull huge ones + toobigs = (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze() culls = culls | toobigs # cull big screen space culls = culls | (self.max_2Dsize > self.config.max_screen_size).squeeze() @@ -507,7 +592,7 @@ class GaussianSplattingModel(Model): viewdirs = means_crop[rend_mask].detach() - camera.camera_to_worlds[..., :3, 3] # (N, 3) viewdirs = viewdirs / viewdirs.norm(dim=-1, keepdim=True) n = min(self.step // self.config.sh_degree_interval,self.config.sh_degree) - # rgbs = eval_sh(n, self.get_colors, viewdirs)[crop_ids] + # rgbs = eval_sh(n, colors_crop[rend_mask,:,:], viewdirs) n_coeffs = num_sh_bases(n) coeffs = colors_crop.permute(0,2,1)[rend_mask,:n_coeffs,:] rgbs = SphericalHarmonics.apply(n, viewdirs,coeffs) @@ -528,19 +613,19 @@ class GaussianSplattingModel(Model): W, background, ) - depth_im = RasterizeGaussians.apply( - self.xys[rend_mask], - depths[rend_mask], - self.radii[rend_mask], - conics[rend_mask], - num_tiles_hit[rend_mask], - depths[:,None], - torch.sigmoid(opacities_crop[rend_mask]), - H, - W, - torch.ones(1,device=self.device)*10, - )[...,0:1] - return {"rgb": rgb[..., :3],'depth':depth_im} + # depth_im = RasterizeGaussians.apply( + # self.xys[rend_mask], + # depths[rend_mask], + # self.radii[rend_mask], + # conics[rend_mask], + # num_tiles_hit[rend_mask], + # depths[rend_mask,None], + # torch.sigmoid(opacities_crop[rend_mask]), + # H, + # W, + # torch.ones(1,device=self.device)*10, + # )[...,0:1] + return {"rgb": rgb}#,'depth':depth_im} def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: """Compute and returns metrics. diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 57e544db3..79e1cd9ca 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -26,7 +26,7 @@ from nerfstudio.utils import colormaps, writer from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter from nerfstudio.viewer.server import viewer_utils from nerfstudio.viewer_beta.utils import CameraState, get_camera - +from nerfstudio.models.gaussian_splatting import GaussianSplattingModel if TYPE_CHECKING: from nerfstudio.viewer_beta.viewer import Viewer @@ -126,7 +126,10 @@ class RenderStateMachine(threading.Thread): with TimeWriter(None, None, write=False) as vis_t: with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext(): - camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb) + if isinstance(self.viewer.get_model(),GaussianSplattingModel): + camera_ray_bundle=None + else: + camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb) self.viewer.get_model().set_crop(obb) self.viewer.get_model().eval() step = self.viewer.step