diff --git a/nerfstudio/cameras/cameras.py b/nerfstudio/cameras/cameras.py index f2020916f..fee7495a8 100644 --- a/nerfstudio/cameras/cameras.py +++ b/nerfstudio/cameras/cameras.py @@ -314,7 +314,7 @@ class Cameras(TensorDataclass): distortion_params_delta: Optional[TensorType["num_rays":..., 6]] = None, keep_shape: Optional[bool] = None, disable_distortion: bool = False, - aabb_box: SceneBox = None, + aabb_box: Optional[SceneBox] = None, ) -> RayBundle: """Generates rays for the given camera indices. diff --git a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index 44e3f5ea0..ec5d95e82 100644 --- a/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx @@ -47,8 +47,11 @@ export function RenderControls() { const crop_enabled = useSelector( (state) => state.renderingState.crop_enabled, ); + console.log('crop_enabled_config panel', crop_enabled); - const crop_bg_color = useSelector((state) => state.renderingState.crop_bg_color); + const crop_bg_color = useSelector( + (state) => state.renderingState.crop_bg_color, + ); const crop_scale = useSelector((state) => state.renderingState.crop_scale); @@ -154,11 +157,12 @@ export function RenderControls() { '2048px': () => setControls({ max_resolution: 2048 }), }), // Enable Crop - crop_options: { + crop_enabled: { label: 'Crop Viewport', value: crop_enabled, hint: 'Crop the viewport to the selected box', onChange: (value) => { + console.log('dispact value', value); dispatch_and_send( websocket, dispatch, @@ -170,7 +174,7 @@ export function RenderControls() { crop_bg_color: { label: '| Background Color', value: crop_bg_color, - render: (get) => get('crop_options'), + render: (get) => get('crop_enabled'), onChange: (v) => { dispatch_and_send( websocket, @@ -186,7 +190,7 @@ export function RenderControls() { min: 0, max: 10, step: 0.05, - render: (get) => get('crop_options'), + render: (get) => get('crop_enabled'), onChange: (v) => { dispatch_and_send( websocket, @@ -202,7 +206,7 @@ export function RenderControls() { min: -10, max: 10, step: 0.05, - render: (get) => get('crop_options'), + render: (get) => get('crop_enabled'), onChange: (v) => { dispatch_and_send( websocket, @@ -251,6 +255,10 @@ export function RenderControls() { setControls({ max_resolution }); setControls({ output_options: outputChoice }); setControls({ colormap_options: colormapChoice }); + setControls({ crop_enabled }); + setControls({ crop_bg_color }); + setControls({ crop_scale }); + setControls({ crop_center }); }, [ setControls, outputOptions, @@ -260,6 +268,10 @@ export function RenderControls() { max_resolution, target_train_util, render_time, + crop_enabled, + crop_bg_color, + crop_scale, + crop_center, display_render_time, ]); diff --git a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx index 27fbf5e36..37e54ad7d 100644 --- a/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx +++ b/nerfstudio/viewer/app/src/modules/SidePanel/CameraPanel/CameraPanel.jsx @@ -63,7 +63,6 @@ function set_camera_position(camera, matrix) { mat.decompose(camera.position, camera.quaternion, camera.scale); } - function RenderTimeSelector(props) { const disabled = props.disabled; const isGlobal = props.isGlobal; @@ -79,7 +78,7 @@ function RenderTimeSelector(props) { if (!isGlobal) { return camera.renderTime; } - camera.renderTime = globalRenderTime + camera.renderTime = globalRenderTime; return globalRenderTime; }; @@ -118,7 +117,7 @@ function RenderTimeSelector(props) { const handleValidation = (e) => { const valueFloat = parseFloat(e.target.value); let valueStr = String(valueFloat); - if (e.target.value >= 0 && e.target.value <= 1){ + if (e.target.value >= 0 && e.target.value <= 1) { setValid(true); if (valueFloat === 1.0) { valueStr = '1.0'; @@ -135,7 +134,7 @@ function RenderTimeSelector(props) { return ( state.renderingState.export_path, - ); + const export_path = useSelector((state) => state.renderingState.export_path); const websocket = useContext(WebSocketContext).socket; const DEFAULT_FOV = 50; @@ -577,7 +574,8 @@ export default function CameraPanel(props) { const [render_modal_open, setRenderModalOpen] = React.useState(false); const [animate, setAnimate] = React.useState(new Set()); const [globalFov, setGlobalFov] = React.useState(DEFAULT_FOV); - const [globalRenderTime, setGlobalRenderTime] = React.useState(DEFAULT_RENDER_TIME); + const [globalRenderTime, setGlobalRenderTime] = + React.useState(DEFAULT_RENDER_TIME); // leva store const cameraPropsStore = useCreateStore(); @@ -603,15 +601,24 @@ export default function CameraPanel(props) { ); const camera_type = useSelector((state) => state.renderingState.camera_type); + const crop_enabled = useSelector( + (state) => state.renderingState.crop_enabled, + ); + const crop_bg_color = useSelector( + (state) => state.renderingState.crop_bg_color, + ); + const crop_center = useSelector((state) => state.renderingState.crop_center); + const crop_scale = useSelector((state) => state.renderingState.crop_scale); + const [display_render_time, set_display_render_time] = React.useState(false); - const receive_temporal_dist = e => { + const receive_temporal_dist = (e) => { const msg = msgpack.decode(new Uint8Array(e.data)); - if (msg.path === "/model/has_temporal_distortion") { - set_display_render_time(msg.data === "true"); - websocket.removeEventListener("message", receive_temporal_dist); + if (msg.path === '/model/has_temporal_distortion') { + set_display_render_time(msg.data === 'true'); + websocket.removeEventListener('message', receive_temporal_dist); } - } + }; websocket.addEventListener('message', receive_temporal_dist); const setRenderHeight = (value) => { @@ -645,6 +652,40 @@ export default function CameraPanel(props) { }); }; + const setCropEnabled = (value) => { + console.log('setting the crop enabled value to: ', value); + dispatch({ + type: 'write', + path: 'renderingState/crop_enabled', + data: value, + }); + console.log('crop enabled value is now: ', crop_enabled); + }; + + const serCropBgColor = (value) => { + dispatch({ + type: 'write', + path: 'renderingState/crop_bg_color', + data: value, + }); + }; + + const setCropCenter = (value) => { + dispatch({ + type: 'write', + path: 'renderingState/crop_center', + data: value, + }); + }; + + const setCropScale = (value) => { + dispatch({ + type: 'write', + path: 'renderingState/crop_scale', + data: value, + }); + }; + const setRenderTime = (value) => { dispatch({ type: 'write', @@ -669,7 +710,7 @@ export default function CameraPanel(props) { if (new_camera_list.length >= 1) { set_camera_position(camera_render, new_camera_list[0].matrix); setFieldOfView(new_camera_list[0].fov); - setRenderTime(new_camera_list[0].renderTime) + setRenderTime(new_camera_list[0].renderTime); set_slider_value(slider_min); } }; @@ -989,6 +1030,15 @@ export default function CameraPanel(props) { }); } + let crop = null; + if (crop_enabled) { + crop = { + crop_bg_color, + crop_center, + crop_scale, + }; + } + // const myData const camera_path_object = { keyframes, @@ -1000,6 +1050,7 @@ export default function CameraPanel(props) { seconds, smoothness_value, is_cycle, + crop, }; return camera_path_object; }; @@ -1066,6 +1117,14 @@ export default function CameraPanel(props) { setCameraProperties(new_properties); setCameras(new_camera_list); reset_slider_render_on_add(new_camera_list); + + if ('crop' in camera_path_object && camera_path_object.crop !== null) { + console.log('crop' in camera_path_object); + setCropEnabled(true); + serCropBgColor(camera_path_object.crop.crop_bg_color); + setCropCenter(camera_path_object.crop.crop_center); + setCropScale(camera_path_object.crop.crop_scale); + } }; const uploadCameraPath = (e) => { @@ -1160,6 +1219,7 @@ export default function CameraPanel(props) { variant="outlined" startIcon={} onClick={export_camera_path} + disabled={cameras.length === 0} > Export Path @@ -1172,6 +1232,7 @@ export default function CameraPanel(props) { size="small" startIcon={} onClick={open_render_modal} + disabled={cameras.length === 0} > Render diff --git a/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx b/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx index d48edfff3..5218b82b8 100644 --- a/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx +++ b/nerfstudio/viewer/app/src/modules/ViewerWindow/ViewerWindow.jsx @@ -18,13 +18,9 @@ const msgpack = require('msgpack-lite'); function CameraToggle() { const dispatch = useDispatch(); - const websocket = useContext(WebSocketContext).socket; const camera_choice = useSelector( (state) => state.renderingState.camera_choice, ); - const crop_enabled = useSelector( - (state) => state.renderingState.crop_enabled, - ); const set_camera_choice = (event: SelectChangeEvent, value: string[]) => { if (value != null) { dispatch({ @@ -32,21 +28,6 @@ function CameraToggle() { path: 'renderingState/camera_choice', data: value, }); - if (websocket.readyState === WebSocket.OPEN) { - const path = 'renderingState/crop_enabled'; - const data = { - type: 'write', - path, - data: false, - }; - if (value === 'Render Camera') { - data.data = false; - } else { - data.data = crop_enabled; - } - const message = msgpack.encode(data); - websocket.send(message); - } } }; diff --git a/scripts/render.py b/scripts/render.py index 357b902ed..4466c5385 100644 --- a/scripts/render.py +++ b/scripts/render.py @@ -11,7 +11,7 @@ import sys from contextlib import ExitStack from dataclasses import dataclass, field from pathlib import Path -from typing import List, Optional +from typing import Any, Dict, List, Optional import mediapy as media import numpy as np @@ -25,10 +25,13 @@ from rich.progress import ( TextColumn, TimeRemainingColumn, ) +from torchtyping import TensorType from typing_extensions import Literal, assert_never from nerfstudio.cameras.camera_paths import get_path_from_json, get_spiral_path 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.eval_utils import eval_setup @@ -42,6 +45,7 @@ def _render_trajectory_video( cameras: Cameras, output_filename: Path, rendered_output_names: List[str], + crop_data: Optional[CropData] = None, rendered_resolution_scaling_factor: float = 1.0, seconds: float = 5.0, output_format: Literal["images", "video"] = "video", @@ -54,6 +58,7 @@ def _render_trajectory_video( cameras: Cameras to render. output_filename: Name of the output file. rendered_output_names: List of outputs to visualise. + crop_data: Crop data to apply to the rendered images. rendered_resolution_scaling_factor: Scaling factor to apply to the camera image resolution. seconds: Length of output video. output_format: How to save output data. @@ -89,9 +94,23 @@ def _render_trajectory_video( with progress: for camera_idx in progress.track(range(cameras.size), description=""): - camera_ray_bundle = cameras.generate_rays(camera_indices=camera_idx) - with torch.no_grad(): - outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + + aabb_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) + + if crop_data is not None: + with renderers.background_color_override_context( + crop_data.background_color.to(pipeline.device) + ), torch.no_grad(): + outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + else: + with torch.no_grad(): + outputs = pipeline.model.get_outputs_for_camera_ray_bundle(camera_ray_bundle) + render_image = [] for rendered_output_name in rendered_output_names: if rendered_output_name not in outputs: @@ -110,8 +129,8 @@ def _render_trajectory_video( media.write_image(output_image_dir / f"{camera_idx:05d}.png", render_image) if output_format == "video": if writer is None: - render_width = int(render_image.shape[1] * rendered_resolution_scaling_factor) - render_height = int(render_image.shape[0] * rendered_resolution_scaling_factor) + render_width = int(render_image.shape[1]) + render_height = int(render_image.shape[0]) writer = stack.enter_context( media.VideoWriter( path=output_filename, @@ -196,28 +215,60 @@ xmlns:GSpherical='http://ns.google.com/videos/1.0/spherical/'> mp4file.close() +@dataclass +class CropData: + """Data for cropping an image.""" + + background_color: TensorType[3] = torch.Tensor([0.0, 0.0, 0.0]) + """background color""" + center: TensorType[3] = torch.Tensor([0.0, 0.0, 0.0]) + """center of the crop""" + scale: TensorType[3] = torch.Tensor([2.0, 2.0, 2.0]) + """scale of the crop""" + + +def get_crop_from_json(camera_json: Dict[str, Any]) -> Optional[CropData]: + """Load crop data from a camera path JSON + + args: + camera_json: camera path data + returns: + Crop data + """ + if "crop" not in camera_json or camera_json["crop"] is None: + return None + + bg_color = camera_json["crop"]["crop_bg_color"] + + 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"]), + ) + + @dataclass class RenderTrajectory: """Load a checkpoint, render a trajectory, and save to a video file.""" - # Path to config YAML file. load_config: Path - # Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis + """Path to config YAML file.""" rendered_output_names: List[str] = field(default_factory=lambda: ["rgb"]) - # Trajectory to render. + """Name of the renderer outputs to use. rgb, depth, etc. concatenates them along y axis""" traj: Literal["spiral", "filename"] = "spiral" - # Scaling factor to apply to the camera image resolution. + """Trajectory to render.""" downscale_factor: int = 1 - # Filename of the camera path to render. + """Scaling factor to apply to the camera image resolution.""" camera_path_filename: Path = Path("camera_path.json") - # Name of the output file. + """Filename of the camera path to render.""" output_path: Path = Path("renders/output.mp4") - # How long the video should be. + """Name of the output file.""" seconds: float = 5.0 - # How to save output data. + """How long the video should be.""" output_format: Literal["images", "video"] = "video" - # Specifies number of rays per chunk during eval. + """How to save output data.""" eval_num_rays_per_chunk: Optional[int] = None + """Specifies number of rays per chunk during eval.""" def main(self) -> None: """Main function.""" @@ -230,6 +281,7 @@ class RenderTrajectory: install_checks.check_ffmpeg_installed() seconds = self.seconds + crop_data = None # TODO(ethan): use camera information from parsing args if self.traj == "spiral": @@ -249,6 +301,7 @@ class RenderTrajectory: camera_type = CameraType.EQUIRECTANGULAR else: camera_type = CameraType.PERSPECTIVE + crop_data = get_crop_from_json(camera_path) camera_path = get_path_from_json(camera_path) else: assert_never(self.traj) @@ -259,6 +312,7 @@ class RenderTrajectory: output_filename=self.output_path, rendered_output_names=self.rendered_output_names, rendered_resolution_scaling_factor=1.0 / self.downscale_factor, + crop_data=crop_data, seconds=seconds, output_format=self.output_format, camera_type=camera_type,