tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Merge branch 'gaussian-splatting' of https://github.com/nerfstudio-project/nerfstudio into gaussian-splatting

Browse Source
This commit is contained in:
Justin Kerr 2023-10-21 14:59:47 -07:00
commit cd119d276d
3 changed files with 37 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
nerfstudio/data/datamanagers/full_images_datamanager.py
View File

@ -246,10 +246,11 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
cached_eval.append(data)
self.eval_dataset.cameras[i].fx = K[0, 0]
self.eval_dataset.cameras[i].fy = K[1, 1]
self.eval_dataset.cameras[i].cx = K[0, 2]
self.eval_dataset.cameras[i].cy = K[1, 2]
self.eval_dataset.cameras.fx[i] = float(K[0, 0])
self.eval_dataset.cameras.fy[i] = float(K[1, 1])
self.eval_dataset.cameras.cx[i] = float(K[0, 2])
self.eval_dataset.cameras.cy[i] = float(K[1, 2])
# print(self.eval_dataset.cameras[i].reshape((1,)))
if cache_images_option == "gpu":
for cache in cached_train:
@ -352,7 +353,10 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
data = self.eval_dataset.get_data(image_idx)
data["image"] = data["image"].to(self.device)
assert len(self.eval_dataset.cameras.shape) == 1, "Assumes single batch dimension"
camera = self.eval_dataset.cameras[image_idx].to(self.device)
camera = self.eval_dataset.cameras[image_idx : image_idx + 1].to(self.device)
# TODO: camera shape torch.Size([]) for eval dataset
# should be torch.Size([1])
print(camera.shape, data.keys())
return camera, data
def next_eval_image(self, step: int) -> Tuple[int, Cameras, Dict]:
@ -368,5 +372,5 @@ class FullImageDatamanager(DataManager, Generic[TDataset]):
data = self.eval_dataset.get_data(image_idx)
data["image"] = data["image"].to(self.device)
assert len(self.eval_dataset.cameras.shape) == 1, "Assumes single batch dimension"
camera = self.eval_dataset.cameras[image_idx].to(self.device)
camera = self.eval_dataset.cameras[image_idx: image_idx + 1].to(self.device)
return image_idx, camera, data

35
nerfstudio/models/gaussian_splatting.py
View File

@ -25,7 +25,7 @@ from nerfstudio.data.scene_box import OrientedBox
from copy import deepcopy
from nerfstudio.cameras.rays import RayBundle
from torch.nn import Parameter
from torchmetrics.image import PeakSignalNoiseRatio
from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
from nerfstudio.cameras.cameras import Cameras
from gsplat._torch_impl import quat_to_rotmat
@ -176,12 +176,13 @@ class GaussianSplattingModel(Model):
self.shs_rest = torch.nn.Parameter(torch.zeros((self.num_points, 3, dim_sh - 1)))
self.opacities = torch.nn.Parameter(torch.logit(0.1 * torch.ones(self.num_points, 1)))
# metrics
self.psnr = PeakSignalNoiseRatio(data_range=1.0)
self.ssim = StructuralSimilarityIndexMeasure()
self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True)
self.step = 0
self.crop_box: Optional[OrientedBox] = None
self.back_color = torch.zeros(3)
@ -488,7 +489,10 @@ class GaussianSplattingModel(Model):
return gps
def _get_downscale_factor(self):
return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0)
if self.training:
return 2 ** max((self.config.num_downscales - self.step // self.config.resolution_schedule), 0)
else:
return 1
def get_outputs(self, camera: Cameras) -> Dict[str, Union[torch.Tensor, List]]:
"""Takes in a Ray Bundle and returns a dictionary of outputs.
@ -506,12 +510,8 @@ class GaussianSplattingModel(Model):
assert camera.shape[0] == 1, "Only one camera at a time"
# dont mutate the input
camera = deepcopy(camera)
if self.training:
#currently relies on the branch vickie/camera-grads
self.camera_optimizer.apply_to_camera(camera)
if self.training:
d = self._get_downscale_factor()
camera.rescale_output_resolution(1 / d)
d = self._get_downscale_factor()
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:
@ -630,8 +630,23 @@ class GaussianSplattingModel(Model):
outputs: the output to compute loss dict to
batch: ground truth batch corresponding to outputs
"""
d = self._get_downscale_factor()
if d > 1:
# use torchvision to resize
import torchvision.transforms.functional as TF
return {}
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"]
metrics_dict = {}
gt_rgb = gt_img.to(self.device) # RGB or RGBA image
# gt_rgb = self.renderer_rgb.blend_background(gt_rgb) # Blend if RGBA
predicted_rgb = outputs["rgb"]
metrics_dict["psnr"] = self.psnr(predicted_rgb, gt_rgb)
self.camera_optimizer.get_metrics_dict(metrics_dict)
return metrics_dict
def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]:
"""Computes and returns the losses dict.

3
nerfstudio/pipelines/base_pipeline.py
View File

@ -341,7 +341,8 @@ class VanillaPipeline(Pipeline):
"""
self.eval()
image_idx, camera_ray_bundle, batch = self.datamanager.next_eval_image(step)
outputs = self.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle)
# TODO ginawu: refactor(?) for gsplat without this hardcoded camera param below
outputs = self.model.get_outputs_for_camera_ray_bundle(None, camera=camera_ray_bundle)
metrics_dict, images_dict = self.model.get_image_metrics_and_images(outputs, batch)
assert "image_idx" not in metrics_dict
metrics_dict["image_idx"] = image_idx