From 50fbb9549e704a614a7ff768109ba5ef6dc79cb8 Mon Sep 17 00:00:00 2001 From: Vickie Ye Date: Tue, 3 Oct 2023 06:24:52 +0000 Subject: [PATCH] adding corrected splitting gradients --- .../models/gaussian_splatting_inhouse.py | 444 ++++++++++-------- 1 file changed, 243 insertions(+), 201 deletions(-) diff --git a/nerfstudio/models/gaussian_splatting_inhouse.py b/nerfstudio/models/gaussian_splatting_inhouse.py index 2d0921513..50abf96fc 100644 --- a/nerfstudio/models/gaussian_splatting_inhouse.py +++ b/nerfstudio/models/gaussian_splatting_inhouse.py @@ -29,9 +29,7 @@ from torchmetrics.image import PeakSignalNoiseRatio from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from diff_rast._torch_impl import compute_sh_color from nerfstudio.cameras.cameras import Cameras -from nerfstudio.engine.callbacks import (TrainingCallback, - TrainingCallbackAttributes, - TrainingCallbackLocation) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.viewer_beta.viewer_elements import ViewerControl from nerfstudio.engine.optimizers import Optimizers from nerfstudio.models.base_model import Model, ModelConfig @@ -44,8 +42,10 @@ from diff_rast.rasterize import RasterizeGaussians from diff_rast.project_gaussians import ProjectGaussians from diff_rast.sh import SphericalHarmonics, num_sh_bases + def get_references_to_object(obj): - return [ref for ref in gc.get_referrers(obj) if isinstance(ref, (list, dict, tuple)) or hasattr(ref, '__dict__')] + return [ref for ref in gc.get_referrers(obj) if isinstance(ref, (list, dict, tuple)) or hasattr(ref, "__dict__")] + def random_quat_tensor(N, **kwargs): u = torch.rand(N, **kwargs) @@ -67,6 +67,7 @@ def identity_quat(N, **kwargs): quat[:, 0] = 1 return quat + def projection_matrix(znear, zfar, fovx, fovy, **kwargs): t = znear * math.tan(0.5 * fovy) b = -t @@ -81,40 +82,42 @@ def projection_matrix(znear, zfar, fovx, fovy, **kwargs): [0.0, 0.0, (f + n) / (f - n), -1.0 * f * n / (f - n)], [0.0, 0.0, 1.0, 0.0], ], - **kwargs + **kwargs, ) @dataclass class GaussianSplattingModelConfig(ModelConfig): """Gaussian Splatting Model Config""" + _target: Type = field(default_factory=lambda: GaussianSplattingModel) - warmup_length:int = 900 + warmup_length: int = 900 """period of steps where refinement is turned off""" - refine_every:int = 100 + refine_every: int = 100 """period of steps where gaussians are culled and densified""" resolution_schedule: int = 300 """training starts at 1/d resolution, every n steps this is doubled""" num_downscales: int = 2 """at the beginning, resolution is 1/2^d, where d is this number""" - cull_alpha_thresh:float = .005 + cull_alpha_thresh: float = 0.005 """threshold of opacity for culling gaussians""" - cull_scale_thresh:float = .5 + cull_scale_thresh: float = 0.5 """threshold of scale for culling gaussians""" - reset_alpha_every:int = 30 + reset_alpha_every: int = 30 """Every this many refinement steps, reset the alpha""" - densify_grad_thresh: float = .0002 + densify_grad_thresh: float = 0.0002 """threshold of positional gradient norm for densifying gaussians""" - densify_size_thresh: float = .01 + densify_size_thresh: float = 0.01 """below this size, gaussians are *duplicated*, otherwise split""" - #set to 3 to effeectively turn off, it seems to make things worse? need to investigate + # set to 3 to effeectively turn off, it seems to make things worse? need to investigate sh_degree_interval: int = 1000 """every n intervals turn on another sh degree""" max_screen_size: int = 700 """maximum screen size of a gaussian, in pixels""" random_init: bool = False """whether to initialize the positions uniformly randomly (not SFM points)""" - ssim_lambda: float = .2 + ssim_lambda: float = 0.2 + class GaussianSplattingModel(Model): """Gaussian Splatting model @@ -124,12 +127,13 @@ class GaussianSplattingModel(Model): """ config: GaussianSplattingModelConfig - def __init__(self,*args,**kwargs): - if 'seed_points' in kwargs: - self.seed_pts = kwargs['seed_points'] + + def __init__(self, *args, **kwargs): + if "seed_points" in kwargs: + self.seed_pts = kwargs["seed_points"] else: self.seed_pts = None - super().__init__(*args,**kwargs) + super().__init__(*args, **kwargs) self.vc = ViewerControl() def populate_modules(self): @@ -137,18 +141,19 @@ class GaussianSplattingModel(Model): if self.seed_pts is not None and not self.config.random_init: # randmeans = torch.rand((30000,3))*8-4 # self.means = torch.nn.Parameter(torch.cat([self.seed_pts[0],randmeans],dim=0)) # (Location, Color) - self.means = torch.nn.Parameter(self.seed_pts[0]) # (Location, Color) + self.means = torch.nn.Parameter(self.seed_pts[0]) # (Location, Color) else: - self.means = torch.nn.Parameter((torch.rand((100000,3))-.5)*2) - self.means_grad_norm = None + self.means = torch.nn.Parameter((torch.rand((100000, 3)) - 0.5) * 2) + self.proj_means = torch.nn.Parameter(torch.zeros(self.num_points, 2)) + self.xys_grad_norm = None # init_scale = torch.log(torch.tensor(.01)).item() # self.scales = torch.nn.Parameter(torch.full((self.num_points,3),init_scale)) self.max_2Dsize = None distances, _ = self.k_nearest_sklearn(self.means.data, 3) distances = torch.from_numpy(distances) - #find the average of the three nearest neighbors for each point and use that as the scale - avg_dist = distances.mean(dim=-1,keepdim=True) - self.scales = torch.nn.Parameter(torch.log(avg_dist.repeat(1,3))) + # find the average of the three nearest neighbors for each point and use that as the scale + avg_dist = distances.mean(dim=-1, keepdim=True) + self.scales = torch.nn.Parameter(torch.log(avg_dist.repeat(1, 3))) # self.scales = torch.nn.Parameter(torch.full((self.num_points,3),init_scale)) self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points)) @@ -162,41 +167,45 @@ class GaussianSplattingModel(Model): # shs[:, 1:, :] = 0.0 # randcolors = torch.rand((30000,dim_sh,3)) # self.colors = torch.nn.Parameter(torch.cat([shs,randcolors])) - fused_color = self.RGB2SH(self.seed_pts[1]/255) + fused_color = self.RGB2SH(self.seed_pts[1] / 255) shs = torch.zeros((fused_color.shape[0], 3, dim_sh)).float().cuda() - shs[:, :3, 0 ] = fused_color + shs[:, :3, 0] = fused_color shs[:, 3:, 1:] = 0.0 - self.colors = torch.nn.Parameter(shs[:,:,0:1]) - self.shs_rest = torch.nn.Parameter(shs[:,:,1:]) + self.colors = torch.nn.Parameter(shs[:, :, 0:1]) + self.shs_rest = torch.nn.Parameter(shs[:, :, 1:]) else: self.colors = torch.nn.Parameter(torch.rand(self.num_points, 1, 3)) self.shs_rest = torch.nn.Parameter(torch.zeros((self.num_points, dim_sh - 1, 3))) - + self.opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1))) # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) from torchmetrics.image import StructuralSimilarityIndexMeasure + self.ssim = StructuralSimilarityIndexMeasure() self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) - self.step=0 - + self.step = 0 + @property def get_colors(self): color = self.colors shs_rest = self.shs_rest return torch.cat((color, shs_rest), dim=2) - - def load_state_dict(self,dict, **kwargs): - #resize the parameters to match the new number of points + + def load_state_dict(self, dict, **kwargs): + # resize the parameters to match the new number of points self.step = 30000 - newp = dict['means'].shape[0] - self.means = torch.nn.Parameter(torch.zeros(newp,3,device=self.device)) - self.scales = torch.nn.Parameter(torch.zeros(newp,3,device=self.device)) - self.quats = torch.nn.Parameter(torch.zeros(newp,4,device=self.device)) - self.colors = torch.nn.Parameter(torch.zeros(self.num_points, 1, 3,device=self.device)) - self.shs_rest = torch.nn.Parameter(torch.zeros(self.num_points, num_sh_bases(self.degree)-1, 3,device=self.device)) - self.opacities = torch.nn.Parameter(torch.zeros(newp,1,device=self.device)) - super().load_state_dict(dict,**kwargs) + newp = dict["means"].shape[0] + self.means = torch.nn.Parameter(torch.zeros(newp, 3, device=self.device)) + self.proj_means = torch.nn.Parameter(torch.zeros(newp, 2, device=self.device)) + self.scales = torch.nn.Parameter(torch.zeros(newp, 3, device=self.device)) + self.quats = torch.nn.Parameter(torch.zeros(newp, 4, device=self.device)) + self.colors = torch.nn.Parameter(torch.zeros(self.num_points, 1, 3, device=self.device)) + self.shs_rest = torch.nn.Parameter( + torch.zeros(self.num_points, num_sh_bases(self.degree) - 1, 3, device=self.device) + ) + self.opacities = torch.nn.Parameter(torch.zeros(newp, 1, device=self.device)) + super().load_state_dict(dict, **kwargs) def k_nearest_sklearn(self, x: torch.Tensor, k: int): """ @@ -208,108 +217,128 @@ class GaussianSplattingModel(Model): x_np = x.cpu().numpy() # Build the nearest neighbors model - nn_model = NearestNeighbors(n_neighbors=k + 1, algorithm='auto', metric='euclidean').fit(x_np) + nn_model = NearestNeighbors(n_neighbors=k + 1, algorithm="auto", metric="euclidean").fit(x_np) # Find the k-nearest neighbors distances, indices = nn_model.kneighbors(x_np) # Exclude the point itself from the result and return return distances[:, 1:].astype(np.float32), indices[:, 1:].astype(np.float32) - + def remove_from_optim(self, optimizer, deleted_mask, new_params): """removes the deleted_mask from the optimizer provided""" assert len(new_params) == 1 assert isinstance(optimizer, torch.optim.Adam), "Only works with Adam" - - param = optimizer.param_groups[0]['params'][0] + + param = optimizer.param_groups[0]["params"][0] param_state = optimizer.state[param] del optimizer.state[param] - + # Modify the state directly without deleting and reassigning. - param_state['exp_avg'] = param_state['exp_avg'][~deleted_mask.squeeze()] - param_state['exp_avg_sq'] = param_state['exp_avg_sq'][~deleted_mask.squeeze()] - + param_state["exp_avg"] = param_state["exp_avg"][~deleted_mask.flatten()] + param_state["exp_avg_sq"] = param_state["exp_avg_sq"][~deleted_mask.flatten()] + # Update the parameter in the optimizer's param group. - del optimizer.param_groups[0]['params'][0] - del optimizer.param_groups[0]['params'] - optimizer.param_groups[0]['params'] = new_params + del optimizer.param_groups[0]["params"][0] + del optimizer.param_groups[0]["params"] + optimizer.param_groups[0]["params"] = new_params optimizer.state[new_params[0]] = param_state - def dup_in_optim(self,optimizer, dup_mask, new_params): + def dup_in_optim(self, optimizer, dup_mask, new_params): """adds the parameters to the optimizer""" - param = optimizer.param_groups[0]['params'][0] + param = optimizer.param_groups[0]["params"][0] param_state = optimizer.state[param] - param_state['exp_avg'] = torch.cat([param_state['exp_avg'], torch.zeros_like(param_state['exp_avg'][dup_mask.squeeze()])],dim=0) - param_state['exp_avg_sq'] = torch.cat([param_state['exp_avg_sq'], torch.zeros_like(param_state['exp_avg_sq'][dup_mask.squeeze()])],dim=0) + param_state["exp_avg"] = torch.cat( + [param_state["exp_avg"], torch.zeros_like(param_state["exp_avg"][dup_mask.flatten()])], dim=0 + ) + param_state["exp_avg_sq"] = torch.cat( + [param_state["exp_avg_sq"], torch.zeros_like(param_state["exp_avg_sq"][dup_mask.flatten()])], dim=0 + ) del optimizer.state[param] optimizer.state[new_params[0]] = param_state - optimizer.param_groups[0]['params'] = new_params + optimizer.param_groups[0]["params"] = new_params del param - def after_train(self,step): + def after_train(self, step): with torch.no_grad(): - #keep track of a moving average of grad norms - visible_mask = self.radii>0 - if self.means_grad_norm is None: - self.means_grad_norm = self.means.grad.detach().norm(dim=-1,keepdim=True) - self.vis_counts = torch.zeros_like(self.means_grad_norm) + # keep track of a moving average of grad norms + visible_mask = (self.radii > 0).flatten() + grads = self.proj_means.grad.detach().norm(dim=-1) + # print(f"grad norm min {grads.min().item()} max {grads.max().item()} mean {grads.mean().item()}") + if self.xys_grad_norm is None: + self.xys_grad_norm = grads + self.vis_counts = torch.zeros_like(self.xys_grad_norm) else: self.vis_counts[visible_mask] = self.vis_counts[visible_mask] + 1 - self.means_grad_norm[visible_mask] = self.means.grad.detach()[visible_mask].norm(dim=-1,keepdim=True) + self.means_grad_norm[visible_mask] + self.xys_grad_norm[visible_mask] = grads[visible_mask] + self.xys_grad_norm[visible_mask] - #update the max screen size + # update the max screen size if self.max_2Dsize is None: self.max_2Dsize = torch.zeros_like(self.radii) newradii = self.radii[visible_mask] - self.max_2Dsize[visible_mask] = torch.maximum(self.max_2Dsize[visible_mask],newradii) + self.max_2Dsize[visible_mask] = torch.maximum(self.max_2Dsize[visible_mask], newradii) - def refinement_after(self, optimizers:Optimizers, step): + def refinement_after(self, optimizers: Optimizers, step): if self.step > self.config.warmup_length: with torch.no_grad(): - #then we densify - avg_grad_norm = self.means_grad_norm / self.vis_counts + # then we densify + avg_grad_norm = self.xys_grad_norm / self.vis_counts high_grads = (avg_grad_norm > self.config.densify_grad_thresh).squeeze() splits = (self.scales.exp().max(dim=-1).values > self.config.densify_size_thresh).squeeze() splits &= high_grads - split_means,split_colors,split_shs,split_opacities,split_scales,split_quats = self.split_gaussians(splits) + ( + split_means, + split_proj, + split_colors, + split_shs, + split_opacities, + split_scales, + split_quats, + ) = self.split_gaussians(splits) dups = (self.scales.exp().max(dim=-1).values <= self.config.densify_size_thresh).squeeze() dups &= high_grads - dup_means,dup_colors,dup_shs,dup_opacities,dup_scales,dup_quats = self.dup_gaussians(dups) - self.means = Parameter(torch.cat([self.means.detach(),split_means, dup_means],dim=0)) - self.colors = Parameter(torch.cat([self.colors.detach(),split_colors, dup_colors],dim=0)) - self.shs_rest = Parameter(torch.cat([self.shs_rest.detach(),split_shs, dup_shs],dim=0)) - self.opacities = Parameter(torch.cat([self.opacities.detach(),split_opacities, dup_opacities],dim=0)) - self.scales = Parameter(torch.cat([self.scales.detach(),split_scales, dup_scales],dim=0)) - self.quats = Parameter(torch.cat([self.quats.detach(),split_quats, dup_quats],dim=0)) - #append zeros to the max_2Dsize tensor - self.max_2Dsize = torch.cat([self.max_2Dsize,torch.zeros_like(split_scales[:,0]),torch.zeros_like(dup_scales[:,0])],dim=0) - + dup_means, dup_proj, dup_colors, dup_shs, dup_opacities, dup_scales, dup_quats = self.dup_gaussians( + dups + ) + self.means = Parameter(torch.cat([self.means.detach(), split_means, dup_means], dim=0)) + self.proj_means = Parameter(torch.cat([self.proj_means.detach(), split_proj, dup_proj], dim=0)) + self.colors = Parameter(torch.cat([self.colors.detach(), split_colors, dup_colors], dim=0)) + self.shs_rest = Parameter(torch.cat([self.shs_rest.detach(), split_shs, dup_shs], dim=0)) + self.opacities = Parameter(torch.cat([self.opacities.detach(), split_opacities, dup_opacities], dim=0)) + self.scales = Parameter(torch.cat([self.scales.detach(), split_scales, dup_scales], dim=0)) + self.quats = Parameter(torch.cat([self.quats.detach(), split_quats, dup_quats], dim=0)) + # append zeros to the max_2Dsize tensor + self.max_2Dsize = torch.cat( + [self.max_2Dsize, torch.zeros_like(split_scales[:, 0]), torch.zeros_like(dup_scales[:, 0])], dim=0 + ) + split_idcs = torch.where(splits)[0] dup_idcs = torch.where(dups)[0] add_idcs = torch.cat([split_idcs, dup_idcs], dim=0) - + param_groups = self.get_param_groups() - for group,param in param_groups.items(): - self.dup_in_optim(optimizers.optimizers[group],add_idcs,param) - #then cull + for group, param in param_groups.items(): + self.dup_in_optim(optimizers.optimizers[group], add_idcs, param) + # then cull deleted_mask = self.cull_gaussians() param_groups = self.get_param_groups() - for group,param in param_groups.items(): - self.remove_from_optim(optimizers.optimizers[group],deleted_mask,param) - + for group, param in param_groups.items(): + self.remove_from_optim(optimizers.optimizers[group], deleted_mask, param) if self.step // self.config.refine_every % self.config.reset_alpha_every == 0: print("Resetting alpha") - reset_value = .01 - self.opacities.data = torch.full_like(self.opacities.data,torch.logit(torch.tensor(reset_value)).item()) - #reset the exp of optimizer - optim = optimizers.optimizers['opacity'] - param = optim.param_groups[0]['params'][0] + reset_value = 0.01 + self.opacities.data = torch.full_like( + self.opacities.data, torch.logit(torch.tensor(reset_value)).item() + ) + # reset the exp of optimizer + optim = optimizers.optimizers["opacity"] + param = optim.param_groups[0]["params"][0] param_state = optim.state[param] - param_state['exp_avg'] = torch.zeros_like(param_state['exp_avg']) - param_state['exp_avg_sq'] = torch.zeros_like(param_state['exp_avg_sq']) - self.means_grad_norm = None + param_state["exp_avg"] = torch.zeros_like(param_state["exp_avg"]) + param_state["exp_avg_sq"] = torch.zeros_like(param_state["exp_avg_sq"]) + self.xys_grad_norm = None self.vis_counts = None self.max_2Dsize = None @@ -318,14 +347,15 @@ class GaussianSplattingModel(Model): This function deletes gaussians with under a certain opacity threshold """ n_bef = self.num_points - #cull transparent ones + # cull transparent ones culls = (torch.sigmoid(self.opacities) < self.config.cull_alpha_thresh).squeeze() - #cull huge ones + # cull huge ones culls = culls | (torch.exp(self.scales).max(dim=-1).values > self.config.cull_scale_thresh).squeeze() - #cull big screen space - if self.step>self.config.refine_every*self.config.reset_alpha_every: + # cull big screen space + if self.step > self.config.refine_every * self.config.reset_alpha_every: culls = culls | (self.max_2Dsize > self.config.max_screen_size).squeeze() self.means = Parameter(self.means[~culls].detach()) + self.proj_means = Parameter(self.proj_means[~culls].detach()) self.scales = Parameter(self.scales[~culls].detach()) self.quats = Parameter(self.quats[~culls].detach()) self.colors = Parameter(self.colors[~culls].detach()) @@ -341,20 +371,21 @@ class GaussianSplattingModel(Model): """ n_splits = split_mask.sum().item() print(f"Splitting {split_mask.sum().item()/self.num_points} gaussians: {n_splits}/{self.num_points}") - #step 1, sample new means - cov_mats = torch.eye(3,device=self.device)[None,...].repeat(n_splits,1,1) + # step 1, sample new means + cov_mats = torch.eye(3, device=self.device)[None, ...].repeat(n_splits, 1, 1) cov3ds = self.cov3d[split_mask] - cov_mats[:,0,0]=cov3ds[:,0] - cov_mats[:,0,1]=cov3ds[:,1] - cov_mats[:,1,0]=cov3ds[:,1] - cov_mats[:,0,2]=cov3ds[:,2] - cov_mats[:,2,0]=cov3ds[:,2] - cov_mats[:,1,1]=cov3ds[:,3] - cov_mats[:,1,2]=cov3ds[:,4] - cov_mats[:,2,1]=cov3ds[:,4] - cov_mats[:,2,2]=cov3ds[:,5] - centered_samples = torch.randn((n_splits,3),device=self.device) - new_means = torch.bmm(cov_mats,centered_samples[...,None]).squeeze() + self.means[split_mask] + cov_mats[:, 0, 0] = cov3ds[:, 0] + cov_mats[:, 0, 1] = cov3ds[:, 1] + cov_mats[:, 1, 0] = cov3ds[:, 1] + cov_mats[:, 0, 2] = cov3ds[:, 2] + cov_mats[:, 2, 0] = cov3ds[:, 2] + cov_mats[:, 1, 1] = cov3ds[:, 3] + cov_mats[:, 1, 2] = cov3ds[:, 4] + cov_mats[:, 2, 1] = cov3ds[:, 4] + cov_mats[:, 2, 2] = cov3ds[:, 5] + centered_samples = torch.randn((n_splits, 3), device=self.device) + new_means = torch.bmm(cov_mats, centered_samples[..., None]).squeeze() + self.means[split_mask] + new_proj = self.proj_means[split_mask] # step 2, sample new colors new_colors = self.colors[split_mask] # step 3, sample new shs @@ -362,46 +393,54 @@ class GaussianSplattingModel(Model): # step 3, sample new opacities new_opacities = self.opacities[split_mask] # step 4, sample new scales - new_scales = torch.logit(torch.exp(self.scales[split_mask])/1.6) - self.scales[split_mask] = torch.logit(torch.exp(self.scales[split_mask])/1.6) + new_scales = torch.logit(torch.exp(self.scales[split_mask]) / 1.6) + self.scales[split_mask] = torch.logit(torch.exp(self.scales[split_mask]) / 1.6) # step 5, sample new quats new_quats = self.quats[split_mask] - return new_means,new_colors,new_shs,new_opacities,new_scales,new_quats + return new_means, new_proj, new_colors, new_shs, new_opacities, new_scales, new_quats - def dup_gaussians(self,dup_mask): + def dup_gaussians(self, dup_mask): """ This function duplicates gaussians that are too small """ n_dups = dup_mask.sum().item() print(f"Would duplicate {dup_mask.sum().item()/self.num_points} gaussians: {n_dups}/{self.num_points}") dup_means = self.means[dup_mask] + self.means.grad[dup_mask] + dup_proj = self.proj_means[dup_mask] dup_colors = self.colors[dup_mask] dup_shs = self.shs_rest[dup_mask] dup_opacities = self.opacities[dup_mask] dup_scales = self.scales[dup_mask] dup_quats = self.quats[dup_mask] - return dup_means,dup_colors,dup_shs,dup_opacities,dup_scales,dup_quats - + return dup_means, dup_proj, dup_colors, dup_shs, dup_opacities, dup_scales, dup_quats + @property def num_points(self): return self.means.shape[0] - + def get_training_callbacks( self, training_callback_attributes: TrainingCallbackAttributes ) -> List[TrainingCallback]: cbs = [] - cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],self.step_cb)) + cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION], self.step_cb)) # The order of these matters - cbs.append(TrainingCallback([TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - self.after_train,)) - cbs.append(TrainingCallback([TrainingCallbackLocation.AFTER_TRAIN_ITERATION], - self.refinement_after, - update_every_num_iters=self.config.refine_every, - args=[training_callback_attributes.optimizers] - )) + cbs.append( + TrainingCallback( + [TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + self.after_train, + ) + ) + cbs.append( + TrainingCallback( + [TrainingCallbackLocation.AFTER_TRAIN_ITERATION], + self.refinement_after, + update_every_num_iters=self.config.refine_every, + args=[training_callback_attributes.optimizers], + ) + ) return cbs - - def step_cb(self,step): + + def step_cb(self, step): self.step = step def get_param_groups(self) -> Dict[str, List[Parameter]]: @@ -420,8 +459,8 @@ class GaussianSplattingModel(Model): } def _get_downscale_factor(self): - return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule),0) - + return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0) + def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]: """Takes in a Ray Bundle and returns a dictionary of outputs. @@ -436,35 +475,36 @@ class GaussianSplattingModel(Model): """ import viser.transforms as vtf import numpy as np - if not isinstance(camera,Cameras): + + if not isinstance(camera, Cameras): print("Called get_outputs with not a camera") return {} - assert camera.shape[0] ==1, "Only one camera at a time" - #dont mutate the input + assert camera.shape[0] == 1, "Only one camera at a time" + # dont mutate the input camera = deepcopy(camera) if self.training: d = self._get_downscale_factor() - camera.rescale_output_resolution(1/d) + camera.rescale_output_resolution(1 / d) - #shift the camera to center of scene looking at center - R = camera.camera_to_worlds[..., :3, :3] # 1 x 3 x 3 - T = camera.camera_to_worlds[..., :3, 3:4] # 1 x 3 x 1 + # shift the camera to center of scene looking at center + R = camera.camera_to_worlds[..., :3, :3] # 1 x 3 x 3 + T = camera.camera_to_worlds[..., :3, 3:4] # 1 x 3 x 1 R = vtf.SO3.from_matrix(R.cpu().squeeze().numpy()) R = R @ vtf.SO3.from_x_radians(np.pi) - R = torch.from_numpy(R.as_matrix()[None,...]).to(self.device,torch.float32) - #vickie pops off here - viewmat = torch.cat([R,T],dim=2) - #add a row of zeros and a 1 to the bottom of the viewmat - viewmat = torch.cat([viewmat,torch.tensor([[[0,0,0,1]]],device=self.device)],dim=1) - #invert it + R = torch.from_numpy(R.as_matrix()[None, ...]).to(self.device, torch.float32) + # vickie pops off here + viewmat = torch.cat([R, T], dim=2) + # add a row of zeros and a 1 to the bottom of the viewmat + viewmat = torch.cat([viewmat, torch.tensor([[[0, 0, 0, 1]]], device=self.device)], dim=1) + # invert it viewmat = torch.inverse(viewmat) - #calculate the FOV of the camera given fx and fy, width and height + # calculate the FOV of the camera given fx and fy, width and height cx = camera.cx.item() cy = camera.cy.item() fovx = 2 * math.atan(camera.width / (2 * camera.fx)) fovy = 2 * math.atan(camera.height / (2 * camera.fy)) W, H = camera.width.item(), camera.height.item() - projmat = projection_matrix(.1,1000,fovx,fovy).to(self.device) + projmat = projection_matrix(0.1, 1000, fovx, fovy).to(self.device) BLOCK_X, BLOCK_Y = 16, 16 tile_bounds = ( (W + BLOCK_X - 1) // BLOCK_X, @@ -475,28 +515,29 @@ class GaussianSplattingModel(Model): background = torch.rand(3, device=self.device) else: background = torch.zeros(3, device=self.device) - xys, depths, self.radii, conics, num_tiles_hit,self.cov3d = ProjectGaussians.apply( + xys, depths, self.radii, conics, num_tiles_hit, self.cov3d = ProjectGaussians.apply( self.means, + self.proj_means, torch.exp(self.scales), 1, self.quats, - viewmat.squeeze()[:3,:], - projmat.squeeze()@viewmat.squeeze(), + viewmat.squeeze()[:3, :], + projmat.squeeze() @ viewmat.squeeze(), camera.fx.item(), camera.fy.item(), H, W, - tile_bounds + tile_bounds, ) if self.degree > 0: viewdirs = self.means - 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.degree) + n = min(self.step // self.config.sh_degree_interval, self.degree) n_bases = num_sh_bases(n) # rgbs = compute_sh_color(viewdirs, self.colors[:,:n_bases]) rgbs = eval_sh(n, self.get_colors, viewdirs) # rgbs = torch.clamp_min(rgbs + 0.5, 0.0) - print("Min", rgbs.min(), "Max", rgbs.max()) + # print("Min", rgbs.min().item(), "Max", rgbs.max().item()) # rgbs = torch.clamp_max(rgbs, 1.0) # rgbs = SphericalHarmonics.apply(n, viewdirs, self.colors[:,:n_bases]) # (N, 3) else: @@ -519,7 +560,6 @@ class GaussianSplattingModel(Model): ) return {"rgb": rgb} - def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: """Compute and returns metrics. @@ -540,24 +580,27 @@ class GaussianSplattingModel(Model): """ d = self._get_downscale_factor() if d > 1: - #use torchvision to resize + # use torchvision to resize import torchvision.transforms.functional as TF - newsize = (batch['image'].shape[0]//d,batch['image'].shape[1]//d) - gt_img = TF.resize(batch['image'].permute(2,0,1),newsize).permute(1,2,0) + + newsize = (batch["image"].shape[0] // d, batch["image"].shape[1] // d) + gt_img = TF.resize(batch["image"].permute(2, 0, 1), newsize).permute(1, 2, 0) else: - gt_img = batch['image'] - Ll1 = torch.nn.functional.l1_loss(gt_img, outputs['rgb']) - simloss = (1-self.ssim(gt_img.permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...]))/2.0 - return {"main_loss": (1-self.config.ssim_lambda)*Ll1 + self.config.ssim_lambda*simloss} + gt_img = batch["image"] + Ll1 = torch.nn.functional.l1_loss(gt_img, outputs["rgb"]) + simloss = (1 - self.ssim(gt_img.permute(2, 0, 1)[None, ...], outputs["rgb"].permute(2, 0, 1)[None, ...])) / 2.0 + return {"main_loss": (1 - self.config.ssim_lambda) * Ll1 + self.config.ssim_lambda * simloss} @torch.no_grad() - def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle, camera:Optional[Cameras]=None) -> Dict[str, torch.Tensor]: + def get_outputs_for_camera_ray_bundle( + self, camera_ray_bundle: RayBundle, camera: Optional[Cameras] = None + ) -> Dict[str, torch.Tensor]: """Takes in camera parameters and computes the output of the model. Args: camera_ray_bundle: ray bundle to calculate outputs over """ - assert camera is not None,'must provide camera to gaussian model' + assert camera is not None, "must provide camera to gaussian model" outs = self.get_outputs(camera.to(self.device)) return outs @@ -599,7 +642,8 @@ class GaussianSplattingModel(Model): def RGB2SH(self, rgb): C0 = 0.28209479177387814 return (rgb - 0.5) / C0 - + + def eval_sh(deg, sh, dirs): """ Evaluate spherical harmonics at unit directions @@ -615,13 +659,7 @@ def eval_sh(deg, sh, dirs): """ C0 = 0.28209479177387814 C1 = 0.4886025119029199 - C2 = [ - 1.0925484305920792, - -1.0925484305920792, - 0.31539156525252005, - -1.0925484305920792, - 0.5462742152960396 - ] + C2 = [1.0925484305920792, -1.0925484305920792, 0.31539156525252005, -1.0925484305920792, 0.5462742152960396] C3 = [ -0.5900435899266435, 2.890611442640554, @@ -629,7 +667,7 @@ def eval_sh(deg, sh, dirs): 0.3731763325901154, -0.4570457994644658, 1.445305721320277, - -0.5900435899266435 + -0.5900435899266435, ] C4 = [ 2.5033429417967046, @@ -641,7 +679,7 @@ def eval_sh(deg, sh, dirs): 0.47308734787878004, -1.7701307697799304, 0.6258357354491761, - ] + ] assert deg <= 4 and deg >= 0 coeff = (deg + 1) ** 2 @@ -650,39 +688,43 @@ def eval_sh(deg, sh, dirs): 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]) + 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]) + 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]) + 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 \ No newline at end of file + 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