Cholesky, optimizations, formatting

nerfstudio-project · Jun 13, 2024 · 99c3680 · 99c3680
1 parent 80d4c99
commit 99c3680
Show file tree

Hide file tree

Showing 6 changed files with 186 additions and 133 deletions.
diff --git a/.prettierignore b/.prettierignore
@@ -0,0 +1,2 @@
+*.mjs
+build/
diff --git a/docs/source/_templates/sidebar/brand.html b/docs/source/_templates/sidebar/brand.html
@@ -1,68 +1,95 @@
-<a class="sidebar-brand{% if logo %} centered{% endif %}" href="{{ pathto(master_doc) }}">
-    {% block brand_content %} {%- if logo_url %}
-    <div class="sidebar-logo-container">
-        <a href="{{ pathto(master_doc) }}"><img class="sidebar-logo" src="{{ logo_url }}" alt="Logo" /></a>
-    </div>
-    {%- endif %} {%- if theme_light_logo and theme_dark_logo %}
-    <div class="sidebar-logo-container" style="margin: .5rem 1em .5rem 0">
-        <img class="sidebar-logo only-light" src="{{ pathto('_static/' + theme_light_logo, 1) }}" alt="logo" />
-        <img class="sidebar-logo only-dark" src="{{ pathto('_static/' + theme_dark_logo, 1) }}" alt="logo" />
-    </div>
-    {%- endif %}
-    <!-- <span class="sidebar-brand-text">{{ project }}</span> -->
+<a
+  class="sidebar-brand{% if logo %} centered{% endif %}"
+  href="{{ pathto(master_doc) }}"
+>
+  {% block brand_content %} {%- if logo_url %}
+  <div class="sidebar-logo-container">
+    <a href="{{ pathto(master_doc) }}"
+      ><img class="sidebar-logo" src="{{ logo_url }}" alt="Logo"
+    /></a>
+  </div>
+  {%- endif %} {%- if theme_light_logo and theme_dark_logo %}
+  <div class="sidebar-logo-container" style="margin: 0.5rem 1em 0.5rem 0">
+    <img
+      class="sidebar-logo only-light"
+      src="{{ pathto('_static/' + theme_light_logo, 1) }}"
+      alt="logo"
+    />
+    <img
+      class="sidebar-logo only-dark"
+      src="{{ pathto('_static/' + theme_dark_logo, 1) }}"
+      alt="logo"
+    />
+  </div>
+  {%- endif %}
+  <!-- <span class="sidebar-brand-text">{{ project }}</span> -->
 
-    {% endblock brand_content %}
+  {% endblock brand_content %}
 </a>
 
 <!-- Dropdown for different versions of the viser docs. Slightly hacky. -->
-<div style="padding: 0 1em;">
-    <script>
-      var viserDocsVersionsPopulated = false;
+<div style="padding: 0 1em">
+  <script>
+    var viserDocsVersionsPopulated = false;
 
-      async function getViserVersionList() {
-        // This index.txt file is written by the docs.yml GitHub action.
-        // https://github.com/nerfstudio-project/viser/blob/main/.github/workflows/docs.yml
-        const response = await fetch(
-          "https://viser.studio/versions/index.txt",
-          { cache: "no-cache" }
-        );
-        return await response.text();
+    async function getViserVersionList() {
+      // This index.txt file is written by the docs.yml GitHub action.
+      // https://github.com/nerfstudio-project/viser/blob/main/.github/workflows/docs.yml
+      const response = await fetch("https://viser.studio/versions/index.txt", {
+        cache: "no-cache",
+      });
+      return await response.text();
+    }
+    async function viserDocsPopulateVersionDropDown() {
+      // Load the version list lazily...
+      if (viserDocsVersionsPopulated) {
+        return;
       }
-      async function viserDocsPopulateVersionDropDown () {
-        // Load the version list lazily...
-        if (viserDocsVersionsPopulated) {
-          return;
-        }
-        viserDocsVersionsPopulated = true;
+      viserDocsVersionsPopulated = true;
 
-        console.log("Populating docs version list!")
-        const versions = (await getViserVersionList()).trim().split("\n").reverse();
-        console.log(versions);
-        let htmlString = "<ul style='margin: 0.5rem 0 0 0'>";
-        htmlString += `<li><a href="https://viser.studio/latest">latest</a></li>`;
-        for (let version of versions) {
-          htmlString += `<li><a href="https://viser.studio/versions/${version}">${version}</a></li>`;
-        }
-
-        htmlString += "</ul>";
-        document.getElementById("viser-version-dropdown").innerHTML = htmlString;
+      console.log("Populating docs version list!");
+      const versions = (await getViserVersionList())
+        .trim()
+        .split("\n")
+        .reverse();
+      console.log(versions);
+      let htmlString = "<ul style='margin: 0.5rem 0 0 0'>";
+      htmlString += `<li><a href="https://viser.studio/latest">latest</a></li>`;
+      for (let version of versions) {
+        htmlString += `<li><a href="https://viser.studio/versions/${version}">${version}</a></li>`;
       }
-    </script>
-    <details
-        style="padding: 0.5rem; background: var(--color-background-primary); border-radius: 0.5rem; border: 1px solid var(--color-sidebar-background-border);"
-        ontoggle="viserDocsPopulateVersionDropDown()"
-    >
-    <summary style="cursor: pointer;"><strong>Version:</strong> <em>{{ version }}</em></summary>
-      <div id="viser-version-dropdown"></div>
-    </details>
-    <!-- End dropdown -->
+
+      htmlString += "</ul>";
+      document.getElementById("viser-version-dropdown").innerHTML = htmlString;
+    }
+  </script>
+  <details
+    style="
+      padding: 0.5rem;
+      background: var(--color-background-primary);
+      border-radius: 0.5rem;
+      border: 1px solid var(--color-sidebar-background-border);
+    "
+    ontoggle="viserDocsPopulateVersionDropDown()"
+  >
+    <summary style="cursor: pointer">
+      <strong>Version:</strong> <em>{{ version }}</em>
+    </summary>
+    <div id="viser-version-dropdown"></div>
+  </details>
+  <!-- End dropdown -->
 </div>
 
 <div style="text-align: left; padding: 1em">
-    <script async defer src="https://buttons.github.io/buttons.js"></script>
-    <a class="github-button" href="https://github.com/nerfstudio-project/viser"
-        data-color-scheme="no-preference: light; light: light; dark: light;" data-size="large" data-show-count="true"
-        aria-label="Download buttons/github-buttons on GitHub">
-        Github
-    </a>
+  <script async defer src="https://buttons.github.io/buttons.js"></script>
+  <a
+    class="github-button"
+    href="https://github.com/nerfstudio-project/viser"
+    data-color-scheme="no-preference: light; light: light; dark: light;"
+    data-size="large"
+    data-show-count="true"
+    aria-label="Download buttons/github-buttons on GitHub"
+  >
+    Github
+  </a>
 </div>
diff --git a/src/viser/_scene_api.py b/src/viser/_scene_api.py
@@ -843,9 +843,11 @@ def add_gaussian_splats(
         assert opacities.shape == (num_gaussians, 1)
         assert covariances.shape == (num_gaussians, 3, 3)
 
-        # Make our covariances more compact!
-        covariances_triu = (
-            covariances.reshape((-1, 9))[:, onp.array([0, 1, 2, 4, 5, 8])]
+        # Get cholesky factor of covariance.
+        cov_cholesky_triu = (
+            onp.linalg.cholesky(covariances)
+            .swapaxes(-1, -2)  # tril => triu
+            .reshape((-1, 9))[:, onp.array([0, 1, 2, 4, 5, 8])]
             .astype(onp.float32)
             .copy()
         )
@@ -855,7 +857,7 @@ def add_gaussian_splats(
                 centers.astype(onp.float32).view(onp.uint8),
                 onp.zeros((num_gaussians, 4), dtype=onp.uint8),
                 # Second texelFetch.
-                covariances_triu.astype(onp.float16).view(onp.uint8),
+                cov_cholesky_triu.astype(onp.float16).view(onp.uint8),
                 _colors_to_uint8(rgbs),
                 _colors_to_uint8(opacities),
             ],

diff --git a/src/viser/client/src/App.tsx b/src/viser/client/src/App.tsx
@@ -48,6 +48,7 @@ import { useDisclosure } from "@mantine/hooks";
 import { rayToViserCoords } from "./WorldTransformUtils";
 import { clickToNDC, clickToOpenCV, isClickValid } from "./ClickUtils";
 import { theme } from "./AppTheme";
+import { GaussianSplatsContext } from "./Splatting/GaussianSplats";
 
 export type ViewerContextContents = {
   // Zustand hooks.
@@ -213,9 +214,13 @@ function ViewerContents() {
               })}
             >
               <Viewer2DCanvas />
-              <ViewerCanvas>
-                <FrameSynchronizedMessageHandler />
-              </ViewerCanvas>
+              <GaussianSplatsContext.Provider
+                value={React.useRef({ numSorting: 0, sortUpdateCallbacks: [] })}
+              >
+                <ViewerCanvas>
+                  <FrameSynchronizedMessageHandler />
+                </ViewerCanvas>
+              </GaussianSplatsContext.Provider>
               {viewer.useGui((state) => state.theme.show_logo) ? (
                 <ViserLogo />
               ) : null}