From e0d8f0426752dc0f97bd332fd6b84e95e8377210 Mon Sep 17 00:00:00 2001 From: Justin Kerr Date: Tue, 19 Sep 2023 16:48:54 -0700 Subject: [PATCH] Render oriented crops in viewer beta (#2447) * add obb to the rendering pipeline * useless line * lint * remove todo * pyright * more pyright --- nerfstudio/engine/trainer.py | 5 +- nerfstudio/scripts/render.py | 33 +++--- nerfstudio/viewer_beta/render_panel.py | 135 +++++++++++++++++++++---- nerfstudio/viewer_beta/viewer.py | 4 +- 4 files changed, 139 insertions(+), 38 deletions(-) diff --git a/nerfstudio/engine/trainer.py b/nerfstudio/engine/trainer.py index 04fa8e05b..b4ccba9ee 100644 --- a/nerfstudio/engine/trainer.py +++ b/nerfstudio/engine/trainer.py @@ -170,10 +170,13 @@ class Trainer: ) banner_messages = [f"Viewer at: {self.viewer_state.viewer_url}"] if self.config.is_viewer_beta_enabled() and self.local_rank == 0: + datapath = self.config.data + if datapath is None: + datapath = self.base_dir self.viewer_state = ViewerBetaState( self.config.viewer, log_filename=viewer_log_path, - datapath=self.base_dir, + datapath=datapath, pipeline=self.pipeline, trainer=self, train_lock=self.train_lock, diff --git a/nerfstudio/scripts/render.py b/nerfstudio/scripts/render.py index 5a3ca78bd..99a7cae01 100644 --- a/nerfstudio/scripts/render.py +++ b/nerfstudio/scripts/render.py @@ -54,7 +54,7 @@ from nerfstudio.cameras.camera_paths import ( ) from nerfstudio.cameras.cameras import Cameras, CameraType from nerfstudio.data.datamanagers.base_datamanager import VanillaDataManager -from nerfstudio.data.scene_box import SceneBox +from nerfstudio.data.scene_box import OrientedBox from nerfstudio.model_components import renderers from nerfstudio.pipelines.base_pipeline import Pipeline from nerfstudio.utils import colormaps, install_checks @@ -127,12 +127,10 @@ def _render_trajectory_video( with progress: for camera_idx in progress.track(range(cameras.size), description=""): - aabb_box = None + obb_box = None if crop_data is not None: - bounding_box_min = crop_data.center - crop_data.scale / 2.0 - bounding_box_max = crop_data.center + crop_data.scale / 2.0 - aabb_box = SceneBox(torch.stack([bounding_box_min, bounding_box_max]).to(pipeline.device)) - camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx, aabb_box=aabb_box) + obb_box = crop_data.obb + camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx, obb_box=obb_box) if crop_data is not None: with renderers.background_color_override_context( @@ -289,10 +287,17 @@ class CropData: background_color: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0]) """background color""" - center: Float[Tensor, "3"] = torch.Tensor([0.0, 0.0, 0.0]) - """center of the crop""" - scale: Float[Tensor, "3"] = torch.Tensor([2.0, 2.0, 2.0]) - """scale of the crop""" + obb: OrientedBox = OrientedBox(R=torch.eye(3), T=torch.zeros(3), S=torch.ones(3) * 2) + """Oriented box representing the crop region""" + + # properties for backwards-compatibility interface + @property + def center(self): + return self.obb.T + + @property + def scale(self): + return self.obb.S def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]: @@ -305,13 +310,13 @@ def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]: """ if "crop" not in camera_json or camera_json["crop"] is None: return None - bg_color = camera_json["crop"]["crop_bg_color"] - + center = camera_json["crop"]["crop_center"] + scale = camera_json["crop"]["crop_scale"] + rot = (0.0, 0.0, 0.0) if "crop_rot" not in camera_json["crop"] else tuple(camera_json["crop"]["crop_rot"]) return CropData( background_color=torch.Tensor([bg_color["r"] / 255.0, bg_color["g"] / 255.0, bg_color["b"] / 255.0]), - center=torch.Tensor(camera_json["crop"]["crop_center"]), - scale=torch.Tensor(camera_json["crop"]["crop_scale"]), + obb=OrientedBox.from_params(center, rot, scale), ) diff --git a/nerfstudio/viewer_beta/render_panel.py b/nerfstudio/viewer_beta/render_panel.py index 02d4d6ce0..7342b0f38 100644 --- a/nerfstudio/viewer_beta/render_panel.py +++ b/nerfstudio/viewer_beta/render_panel.py @@ -13,17 +13,19 @@ # limitations under the License. from __future__ import annotations - +from pathlib import Path import colorsys import dataclasses import threading import time from typing import Dict, List, Optional, Tuple - +import datetime +from nerfstudio.viewer_beta.control_panel import ControlPanel import numpy as onp import splines import splines.quaternion import viser +import json import viser.transforms as tf @@ -256,7 +258,11 @@ class CameraPath: ) -def populate_render_tab(server: viser.ViserServer) -> None: +def populate_render_tab( + server: viser.ViserServer, config_path: Path, datapath: Path, control_panel: ControlPanel +) -> None: + from nerfstudio.viewer_beta.viewer import VISER_NERFSTUDIO_SCALE_RATIO + fov_degrees = server.add_gui_slider( "FOV", initial_value=90.0, @@ -528,6 +534,11 @@ def populate_render_tab(server: viser.ViserServer) -> None: play_button.visible = True pause_button.visible = False + # set the initial value to the current date-time string + now = datetime.datetime.now() + render_name_text = server.add_gui_text( + "Render Name", initial_value=now.strftime("%Y-%m-%d-%H-%M-%S"), hint="Name of the render" + ) render_button = server.add_gui_button( "Generate Command", color="green", @@ -537,10 +548,107 @@ def populate_render_tab(server: viser.ViserServer) -> None: @render_button.on_click def _(event: viser.GuiEvent) -> None: - """TODO: write the render JSON and show the render command.""" assert event.client is not None - with event.client.add_gui_modal("TODO") as modal: - event.client.add_gui_markdown("TODO") + num_frames = int(framerate_slider.value * duration_number.value) + json_data = {} + # json data has the properties: + # keyframes: list of keyframes with + # matrix : flattened 4x4 matrix + # fov: float in degrees + # aspect: float + # camera_type: string of camera type + # render_height: int + # render_width: int + # fps: int + # seconds: float + # is_cycle: bool + # smoothness_value: float + # camera_path: list of frames with properties + # camera_to_world: flattened 4x4 matrix + # fov: float in degrees + # aspect: float + # first populate the keyframes: + keyframes = [] + for keyframe, dummy in camera_path._keyframes.values(): + pose = tf.SE3.from_rotation_and_translation( + tf.SO3(keyframe.wxyz) @ tf.SO3.from_x_radians(onp.pi), + keyframe.position / VISER_NERFSTUDIO_SCALE_RATIO, + ) + keyframes.append( + { + "matrix": pose.as_matrix().flatten().tolist(), + "fov": onp.rad2deg(keyframe.override_fov_value) + if keyframe.override_fov_enabled + else fov_degrees.value, + "aspect": keyframe.aspect, + } + ) + json_data["keyframes"] = keyframes + json_data["camera_type"] = "perspective" + json_data["render_height"] = resolution.value[1] + json_data["render_width"] = resolution.value[0] + json_data["fps"] = framerate_slider.value + json_data["seconds"] = duration_number.value + json_data["is_cycle"] = loop.value + json_data["smoothness_value"] = smoothness.value + # now populate the camera path: + camera_path_list = [] + for i in range(num_frames): + maybe_pose_and_fov = camera_path.interpolate_pose_and_fov(i / num_frames) + if maybe_pose_and_fov is None: + return + pose, fov = maybe_pose_and_fov + # rotate the axis of the camera 180 about x axis + pose = tf.SE3.from_rotation_and_translation( + pose.rotation() @ tf.SO3.from_x_radians(onp.pi), + pose.translation() / VISER_NERFSTUDIO_SCALE_RATIO, + ) + camera_path_list.append( + { + "camera_to_world": pose.as_matrix().flatten().tolist(), + "fov": onp.rad2deg(fov), + "aspect": resolution.value[0] / resolution.value[1], + } + ) + json_data["camera_path"] = camera_path_list + # finally add crop data if crop is enabled + if control_panel.crop_viewport: + obb = control_panel.crop_obb + rpy = tf.SO3.from_matrix(obb.R.numpy()).as_rpy_radians() + color = control_panel.background_color + json_data["crop"] = { + "crop_center": obb.T.tolist(), + "crop_scale": obb.S.tolist(), + "crop_rot": [rpy.roll, rpy.pitch, rpy.yaw], + "crop_bg_color": {"r": color[0], "g": color[1], "b": color[2]}, + } + + # now write the json file + json_outfile = datapath / "camera_paths" / f"{render_name_text.value}.json" + with open(json_outfile.absolute(), "w") as outfile: + json.dump(json_data, outfile) + # now show the command + with event.client.add_gui_modal("Render Command") as modal: + dataname = datapath.name + command = " ".join( + [ + "ns-render camera-path", + f"--load-config {config_path}", + f"--camera-path-filename {json_outfile.absolute()}", + f"--output-path renders/{dataname}/{render_name_text.value}.mp4", + ] + ) + event.client.add_gui_markdown( + "\n".join( + [ + "To render the trajectory, run the following from the command line:", + "", + "```", + command, + "```", + ] + ) + ) close_button = event.client.add_gui_button("Close") @close_button.on_click @@ -551,18 +659,3 @@ def populate_render_tab(server: viser.ViserServer) -> None: camera_path.default_fov = fov_degrees.value / 180.0 * onp.pi transform_controls: List[viser.SceneNodeHandle] = [] - - -def main() -> None: - """Launch a GUI with just the render panel, for development purposes.""" - server = viser.ViserServer() - server.configure_theme(dark_mode=True, control_layout="collapsible") - server.world_axes.visible = True - - populate_render_tab(server) - while True: - time.sleep(10.0) - - -if __name__ == "__main__": - main() diff --git a/nerfstudio/viewer_beta/viewer.py b/nerfstudio/viewer_beta/viewer.py index d2a53b198..6e6e02a29 100644 --- a/nerfstudio/viewer_beta/viewer.py +++ b/nerfstudio/viewer_beta/viewer.py @@ -154,11 +154,11 @@ class Viewer: self._toggle_training_state, self.set_camera_visibility, ) + config_path = self.log_filename.parents[0] / "config.yml" with tabs.add_tab("Render", viser.Icon.CAMERA): - populate_render_tab(self.viser_server) + populate_render_tab(self.viser_server, config_path, self.datapath, self.control_panel) with tabs.add_tab("Export", viser.Icon.PACKAGE_EXPORT): - config_path = self.log_filename.parents[0] / "config.yml" populate_export_tab(self.viser_server, self.control_panel, config_path) def nested_folder_install(folder_labels: List[str], element: ViewerElement):