mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 10:34:29 +08:00
Merge branch 'justin/inhouse' of https://github.com/jake-austin/gaussian-splatting-nerfstudio into justin/inhouse
This commit is contained in:
AdamRashid96
commit
8dcaa09828
@ -590,11 +590,11 @@ method_configs["gaussian-splatting"] = TrainerConfig(
|
||||
method_name="gaussian-splatting",
|
||||
steps_per_eval_image=10,
|
||||
steps_per_eval_batch=10,
|
||||
steps_per_save=2000000,#saving is broken
|
||||
steps_per_save=2000,
|
||||
steps_per_eval_all_images=1000000, # set to a very large model so we don't eval with all images
|
||||
max_num_iterations=3000000,
|
||||
max_num_iterations=30000,
|
||||
mixed_precision=False,
|
||||
gradient_accumulation_steps=4,
|
||||
gradient_accumulation_steps=1,
|
||||
pipeline=VanillaPipelineConfig(
|
||||
datamanager=FullImageDatamanagerConfig(
|
||||
dataparser=ColmapDataParserConfig(load_3D_points=True),
|
||||
|
||||
@ -164,7 +164,8 @@ class Model(nn.Module):
|
||||
"""
|
||||
|
||||
@torch.no_grad()
|
||||
def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: Union[RayBundle, Cameras]) -> 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:
|
||||
|
||||
@ -22,6 +22,7 @@ from dataclasses import dataclass, field
|
||||
from typing import Any, Dict, List, Literal, Optional, Tuple, Type, Union
|
||||
import torch
|
||||
from copy import deepcopy
|
||||
from nerfstudio.cameras.rays import RayBundle
|
||||
from torch.nn import Parameter
|
||||
from torchmetrics.functional import structural_similarity_index_measure
|
||||
from torchmetrics.image import PeakSignalNoiseRatio
|
||||
@ -88,11 +89,11 @@ def projection_matrix(znear, zfar, fovx, fovy, **kwargs):
|
||||
class GaussianSplattingModelConfig(ModelConfig):
|
||||
"""Gaussian Splatting Model Config"""
|
||||
_target: Type = field(default_factory=lambda: GaussianSplattingModel)
|
||||
warmup_length:int = 1500
|
||||
warmup_length:int = 900
|
||||
"""period of steps where refinement is turned off"""
|
||||
refine_every:int = 100
|
||||
"""period of steps where gaussians are culled and densified"""
|
||||
resolution_schedule: int = 500
|
||||
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"""
|
||||
@ -107,10 +108,13 @@ class GaussianSplattingModelConfig(ModelConfig):
|
||||
densify_size_thresh: float = .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
|
||||
sh_degree_interval: int = 3
|
||||
sh_degree_interval: int = 1000
|
||||
"""every n intervals turn on another sh degree"""
|
||||
max_screen_size: int = 200
|
||||
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
|
||||
|
||||
class GaussianSplattingModel(Model):
|
||||
"""Gaussian Splatting model
|
||||
@ -121,18 +125,25 @@ class GaussianSplattingModel(Model):
|
||||
|
||||
config: GaussianSplattingModelConfig
|
||||
def __init__(self,*args,**kwargs):
|
||||
self.seed_pts = kwargs['seed_points']
|
||||
if 'seed_points' in kwargs:
|
||||
self.seed_pts = kwargs['seed_points']
|
||||
else:
|
||||
self.seed_pts = None
|
||||
super().__init__(*args,**kwargs)
|
||||
self.vc = ViewerControl()
|
||||
|
||||
def populate_modules(self):
|
||||
# TODO (jake-austin): clean this up, this is transplanted code across all the implementation functions
|
||||
self.means = torch.nn.Parameter(self.seed_pts[0]) # (Location, Color)
|
||||
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)
|
||||
else:
|
||||
self.means = torch.nn.Parameter((torch.rand((100000,3))-.5)*2)
|
||||
self.means_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.seed_pts[0], 3)
|
||||
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)
|
||||
@ -142,49 +153,42 @@ class GaussianSplattingModel(Model):
|
||||
self.quats = torch.nn.Parameter(random_quat_tensor(self.num_points))
|
||||
self.degree = 4
|
||||
dim_sh = num_sh_bases(self.degree)
|
||||
# self.colors = torch.nn.Parameter(torch.rand(self.num_points, dim_sh, 3))
|
||||
|
||||
fused_color = self.RGB2SH(self.seed_pts[1]/255)
|
||||
shs = torch.zeros((fused_color.shape[0], dim_sh, 3)).float().cuda()
|
||||
shs[:, 0, :] = fused_color
|
||||
shs[:, 1:, :] = 0.0
|
||||
self.colors = torch.nn.Parameter(shs)
|
||||
if self.seed_pts is not None and not self.config.random_init:
|
||||
fused_color = self.RGB2SH(self.seed_pts[1]/255)
|
||||
shs = torch.zeros((fused_color.shape[0], dim_sh, 3))
|
||||
shs[:, 0, :] = fused_color
|
||||
shs[:, 1:, :] = 0.0
|
||||
randcolors = torch.rand((30000,dim_sh,3))
|
||||
self.colors = torch.nn.Parameter(torch.cat([shs,randcolors]))
|
||||
else:
|
||||
self.colors = torch.nn.Parameter(torch.rand(self.num_points, dim_sh, 3))
|
||||
|
||||
self.opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1)))
|
||||
self.opacities = torch.nn.Parameter(torch.logit(0.01 * 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
|
||||
|
||||
def k_nearest_faiss(self, x, k):
|
||||
"""
|
||||
Find k-nearest neighbors using FAISS.
|
||||
x: The data matrix (NumPy array) of shape [num_samples, num_features]
|
||||
k: The number of neighbors to retrieve
|
||||
Returns: A tuple (distances, indices) where both are Numpy arrays of shape [num_samples, k]
|
||||
"""
|
||||
num_samples = x.shape[0]
|
||||
|
||||
# Build the index
|
||||
index = faiss.IndexFlatL2(x.shape[1]) # Here, using a flat (brute force) index for L2 distance
|
||||
index.add(x)
|
||||
|
||||
# Search the index
|
||||
distances, indices = index.search(x, k + 1) # +1 because the point itself is always returned
|
||||
|
||||
# Exclude the point itself from the result and return
|
||||
return distances[:, 1:], indices[:, 1:]
|
||||
|
||||
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, num_sh_bases(self.degree), 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):
|
||||
"""
|
||||
Find k-nearest neighbors using sklearn's NearestNeighbors.
|
||||
x: The data tensor of shape [num_samples, num_features]
|
||||
k: The number of neighbors to retrieve
|
||||
"""
|
||||
num_samples = x.size(0)
|
||||
|
||||
# Convert tensor to numpy array
|
||||
x_np = x.cpu().numpy()
|
||||
|
||||
@ -196,7 +200,6 @@ class GaussianSplattingModel(Model):
|
||||
|
||||
# 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"""
|
||||
@ -217,7 +220,6 @@ class GaussianSplattingModel(Model):
|
||||
optimizer.param_groups[0]['params'] = new_params
|
||||
optimizer.state[new_params[0]] = param_state
|
||||
|
||||
|
||||
def dup_in_optim(self,optimizer, dup_mask, new_params):
|
||||
"""adds the parameters to the optimizer"""
|
||||
param = optimizer.param_groups[0]['params'][0]
|
||||
@ -246,19 +248,6 @@ class GaussianSplattingModel(Model):
|
||||
newradii = self.radii[visible_mask]
|
||||
self.max_2Dsize[visible_mask] = torch.maximum(self.max_2Dsize[visible_mask],newradii)
|
||||
|
||||
|
||||
def refinement_before(self, optimizers:Optimizers, step):
|
||||
print("Inside refinement before")
|
||||
if self.step > self.config.warmup_length:
|
||||
with torch.no_grad():
|
||||
# do all the refinement stuff here
|
||||
#first we cull gaussians
|
||||
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)
|
||||
|
||||
|
||||
def refinement_after(self, optimizers:Optimizers, step):
|
||||
if self.step > self.config.warmup_length:
|
||||
with torch.no_grad():
|
||||
@ -277,6 +266,8 @@ class GaussianSplattingModel(Model):
|
||||
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]
|
||||
@ -285,6 +276,12 @@ class GaussianSplattingModel(Model):
|
||||
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
|
||||
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)
|
||||
|
||||
|
||||
if self.step // self.config.refine_every % self.config.reset_alpha_every == 0:
|
||||
print("Resetting alpha")
|
||||
@ -312,9 +309,6 @@ class GaussianSplattingModel(Model):
|
||||
#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_grad_norm = self.means_grad_norm[~culls].detach()
|
||||
self.max_2Dsize = self.max_2Dsize[~culls].detach()
|
||||
self.vis_counts = self.vis_counts[~culls].detach()
|
||||
self.means = Parameter(self.means[~culls].detach())
|
||||
self.scales = Parameter(self.scales[~culls].detach())
|
||||
self.quats = Parameter(self.quats[~culls].detach())
|
||||
@ -361,7 +355,7 @@ class GaussianSplattingModel(Model):
|
||||
"""
|
||||
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]
|
||||
dup_means = self.means[dup_mask] + self.means.grad[dup_mask]
|
||||
dup_colors = self.colors[dup_mask]
|
||||
dup_opacities = self.opacities[dup_mask]
|
||||
dup_scales = self.scales[dup_mask]
|
||||
@ -377,11 +371,6 @@ class GaussianSplattingModel(Model):
|
||||
) -> List[TrainingCallback]:
|
||||
cbs = []
|
||||
cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],self.step_cb))
|
||||
cbs.append(TrainingCallback([TrainingCallbackLocation.BEFORE_TRAIN_ITERATION],
|
||||
self.refinement_before,
|
||||
update_every_num_iters=self.config.refine_every,
|
||||
args=[training_callback_attributes.optimizers]
|
||||
))
|
||||
# The order of these matters
|
||||
cbs.append(TrainingCallback([TrainingCallbackLocation.AFTER_TRAIN_ITERATION],
|
||||
self.after_train,))
|
||||
@ -454,7 +443,7 @@ class GaussianSplattingModel(Model):
|
||||
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(.0001,1000,fovx,fovy).to(self.device)
|
||||
projmat = projection_matrix(.1,1000,fovx,fovy).to(self.device)
|
||||
BLOCK_X, BLOCK_Y = 16, 16
|
||||
tile_bounds = (
|
||||
(W + BLOCK_X - 1) // BLOCK_X,
|
||||
@ -535,19 +524,18 @@ class GaussianSplattingModel(Model):
|
||||
else:
|
||||
gt_img = batch['image']
|
||||
Ll1 = torch.nn.functional.l1_loss(gt_img, outputs['rgb'])
|
||||
# This simloss makes the results look weird, removing for now
|
||||
# simloss = self.ssim(gt_img.permute(2,0,1)[None,...], outputs['rgb'].permute(2,0,1)[None,...])
|
||||
return {"main_loss": Ll1}
|
||||
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: Cameras) -> 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
|
||||
"""
|
||||
outs = self.get_outputs(camera_ray_bundle.to(self.device))
|
||||
outs["rgb"] = torch.clamp(outs["rgb"], 0.0, 1.0)
|
||||
assert camera is not None,'must provide camera to gaussian model'
|
||||
outs = self.get_outputs(camera.to(self.device))
|
||||
return outs
|
||||
|
||||
def get_image_metrics_and_images(
|
||||
|
||||
@ -136,10 +136,10 @@ 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)
|
||||
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)
|
||||
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:
|
||||
|
||||
@ -104,7 +104,7 @@ def _start_viewer(config: TrainerConfig, pipeline: Pipeline, step: int):
|
||||
viewer_state = ViewerBetaState(
|
||||
config.viewer,
|
||||
log_filename=viewer_log_path,
|
||||
datapath=base_dir,
|
||||
datapath=pipeline.datamanager.get_datapath(),
|
||||
pipeline=pipeline,
|
||||
share=config.viewer.make_share_url,
|
||||
)
|
||||
|
||||
@ -142,10 +142,10 @@ 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(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(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
|
||||
|
||||
Loading…
Reference in New Issue
Block a user