diff --git a/nerfactory/utils/writer.py b/nerfactory/utils/writer.py index c1a736032..810118fa8 100644 --- a/nerfactory/utils/writer.py +++ b/nerfactory/utils/writer.py @@ -116,6 +116,7 @@ def put_time(name: str, duration: float, step: int, avg_over_steps: bool = True, remain_iter = GLOBAL_BUFFER["max_iter"] - step remain_time = remain_iter * GLOBAL_BUFFER["events"][name]["avg"] put_scalar(EventName.ETA, remain_time, step) + GLOBAL_BUFFER["viewer/train_eta"] = _format_time(remain_time) @check_main_thread @@ -152,6 +153,7 @@ def setup_event_writers(config: LoggingConfig, max_iter: int) -> None: GLOBAL_BUFFER["max_buffer_size"] = config.max_buffer_size GLOBAL_BUFFER["steps_per_log"] = config.steps_per_log GLOBAL_BUFFER["events"] = {} + GLOBAL_BUFFER["viewer/train_eta"] = "Paused" class Writer: diff --git a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx b/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx index ccf97f1d7..5742e7ad1 100644 --- a/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx +++ b/nerfactory/viewer/app/src/modules/ConfigPanel/ConfigPanel.jsx @@ -31,6 +31,12 @@ export function RenderControls() { const field_of_view = useSelector( (state) => state.renderingState.field_of_view, ); + let eval_fps = useSelector( + (state) => state.renderingState.eval_fps, + ); + let train_eta = useSelector( + (state) => state.renderingState.train_eta, + ); const dispatch = useDispatch(); @@ -111,18 +117,8 @@ export function RenderControls() { }; const [, setControls] = useControls( - 'Rendering State', + 'Rendering Controls', () => ({ - // WebSocket isConnected - // button does nothing except be an indicator - 'WebSocket Connected': button(() => {}, { - disabled: !isWebsocketConnected, - }), - // webRTC isConnected - // button does nothing except be an indicator - 'WebRTC Connected': button(() => {}, { - disabled: !isWebrtcConnected, - }), // isTraining 'Pause Training': button( () => { @@ -220,8 +216,6 @@ export function RenderControls() { }, }), [ - isWebsocketConnected, - isWebrtcConnected, isTraining, outputOptions, outputChoice, @@ -231,6 +225,38 @@ export function RenderControls() { ], ); + const [, setState] = useControls( + 'Rendering Monitor', + () => ({ + // WebSocket isConnected + // button does nothing except be an indicator + 'WebSocket Connected': button(() => {}, { + disabled: !isWebsocketConnected, + }), + // webRTC isConnected + // button does nothing except be an indicator + 'WebRTC Connected': button(() => {}, { + disabled: !isWebrtcConnected, + }), + eval_fps: { + label: 'Eval FPS', + value: eval_fps, + disabled: true, + }, + train_eta: { + label: 'Train ETA', + value: train_eta, + disabled: true, + }, + }), + [ + isWebsocketConnected, + isWebrtcConnected, + eval_fps, + train_eta, + ], + ); + useEffect(() => { setControls({ min_resolution }); setControls({ max_resolution }); @@ -238,8 +264,6 @@ export function RenderControls() { setControls({ 'Camera FoV': field_of_view }); }, [ setControls, - isWebsocketConnected, - isWebrtcConnected, isTraining, outputOptions, outputChoice, @@ -248,5 +272,28 @@ export function RenderControls() { field_of_view, ]); + useEffect(() => {}, [ + setState, + isWebsocketConnected, + isWebrtcConnected, + ]); + + useEffect(() => { + websocket.addEventListener('message', (originalCmd) => { + // set the remote description when the offer is received + const cmd = msgpack.decode(new Uint8Array(originalCmd.data)); + if (cmd.path === '/renderingState/eval_fps') { + eval_fps = cmd.data; + console.log("setting state", eval_fps); + setState({ eval_fps }); + } + if (cmd.path === '/renderingState/train_eta') { + train_eta = cmd.data; + console.log("setting state", train_eta) + setState({ train_eta }); + } + }); + }, []); + return null; } diff --git a/nerfactory/viewer/app/src/reducer.js b/nerfactory/viewer/app/src/reducer.js index d9d2745aa..ae3b0bab6 100644 --- a/nerfactory/viewer/app/src/reducer.js +++ b/nerfactory/viewer/app/src/reducer.js @@ -18,6 +18,8 @@ const initialState = { minResolution: 50, maxResolution: 512, field_of_view: 80, + eval_fps: 'Paused', + train_eta: 'Paused', }, // the scene state sceneState: { diff --git a/nerfactory/viewer/server/viewer_utils.py b/nerfactory/viewer/server/viewer_utils.py index fa3aa62de..235b97534 100644 --- a/nerfactory/viewer/server/viewer_utils.py +++ b/nerfactory/viewer/server/viewer_utils.py @@ -32,6 +32,7 @@ from nerfactory.graphs.base import Graph from nerfactory.utils import profiler from nerfactory.utils.config import ViewerConfig from nerfactory.utils.decorators import check_visualizer_enabled, decorate_all +from nerfactory.utils.writer import GLOBAL_BUFFER from nerfactory.viewer.server.utils import get_intrinsics_matrix_and_camera_to_world_h from nerfactory.viewer.server.visualizer import Viewer @@ -279,6 +280,16 @@ class VisualizerState: image = (image_output).astype("uint8") self.vis.set_image(image) + def _update_viewer_stats(self, render_time, image_height) -> None: + is_training = self.vis["renderingState/isTraining"].read() + if is_training is None or is_training: + eval_fps = f"{1 / render_time:.2f} fps at {image_height} res" + self.vis["renderingState/eval_fps"].write(eval_fps) + self.vis["renderingState/train_eta"].write(GLOBAL_BUFFER["viewer/train_eta"]) + else: + self.vis["renderingState/eval_fps"].write("Paused") + self.vis["renderingState/train_eta"].write("Paused") + @profiler.time_function def _render_image_in_viewer(self, graph: Graph) -> None: """ @@ -336,21 +347,25 @@ class VisualizerState: camera_ray_bundle.num_rays_per_chunk = self.config.num_rays_per_chunk graph.eval() + check_thread = CheckThread(state=self) render_thread = RenderThread(state=self, graph=graph, camera_ray_bundle=camera_ray_bundle) + start_time = time.time() check_thread.start() render_thread.start() - try: render_thread.join() check_thread.join() except Exception: # pylint: disable=broad-except pass + render_duration = time.time() - start_time graph.train() + outputs = render_thread.vis_outputs if outputs is not None: self._send_output_to_viewer(outputs) + self._update_viewer_stats(render_duration, image_height) def get_default_vis() -> Viewer: