diff --git a/nerfstudio/viewer/render_state_machine.py b/nerfstudio/viewer/render_state_machine.py index b9d43a782..05bd658bc 100644 --- a/nerfstudio/viewer/render_state_machine.py +++ b/nerfstudio/viewer/render_state_machine.py @@ -17,18 +17,21 @@ from __future__ import annotations import contextlib import threading +import time from dataclasses import dataclass from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple, get_args import numpy as np import torch +import torch.nn.functional as F from nerfstudio.cameras.cameras import Cameras from nerfstudio.model_components.renderers import background_color_override_context from nerfstudio.models.gaussian_splatting import GaussianSplattingModel from nerfstudio.utils import colormaps, writer from nerfstudio.utils.writer import GLOBAL_BUFFER, EventName, TimeWriter -from nerfstudio.viewer_legacy.server import viewer_utils from nerfstudio.viewer.utils import CameraState, get_camera +from nerfstudio.viewer_legacy.server import viewer_utils + from viser import ClientHandle if TYPE_CHECKING: @@ -166,11 +169,21 @@ class RenderStateMachine(threading.Thread): num_rays = (camera.height * camera.width).item() if self.viewer.control_panel.layer_depth: if isinstance(self.viewer.get_model(), GaussianSplattingModel): - # TODO: sending depth at high resolution lags the network a lot, figure out how to do this more efficiently - # outputs["gl_z_buf_depth"] = outputs["depth"] - pass + # Gaussians render much faster than we can send depth images, so we do some downsampling. + assert len(outputs["depth"].shape) == 3 + assert outputs["depth"].shape[-1] == 1 + + desired_depth_pixels = {"low_move": 128, "low_static": 128, "high": 512}[self.state] ** 2 + current_depth_pixels = outputs["depth"].shape[0] * outputs["depth"].shape[1] + scale = min(desired_depth_pixels / current_depth_pixels, 1.0) + + outputs["gl_z_buf_depth"] = F.interpolate( + outputs["depth"].squeeze(dim=-1)[None, None, ...], + size=(int(outputs["depth"].shape[0] * scale), int(outputs["depth"].shape[1] * scale)), + mode="bilinear", + )[0, 0, :, :, None] else: - # convert to z_depth if depth compositing is enabled + # Convert to z_depth if depth compositing is enabled. R = camera.camera_to_worlds[0, 0:3, 0:3].T camera_ray_bundle = camera.generate_rays(camera_indices=0, obb_box=obb) pts = camera_ray_bundle.directions * outputs["depth"] @@ -186,6 +199,9 @@ class RenderStateMachine(threading.Thread): def run(self): """Main loop for the render thread""" while self.running: + if not self.viewer.ready: + time.sleep(0.1) + continue if not self.render_trigger.wait(0.2): # if we haven't received a trigger in a while, send a static action self.action(RenderAction(action="static", camera_state=self.viewer.get_camera_state(self.client)))