Render with colormaps (#1893)

* render with colormaps

* black

scripts/render.py

@@ -40,7 +40,7 @@ from nerfstudio.cameras.cameras import Cameras, CameraType
 from nerfstudio.data.scene_box import SceneBox
 from nerfstudio.model_components import renderers
 from nerfstudio.pipelines.base_pipeline import Pipeline
-from nerfstudio.utils import install_checks
+from nerfstudio.utils import colormaps, install_checks
 from nerfstudio.utils.eval_utils import eval_setup
 from nerfstudio.utils.rich_utils import ItersPerSecColumn
 
@@ -57,6 +57,11 @@ def _render_trajectory_video(
     seconds: float = 5.0,
     output_format: Literal["images", "video"] = "video",
     camera_type: CameraType = CameraType.PERSPECTIVE,
+    colormap: colormaps.Colormaps = "turbo",
+    normalize: bool = True,
+    colormap_min: float = 0.0,
+    colormap_max: float = 1.0,
+    invert: bool = False,
 ) -> None:
     """Helper function to create a video of the spiral trajectory.
 
@@ -70,6 +75,11 @@ def _render_trajectory_video(
         seconds: Length of output video.
         output_format: How to save output data.
         camera_type: Camera projection format type.
+        colormap: Colormap to use if single channel output.
+        normalize: Whether to normalize the output images, only works with colormaps.
+        colormap_min: Minimum colormap value.
+        colormap_max: Maximum colormap value.
+        invert: Whether to invert the colormap.
     """
     CONSOLE.print("[bold green]Creating trajectory " + output_format)
     cameras.rescale_output_resolution(rendered_resolution_scaling_factor)
@@ -126,9 +136,19 @@ def _render_trajectory_video(
                             f"Please set --rendered_output_name to one of: {outputs.keys()}", justify="center"
                         )
                         sys.exit(1)
-                    output_image = outputs[rendered_output_name].cpu().numpy()
-                    if output_image.shape[-1] == 1:
-                        output_image = np.concatenate((output_image,) * 3, axis=-1)
+                    output_image = outputs[rendered_output_name]
+                    output_image = (
+                        colormaps.apply_colormap(
+                            image=output_image,
+                            colormap=colormap,
+                            normalize=normalize,
+                            colormap_min=colormap_min,
+                            colormap_max=colormap_max,
+                            invert=invert,
+                        )
+                        .cpu()
+                        .numpy()
+                    )
                     render_image.append(output_image)
                 render_image = np.concatenate(render_image, axis=1)
                 if output_format == "images":
@@ -293,6 +313,15 @@ class RenderTrajectory:
     """Number of interpolation steps between eval dataset cameras."""
     eval_num_rays_per_chunk: Optional[int] = None
     """Specifies number of rays per chunk during eval."""
+    colormap: colormaps.Colormaps = "turbo"
+    normalize: bool = False
+    """Whether to normalize the output if using a colormap."""
+    colormap_min: float = 0.0
+    """Minimum value of the colormap."""
+    colormap_max: float = 1.0
+    """Maximum value of the colormap."""
+    invert: bool = False
+    """Whether to invert the colormap."""
 
     def main(self) -> None:
         """Main function."""
@@ -345,6 +374,11 @@ class RenderTrajectory:
             seconds=seconds,
             output_format=self.output_format,
             camera_type=camera_type,
+            colormap=self.colormap,
+            normalize=self.normalize,
+            colormap_min=self.colormap_min,
+            colormap_max=self.colormap_max,
+            invert=self.invert,
         )