tink/nerfstudio
4
0 Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-14 11:54:01 +08:00
Code Activity

lint

Browse Source
This commit is contained in:
Justin Kerr committed 2023-10-13 00:22:27 -07:00
1 parent d8b76fea4d
commit f4cc901e08
11 files changed
+231 -160
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

nerfstudio/configs/method_configs.py
+10 -10
View File
@@ -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,
nerfstudio/data/datamanagers/full_images_datamanager.py
+52 -29
View File
@@ -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))]
nerfstudio/models/base_model.py
+4 -4
View File
@@ -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:
nerfstudio/models/gaussian_splatting.py
+121 -84
View File
@@ -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.stop_split_at:
#then we densify
avg_grad_norm = (self.xys_grad_norm / self.vis_counts) * 0.5* max(self.last_size[0],self.last_size[1])
if self.step < self.config.stop_split_at:
# then we densify
avg_grad_norm = (
(self.xys_grad_norm / self.vis_counts) * 0.5 * max(self.last_size[0], self.last_size[1])
)
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
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(
nerfstudio/pipelines/base_pipeline.py
+6 -8
View File
@@ -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)
nerfstudio/scripts/render.py
+6 -2
View File
@@ -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:
nerfstudio/utils/gaussian_utils.py
+6 -9
View File
@@ -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)
nerfstudio/viewer/server/render_state_machine.py
+3 -4
View File
@@ -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
nerfstudio/viewer_beta/control_panel.py
+1 -1
View File
@@ -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]):
"""
nerfstudio/viewer_beta/render_state_machine.py
+21 -8
View File
@@ -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
nerfstudio/viewer_beta/utils.py
+1 -1
View File
@@ -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