Sogs load tweaks (#7667)

Author: slimbuck

src/framework/parsers/sogs.js

@@ -8,14 +8,6 @@ import { GSplatSogsData } from '../../scene/gsplat/gsplat-sogs-data.js';
  * @import { ResourceHandlerCallback } from '../handlers/handler.js'
  */
 
-const readImageDataAsync = (texture) => {
-    return texture.read(0, 0, texture.width, texture.height, {
-        mipLevel: 0,
-        face: 0,
-        immediate: true
-    });
-};
-
 class SogsParser {
     /** @type {AppBase} */
     app;
@@ -66,18 +58,10 @@ class SogsParser {
         // wait for all textures to complete loading
         await Promise.allSettled(promises);
 
-        // sh palette has 64 sh entries per row. use width to calculate number of bands
-        const widths = {
-            192: 1,     // 64 * 3
-            512: 2,     // 64 * 8
-            960: 3      // 64 * 15
-        };
-
         // construct the gsplat resource
         const data = new GSplatSogsData();
         data.meta = meta;
         data.numSplats = meta.means.shape[0];
-        data.shBands = widths[textures.shN?.[0]?.resource?.width] ?? 0;
         data.means_l = textures.means[0].resource;
         data.means_u = textures.means[1].resource;
         data.quats = textures.quats[0].resource;
@@ -87,27 +71,10 @@ class SogsParser {
         data.sh_labels = textures.shN?.[1]?.resource;
 
         if (asset.data?.reorder ?? true) {
-            // copy back means_l and means_u data from gpu so cpu reorder has access to it
-            data.means_l._levels[0] = await readImageDataAsync(data.means_l);
-            data.means_u._levels[0] = await readImageDataAsync(data.means_u);
-            data.reorderData();
+            await data.reorderData();
         }
 
-        let resource;
-        if (asset.data?.decompress) {
-            // copy back gpu texture data so cpu iterator has access to it
-            const { means_l, means_u, quats, scales, sh0, sh_labels, sh_centroids } = data;
-            means_l._levels[0] = await readImageDataAsync(means_l);
-            means_u._levels[0] = await readImageDataAsync(means_u);
-            quats._levels[0] = await readImageDataAsync(quats);
-            scales._levels[0] = await readImageDataAsync(scales);
-            sh0._levels[0] = await readImageDataAsync(sh0);
-            sh_labels._levels[0] = await readImageDataAsync(sh_labels);
-            sh_centroids._levels[0] = await readImageDataAsync(sh_centroids);
-            resource = new GSplatResource(this.app, data.decompress(), []);
-        } else {
-            resource = new GSplatResource(this.app, data, []);
-        }
+        const resource = new GSplatResource(this.app, (asset.data?.decompress) ? (await data.decompress()) : data, []);
 
         callback(null, resource);
     }

src/index.js

@@ -236,6 +236,7 @@ export { Sky } from './scene/skybox/sky.js';
 export { GSplatData } from './scene/gsplat/gsplat-data.js';
 export { GSplat } from './scene/gsplat/gsplat.js';
 export { GSplatInstance } from './scene/gsplat/gsplat-instance.js';
+export { GSplatSogsData } from './scene/gsplat/gsplat-sogs-data.js';
 
 // FRAMEWORK
 export * from './framework/constants.js';

src/scene/gsplat/gsplat-sogs-data.js

@@ -38,6 +38,14 @@ const reorderFS = /* glsl */`
     }
 `;
 
+const readImageDataAsync = (texture) => {
+    return texture.read(0, 0, texture.width, texture.height, {
+        mipLevel: 0,
+        face: 0,
+        immediate: true
+    });
+};
+
 const resolve = (scope, values) => {
     for (const key in values) {
         scope.resolve(key).setValue(values[key]);
@@ -129,8 +137,6 @@ class GSplatSogsData {
 
     numSplats;
 
-    shBands;
-
     means_l;
 
     means_u;
@@ -192,7 +198,17 @@ class GSplatSogsData {
         return true;
     }
 
-    decompress() {
+    get shBands() {
+        // sh palette has 64 sh entries per row. use width to calculate number of bands
+        const widths = {
+            192: 1,     // 64 * 3
+            512: 2,     // 64 * 8
+            960: 3      // 64 * 15
+        };
+        return widths[this.sh_centroids?.resource?.width] ?? 0;
+    }
+
+    async decompress() {
         const members = [
             'x', 'y', 'z',
             'f_dc_0', 'f_dc_1', 'f_dc_2', 'opacity',
@@ -202,6 +218,16 @@ class GSplatSogsData {
 
         const { shBands } = this;
 
+        // copy back gpu texture data so cpu iterator has access to it
+        const { means_l, means_u, quats, scales, sh0, sh_labels, sh_centroids } = this;
+        means_l._levels[0] = await readImageDataAsync(means_l);
+        means_u._levels[0] = await readImageDataAsync(means_u);
+        quats._levels[0] = await readImageDataAsync(quats);
+        scales._levels[0] = await readImageDataAsync(scales);
+        sh0._levels[0] = await readImageDataAsync(sh0);
+        sh_labels._levels[0] = await readImageDataAsync(sh_labels);
+        sh_centroids._levels[0] = await readImageDataAsync(sh_centroids);
+
         // allocate spherical harmonics data
         if (shBands > 0) {
             const shMembers = [];
@@ -362,23 +388,25 @@ class GSplatSogsData {
         return order;
     }
 
-    reorderData() {
-        if (!this.orderTexture) {
-            const { device, height, width } = this.means_l;
-
-            this.orderTexture = new Texture(device, {
-                name: 'orderTexture',
-                width,
-                height,
-                format: PIXELFORMAT_R32U,
-                mipmaps: false,
-                levels: [this.calcMortonOrder()]
-            });
+    async reorderData() {
+        const { device, height, width } = this.means_l;
 
-            device.on('devicerestored', () => {
-                this.reorderGpuMemory();
-            });
-        }
+        // copy back means_l and means_u data from gpu so cpu reorder has access to it
+        this.means_l._levels[0] = await readImageDataAsync(this.means_l);
+        this.means_u._levels[0] = await readImageDataAsync(this.means_u);
+
+        this.orderTexture = new Texture(device, {
+            name: 'orderTexture',
+            width,
+            height,
+            format: PIXELFORMAT_R32U,
+            mipmaps: false,
+            levels: [this.calcMortonOrder()]
+        });
+
+        device.on('devicerestored', () => {
+            this.reorderGpuMemory();
+        });
 
         this.reorderGpuMemory();
     }