From 00b81cce28919ded60e21c58885df497c83ec2a7 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Fri, 4 Aug 2023 15:20:19 +0000 Subject: [PATCH] reverting changes --- nerfstudio/models/nerfacto.py | 7 +++--- nerfstudio/scripts/render.py | 6 ----- nerfstudio/utils/tensor_dataclass.py | 34 +++++----------------------- 3 files changed, 10 insertions(+), 37 deletions(-) diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py index 123b3d152..1122d5c20 100644 --- a/nerfstudio/models/nerfacto.py +++ b/nerfstudio/models/nerfacto.py @@ -26,6 +26,7 @@ import torch from torch.nn import Parameter from torchmetrics.functional import structural_similarity_index_measure from torchmetrics.image import PeakSignalNoiseRatio +from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity from nerfstudio.cameras.rays import RayBundle, RaySamples from nerfstudio.engine.callbacks import TrainingCallback, TrainingCallbackAttributes, TrainingCallbackLocation @@ -232,7 +233,7 @@ class NerfactoModel(Model): # metrics self.psnr = PeakSignalNoiseRatio(data_range=1.0) self.ssim = structural_similarity_index_measure - # self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) + self.lpips = LearnedPerceptualImagePatchSimilarity(normalize=True) self.step = 0 def get_param_groups(self) -> Dict[str, List[Parameter]]: @@ -387,11 +388,11 @@ class NerfactoModel(Model): psnr = self.psnr(gt_rgb, predicted_rgb) ssim = self.ssim(gt_rgb, predicted_rgb) - # lpips = self.lpips(gt_rgb, predicted_rgb) + lpips = self.lpips(gt_rgb, predicted_rgb) # all of these metrics will be logged as scalars metrics_dict = {"psnr": float(psnr.item()), "ssim": float(ssim)} # type: ignore - # metrics_dict["lpips"] = float(lpips) + metrics_dict["lpips"] = float(lpips) images_dict = {"img": combined_rgb, "accumulation": combined_acc, "depth": combined_depth} diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 77c9b4f52..5a3ca78bd 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -497,8 +497,6 @@ class RenderInterpolated(BaseRender): """Frame rate of the output video.""" output_format: Literal["images", "video"] = "video" """How to save output data.""" - loop: bool = False - """Whether to start and end with the first camera.""" def main(self) -> None: """Main function.""" @@ -517,10 +515,6 @@ class RenderInterpolated(BaseRender): assert pipeline.datamanager.train_dataset is not None cameras = pipeline.datamanager.train_dataset.cameras - if self.loop: - # This will make the camera path start and end with the same camera - cameras = cameras.cat(cameras[:1]) - seconds = self.interpolation_steps * len(cameras) / self.frame_rate camera_path = get_interpolated_camera_path( cameras=cameras, diff --git a/nerfstudio/utils/tensor_dataclass.py b/nerfstudio/utils/tensor_dataclass.py index 7d2147332..0ebe7b5ce 100644 --- a/nerfstudio/utils/tensor_dataclass.py +++ b/nerfstudio/utils/tensor_dataclass.py @@ -145,12 +145,12 @@ class TensorDataclass: def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT: if isinstance(indices, (torch.Tensor)): - return self._apply_fn_to_fields(lambda field_name, x: x[indices]) + return self._apply_fn_to_fields(lambda x: x[indices]) if isinstance(indices, (int, slice, type(Ellipsis))): indices = (indices,) assert isinstance(indices, tuple) - def tensor_fn(field_name, x): + def tensor_fn(x): return x[indices + (slice(None),)] def dataclass_fn(x): @@ -195,28 +195,6 @@ class TensorDataclass: """Returns the number of dimensions of the tensor dataclass.""" return len(self._shape) - def cat(self: TensorDataclassT, other: TensorDataclassT, dim: int = 0) -> TensorDataclassT: - """Returns a new TensorDataclass with the other TensorDataclass's values concatenated along the specified dimension. - - Args: - other: The other TensorDataclass to concatenate with. - dim: The dimension along which to concatenate the tensors. - - Returns: - A new TensorDataclass with the tensors concatenated along the specified dimension. - """ - if not isinstance(other, TensorDataclass): - raise TypeError("Expected 'other' to be a TensorDataclass.") - - def tensor_fn(field_name, x): - return torch.cat([x, getattr(other, field_name)], dim=dim) - - dataclass_fn = None - - custom_tensor_dims_fn = None - - return self._apply_fn_to_fields(tensor_fn, dataclass_fn, custom_tensor_dims_fn=custom_tensor_dims_fn) - def reshape(self: TensorDataclassT, shape: Tuple[int, ...]) -> TensorDataclassT: """Returns a new TensorDataclass with the same data but with a new shape. @@ -231,7 +209,7 @@ class TensorDataclass: if isinstance(shape, int): shape = (shape,) - def tensor_fn(field_name, x): + def tensor_fn(x): return x.reshape((*shape, x.shape[-1])) def dataclass_fn(x): @@ -269,7 +247,7 @@ class TensorDataclass: return v.broadcast_to((*shape, *v.shape[-custom_dims:])) return self._apply_fn_to_fields( - lambda field_name, x: x.broadcast_to((*shape, x.shape[-1])), custom_tensor_dims_fn=custom_tensor_dims_fn + lambda x: x.broadcast_to((*shape, x.shape[-1])), custom_tensor_dims_fn=custom_tensor_dims_fn ) def to(self: TensorDataclassT, device) -> TensorDataclassT: @@ -281,7 +259,7 @@ class TensorDataclass: Returns: A new TensorDataclass with the same data but on the specified device. """ - return self._apply_fn_to_fields(lambda field_name, x: x.to(device)) + return self._apply_fn_to_fields(lambda x: x.to(device)) def _apply_fn_to_fields( self: TensorDataclassT, @@ -352,7 +330,7 @@ class TensorDataclass: ): new_dict[f] = custom_tensor_dims_fn(f, v) elif isinstance(v, (torch.Tensor, TensorDataclass)): - new_dict[f] = fn(f, v) + new_dict[f] = fn(v) elif isinstance(v, Dict): new_dict[f] = self._apply_fn_to_dict(v, fn, dataclass_fn) else: