(gaussians) Fix flicker when parents have unmountWhenInvisible=true

nerfstudio-project · Jul 26, 2024 · edc7b82 · edc7b82
1 parent 203d92a
commit edc7b82
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 43 deletions.
diff --git a/src/viser/client/src/ControlPanel/SceneTreeTable.tsx b/src/viser/client/src/ControlPanel/SceneTreeTable.tsx
@@ -50,7 +50,6 @@ const SceneTreeTableRow = React.memo(function SceneTreeTableRow(props: {
   function setOverrideVisibility(name: string, visible: boolean | undefined) {
     const attr = viewer.nodeAttributesFromName.current;
     attr[name]!.overrideVisibility = visible;
-    console.log(name, visible);
     rerenderTable();
   }
   const setLabelVisibility = viewer.useSceneTree(

diff --git a/src/viser/client/src/SceneTree.tsx b/src/viser/client/src/SceneTree.tsx
@@ -205,53 +205,59 @@ export function SceneNodeThreeObject(props: {
 
   // Update attributes on a per-frame basis. Currently does redundant work,
   // although this shouldn't be a bottleneck.
-  useFrame(() => {
-    const attrs = viewer.nodeAttributesFromName.current[props.name];
-    everyFrameCallback && everyFrameCallback();
+  useFrame(
+    () => {
+      const attrs = viewer.nodeAttributesFromName.current[props.name];
+      everyFrameCallback && everyFrameCallback();
 
-    // Unmount when invisible.
-    // Examples: <Html /> components, PivotControls.
-    //
-    // This is a workaround for situations where just setting `visible` doesn't
-    // work (like <Html />), or to prevent invisible elements from being
-    // interacted with (<PivotControls />).
-    //
-    // https://github.com/pmndrs/drei/issues/1323
-    if (unmountWhenInvisible) {
-      const displayed = isDisplayed();
-      if (displayed && unmount) {
-        setUnmount(false);
-      }
-      if (!displayed && !unmount) {
-        setUnmount(true);
+      // Unmount when invisible.
+      // Examples: <Html /> components, PivotControls.
+      //
+      // This is a workaround for situations where just setting `visible` doesn't
+      // work (like <Html />), or to prevent invisible elements from being
+      // interacted with (<PivotControls />).
+      //
+      // https://github.com/pmndrs/drei/issues/1323
+      if (unmountWhenInvisible) {
+        const displayed = isDisplayed();
+        if (displayed && unmount) {
+          if (obj !== null) obj.visible = false;
+          setUnmount(false);
+        }
+        if (!displayed && !unmount) {
+          setUnmount(true);
+        }
       }
-    }
 
-    if (obj === null) return;
-    if (attrs === undefined) return;
+      if (obj === null) return;
+      if (attrs === undefined) return;
 
-    const visibility =
-      (attrs?.overrideVisibility === undefined
-        ? attrs?.visibility
-        : attrs.overrideVisibility) ?? true;
-    obj.visible = visibility;
+      const visibility =
+        (attrs?.overrideVisibility === undefined
+          ? attrs?.visibility
+          : attrs.overrideVisibility) ?? true;
+      obj.visible = visibility;
 
-    if (attrs.poseUpdateState == "needsUpdate") {
-      attrs.poseUpdateState = "updated";
-      const wxyz = attrs.wxyz;
-      if (wxyz !== undefined) {
-        obj.quaternion.set(wxyz[1], wxyz[2], wxyz[3], wxyz[0]);
-      }
-      const position = attrs.position;
-      if (position !== undefined) {
-        obj.position.set(position[0], position[1], position[2]);
-      }
+      if (attrs.poseUpdateState == "needsUpdate") {
+        attrs.poseUpdateState = "updated";
+        const wxyz = attrs.wxyz;
+        if (wxyz !== undefined) {
+          obj.quaternion.set(wxyz[1], wxyz[2], wxyz[3], wxyz[0]);
+        }
+        const position = attrs.position;
+        if (position !== undefined) {
+          obj.position.set(position[0], position[1], position[2]);
+        }
 
-      // Update matrices if necessary. This is necessary for PivotControls.
-      if (!obj.matrixAutoUpdate) obj.updateMatrix();
-      if (!obj.matrixWorldAutoUpdate) obj.updateMatrixWorld();
-    }
-  });
+        // Update matrices if necessary. This is necessary for PivotControls.
+        if (!obj.matrixAutoUpdate) obj.updateMatrix();
+        if (!obj.matrixWorldAutoUpdate) obj.updateMatrixWorld();
+      }
+    },
+    // Other useFrame hooks may depend on transforms + visibility. So it's best
+    // to call this hook early.
+    -10000,
+  );
 
   // Clicking logic.
   const sendClicksThrottled = useThrottledMessageSender(50);

diff --git a/src/viser/client/src/Splatting/GaussianSplats.tsx b/src/viser/client/src/Splatting/GaussianSplats.tsx
@@ -249,6 +249,7 @@ export default function GlobalGaussianSplats() {
   const prevRowMajorT_camera_groups = meshProps.rowMajorT_camera_groups
     .slice()
     .fill(0);
+  let prevVisible = false;
   useFrame((state, delta) => {
     const mesh = meshRef.current;
     if (mesh === null || sortWorker === null) return;
@@ -305,11 +306,17 @@ export default function GlobalGaussianSplats() {
           visible = visible && ancestor.visible;
         });
       }
-      if (!visible) {
+      if (!(visible && prevVisible)) {
         meshProps.rowMajorT_camera_groups[sortedGroupIndex * 12 + 3] = 1e10;
         meshProps.rowMajorT_camera_groups[sortedGroupIndex * 12 + 7] = 1e10;
         meshProps.rowMajorT_camera_groups[sortedGroupIndex * 12 + 11] = 1e10;
       }
+
+      // If the parent has unmountWhenInvisible=true, the first frame after
+      // showing a hidden parent can have visible=true with an incorrect
+      // matrixWorld transform. There might be a better fix, but `prevVisible`
+      // is an easy workaround for this.
+      prevVisible = visible;
     }
 
     if (