Merge pull request #63 from ethanweber/tancik/improve_depth_vis

Improve depth vis

mattport/nerf/graph/vanilla_nerf.py

@@ -16,7 +16,7 @@ from mattport.nerf.field_modules.mlp import MLP
 from mattport.nerf.field_modules.ray_generator import RayGenerator
 from mattport.nerf.graph.base import Graph
 from mattport.nerf.loss import MSELoss
-from mattport.nerf.renderers import AccumulationRenderer, DisparityRenderer, RGBRenderer
+from mattport.nerf.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer
 from mattport.nerf.sampler import PDFSampler, UniformSampler
 from mattport.structures import colors
 from mattport.structures.rays import RaySamples
@@ -114,7 +114,7 @@ class NeRFGraph(Graph):
         # renderers
         self.renderer_rgb = RGBRenderer(background_color=colors.WHITE)
         self.renderer_accumulation = AccumulationRenderer()
-        self.renderer_disparity = DisparityRenderer()
+        self.renderer_depth = DepthRenderer()
 
         # losses
         self.rgb_loss = MSELoss()
@@ -148,7 +148,7 @@ class NeRFGraph(Graph):
             weights=coarse_weights,
         )  # RendererOutputs
         coarse_renderer_accumulation = self.renderer_accumulation(coarse_weights)  # RendererOutputs
-        coarse_renderer_disparity = self.renderer_disparity(coarse_weights, uniform_ray_samples.ts)
+        coarse_renderer_depth = self.renderer_depth(coarse_weights, uniform_ray_samples.ts)
 
         # fine network:
         pdf_ray_samples = self.sampler_pdf(uniform_ray_samples, coarse_weights)  # RaySamples
@@ -161,9 +161,21 @@ class NeRFGraph(Graph):
             weights=fine_weights,
         )  # RendererOutputs
         fine_renderer_accumulation = self.renderer_accumulation(fine_weights)  # RendererOutputs
-        fine_renderer_disparity = self.renderer_disparity(fine_weights, pdf_ray_samples.ts)
+        fine_renderer_depth = self.renderer_depth(fine_weights, pdf_ray_samples.ts)
 
         # TODO refactor this into "vis" section. Doesn't need to be run during training.
+        coarse_renderer_depth = visualization.apply_depth_colormap(
+            coarse_renderer_depth.depth,
+            accumulation=coarse_renderer_accumulation.accumulation,
+            near_plane=self.near_plane,
+            far_plane=self.far_plane,
+        )
+        fine_renderer_depth = visualization.apply_depth_colormap(
+            fine_renderer_depth.depth,
+            accumulation=fine_renderer_accumulation.accumulation,
+            near_plane=self.near_plane,
+            far_plane=self.far_plane,
+        )
         coarse_renderer_accumulation = visualization.apply_colormap(coarse_renderer_accumulation.accumulation)
         fine_renderer_accumulation = visualization.apply_colormap(fine_renderer_accumulation.accumulation)
 
@@ -173,8 +185,8 @@ class NeRFGraph(Graph):
             "rgb_fine": fine_renderer_outputs.rgb,
             "accumulation_coarse": coarse_renderer_accumulation,
             "accumulation_fine": fine_renderer_accumulation,
-            "disparity_coarse": coarse_renderer_disparity.disparity,
-            "disparity_fine": fine_renderer_disparity.disparity,
+            "depth_coarse": coarse_renderer_depth,
+            "depth_fine": fine_renderer_depth,
         }
         return outputs
 

mattport/nerf/renderers.py

@@ -16,7 +16,7 @@ class RendererOutputs:
     rgb: TensorType["num_rays", 3] = None
     density: TensorType["num_rays", 1] = None
     accumulation: TensorType["num_rays", 1] = None
-    disparity: TensorType["num_rays", 1] = None
+    depth: TensorType["num_rays", 1] = None
 
 
 class RGBRenderer(nn.Module):
@@ -82,7 +82,7 @@ class AccumulationRenderer(nn.Module):
         return renderer_outputs
 
 
-class DisparityRenderer(nn.Module):
+class DepthRenderer(nn.Module):
     """Calcualte depth along ray."""
 
     def __init__(self, method: str = "expected") -> None:
@@ -103,14 +103,15 @@ class DisparityRenderer(nn.Module):
             ts (TensorType[..., "num_samples"]): Sample locations along rays
 
         Returns:
-            RendererOutputs: Outputs with disparity values.
+            RendererOutputs: Outputs with depth values.
         """
 
         if self.method == "expected":
-            depth = torch.sum(weights * ts, dim=-1)
             eps = 1e-10
-            disparity = 1.0 / torch.max(eps * torch.ones_like(depth), depth / torch.sum(weights, -1))
+            depth = torch.sum(weights * ts, dim=-1) / (torch.sum(weights, -1) + eps)
 
-            return RendererOutputs(disparity=disparity)
+            depth = torch.clip(depth, ts[..., 0], ts[..., -1])
+
+            return RendererOutputs(depth=depth[..., None])
 
         raise NotImplementedError(f"Method {self.method} not implemented")

mattport/nerf/trainer.py

@@ -259,8 +259,8 @@ class Trainer:
             rgb_fine = []
             accumulation_coarse = []
             accumulation_fine = []
-            disparity_coarse = []
-            disparity_fine = []
+            depth_coarse = []
+            depth_fine = []
             for i in range(0, num_rays, chunk_size):
                 ray_indices = all_ray_indices[i : i + chunk_size].to(self.device)
                 graph_outputs = self.graph(ray_indices)
@@ -268,14 +268,14 @@ class Trainer:
                 rgb_fine.append(graph_outputs["rgb_fine"])
                 accumulation_coarse.append(graph_outputs["accumulation_coarse"])
                 accumulation_fine.append(graph_outputs["accumulation_fine"])
-                disparity_coarse.append(graph_outputs["disparity_coarse"])
-                disparity_fine.append(graph_outputs["disparity_fine"])
+                depth_coarse.append(graph_outputs["depth_coarse"])
+                depth_fine.append(graph_outputs["depth_fine"])
             rgb_coarse = torch.cat(rgb_coarse).view(image_height, image_width, 3).detach().cpu()
             rgb_fine = torch.cat(rgb_fine).view(image_height, image_width, 3).detach().cpu()
             accumulation_coarse = torch.cat(accumulation_coarse).view(image_height, image_width, 3).detach().cpu()
             accumulation_fine = torch.cat(accumulation_fine).view(image_height, image_width, 3).detach().cpu()
-            disparity_coarse = torch.cat(disparity_coarse).view(image_height, image_width, 1).detach().cpu()
-            disparity_fine = torch.cat(disparity_fine).view(image_height, image_width, 1).detach().cpu()
+            depth_coarse = torch.cat(depth_coarse).view(image_height, image_width, 3).detach().cpu()
+            depth_fine = torch.cat(depth_fine).view(image_height, image_width, 3).detach().cpu()
 
         combined_image = torch.cat([image, rgb_coarse, rgb_fine], dim=1)
         writer.write_event(
@@ -287,10 +287,8 @@ class Trainer:
             {"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_accumulation"}
         )
 
-        combined_image = torch.cat([disparity_coarse, disparity_fine], dim=1)
-        writer.write_event(
-            {"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_disparity"}
-        )
+        combined_image = torch.cat([depth_coarse, depth_fine], dim=1)
+        writer.write_event({"name": f"image_idx_{image_idx}", "x": combined_image, "step": step, "group": "val_depth"})
 
         coarse_psnr = get_psnr(image, rgb_coarse)
         writer.write_event(

mattport/utils/visualization.py

@@ -1,5 +1,6 @@
 """ Helper functions for visualizing outputs """
 
+from typing import Optional
 import torch
 from matplotlib import cm
 from torchtyping import TensorType
@@ -19,6 +20,45 @@ def apply_colormap(image: TensorType[..., 1], cmap="viridis") -> TensorType[...,
     colormap = cm.get_cmap(cmap)
     colormap = torch.tensor(colormap.colors).to(image.device)
 
+    image2 = (image * 255).long()
+
+    if torch.min(image2) < 0 or torch.max(image2) > 255:
+        print("hi")
+
     image = (image * 255).long()
 
-    return colormap[image]
+    return colormap[image[..., 0]]
+
+
+def apply_depth_colormap(
+    depth: TensorType[..., 1],
+    accumulation: Optional[TensorType[..., 1]] = None,
+    near_plane: Optional[float] = None,
+    far_plane: Optional[float] = None,
+    cmap="turbo",
+) -> TensorType[..., 3]:
+    """Converts a depth image to color for easier analysis.
+
+    Args:
+        depth (TensorType[..., 1]): Depth image.
+        accumulation (TensorType[..., 1], optional): Ray accumulation used for masking vis. Defaults to None.
+        near_plane (float, optional): Closest depth to consider. If None, use min image value. Defaults to None.
+        far_plane (float, optional): Furthest depth to consider. If None, use max image value. Defaults to None.
+        cmap (str, optional): Colormap to apply. Defaults to "turbo".
+
+    Returns:
+        TensorType[..., 3]: Colored depth image
+    """
+
+    near_plane = near_plane or torch.min(depth)
+    far_plane = far_plane or torch.max(depth)
+
+    depth = (depth - near_plane) / (far_plane - near_plane + 1e-10)
+    depth = torch.clip(depth, 0, 1)
+
+    colored_image = apply_colormap(depth, cmap=cmap)
+
+    if accumulation is not None:
+        colored_image = colored_image * accumulation + (1 - accumulation)
+
+    return colored_image

mattport/utils/writer.py

@@ -27,7 +27,6 @@ def write_event(args: Dict[str, Any]):
     assert comms.is_main_process(), "Writing out with process other than main"
     for writer in EVENT_WRITERS:
         if "x" in args:
-            print(args["x"].shape)
             writer.write_image(**args)
         elif "scalar" in args:
             writer.write_scalar(**args)