updated render and metrics files

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

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

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)
 

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()