From 8146ea1ed4616c77895b089355d8f8d94c518281 Mon Sep 17 00:00:00 2001 From: Ethan Weber Date: Wed, 12 Apr 2023 22:48:40 +0000 Subject: [PATCH] updated render and metrics files --- nerfstudio/utils/lpips_utils.py | 2 +- nerfstudio/utils/metrics.py | 22 +++++++++++++--------- scripts/render.py | 14 +++++++------- tests/utils/test_metrics.py | 9 ++------- 4 files changed, 23 insertions(+), 24 deletions(-) diff --git a/nerfstudio/utils/lpips_utils.py b/nerfstudio/utils/lpips_utils.py index dd4d9b50d..a8b118d4d 100644 --- a/nerfstudio/utils/lpips_utils.py +++ b/nerfstudio/utils/lpips_utils.py @@ -23,7 +23,7 @@ from torch import Tensor from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity -def _valid_img(img: Tensor, normalize: bool) -> bool: +def _valid_img(img: Tensor, normalize: bool): """check that input is a valid image to the network.""" value_check = img.max() <= 1.0 and img.min() >= 0.0 if normalize else img.min() >= -1 return img.ndim == 4 and img.shape[1] == 3 and value_check diff --git a/nerfstudio/utils/metrics.py b/nerfstudio/utils/metrics.py index 3fab13f0e..0331e0c41 100644 --- a/nerfstudio/utils/metrics.py +++ b/nerfstudio/utils/metrics.py @@ -20,7 +20,7 @@ from abc import abstractmethod from typing import Optional import torch -from torch import Tensor, nn +from torch import nn from torchmetrics.functional import ( peak_signal_noise_ratio, structural_similarity_index_measure, @@ -50,7 +50,7 @@ class ImageMetricModule(nn.Module): preds: TensorType["bs", 3, "H", "W"], target: TensorType["bs", 3, "H", "W"], mask: Optional[TensorType["bs", 1, "H", "W"]] = None, - ) -> Tensor: + ) -> TensorType["bs"]: """Computes the metric. Args: @@ -59,7 +59,7 @@ class ImageMetricModule(nn.Module): mask: Mask to use to only compute the metrics where the mask is True. Returns: - TensorType["bs"]: Metric value. + Metric value. """ @@ -71,7 +71,7 @@ class PSNRModule(ImageMetricModule): preds: TensorType["bs", 3, "H", "W"], target: TensorType["bs", 3, "H", "W"], mask: Optional[TensorType["bs", 1, "H", "W"]] = None, - ) -> Tensor: + ) -> TensorType["bs"]: bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3] hw = h * w @@ -86,8 +86,10 @@ class PSNRModule(ImageMetricModule): # the masked version psnr_reshaped = psnr_image.view(bs, hw) + psnr_reshaped = torch.nan_to_num(psnr_reshaped, nan=0.0, posinf=0.0, neginf=0.0) mask_reshaped = mask.view(bs, hw) - psnr = (psnr_reshaped * mask_reshaped).sum(-1) / hw + den = mask_reshaped.sum(-1, keepdim=True) + psnr = (psnr_reshaped * mask_reshaped / den).sum(-1) return psnr @@ -99,7 +101,7 @@ class SSIMModule(ImageMetricModule): preds: TensorType["bs", 3, "H", "W"], target: TensorType["bs", 3, "H", "W"], mask: Optional[TensorType["bs", 1, "H", "W"]] = None, - ) -> Tensor: + ) -> TensorType["bs"]: bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3] hw = h * w @@ -118,7 +120,8 @@ class SSIMModule(ImageMetricModule): # the masked version ssim_reshaped = ssim_image.view(bs, hw) mask_reshaped = mask.view(bs, hw) - ssim = (ssim_reshaped * mask_reshaped).sum(-1) / hw + den = mask_reshaped.sum(-1, keepdim=True) + ssim = (ssim_reshaped * mask_reshaped / den).sum(-1) return ssim @@ -134,7 +137,7 @@ class LPIPSModule(ImageMetricModule): preds: TensorType["bs", 3, "H", "W"], target: TensorType["bs", 3, "H", "W"], mask: Optional[TensorType["bs", 1, "H", "W"]] = None, - ) -> Tensor: + ) -> TensorType["bs"]: bs, h, w = preds.shape[0], preds.shape[2], preds.shape[3] hw = h * w @@ -152,5 +155,6 @@ class LPIPSModule(ImageMetricModule): # the masked version lpips_reshaped = lpips_image.view(bs, hw) mask_reshaped = mask.view(bs, hw) - lpips = (lpips_reshaped * mask_reshaped).sum(-1) / hw + den = mask_reshaped.sum(-1, keepdim=True) + lpips = (lpips_reshaped * mask_reshaped / den).sum(-1) return lpips diff --git a/scripts/render.py b/scripts/render.py index 2d59a0735..453491639 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -58,7 +58,7 @@ def _render_trajectory_video( depth_far_plane: Optional[float] = None, using_eval_images: bool = False, ) -> None: - """Helper function to create a video of the spiral trajectory. + """Helper function to create a video from cameras. Args: pipeline: Pipeline to evaluate with. @@ -73,7 +73,6 @@ def _render_trajectory_video( depth_near_plane: Closest depth to consider. If None, use min image value. depth_far_plane: Furthest depth to consider. If None, use max image value. using_eval_images: Whether the eval images are being used or not. - depth_method: Method to use for depth rendering. If None, use model default. """ # pylint: disable=too-many-statements CONSOLE.print("[bold green]Creating trajectory " + output_format) @@ -142,10 +141,10 @@ def _render_trajectory_video( far_plane=depth_far_plane, ) if rendered_output_name == "visibility_median_count": - # NOTE(ethan): we cap the number of visible images to 255 here + # NOTE: we cap the number of visible images to 255 by using torch.uint8 outputs[rendered_output_name] = outputs[rendered_output_name].clamp(0, 255).to(torch.uint8) if output_format == "video": - # rescale to [0, 1] is needed to concatenate into a video format (see below) + # rescaling to [0, 1] is needed to concatenate into a video output format (see below) outputs[rendered_output_name] = outputs[rendered_output_name].float() / 255.0 output_image = outputs[rendered_output_name].cpu().numpy() if output_image.shape[-1] == 1: @@ -165,11 +164,12 @@ def _render_trajectory_video( names = rendered_output_names if output_format == "images": - # for images, we keep outputs in different folders + # for images, we write outputs to separate folders for name, im in zip(names, render_image): folder = output_image_dir / name folder.mkdir(parents=True, exist_ok=True) media.write_image(folder / f"{camera_idx:05d}.png", im) + # save depth_raw, which doesn't have the colormap applied if depth_raw is not None: folder = output_image_dir / "depth_raw" folder.mkdir(parents=True, exist_ok=True) @@ -370,9 +370,9 @@ class RenderTrajectory: cameras=pipeline.datamanager.eval_dataloader.cameras, steps=self.interpolation_steps ) elif self.traj == "eval-images": - camera_type = CameraType.PERSPECTIVE - # TODO(ethan): support any camera type, reading from the cameras object camera_path = pipeline.datamanager.eval_dataloader.cameras + camera_type = pipeline.datamanager.eval_dataloader.cameras.camera_type[0] + # TODO: support more than one camera type else: assert_never(self.traj) diff --git a/tests/utils/test_metrics.py b/tests/utils/test_metrics.py index 599069924..e5a49c708 100644 --- a/tests/utils/test_metrics.py +++ b/tests/utils/test_metrics.py @@ -3,7 +3,7 @@ Test metrics """ import torch -from nerfstudio.utils.metrics import PSNRModule, SSIMModule, LPIPSModule +from nerfstudio.utils.metrics import LPIPSModule, PSNRModule, SSIMModule def test_psnr(): @@ -55,14 +55,9 @@ def test_lpips(): # test non-masked version lpips = lpips_module(preds=pred, target=target) + print(lpips) assert lpips.shape == (bs,) # test masked version lpips = lpips_module(preds=pred, target=target, mask=mask) assert lpips.shape == (bs,) - - -if __name__ == "__main__": - test_psnr() - test_ssim() - test_lpips()