Add split screen (#1898)

* Add split screen

* bug fixes

nerfstudio/scripts/render.py

@@ -71,11 +71,7 @@ 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,
+    colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions(),
 ) -> None:
     """Helper function to create a video of the spiral trajectory.
 
@@ -89,11 +85,7 @@ 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.
+        colormap_options: Options for colormap.
     """
     CONSOLE.print("[bold green]Creating trajectory " + output_format)
     cameras.rescale_output_resolution(rendered_resolution_scaling_factor)
@@ -154,11 +146,7 @@ def _render_trajectory_video(
                     output_image = (
                         colormaps.apply_colormap(
                             image=output_image,
-                            colormap=colormap,
-                            normalize=normalize,
-                            colormap_min=colormap_min,
-                            colormap_max=colormap_max,
-                            invert=invert,
+                            colormap_options=colormap_options,
                         )
                         .cpu()
                         .numpy()
@@ -327,15 +315,8 @@ 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."""
+    colormap_options: colormaps.ColormapOptions = colormaps.ColormapOptions()
+    """Colormap options."""
 
     def main(self) -> None:
         """Main function."""
@@ -388,11 +369,7 @@ 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,
+            colormap_options=self.colormap_options,
         )
 
 

nerfstudio/utils/colormaps.py

@@ -14,6 +14,7 @@
 
 """ Helper functions for visualizing outputs """
 
+from dataclasses import dataclass
 from typing import Literal, Optional
 
 import matplotlib
@@ -25,13 +26,25 @@ from nerfstudio.utils import colors
 Colormaps = Literal["default", "turbo", "viridis", "magma", "inferno", "cividis"]
 
 
+@dataclass
+class ColormapOptions:
+    """Options for colormap"""
+
+    colormap: Colormaps = "default"
+    """ The colormap to use """
+    normalize: bool = False
+    """ Whether to normalize the input tensor image """
+    colormap_min: float = 0
+    """ Minimum value for the output colormap """
+    colormap_max: float = 1
+    """ Maximum value for the output colormap """
+    invert: bool = False
+    """ Whether to invert the output colormap """
+
+
 def apply_colormap(
     image: TensorType["bs":..., 1],
-    colormap: Colormaps = "turbo",
-    normalize: bool = False,
-    colormap_min: float = 0,
-    colormap_max: float = 1,
-    invert: bool = False,
+    colormap_options: ColormapOptions = ColormapOptions(),
     eps: float = 1e-9,
 ) -> TensorType["bs":..., "rgb":3]:
     """
@@ -39,11 +52,6 @@ def apply_colormap(
 
     Args:
         image: Input tensor image.
-        colormap: Colormap to apply.
-        normalize: Whether to normalize the input tensor image.
-        colormap_min: Minimum value for the output colormap.
-        colormap_max: Maximum value for the output colormap.
-        invert: Whether to invert the output colormap.
         eps: Epsilon value for numerical stability.
 
     Returns:
@@ -57,14 +65,16 @@ def apply_colormap(
     # rendering depth outputs
     if image.shape[-1] == 1 and torch.is_floating_point(image):
         output = image
-        if normalize:
+        if colormap_options.normalize:
             output = output - torch.min(output)
             output = output / (torch.max(output) + eps)
-        output = output * (colormap_max - colormap_min) + colormap_min
+        output = (
+            output * (colormap_options.colormap_max - colormap_options.colormap_min) + colormap_options.colormap_min
+        )
         output = torch.clip(output, 0, 1)
-        if invert:
+        if colormap_options.invert:
             output = 1 - output
-        return apply_float_colormap(output, colormap=colormap)
+        return apply_float_colormap(output, colormap=colormap_options.colormap)
 
     # rendering boolean outputs
     if image.dtype == torch.bool:
@@ -104,7 +114,7 @@ def apply_depth_colormap(
     accumulation: Optional[TensorType["bs":..., 1]] = None,
     near_plane: Optional[float] = None,
     far_plane: Optional[float] = None,
-    colormap: Colormaps = "turbo",
+    colormap_options: ColormapOptions = ColormapOptions(),
 ) -> TensorType["bs":..., "rgb":3]:
     """Converts a depth image to color for easier analysis.
 
@@ -126,7 +136,7 @@ def apply_depth_colormap(
     depth = torch.clip(depth, 0, 1)
     # depth = torch.nan_to_num(depth, nan=0.0) # TODO(ethan): remove this
 
-    colored_image = apply_colormap(depth, colormap=colormap)
+    colored_image = apply_colormap(depth, colormap_options=colormap_options)
 
     if accumulation is not None:
         colored_image = colored_image * accumulation + (1 - accumulation)

nerfstudio/viewer/app/src/themes/leva_theme.json

@@ -20,7 +20,7 @@
   "space": {
     "sm": "6px",
     "md": "5px",
-    "rowGap": "6px",
+    "rowGap": "3px",
     "colGap": "7px"
   },
   "fontSizes": {
@@ -31,15 +31,15 @@
     "controlWidth": "170px",
     "scrubberWidth": "8px",
     "scrubberHeight": "16px",
-    "rowHeight": "24px",
-    "folderHeight": "20px",
+    "rowHeight": "22px",
+    "folderHeight": "10px",
     "checkboxSize": "16px",
     "joystickWidth": "100px",
     "joystickHeight": "100px",
     "colorPickerWidth": "160px",
     "colorPickerHeight": "100px",
     "monitorHeight": "60px",
-    "titleBarHeight": "39px"
+    "titleBarHeight": "20px"
   },
   "borderWidths": {
     "root": "0px",

nerfstudio/viewer/server/control_panel.py

@@ -18,7 +18,7 @@ from typing import Callable, DefaultDict, List, Tuple, get_args
 
 import torch
 
-from nerfstudio.utils.colormaps import Colormaps
+from nerfstudio.utils.colormaps import ColormapOptions, Colormaps
 from nerfstudio.viewer.server.viewer_elements import (
     ViewerCheckbox,
     ViewerDropdown,
@@ -42,8 +42,17 @@ class ControlPanel:
         update_output_cb: a callback that will be called when the user changes the output render
     """
 
-    def __init__(self, time_enabled: bool, rerender_cb: Callable, crop_update_cb: Callable, update_output_cb: Callable):
+    def __init__(
+        self,
+        viser_server: ViserServer,
+        time_enabled: bool,
+        rerender_cb: Callable,
+        crop_update_cb: Callable,
+        update_output_cb: Callable,
+        update_split_output_cb: Callable,
+    ):
         # elements holds a mapping from tag: [elements]
+        self.viser_server = viser_server
         self._elements_by_tag: DefaultDict[str, List[ViewerElement]] = defaultdict(lambda: [])
 
         self._train_speed = ViewerDropdown(
@@ -56,39 +65,74 @@ class ControlPanel:
             cb_hook=lambda han: [self.update_control_panel(), update_output_cb(han), rerender_cb(han)],
         )
         self._colormap = ViewerDropdown[Colormaps]("Colormap", "default", ["default"], cb_hook=rerender_cb)
-        self._invert = ViewerCheckbox("| Invert", False, cb_hook=rerender_cb)
-        self._normalize = ViewerCheckbox("| Normalize", False, cb_hook=rerender_cb)
-        self._min = ViewerNumber("| Min", 0.0, cb_hook=rerender_cb)
-        self._max = ViewerNumber("| Max", 1.0, cb_hook=rerender_cb)
+        self._invert = ViewerCheckbox("Invert", False, cb_hook=rerender_cb)
+        self._normalize = ViewerCheckbox("Normalize", False, cb_hook=rerender_cb)
+        self._min = ViewerNumber("Min", 0.0, cb_hook=rerender_cb)
+        self._max = ViewerNumber("Max", 1.0, cb_hook=rerender_cb)
+
+        self._split = ViewerCheckbox(
+            "Enable", False, cb_hook=lambda han: [self.update_control_panel(), rerender_cb(han)]
+        )
+        self._split_percentage = ViewerSlider("Split Percentage", 0.5, 0.0, 1.0, 0.01, cb_hook=rerender_cb)
+        self._split_output_render = ViewerDropdown(
+            "Output Render Split",
+            "not set",
+            ["not set"],
+            cb_hook=lambda han: [self.update_control_panel(), update_split_output_cb(han), rerender_cb(han)],
+        )
+        # Hack: spaces are after at the end of the names to make them unique
+        self._split_colormap = ViewerDropdown[Colormaps]("Colormap ", "default", ["default"], cb_hook=rerender_cb)
+        self._split_invert = ViewerCheckbox("Invert ", False, cb_hook=rerender_cb)
+        self._split_normalize = ViewerCheckbox("Normalize ", True, cb_hook=rerender_cb)
+        self._split_min = ViewerNumber("Min ", 0.0, cb_hook=rerender_cb)
+        self._split_max = ViewerNumber("Max ", 1.0, cb_hook=rerender_cb)
+
         self._train_util = ViewerSlider("Train Util", 0.85, 0, 1, 0.05)
         self._max_res = ViewerSlider("Max Res", 512, 64, 2048, 100, cb_hook=rerender_cb)
         self._crop_viewport = ViewerCheckbox(
-            "Crop Viewport",
+            "Enable ",
             False,
             cb_hook=lambda han: [self.update_control_panel(), crop_update_cb(han), rerender_cb(han)],
         )
-        self._background_color = ViewerRGB("| Background color", (38, 42, 55), cb_hook=crop_update_cb)
-        self._crop_min = ViewerVec3("| Crop Min", (-1, -1, -1), 0.05, cb_hook=crop_update_cb)
-        self._crop_max = ViewerVec3("| Crop Max", (1, 1, 1), 0.05, cb_hook=crop_update_cb)
+        self._background_color = ViewerRGB("Background color", (38, 42, 55), cb_hook=crop_update_cb)
+        self._crop_min = ViewerVec3("Crop Min", (-1, -1, -1), 0.05, cb_hook=crop_update_cb)
+        self._crop_max = ViewerVec3("Crop Max", (1, 1, 1), 0.05, cb_hook=crop_update_cb)
         self._time = ViewerSlider("Time", 0.0, 0.0, 1.0, 0.01, cb_hook=rerender_cb)
         self._time_enabled = time_enabled
 
         self.add_element(self._train_speed)
-        self.add_element(self._output_render)
-        self.add_element(self._colormap)
-        # colormap options
-        self.add_element(self._invert, additional_tags=("colormap",))
-        self.add_element(self._normalize, additional_tags=("colormap",))
-        self.add_element(self._min, additional_tags=("colormap",))
-        self.add_element(self._max, additional_tags=("colormap",))
-
         self.add_element(self._train_util)
-        self.add_element(self._max_res)
-        self.add_element(self._crop_viewport)
-        # Crop options
-        self.add_element(self._background_color, additional_tags=("crop",))
-        self.add_element(self._crop_min, additional_tags=("crop",))
-        self.add_element(self._crop_max, additional_tags=("crop",))
+        with self.viser_server.gui_folder("Render Options"):
+            self.add_element(self._max_res)
+            self.add_element(self._output_render)
+            self.add_element(self._colormap)
+            # colormap options
+            with self.viser_server.gui_folder(" "):
+                self.add_element(self._invert, additional_tags=("colormap",))
+                self.add_element(self._normalize, additional_tags=("colormap",))
+                self.add_element(self._min, additional_tags=("colormap",))
+                self.add_element(self._max, additional_tags=("colormap",))
+
+        # split options
+        with self.viser_server.gui_folder("Split Screen"):
+            self.add_element(self._split)
+
+            self.add_element(self._split_percentage, additional_tags=("split",))
+            self.add_element(self._split_output_render, additional_tags=("split",))
+            self.add_element(self._split_colormap, additional_tags=("split",))
+            with self.viser_server.gui_folder("  "):
+                self.add_element(self._split_invert, additional_tags=("split_colormap",))
+                self.add_element(self._split_normalize, additional_tags=("split_colormap",))
+                self.add_element(self._split_min, additional_tags=("split_colormap",))
+                self.add_element(self._split_max, additional_tags=("split_colormap",))
+
+        with self.viser_server.gui_folder("Crop Viewport"):
+            self.add_element(self._crop_viewport)
+
+            # Crop options
+            self.add_element(self._background_color, additional_tags=("crop",))
+            self.add_element(self._crop_min, additional_tags=("crop",))
+            self.add_element(self._crop_max, additional_tags=("crop",))
 
         self.add_element(self._time, additional_tags=("time",))
 
@@ -104,22 +148,14 @@ class ControlPanel:
             self._train_util.value = 0.5
             self._max_res.value = 1024
 
-    def install(self, viser_server: ViserServer) -> None:
-        """Installs the control panel on the viser server
-
-        Args:
-            viser_server: the viser server
-        """
-        for e in self._elements_by_tag["all"]:
-            e.install(viser_server)
-        self.update_control_panel()
-
     def update_output_options(self, new_options: List[str]):
         """
         Args:
             new_options: a list of new output options
         """
         self._output_render.set_options(new_options)
+        self._split_output_render.set_options(new_options)
+        self._split_output_render.value = new_options[-1]
 
     def add_element(self, e: ViewerElement, additional_tags: Tuple[str] = tuple()) -> None:
         """Adds an element to the control panel
@@ -131,6 +167,7 @@ class ControlPanel:
         self._elements_by_tag["all"].append(e)
         for t in additional_tags:
             self._elements_by_tag[t].append(e)
+        e.install(self.viser_server)
 
     def update_control_panel(self) -> None:
         """
@@ -139,9 +176,15 @@ class ControlPanel:
         self._colormap.set_disabled(self.output_render == "rgb")
         for e in self._elements_by_tag["colormap"]:
             e.set_hidden(self.output_render == "rgb")
+        for e in self._elements_by_tag["split_colormap"]:
+            e.set_hidden(not self._split.value or self.split_output_render == "rgb")
         for e in self._elements_by_tag["crop"]:
             e.set_hidden(not self.crop_viewport)
         self._time.set_hidden(not self._time_enabled)
+        self._split_percentage.set_hidden(not self._split.value)
+        self._split_output_render.set_hidden(not self._split.value)
+        self._split_colormap.set_hidden(not self._split.value)
+        self._split_colormap.set_disabled(self.split_output_render == "rgb")
 
     def update_colormap_options(self, dimensions: int, dtype: type) -> None:
         """update the colormap options based on the current render
@@ -150,13 +193,16 @@ class ControlPanel:
             dimensions: the number of dimensions of the render
             dtype: the data type of the render
         """
-        colormap_options: List[Colormaps] = []
-        if dimensions == 3:
-            colormap_options = ["default"]
-        if dimensions == 1 and dtype == torch.float:
-            colormap_options = [c for c in list(get_args(Colormaps)) if c != "default"]
+        self._colormap.set_options(_get_colormap_options(dimensions, dtype))
 
-        self._colormap.set_options(colormap_options)
+    def update_split_colormap_options(self, dimensions: int, dtype: type) -> None:
+        """update the split colormap options based on the current render
+
+        Args:
+            dimensions: the number of dimensions of the render
+            dtype: the data type of the render
+        """
+        self._split_colormap.set_options(_get_colormap_options(dimensions, dtype))
 
     @property
     def train_speed(self) -> str:
@@ -169,30 +215,19 @@ class ControlPanel:
         return self._output_render.value
 
     @property
-    def colormap(self) -> Colormaps:
-        """Returns the current colormap"""
-        colormap = self._colormap.value
-        return colormap
+    def split_output_render(self) -> str:
+        """Returns the current output for the split render"""
+        return self._split_output_render.value
 
     @property
-    def colormap_invert(self) -> bool:
-        """Returns the current colormap invert setting"""
-        return self._invert.value
+    def split(self) -> bool:
+        """Returns whether the split is enabled"""
+        return self._split.value
 
     @property
-    def colormap_normalize(self) -> bool:
-        """Returns the current colormap normalize setting"""
-        return self._normalize.value
-
-    @property
-    def colormap_min(self) -> float:
-        """Returns the current colormap min setting"""
-        return self._min.value
-
-    @property
-    def colormap_max(self) -> float:
-        """Returns the current colormap max setting"""
-        return self._max.value
+    def split_percentage(self) -> float:
+        """Returns the percentage of the screen to split"""
+        return self._split_percentage.value
 
     @property
     def train_util(self) -> float:
@@ -253,3 +288,44 @@ class ControlPanel:
     def time(self, value: float):
         """Sets the background color"""
         self._time.value = value
+
+    @property
+    def colormap_options(self) -> ColormapOptions:
+        """Returns the current colormap options"""
+        return ColormapOptions(
+            colormap=self._colormap.value,
+            normalize=self._normalize.value,
+            colormap_min=self._min.value,
+            colormap_max=self._max.value,
+            invert=self._invert.value,
+        )
+
+    @property
+    def split_colormap_options(self) -> ColormapOptions:
+        """Returns the current colormap options"""
+        return ColormapOptions(
+            colormap=self._split_colormap.value,
+            normalize=self._split_normalize.value,
+            colormap_min=self._split_min.value,
+            colormap_max=self._split_max.value,
+            invert=self._split_invert.value,
+        )
+
+
+def _get_colormap_options(dimensions: int, dtype: type) -> List[Colormaps]:
+    """
+    Given the number of dimensions and data type, returns a list of available colormap options
+    to use with the visualize() function.
+
+    Args:
+        dimensions: the number of dimensions of the render
+        dtype: the data type of the render
+    Returns:
+        a list of available colormap options
+    """
+    colormap_options: List[Colormaps] = []
+    if dimensions == 3:
+        colormap_options = ["default"]
+    if dimensions == 1 and dtype == torch.float:
+        colormap_options = [c for c in list(get_args(Colormaps)) if c != "default"]
+    return colormap_options

nerfstudio/viewer/server/render_state_machine.py

@@ -240,12 +240,25 @@ class RenderStateMachine(threading.Thread):
         )
         selected_output = colormaps.apply_colormap(
             image=outputs[self.viewer.control_panel.output_render],
-            colormap=self.viewer.control_panel.colormap,
-            normalize=self.viewer.control_panel.colormap_normalize,
-            colormap_min=self.viewer.control_panel.colormap_min,
-            colormap_max=self.viewer.control_panel.colormap_max,
-            invert=self.viewer.control_panel.colormap_invert,
+            colormap_options=self.viewer.control_panel.colormap_options,
         )
+
+        if self.viewer.control_panel.split:
+            split_output_render = self.viewer.control_panel.split_output_render
+            self.viewer.update_colormap_options(
+                dimensions=outputs[split_output_render].shape[-1], dtype=outputs[split_output_render].dtype
+            )
+            split_output = colormaps.apply_colormap(
+                image=outputs[self.viewer.control_panel.split_output_render],
+                colormap_options=self.viewer.control_panel.split_colormap_options,
+            )
+            split_index = min(
+                int(self.viewer.control_panel.split_percentage * selected_output.shape[1]),
+                selected_output.shape[1] - 1,
+            )
+            selected_output = torch.cat([selected_output[:, :split_index], split_output[:, split_index:]], dim=1)
+            selected_output[:, split_index] = torch.tensor([0.133, 0.157, 0.192], device=selected_output.device)
+
         selected_output = (selected_output * 255).type(torch.uint8)
 
         self.viewer.viser_server.set_background_image(

nerfstudio/viewer/server/viewer_state.py

@@ -115,6 +115,7 @@ class ViewerState:
 
         # viewer specific variables
         self.output_type_changed = True
+        self.output_split_type_changed = True
         self.step = 0
         self.train_btn_state: Literal["training", "paused", "completed"] = "training"
         self._prev_train_state: Literal["training", "paused", "completed"] = "training"
@@ -134,9 +135,13 @@ class ViewerState:
             self.viser_server.register_handler(TimeConditionMessage, self._handle_time_condition_message)
 
         self.control_panel = ControlPanel(
-            self.include_time, self._interrupt_render, self._crop_params_update, self._output_type_change
+            self.viser_server,
+            self.include_time,
+            self._interrupt_render,
+            self._crop_params_update,
+            self._output_type_change,
+            self._output_split_type_change,
         )
-        self.control_panel.install(self.viser_server)
 
         def nested_folder_install(folder_labels: List[str], element: ViewerElement):
             if len(folder_labels) == 0:
@@ -163,6 +168,9 @@ class ViewerState:
     def _output_type_change(self, _):
         self.output_type_changed = True
 
+    def _output_split_type_change(self, _):
+        self.output_split_type_changed = True
+
     def _interrupt_render(self, _) -> None:
         """Interrupt current render."""
         if self.camera_message is not None:
@@ -350,6 +358,17 @@ class ViewerState:
             self.control_panel.update_colormap_options(dimensions, dtype)
             self.output_type_changed = False
 
+    def update_split_colormap_options(self, dimensions: int, dtype: type) -> None:
+        """update the colormap options based on the current render
+
+        Args:
+            dimensions: the number of dimensions of the render
+            dtype: the data type of the render
+        """
+        if self.output_split_type_changed:
+            self.control_panel.update_split_colormap_options(dimensions, dtype)
+            self.output_split_type_changed = False
+
     def get_model(self) -> Model:
         """Returns the model."""
         return self.pipeline.model