diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index f3f6b6789..ae6937380 100644 --- a/nerfstudio/scripts/render.py +++ b/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, ) diff --git a/nerfstudio/utils/colormaps.py b/nerfstudio/utils/colormaps.py index c714f88db..a9eb6fc8b 100644 --- a/nerfstudio/utils/colormaps.py +++ b/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) diff --git a/nerfstudio/viewer/app/src/themes/leva_theme.json b/nerfstudio/viewer/app/src/themes/leva_theme.json index 9a0969ef5..db09d14c4 100644 --- a/nerfstudio/viewer/app/src/themes/leva_theme.json +++ b/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", diff --git a/nerfstudio/viewer/server/control_panel.py b/nerfstudio/viewer/server/control_panel.py index 74f305035..7ee539fa6 100644 --- a/nerfstudio/viewer/server/control_panel.py +++ b/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 diff --git a/nerfstudio/viewer/server/render_state_machine.py b/nerfstudio/viewer/server/render_state_machine.py index f3af3e3fe..1001b4f18 100644 --- a/nerfstudio/viewer/server/render_state_machine.py +++ b/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( diff --git a/nerfstudio/viewer/server/viewer_state.py b/nerfstudio/viewer/server/viewer_state.py index 0d37ee902..513bc480c 100644 --- a/nerfstudio/viewer/server/viewer_state.py +++ b/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