From f4cc901e08e801c9d921f33292a5d59c35b7e18e Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Fri, 13 Oct 2023 00:22:27 -0700 Subject: [PATCH] lint --- nerfstudio/configs/method_configs.py | 20 +- .../datamanagers/full_images_datamanager.py | 81 ++++--- nerfstudio/models/base_model.py | 8 +- nerfstudio/models/gaussian_splatting.py | 205 +++++++++++------- nerfstudio/pipelines/base_pipeline.py | 14 +- nerfstudio/scripts/render.py | 8 +- nerfstudio/utils/gaussian_utils.py | 15 +- .../viewer/server/render_state_machine.py | 7 +- nerfstudio/viewer_beta/control_panel.py | 2 +- .../viewer_beta/render_state_machine.py | 29 ++- nerfstudio/viewer_beta/utils.py | 2 +- 11 files changed, 231 insertions(+), 160 deletions(-) diff --git a/nerfstudio/configs/method_configs.py b/nerfstudio/configs/method_configs.py index 22448ae21..8e57546ba 100644 --- a/nerfstudio/configs/method_configs.py +++ b/nerfstudio/configs/method_configs.py @@ -49,14 +49,14 @@ from nerfstudio.data.datasets.depth_dataset import DepthDataset from nerfstudio.data.datasets.sdf_dataset import SDFDataset from nerfstudio.data.datasets.semantic_dataset import SemanticDataset from nerfstudio.data.pixel_samplers import PairPixelSamplerConfig -from nerfstudio.engine.optimizers import (AdamOptimizerConfig, - RAdamOptimizerConfig) -from nerfstudio.engine.schedulers import (CosineDecaySchedulerConfig, - ExponentialDecaySchedulerConfig, - MultiStepSchedulerConfig) +from nerfstudio.engine.optimizers import AdamOptimizerConfig, RAdamOptimizerConfig +from nerfstudio.engine.schedulers import ( + CosineDecaySchedulerConfig, + ExponentialDecaySchedulerConfig, + MultiStepSchedulerConfig, +) from nerfstudio.engine.trainer import TrainerConfig -from nerfstudio.field_components.temporal_distortions import \ - TemporalDistortionKind +from nerfstudio.field_components.temporal_distortions import TemporalDistortionKind from nerfstudio.fields.sdf_field import SDFFieldConfig from nerfstudio.models.depth_nerfacto import DepthNerfactoModelConfig from nerfstudio.models.gaussian_splatting import GaussianSplattingModelConfig @@ -629,9 +629,9 @@ method_configs["gaussian-splatting"] = TrainerConfig( ), }, "shs": { - "optimizer": AdamOptimizerConfig(lr=2.5e-3/20, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=2.5e-3 / 20, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( - lr_final=1e-3/20, + lr_final=1e-3 / 20, max_steps=30000, ), }, @@ -647,7 +647,7 @@ method_configs["gaussian-splatting"] = TrainerConfig( ), }, "rotation": { - "optimizer": AdamOptimizerConfig(lr=.001, eps=1e-15), + "optimizer": AdamOptimizerConfig(lr=0.001, eps=1e-15), "scheduler": ExponentialDecaySchedulerConfig( lr_final=1e-4, max_steps=15000, diff --git a/nerfstudio/data/datamanagers/full_images_datamanager.py b/nerfstudio/data/datamanagers/full_images_datamanager.py index 144581162..2407f6727 100644 --- a/nerfstudio/data/datamanagers/full_images_datamanager.py +++ b/nerfstudio/data/datamanagers/full_images_datamanager.py @@ -25,28 +25,25 @@ import random from dataclasses import dataclass, field from functools import cached_property from pathlib import Path -from typing import (Dict, ForwardRef, Generic, List, Literal, - Optional, Tuple, Type, Union, cast, get_args, get_origin) +from typing import Dict, ForwardRef, Generic, List, Literal, Optional, Tuple, Type, Union, cast, get_args, get_origin import cv2 import numpy as np import torch from copy import deepcopy -from rich.progress import Console from torch.nn import Parameter from tqdm import tqdm from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.configs.dataparser_configs import AnnotatedDataParserUnion -from nerfstudio.data.datamanagers.base_datamanager import (DataManager, - DataManagerConfig,TDataset) +from nerfstudio.data.datamanagers.base_datamanager import DataManager, DataManagerConfig, TDataset from nerfstudio.data.dataparsers.base_dataparser import DataparserOutputs -from nerfstudio.data.dataparsers.nerfstudio_dataparser import \ - NerfstudioDataParserConfig +from nerfstudio.data.dataparsers.nerfstudio_dataparser import NerfstudioDataParserConfig from nerfstudio.data.datasets.base_dataset import InputDataset from nerfstudio.utils.misc import get_orig_class from nerfstudio.utils.rich_utils import CONSOLE + @dataclass class FullImageDatamanagerConfig(DataManagerConfig): _target: Type = field(default_factory=lambda: FullImageDatamanager) @@ -65,6 +62,7 @@ class FullImageDatamanagerConfig(DataManagerConfig): cache_images: Literal["no-cache", "cpu", "gpu"] = "cpu" """Whether to cache images in memory. If "numpy", caches as numpy arrays, if "torch", caches as torch tensors.""" + class FullImageDatamanager(DataManager, Generic[TDataset]): """ A datamanager that outputs full images and cameras instead of raybundles. This makes the @@ -101,7 +99,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): if test_mode == "inference": self.dataparser.downscale_factor = 1 # Avoid opening images self.includes_time = self.dataparser.includes_time - + self.train_dataparser_outputs: DataparserOutputs = self.dataparser.get_dataparser_outputs(split="train") self.train_dataset = self.create_train_dataset() self.eval_dataset = self.create_eval_dataset() @@ -119,7 +117,6 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): super().__init__() - def cache_images(self, cache_images_option): cached_train = [] CONSOLE.log("Caching / undistorting train images") @@ -132,10 +129,23 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): image = data["image"].numpy() if camera.camera_type.item() == CameraType.PERSPECTIVE.value: - distortion_params = np.array([distortion_params[0], distortion_params[1], distortion_params[4], distortion_params[5], distortion_params[2], distortion_params[3], 0, 0]) - image = cv2.undistort(image, K, distortion_params, None, K) # Should update K in-place + distortion_params = np.array( + [ + distortion_params[0], + distortion_params[1], + distortion_params[4], + distortion_params[5], + distortion_params[2], + distortion_params[3], + 0, + 0, + ] + ) + image = cv2.undistort(image, K, distortion_params, None, K) # Should update K in-place elif camera.camera_type.item() == CameraType.FISHEYE.value: - distortion_params = np.array([distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]]) + distortion_params = np.array( + [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]] + ) image = cv2.fisheye.undistortImage(image, K, distortion_params, None, K) else: raise NotImplementedError("Only perspective and fisheye cameras are supported") @@ -169,10 +179,23 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): image = data["image"].numpy() if camera.camera_type.item() == CameraType.PERSPECTIVE.value: - distortion_params = np.array([distortion_params[0], distortion_params[1], distortion_params[4], distortion_params[5], distortion_params[2], distortion_params[3], 0, 0]) - image = cv2.undistort(image, K, distortion_params, None, K) # Should update K in-place + distortion_params = np.array( + [ + distortion_params[0], + distortion_params[1], + distortion_params[4], + distortion_params[5], + distortion_params[2], + distortion_params[3], + 0, + 0, + ] + ) + image = cv2.undistort(image, K, distortion_params, None, K) # Should update K in-place elif camera.camera_type.item() == CameraType.FISHEYE.value: - distortion_params = np.array([distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]]) + distortion_params = np.array( + [distortion_params[0], distortion_params[1], distortion_params[2], distortion_params[3]] + ) image = cv2.fisheye.undistortImage(image, K, distortion_params, None, K) else: raise NotImplementedError("Only perspective and fisheye cameras are supported") @@ -245,7 +268,7 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): if issubclass(value, InputDataset): return cast(Type[TDataset], value) return default - + def get_datapath(self) -> Path: return self.config.dataparser.data @@ -261,32 +284,32 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): A list of dictionaries containing the data manager's param groups. """ return {} - + def get_train_rays_per_batch(self): # TODO: fix this to be the resolution of the last image rendered - return 800*800 - + return 800 * 800 + def next_train(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next training batch - + Returns a Camera instead of raybundle""" - image_idx = self.train_unseen_cameras.pop(random.randint(0, len(self.train_unseen_cameras)-1)) + image_idx = self.train_unseen_cameras.pop(random.randint(0, len(self.train_unseen_cameras) - 1)) # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.train_unseen_cameras) == 0: self.train_unseen_cameras = [i for i in range(len(self.train_dataset))] - + data = deepcopy(self.cached_train[image_idx]) data["image"] = data["image"].to(self.device) - + assert len(self.train_dataset.cameras.shape) == 1, "Assumes single batch dimension" - camera = self.train_dataset.cameras[image_idx:image_idx+1].to(self.device) + camera = self.train_dataset.cameras[image_idx : image_idx + 1].to(self.device) return camera, data def next_eval(self, step: int) -> Tuple[Cameras, Dict]: """Returns the next evaluation batch - + Returns a Camera instead of raybundle""" - image_idx = self.eval_unseen_cameras.pop(random.randint(0, len(self.eval_unseen_cameras)-1)) + image_idx = self.eval_unseen_cameras.pop(random.randint(0, len(self.eval_unseen_cameras) - 1)) # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.eval_unseen_cameras) == 0: self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] @@ -298,11 +321,11 @@ class FullImageDatamanager(DataManager, Generic[TDataset]): def next_eval_image(self, step: int) -> Tuple[int, Cameras, Dict]: """Returns the next evaluation batch - + Returns a Camera instead of raybundle - + TODO: Make sure this logic is consistent with the vanilladatamanager""" - image_idx = self.eval_unseen_cameras.pop(random.randint(0, len(self.eval_unseen_cameras)-1)) + image_idx = self.eval_unseen_cameras.pop(random.randint(0, len(self.eval_unseen_cameras) - 1)) # Make sure to re-populate the unseen cameras list if we have exhausted it if len(self.eval_unseen_cameras) == 0: self.eval_unseen_cameras = [i for i in range(len(self.eval_dataset))] diff --git a/nerfstudio/models/base_model.py b/nerfstudio/models/base_model.py index 9f0ae77d9..3d52f3f23 100644 --- a/nerfstudio/models/base_model.py +++ b/nerfstudio/models/base_model.py @@ -32,8 +32,7 @@ from nerfstudio.cameras.rays import RayBundle from nerfstudio.configs.base_config import InstantiateConfig from nerfstudio.configs.config_utils import to_immutable_dict from nerfstudio.data.scene_box import SceneBox -from nerfstudio.engine.callbacks import (TrainingCallback, - TrainingCallbackAttributes) +from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes from nerfstudio.model_components.scene_colliders import NearFarCollider @@ -164,8 +163,9 @@ class Model(nn.Module): """ @torch.no_grad() - def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: Union[RayBundle, Cameras], - camera:Optional[Cameras] = None) -> Dict[str, torch.Tensor]: + def get_outputs_for_camera_ray_bundle( + self, camera_ray_bundle: Union[RayBundle, Cameras], camera: Optional[Cameras] = None + ) -> Dict[str, torch.Tensor]: """Takes in camera parameters and computes the output of the model. TODO (jake-austin): Make 2 cases for when this is a ray or camera based model Args: diff --git a/nerfstudio/models/gaussian_splatting.py b/nerfstudio/models/gaussian_splatting.py index 3226d1c9b..981238fa7 100644 --- a/nerfstudio/models/gaussian_splatting.py +++ b/nerfstudio/models/gaussian_splatting.py @@ -28,6 +28,7 @@ from torch.nn import Parameter from torchmetrics.image import PeakSignalNoiseRatio from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.cameras import Cameras +from gsplat._torch_impl import quat_to_rotmat from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation from nerfstudio.engine.optimizers import Optimizers from nerfstudio.models.base_model import Model, ModelConfig @@ -41,6 +42,7 @@ from gsplat.rasterize import RasterizeGaussians from gsplat.project_gaussians import ProjectGaussians from gsplat.sh import SphericalHarmonics, num_sh_bases + def eval_sh(deg, sh, dirs): """ Evaluate spherical harmonics at unit directions @@ -126,6 +128,7 @@ def eval_sh(deg, sh, dirs): ) return result + def random_quat_tensor(N, **kwargs): u = torch.rand(N, **kwargs) v = torch.rand(N, **kwargs) @@ -140,10 +143,12 @@ def random_quat_tensor(N, **kwargs): dim=-1, ) + def RGB2SH(rgb): C0 = 0.28209479177387814 return (rgb - 0.5) / C0 + def projection_matrix(znear, zfar, fovx, fovy, **kwargs): t = znear * math.tan(0.5 * fovy) b = -t @@ -167,7 +172,7 @@ class GaussianSplattingModelConfig(ModelConfig): """Gaussian Splatting Model Config""" _target: Type = field(default_factory=lambda: GaussianSplattingModel) - warmup_length:int = 1000 + warmup_length: int = 1000 """period of steps where refinement is turned off""" refine_every: int = 100 """period of steps where gaussians are culled and densified""" @@ -191,13 +196,14 @@ class GaussianSplattingModelConfig(ModelConfig): """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 """weight of ssim loss""" stop_split_at: int = 25000 """stop splitting at this step""" sh_degree: int = 4 """maximum degree of spherical harmonics to use""" + class GaussianSplattingModel(Model): """Gaussian Splatting model @@ -249,9 +255,9 @@ class GaussianSplattingModel(Model): self.ssim = StructuralSimilarityIndexMeasure() self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) self.step = 0 - self.crop_box:Optional[OrientedBox] = None + self.crop_box: Optional[OrientedBox] = None self.back_color = torch.zeros(3) - + @property def get_colors(self): color = self.colors @@ -261,14 +267,16 @@ class GaussianSplattingModel(Model): 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, 3, 1,device=self.device)) - self.shs_rest = torch.nn.Parameter(torch.zeros(self.num_points, 3, num_sh_bases(self.config.sh_degree)-1,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.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, 3, 1, device=self.device)) + self.shs_rest = torch.nn.Parameter( + torch.zeros(self.num_points, 3, num_sh_bases(self.config.sh_degree) - 1, 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): """ @@ -298,22 +306,31 @@ class GaussianSplattingModel(Model): del optimizer.state[param] # Modify the state directly without deleting and reassigning. - param_state['exp_avg'] = param_state['exp_avg'][~deleted_mask] - param_state['exp_avg_sq'] = param_state['exp_avg_sq'][~deleted_mask] - + param_state["exp_avg"] = param_state["exp_avg"][~deleted_mask] + param_state["exp_avg_sq"] = param_state["exp_avg_sq"][~deleted_mask] + # 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 optimizer.state[new_params[0]] = param_state - def dup_in_optim(self,optimizer, dup_mask, new_params,n=2): + def dup_in_optim(self, optimizer, dup_mask, new_params, n=2): """adds the parameters to the optimizer""" param = optimizer.param_groups[0]["params"][0] param_state = optimizer.state[param] - repeat_dims = (n,) + tuple(1 for _ in range(param_state['exp_avg'].dim() - 1)) - param_state['exp_avg'] = torch.cat([param_state['exp_avg'], torch.zeros_like(param_state['exp_avg'][dup_mask.squeeze()]).repeat(*repeat_dims)],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()]).repeat(*repeat_dims)],dim=0) + repeat_dims = (n,) + tuple(1 for _ in range(param_state["exp_avg"].dim() - 1)) + param_state["exp_avg"] = torch.cat( + [param_state["exp_avg"], torch.zeros_like(param_state["exp_avg"][dup_mask.squeeze()]).repeat(*repeat_dims)], + 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()]).repeat(*repeat_dims), + ], + dim=0, + ) del optimizer.state[param] optimizer.state[new_params[0]] = param_state optimizer.param_groups[0]["params"] = new_params @@ -323,7 +340,7 @@ class GaussianSplattingModel(Model): with torch.no_grad(): # keep track of a moving average of grad norms visible_mask = (self.radii > 0).flatten() - grads = self.xys.grad.detach().norm(dim=-1)#TODO fill in + grads = self.xys.grad.detach().norm(dim=-1) # TODO fill in # print(f"grad norm min {grads.min().item()} max {grads.max().item()} mean {grads.mean().item()} size {grads.shape}") if self.xys_grad_norm is None: self.xys_grad_norm = grads @@ -337,47 +354,65 @@ class GaussianSplattingModel(Model): 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) - - def set_crop(self,crop_box:OrientedBox): + + def set_crop(self, crop_box: OrientedBox): self.crop_box = crop_box + def set_background(self, back_color: torch.Tensor): + assert back_color.shape == (3,) + self.back_color = back_color + def refinement_after(self, optimizers: Optimizers, step): if self.step > self.config.warmup_length: with torch.no_grad(): - if self.step self.config.densify_grad_thresh).squeeze() splits = (self.scales.exp().max(dim=-1).values > self.config.densify_size_thresh).squeeze() splits &= high_grads - nsamps=2 - split_means,split_colors,split_shs,split_opacities,split_scales,split_quats = self.split_gaussians(splits,nsamps) + nsamps = 2 + ( + split_means, + split_colors, + split_shs, + split_opacities, + split_scales, + split_quats, + ) = self.split_gaussians(splits, nsamps) 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_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, + ) split_idcs = torch.where(splits)[0] param_groups = self.get_param_groups() - for group,param in param_groups.items(): - self.dup_in_optim(optimizers.optimizers[group],split_idcs,param,n=nsamps) + for group, param in param_groups.items(): + self.dup_in_optim(optimizers.optimizers[group], split_idcs, param, n=nsamps) dup_idcs = torch.where(dups)[0] - + param_groups = self.get_param_groups() - for group,param in param_groups.items(): - self.dup_in_optim(optimizers.optimizers[group],dup_idcs,param,1) - - #only cull if we've seen every image since opacity reset - reset_interval = (self.config.reset_alpha_every * self.config.refine_every) + for group, param in param_groups.items(): + self.dup_in_optim(optimizers.optimizers[group], dup_idcs, param, 1) + + # only cull if we've seen every image since opacity reset + reset_interval = self.config.reset_alpha_every * self.config.refine_every if self.step % reset_interval > self.num_train_data: - #then cull + # then cull deleted_mask = self.cull_gaussians() param_groups = self.get_param_groups() for group, param in param_groups.items(): @@ -421,38 +456,40 @@ class GaussianSplattingModel(Model): print(f"Culled {n_bef - self.num_points} gaussians") return culls - def split_gaussians(self, split_mask,samps): + def split_gaussians(self, split_mask, samps): """ This function splits gaussians that are too large """ - from gsplat._torch_impl import quat_to_rotmat + n_splits = split_mask.sum().item() print(f"Splitting {split_mask.sum().item()/self.num_points} gaussians: {n_splits}/{self.num_points}") - centered_samples = torch.randn((samps*n_splits,3),device=self.device) # Nx3 of axis-aligned scales - scaled_samples = torch.exp(self.scales[split_mask].repeat(samps,1)) * centered_samples # how these scales are rotated - quats = self.quats[split_mask]/self.quats[split_mask].norm(dim=-1,keepdim=True) #normalize them first - rots = quat_to_rotmat(quats.repeat(samps,1)) # how these scales are rotated - rotated_samples = torch.bmm(rots,scaled_samples[...,None]).squeeze() - new_means = rotated_samples + self.means[split_mask].repeat(samps,1) + centered_samples = torch.randn((samps * n_splits, 3), device=self.device) # Nx3 of axis-aligned scales + scaled_samples = ( + torch.exp(self.scales[split_mask].repeat(samps, 1)) * centered_samples + ) # how these scales are rotated + quats = self.quats[split_mask] / self.quats[split_mask].norm(dim=-1, keepdim=True) # normalize them first + rots = quat_to_rotmat(quats.repeat(samps, 1)) # how these scales are rotated + rotated_samples = torch.bmm(rots, scaled_samples[..., None]).squeeze() + new_means = rotated_samples + self.means[split_mask].repeat(samps, 1) # step 2, sample new colors - new_colors = self.colors[split_mask].repeat(samps,1,1) + new_colors = self.colors[split_mask].repeat(samps, 1, 1) # step 3, sample new shs - new_shs = self.shs_rest[split_mask].repeat(samps,1,1) + new_shs = self.shs_rest[split_mask].repeat(samps, 1, 1) # step 3, sample new opacities - new_opacities = self.opacities[split_mask].repeat(samps,1) + new_opacities = self.opacities[split_mask].repeat(samps, 1) # step 4, sample new scales - new_scales = torch.log(torch.exp(self.scales[split_mask])/1.6).repeat(samps,1) - self.scales[split_mask] = torch.log(torch.exp(self.scales[split_mask])/1.6) + new_scales = torch.log(torch.exp(self.scales[split_mask]) / 1.6).repeat(samps, 1) + self.scales[split_mask] = torch.log(torch.exp(self.scales[split_mask]) / 1.6) # step 5, sample new quats - new_quats = self.quats[split_mask].repeat(samps,1) - return new_means,new_colors,new_shs,new_opacities,new_scales,new_quats + new_quats = self.quats[split_mask].repeat(samps, 1) + return new_means, new_colors, new_shs, new_opacities, new_scales, new_quats 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}") + print(f"Duplicating {dup_mask.sum().item()/self.num_points} gaussians: {n_dups}/{self.num_points}") dup_means = self.means[dup_mask] dup_colors = self.colors[dup_mask] dup_shs = self.shs_rest[dup_mask] @@ -526,16 +563,16 @@ class GaussianSplattingModel(Model): camera = deepcopy(camera) if self.training: d = self._get_downscale_factor() - camera.rescale_output_resolution(1/d) + camera.rescale_output_resolution(1 / d) if self.crop_box is not None and not self.training: crop_ids = self.crop_box.within(self.means).squeeze() - if crop_ids.sum()==0: - return {'rgb':torch.full((camera.height.item(),camera.width.item(),3),.5,device=self.device)} + if crop_ids.sum() == 0: + return {"rgb": torch.full((camera.height.item(), camera.width.item(), 3), 0.5, device=self.device)} else: - crop_ids = torch.ones_like(self.means[:,0],dtype=torch.bool) - #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 + crop_ids = torch.ones_like(self.means[:, 0], dtype=torch.bool) + # 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) @@ -561,8 +598,8 @@ class GaussianSplattingModel(Model): if self.training: background = torch.rand(3, device=self.device) else: - background = torch.zeros(3, device=self.device) - + background = self.back_color + opacities_crop = self.opacities[crop_ids] means_crop = self.means[crop_ids] colors_crop = self.get_colors[crop_ids] @@ -581,19 +618,19 @@ class GaussianSplattingModel(Model): W, tile_bounds, ) - #Important to allow xys grads to populate properly + # Important to allow xys grads to populate properly if self.training: self.xys.retain_grad() - rend_mask = self.radii>0 + rend_mask = self.radii > 0 if self.config.sh_degree > 0: 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) + n = min(self.step // self.config.sh_degree_interval, self.config.sh_degree) # 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) - rgbs = torch.clamp(rgbs + 0.5,0.0,1.0) + coeffs = colors_crop.permute(0, 2, 1)[rend_mask, :n_coeffs, :] + rgbs = SphericalHarmonics.apply(n, viewdirs, coeffs) + rgbs = torch.clamp(rgbs + 0.5, 0.0, 1.0) else: rgbs = self.get_colors.squeeze()[rend_mask] # (N, 3) rgbs = torch.sigmoid(rgbs) @@ -615,13 +652,13 @@ class GaussianSplattingModel(Model): self.radii[rend_mask], conics[rend_mask], num_tiles_hit[rend_mask], - depths[rend_mask,None].repeat(1,3), + depths[rend_mask, None].repeat(1, 3), torch.sigmoid(opacities_crop[rend_mask]), H, W, - torch.ones(3,device=self.device)*10, - )[...,0:1] - return {"rgb": rgb,'depth':depth_im} + torch.ones(3, 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. @@ -649,10 +686,10 @@ class GaussianSplattingModel(Model): 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.abs(gt_img- outputs['rgb']).mean() - simloss = (1-self.ssim(gt_img.permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...])) - return {"main_loss": (1-self.config.ssim_lambda)*Ll1 + self.config.ssim_lambda*simloss} + gt_img = batch["image"] + Ll1 = torch.abs(gt_img - outputs["rgb"]).mean() + simloss = 1 - self.ssim(gt_img.permute(2, 0, 1)[None, ...], outputs["rgb"].permute(2, 0, 1)[None, ...]) + return {"main_loss": (1 - self.config.ssim_lambda) * Ll1 + self.config.ssim_lambda * simloss} @torch.no_grad() def get_outputs_for_camera_ray_bundle( diff --git a/nerfstudio/pipelines/base_pipeline.py b/nerfstudio/pipelines/base_pipeline.py index b3f27b91d..a520f563c 100644 --- a/nerfstudio/pipelines/base_pipeline.py +++ b/nerfstudio/pipelines/base_pipeline.py @@ -22,14 +22,12 @@ from abc import abstractmethod from dataclasses import dataclass, field from pathlib import Path from time import time -from typing import (Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, - Union, cast) +from typing import Any, Dict, List, Literal, Mapping, Optional, Tuple, Type, Union, cast import torch import torch.distributed as dist from PIL import Image -from rich.progress import (BarColumn, MofNCompleteColumn, Progress, TextColumn, - TimeElapsedColumn) +from rich.progress import BarColumn, MofNCompleteColumn, Progress, TextColumn, TimeElapsedColumn from torch import nn from torch.cuda.amp.grad_scaler import GradScaler from torch.nn import Parameter @@ -261,9 +259,9 @@ class VanillaPipeline(Pipeline): self.datamanager: DataManager = config.datamanager.setup( device=device, test_mode=test_mode, world_size=world_size, local_rank=local_rank ) - #TODO make this work with normal nerfacto too - pts = self.datamanager.train_dataparser_outputs.metadata['points3D_xyz'] - pts_rgb = self.datamanager.train_dataparser_outputs.metadata['points3D_rgb'] + # TODO make this work with normal nerfacto too + pts = self.datamanager.train_dataparser_outputs.metadata["points3D_xyz"] + pts_rgb = self.datamanager.train_dataparser_outputs.metadata["points3D_rgb"] self.datamanager.to(device) # TODO(ethan): get rid of scene_bounds from the model assert self.datamanager.train_dataset is not None, "Missing input dataset" @@ -274,7 +272,7 @@ class VanillaPipeline(Pipeline): metadata=self.datamanager.train_dataset.metadata, device=device, grad_scaler=grad_scaler, - seed_points = (pts,pts_rgb) + seed_points=(pts, pts_rgb), ) self.model.to(device) diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 3ce172389..f20c4c7b2 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -140,10 +140,14 @@ def _render_trajectory_video( with renderers.background_color_override_context( crop_data.background_color.to(pipeline.device) ), torch.no_grad(): - outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=cameras[camera_idx:camera_idx+1]) + outputs = pipeline.model.get_outputs_for_camera_ray_bundle( + camera_ray_bundle, camera=cameras[camera_idx : camera_idx + 1] + ) else: with torch.no_grad(): - outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=cameras[camera_idx:camera_idx+1]) + outputs = pipeline.model.get_outputs_for_camera_ray_bundle( + camera_ray_bundle, camera=cameras[camera_idx : camera_idx + 1] + ) render_image = [] for rendered_output_name in rendered_output_names: diff --git a/nerfstudio/utils/gaussian_utils.py b/nerfstudio/utils/gaussian_utils.py index b7f2a966f..48e4ea859 100644 --- a/nerfstudio/utils/gaussian_utils.py +++ b/nerfstudio/utils/gaussian_utils.py @@ -3,7 +3,7 @@ # # GRAPHDECO research group, https://team.inria.fr/graphdeco # # All rights reserved. # # -# # This software is free for non-commercial, research and evaluation use +# # This software is free for non-commercial, research and evaluation use # # under the terms of the LICENSE.md file. # # # # For inquiries contact george.drettakis@inria.fr @@ -107,13 +107,13 @@ # def render_from_dict(data : dict, bg_color : torch.Tensor, scaling_modifier = 1.0, override_color = None): # """ -# Render the scene. - +# Render the scene. + # Background tensor (bg_color) must be on GPU! # """ # for k in ['active_sh_degree', 'max_sh_degree', 'xyz', 'scaling', 'rotation', 'features', 'opacity', 'FoVx', 'FoVy', 'image_height', 'image_width', 'world_view_transform', 'full_proj_transform', 'camera_center']: # assert k in data, f'Required key {k} not in {data.keys()}' - + # # Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means # screenspace_points = torch.zeros_like(data['xyz'], dtype=data['xyz'].dtype, requires_grad=True, device="cuda") + 0 # try: @@ -167,7 +167,7 @@ # colors_precomp = override_color -# # Rasterize visible Gaussians to image, obtain their radii (on screen). +# # Rasterize visible Gaussians to image, obtain their radii (on screen). # rendered_image, radii = rasterizer( # means3D = means3D, # means2D = means2D, @@ -295,7 +295,7 @@ # class MiniCam: # def __init__(self, width, height, fovy, fovx, znear, zfar, world_view_transform, full_proj_transform): # self.image_width = width -# self.image_height = height +# self.image_height = height # self.FoVy = fovy # self.FoVx = fovx # self.znear = znear @@ -306,8 +306,6 @@ # self.camera_center = view_inv[3][:3] - - # from math import exp # import torch @@ -362,4 +360,3 @@ # return ssim_map.mean() # else: # return ssim_map.mean(1).mean(1).mean(1) - diff --git a/nerfstudio/viewer/server/render_state_machine.py b/nerfstudio/viewer/server/render_state_machine.py index b02dfae9d..f0b60d5b7 100644 --- a/nerfstudio/viewer/server/render_state_machine.py +++ b/nerfstudio/viewer/server/render_state_machine.py @@ -23,8 +23,7 @@ from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, get_args import torch from nerfstudio.cameras.cameras import Cameras -from nerfstudio.model_components.renderers import \ - background_color_override_context +from nerfstudio.model_components.renderers import background_color_override_context from nerfstudio.models.gaussian_splatting import GaussianSplattingModel from nerfstudio.utils import colormaps, writer from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter @@ -129,8 +128,8 @@ class RenderStateMachine(threading.Thread): assert camera is not None, "render called before viewer connected" with self.viewer.train_lock if self.viewer.train_lock is not None else contextlib.nullcontext(): - #TODO jake-austin: Make this check whether the model inherits from a camera based model or a ray based model - #TODO Zhuoyang: First made some dummy judgements, need to be fixed later + # TODO jake-austin: Make this check whether the model inherits from a camera based model or a ray based model + # TODO Zhuoyang: First made some dummy judgements, need to be fixed later isGaussianSplattingModel = isinstance(self.viewer.get_model(), GaussianSplattingModel) if isGaussianSplattingModel: pass diff --git a/nerfstudio/viewer_beta/control_panel.py b/nerfstudio/viewer_beta/control_panel.py index b2c1c8ebb..a5551ca3f 100644 --- a/nerfstudio/viewer_beta/control_panel.py +++ b/nerfstudio/viewer_beta/control_panel.py @@ -270,7 +270,7 @@ class ControlPanel: step: the train step to set the model to """ with self.viser_server.atomic(), self.stat_folder: - self.markdown.content=f"Step: {step}" + self.markdown.content = f"Step: {step}" def update_output_options(self, new_options: List[str]): """ diff --git a/nerfstudio/viewer_beta/render_state_machine.py b/nerfstudio/viewer_beta/render_state_machine.py index 9283b0e0a..cb6338942 100644 --- a/nerfstudio/viewer_beta/render_state_machine.py +++ b/nerfstudio/viewer_beta/render_state_machine.py @@ -27,6 +27,7 @@ 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,9 +127,15 @@ 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(): - if isinstance(self.viewer.get_model(),GaussianSplattingModel): - camera_ray_bundle=None + if isinstance(self.viewer.get_model(), GaussianSplattingModel): + camera_ray_bundle = None self.viewer.get_model().set_crop(obb) + color = self.viewer.control_panel.background_color + background_color = torch.tensor( + [color[0] / 255.0, color[1] / 255.0, color[2] / 255.0], + device=self.viewer.get_model().device, + ) + self.viewer.get_model().set_background(background_color) else: camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb) self.viewer.get_model().eval() @@ -146,21 +153,27 @@ class RenderStateMachine(threading.Thread): with background_color_override_context( background_color ), torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt): - outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=camera) + outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle( + camera_ray_bundle, camera=camera + ) else: with torch.no_grad(), viewer_utils.SetTrace(self.check_interrupt): - outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle(camera_ray_bundle,camera=camera) + outputs = self.viewer.get_model().get_outputs_for_camera_ray_bundle( + camera_ray_bundle, camera=camera + ) except viewer_utils.IOChangeException: self.viewer.get_model().train() raise self.viewer.get_model().train() - if isinstance(self.viewer.get_model(),GaussianSplattingModel): - num_rays = image_height*image_width + if isinstance(self.viewer.get_model(), GaussianSplattingModel): + num_rays = image_height * image_width else: num_rays = len(camera_ray_bundle) if self.viewer.control_panel.layer_depth: - if isinstance(self.viewer.get_model(),GaussianSplattingModel): - outputs["gl_z_buf_depth"] = outputs['depth'] + if isinstance(self.viewer.get_model(), GaussianSplattingModel): + # TODO: sending depth at high resolution lags the network a lot, figure out how to do this more efficiently + # outputs["gl_z_buf_depth"] = outputs["depth"] + pass else: # convert to z_depth if depth compositing is enabled R = camera.camera_to_worlds[0:3, 0:3].T diff --git a/nerfstudio/viewer_beta/utils.py b/nerfstudio/viewer_beta/utils.py index 429edcb05..3614fce41 100644 --- a/nerfstudio/viewer_beta/utils.py +++ b/nerfstudio/viewer_beta/utils.py @@ -65,7 +65,7 @@ def get_camera( cx=pp_w, cy=pp_h, camera_type=camera_type, - camera_to_worlds=camera_state.c2w.to(torch.float32)[None,...], + camera_to_worlds=camera_state.c2w.to(torch.float32)[None, ...], times=torch.tensor([0.0], dtype=torch.float32), ) return camera