tink/nerfstudio
4
0
Fork 0
mirror of https://github.com/nerfstudio-project/nerfstudio.git synced 2026-11-11 02:24:38 +08:00
Code

Viewer depth compositing for Gaussian Splatting models (#2744)

Browse Source 
* Viewer depth compositing for Gaussian Splatting models

* Tune
This commit is contained in:
Brent Yi 2024-01-10 20:04:00 -08:00 committed by GitHub
parent e7e713e900
commit 3b6ca91dbc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 21 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
nerfstudio/viewer/render_state_machine.py
View File

@ -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)))