diff --git a/src/viser/_scene_handles.py b/src/viser/_scene_handles.py index e017883a8..09dcdf5e8 100644 --- a/src/viser/_scene_handles.py +++ b/src/viser/_scene_handles.py @@ -26,7 +26,7 @@ from typing_extensions import Self, deprecated, override from . import _messages -from ._assignable_props_api import AssignablePropsBase +from ._assignable_props_api import AssignablePropsBase, colors_to_uint8 from .infra._infra import ( WebsockClientConnection, WebsockServer, @@ -1235,6 +1235,70 @@ def _ensure_buffer_size(self, num_gaussians: int) -> None: self.buffer = new_buffer + @staticmethod + def _pack_centers(buffer: np.ndarray, centers: np.ndarray) -> None: + buffer[:, 0:3] = np.ascontiguousarray(centers, dtype=np.float32).view(np.uint32) + + @staticmethod + def _pack_covariances(buffer: np.ndarray, covariances: np.ndarray) -> None: + # Extract upper-triangular terms: indices [0,1,2,4,5,8] from flattened 3x3. + cov_triu = covariances.reshape((-1, 9))[:, np.array([0, 1, 2, 4, 5, 8])] + cov_triu_f16 = np.ascontiguousarray(cov_triu, dtype=np.float16) + buffer[:, 4:7] = cov_triu_f16.view(np.uint32) + + @staticmethod + def _pack_rgba( + buffer: np.ndarray, + rgbs: np.ndarray | None = None, + opacities: np.ndarray | None = None, + ) -> None: + rgba = buffer[:, 7:8].view(np.uint8).reshape(-1, 4) + if rgbs is not None: + rgba[:, :3] = colors_to_uint8(rgbs) + if opacities is not None: + rgba[:, 3:4] = colors_to_uint8(opacities) + buffer[:, 7:8] = rgba.view(np.uint32) + + def set_gaussians( + self, + centers: np.ndarray, + covariances: np.ndarray, + rgbs: np.ndarray, + opacities: np.ndarray, + ) -> None: + """Atomically update all Gaussian attributes, including count changes. + + This is the preferred fast path when per-frame updates may change the + number of Gaussians: all attributes land in a single buffer update, + preventing transient mixed-state frames from sequential property + assignments (centers/covariances/rgbs/opacities one-by-one). A call + that leaves the buffer numerically unchanged sends no message. + """ + assert centers.ndim == 2 and centers.shape[1] == 3, ( + f"centers must have shape (N, 3), got {centers.shape}" + ) + num_gaussians = centers.shape[0] + assert covariances.ndim == 3 and covariances.shape == (num_gaussians, 3, 3), ( + f"covariances must have shape ({num_gaussians}, 3, 3), got {covariances.shape}" + ) + assert rgbs.ndim == 2 and rgbs.shape == (num_gaussians, 3), ( + f"rgbs must have shape ({num_gaussians}, 3), got {rgbs.shape}" + ) + assert opacities.ndim == 2 and opacities.shape == (num_gaussians, 1), ( + f"opacities must have shape ({num_gaussians}, 1), got {opacities.shape}" + ) + + # Assemble the full buffer locally, then store it with a single + # property assignment: props_setattr rejects writes to removed handles, + # resizes or copies in place as needed, and queues exactly one private + # snapshot for the wire. Routing a resize through _ensure_buffer_size + # here would queue an extra all-default buffer message first. + buffer = np.zeros((num_gaussians, 8), dtype=np.uint32) + self._pack_centers(buffer, centers) + self._pack_covariances(buffer, covariances) + self._pack_rgba(buffer, rgbs=rgbs, opacities=opacities) + self.buffer = buffer + @property def centers(self) -> npt.NDArray[np.float32]: """Centers of the Gaussians. Shape: (N, 3). Synchronized automatically when assigned.""" @@ -1246,7 +1310,7 @@ def centers(self, centers: np.ndarray) -> None: f"centers must have shape (N, 3), got {centers.shape}" ) self._ensure_buffer_size(centers.shape[0]) - self.buffer[:, 0:3] = centers.astype(np.float32).view(np.uint32) + self._pack_centers(self.buffer, centers) # Queue a private snapshot: the stored buffer is mutated in place by # later sub-property assignments, possibly while the event loop is still # serializing this message. Matches the guard in props_setattr. @@ -1260,15 +1324,11 @@ def rgbs(self) -> npt.NDArray[np.uint8]: @rgbs.setter def rgbs(self, rgbs: np.ndarray) -> None: - from ._assignable_props_api import colors_to_uint8 - assert rgbs.ndim == 2 and rgbs.shape[1] == 3, ( f"rgbs must have shape (N, 3), got {rgbs.shape}" ) self._ensure_buffer_size(rgbs.shape[0]) - rgba = self.buffer[:, 7:8].view(np.uint8).reshape(-1, 4) - rgba[:, :3] = colors_to_uint8(rgbs) - self.buffer[:, 7:8] = rgba.view(np.uint32) + self._pack_rgba(self.buffer, rgbs=rgbs) self._queue_update("buffer", self.buffer.copy()) @property @@ -1280,15 +1340,11 @@ def opacities(self) -> npt.NDArray[np.uint8]: @opacities.setter def opacities(self, opacities: np.ndarray) -> None: - from ._assignable_props_api import colors_to_uint8 - assert opacities.ndim == 2 and opacities.shape[1] == 1, ( f"opacities must have shape (N, 1), got {opacities.shape}" ) self._ensure_buffer_size(opacities.shape[0]) - rgba = self.buffer[:, 7:8].view(np.uint8).reshape(-1, 4) - rgba[:, 3:4] = colors_to_uint8(opacities) - self.buffer[:, 7:8] = rgba.view(np.uint32) + self._pack_rgba(self.buffer, opacities=opacities) self._queue_update("buffer", self.buffer.copy()) @property @@ -1317,10 +1373,7 @@ def covariances(self, covariances: np.ndarray) -> None: f"covariances must have shape (N, 3, 3), got {covariances.shape}" ) self._ensure_buffer_size(covariances.shape[0]) - # Extract upper-triangular terms: indices [0,1,2,4,5,8] from flattened 3x3. - cov_triu = covariances.reshape((-1, 9))[:, np.array([0, 1, 2, 4, 5, 8])] - cov_triu_f16 = cov_triu.astype(np.float16) - self.buffer[:, 4:7] = np.ascontiguousarray(cov_triu_f16).view(np.uint32) + self._pack_covariances(self.buffer, covariances) self._queue_update("buffer", self.buffer.copy()) diff --git a/tests/test_handle_lifecycle_bugs.py b/tests/test_handle_lifecycle_bugs.py index 378fc11b1..9ba606a80 100644 --- a/tests/test_handle_lifecycle_bugs.py +++ b/tests/test_handle_lifecycle_bugs.py @@ -448,6 +448,107 @@ def latest_buffer_update() -> np.ndarray: assert np.array_equal(queued, before) +def _add_test_splats(server: viser.ViserServer, n: int): + return server.scene.add_gaussian_splats( + "/splats", + centers=np.zeros((n, 3), np.float32), + covariances=np.tile(np.eye(3, dtype=np.float32) * 0.01, (n, 1, 1)), + rgbs=np.zeros((n, 3), np.uint8), + opacities=np.ones((n, 1), np.float32), + ) + + +def test_set_gaussians_round_trip() -> None: + """set_gaussians() must store values readable back through the per-property + getters, including when the Gaussian count changes.""" + with _server() as server: + s = _add_test_splats(server, n=2) + + m = 4 # Different count from the initial 2. + centers = np.arange(m * 3, dtype=np.float32).reshape(m, 3) + covariances = np.tile(np.eye(3, dtype=np.float32) * 0.25, (m, 1, 1)) + rgbs = np.arange(m * 3, dtype=np.uint8).reshape(m, 3) + opacities = np.linspace(0.0, 1.0, m, dtype=np.float32).reshape(m, 1) + s.set_gaussians(centers, covariances, rgbs, opacities) + + assert np.array_equal(s.centers, centers) + assert np.array_equal(s.rgbs, rgbs) + # Opacities are quantized to uint8 on the way in. + assert np.array_equal( + s.opacities, np.clip(opacities * 255.0, 0, 255).astype(np.uint8) + ) + # Covariances round-trip through float16. + np.testing.assert_allclose(s.covariances, covariances, rtol=1e-3) + + # Zero Gaussians is a valid state. + s.set_gaussians( + np.zeros((0, 3), np.float32), + np.zeros((0, 3, 3), np.float32), + np.zeros((0, 3), np.uint8), + np.zeros((0, 1), np.float32), + ) + assert s.centers.shape == (0, 3) + + +def test_set_gaussians_count_change_queues_single_update() -> None: + """A count-changing set_gaussians() must push exactly one buffer update: an + intermediate resize message would flash an all-default buffer on clients + whenever the ~60 Hz flush lands between the two pushes.""" + with _server() as server: + s = _add_test_splats(server, n=2) + + interface = s._impl.api._websock_interface + original_queue_message = interface.queue_message + buffer_updates: list[np.ndarray] = [] + + def spy(message) -> None: + if ( + type(message).__name__ == "SceneNodeUpdateMessage" + and "buffer" in message.updates + ): + buffer_updates.append(message.updates["buffer"]) + original_queue_message(message) + + interface.queue_message = spy # type: ignore[method-assign] + try: + m = 5 + s.set_gaussians( + np.ones((m, 3), np.float32), + np.tile(np.eye(3, dtype=np.float32) * 0.01, (m, 1, 1)), + np.zeros((m, 3), np.uint8), + np.ones((m, 1), np.float32), + ) + finally: + interface.queue_message = original_queue_message # type: ignore[method-assign] + + assert len(buffer_updates) == 1 + # The queued snapshot must alias neither the live server-owned buffer + # (mutated in place by later sub-property writes) nor anything the + # caller can touch. + assert buffer_updates[0] is not s._impl.props.buffer + queued_before = buffer_updates[0].copy() + s.centers = np.full((m, 3), 7.0, np.float32) + assert np.array_equal(buffer_updates[0], queued_before) + + +def test_set_gaussians_rejects_removed_handle() -> None: + """set_gaussians() on a removed handle must raise like property assignment + does, instead of queuing updates that resurrect the node on clients.""" + with _server() as server: + n = 3 + s = _add_test_splats(server, n=n) + s.remove() + + for count in (n, n + 1): # Same-count and resizing calls alike. + with pytest.raises(RuntimeError): + s.set_gaussians( + np.zeros((count, 3), np.float32), + np.tile(np.eye(3, dtype=np.float32) * 0.01, (count, 1, 1)), + np.zeros((count, 3), np.uint8), + np.ones((count, 1), np.float32), + ) + + def test_upload_part_after_button_removal_still_acks_and_completes() -> None: """A FileTransferPart arriving after button.remove() mid-upload must not assert (the handle is legitimately gone from the registry): parts keep