Working global sorting + SIMD

nerfstudio-project · Jul 23, 2024 · 620bec4 · 620bec4
1 parent a0d26e0
commit 620bec4
Show file tree

Hide file tree

Showing 9 changed files with 279 additions and 99 deletions.
diff --git a/src/viser/client/src/App.tsx b/src/viser/client/src/App.tsx
@@ -45,9 +45,13 @@ import { useDisclosure } from "@mantine/hooks";
 import { rayToViserCoords } from "./WorldTransformUtils";
 import { ndcFromPointerXy, opencvXyFromPointerXy } from "./ClickUtils";
 import { theme } from "./AppTheme";
-import { GaussianSplatsContext } from "./Splatting/GaussianSplats";
+import {
+  GaussianSplatsContext,
+  useGaussianSplatStore,
+} from "./Splatting/SplatContext";
 import { FrameSynchronizedMessageHandler } from "./MessageHandler";
 import { PlaybackFromFile } from "./FilePlayback";
+import GlobalGaussianSplats from "./Splatting/GaussianSplats";
 
 export type ViewerContextContents = {
   messageSource: "websocket" | "file_playback";
@@ -259,10 +263,9 @@ function ViewerContents({ children }: { children: React.ReactNode }) {
               })}
             >
               <Viewer2DCanvas />
-              <GaussianSplatsContext.Provider
-                value={React.useRef({ numSorting: 0, sortUpdateCallbacks: [] })}
-              >
+              <GaussianSplatsContext.Provider value={useGaussianSplatStore()}>
                 <ViewerCanvas>
+                  <GlobalGaussianSplats />
                   <FrameSynchronizedMessageHandler />
                 </ViewerCanvas>
               </GaussianSplatsContext.Provider>

diff --git a/src/viser/client/src/MessageHandler.tsx b/src/viser/client/src/MessageHandler.tsx
@@ -29,7 +29,7 @@ import GeneratedGuiContainer from "./ControlPanel/Generated";
 import { Paper, Progress } from "@mantine/core";
 import { IconCheck } from "@tabler/icons-react";
 import { computeT_threeworld_world } from "./WorldTransformUtils";
-import GaussianSplats from "./Splatting/GaussianSplats";
+import { GaussianSplatsContext } from "./Splatting/SplatContext";
 
 /** Convert raw RGB color buffers to linear color buffers. **/
 function threeColorBufferFromUint8Buffer(colors: ArrayBuffer) {
@@ -51,6 +51,10 @@ function useMessageHandler() {
   const viewer = useContext(ViewerContext)!;
   const ContextBridge = useContextBridge();
 
+  const splatContext = useContext(GaussianSplatsContext)!;
+  const setGaussianBuffer = splatContext((state) => state.setBuffer);
+  const removeGaussianBuffer = splatContext((state) => state.removeBuffer);
+
   // We could reduce the redundancy here if we wanted to.
   // https://github.com/nerfstudio-project/viser/issues/39
   const removeSceneNode = viewer.useSceneTree((state) => state.removeSceneNode);
@@ -966,23 +970,35 @@ function useMessageHandler() {
         return;
       }
       case "GaussianSplatsMessage": {
+        // <GaussianSplats
+        //   buffers={{
+        //     buffer: new Uint32Array(
+        //       message.buffer.buffer.slice(
+        //         message.buffer.byteOffset,
+        //         message.buffer.byteOffset + message.buffer.byteLength,
+        //       ),
+        //     ),
+        //   }}
+        // />
+        setGaussianBuffer(
+          message.name,
+          new Uint32Array(
+            message.buffer.buffer.slice(
+              message.buffer.byteOffset,
+              message.buffer.byteOffset + message.buffer.byteLength,
+            ),
+          ),
+        );
         addSceneNodeMakeParents(
-          new SceneNode<THREE.Group>(message.name, (ref) => {
-            return (
-              <group ref={ref}>
-                <GaussianSplats
-                  buffers={{
-                    buffer: new Uint32Array(
-                      message.buffer.buffer.slice(
-                        message.buffer.byteOffset,
-                        message.buffer.byteOffset + message.buffer.byteLength,
-                      ),
-                    ),
-                  }}
-                />
-              </group>
-            );
-          }),
+          new SceneNode<THREE.Group>(
+            message.name,
+            (ref) => {
+              return <group ref={ref}></group>;
+            },
+            () => {
+              removeGaussianBuffer(message.name);
+            },
+          ),
         );
         return;
       }