Introduce SplatUtils.mergeSplats and solve envMap global sorting

Author: mrxz

examples/envmap/index.html

@@ -40,7 +40,7 @@
   </script>
   <script type="module">
     import * as THREE from "three";
-    import { Splat, SplatLoader } from "@sparkjsdev/spark";
+    import { Splat, SplatLoader, SplatUtils } from "@sparkjsdev/spark";
     import { GLTFLoader } from "three/addons/loaders/GLTFLoader.js";
     import { getAssetFileURL } from "/examples/js/get-asset-url.js";
 
@@ -61,18 +61,22 @@
 
     const splatLoader = new SplatLoader();
     const splatURL = await getAssetFileURL("fireplace.spz");
+
     const background = await splatLoader.loadAsync(splatURL);
     background.quaternion.set(1, 0, 0, 0);
     background.position.set(0.5, 0, -1);
     background.scale.setScalar(0.5);
-    scene.add(background);
 
     const background2 = background.clone();
     background2.quaternion.set(1, 0, 0, 0);
     background2.rotation.y = Math.PI;
     background2.position.set(-0.5, 0, 0.0);
     background2.scale.setScalar(0.5);
-    scene.add(background2); // FIXME: Combine/merge splats before rendering
+
+    const combinedBackground = SplatUtils.mergeSplats([ background, background2 ]);
+    scene.add(combinedBackground);
+    combinedBackground.sorter.sortRadial = true;
+    combinedBackground.sorter.sort360 = true;
 
     // Add rubber duck
     const gltfLoader = new GLTFLoader();
@@ -82,15 +86,9 @@
     duck.position.set(0, 0.45, -0.4);
     scene.add(duck);
 
-    background2.sorter.sortRadial = true;
-    background2.sorter.sort360 = true;
-    Promise.all([
-      background.sortFor(renderer, camera),
-      background2.sortFor(renderer, camera),
-    ]).then(() => {
+    combinedBackground.sortFor(renderer, camera).then(() => {
       // Disable sorting while rendering the env map
-      background.autoSort = false;
-      background2.autoSort = false;
+      combinedBackground.autoSort = false;
 
       const sceneRT = new THREE.PMREMGenerator(renderer).fromScene(scene);
       for (let obj of duck.children) {
@@ -100,8 +98,7 @@
         obj.material.roughness = 0.02;
       }
 
-      background.autoSort = true;
-      background2.autoSort = true;
+      combinedBackground.autoSort = true;
     });
 
     let lastTime;

src/Splat.ts

@@ -48,6 +48,42 @@ export interface SplatData {
   ) => void;
 }
 
+export type SplatCallback = (
+  i: number,
+  x: number,
+  y: number,
+  z: number,
+  scaleX: number,
+  scaleY: number,
+  scaleZ: number,
+  quatX: number,
+  quatY: number,
+  quatZ: number,
+  quatW: number,
+  opacity: number,
+  r: number,
+  g: number,
+  b: number,
+  sh?: ArrayLike<number>,
+) => void;
+
+/**
+ * Extended SplatData interface that allows the splat properties
+ * to be decoded and read back.
+ */
+export interface IterableSplatData extends SplatData {
+  iterateSplats: (callback: SplatCallback) => void;
+}
+
+export type SortContext = {
+  lastOriginToCamera: THREE.Matrix4;
+  sortJob: Promise<SplatOrdering> | null;
+  ordering: Uint32Array;
+  pendingOrdering: Uint32Array;
+  activeSplats: number;
+  orderingId: number;
+};
+
 /**
  * Object representing a collection of Gaussian Splats in a scene.
  */
@@ -76,19 +112,9 @@ export class Splat extends THREE.Mesh<SplatGeometry, THREE.ShaderMaterial> {
   sorter: SplatSorter = new ReadbackSplatSorter();
   /**
    * Mapping from camera to sort context.
-   * This allows multiple ca
+   * This allows multiple viewpoints from different cameras.
    */
-  private sortContext: WeakMap<
-    THREE.Camera,
-    {
-      lastOriginToCamera: THREE.Matrix4;
-      sortJob: Promise<SplatOrdering> | null;
-      ordering: Uint32Array;
-      pendingOrdering: Uint32Array;
-      activeSplats: number;
-      orderingId: number;
-    }
-  > = new WeakMap();
+  private sortContext: WeakMap<THREE.Camera, SortContext> = new WeakMap();
   /**
    * Id of the current ordering used by the SplatGeometry.
    */
@@ -360,15 +386,7 @@ export class Splat extends THREE.Mesh<SplatGeometry, THREE.ShaderMaterial> {
         renderer,
         context.pendingOrdering,
       );
-      context.sortJob.then((result) => {
-        context.sortJob = null;
-        // Swap ordering arrays
-        context.pendingOrdering = context.ordering;
-        context.ordering = result.ordering;
-
-        context.activeSplats = result.activeSplats;
-        context.orderingId = globalOrderingId++;
-      });
+      context.sortJob.then((result) => this.onSortComplete(context, result));
     } else if (needsSort) {
       console.log("Suppressing sort as one is in progress...");
     }
@@ -383,6 +401,16 @@ export class Splat extends THREE.Mesh<SplatGeometry, THREE.ShaderMaterial> {
     }
   }
 
+  protected onSortComplete(context: SortContext, result: SplatOrdering) {
+    context.sortJob = null;
+    // Swap ordering arrays
+    context.pendingOrdering = context.ordering;
+    context.ordering = result.ordering;
+
+    context.activeSplats = result.activeSplats;
+    context.orderingId = globalOrderingId++;
+  }
+
   setShaderHooks(hooks: ShaderHooks | null): ShaderHooks | null {
     const previousShaderHooks = this.shaderHooks;
     this.shaderHooks = hooks;

src/SplatUtils.ts

@@ -0,0 +1,111 @@
+import * as THREE from "three";
+import { type IterableSplatData, Splat } from "./Splat";
+import { DefaultSplatEncoding, type SplatEncoder } from "./encoding/encoder";
+
+const tempCenter = new THREE.Vector3();
+const tempScales = new THREE.Vector3();
+const tempQuat = new THREE.Quaternion();
+
+/**
+ * Combines multiple Splat objects into a single Splat object. The individual
+ * world transforms are applied to the individual splats. Each splat must
+ * have the same number of spherical harmonics and the underlying SplatData
+ * must be iterable.
+ * @param splats The splats to combine
+ * @param options Additional options
+ * @returns The combined splat
+ */
+export function mergeSplats<T>(
+  splats: Array<Splat>,
+  options?: {
+    splatEncoder?: SplatEncoder<T> | (() => SplatEncoder<T>);
+  },
+): Splat | null {
+  const numSh = splats[0].splatData.numSh;
+  const splatEncoderFactory =
+    options?.splatEncoder ?? DefaultSplatEncoding.createSplatEncoder;
+  const splatEncoder =
+    typeof splatEncoderFactory === "function"
+      ? splatEncoderFactory()
+      : splatEncoderFactory;
+
+  // Sum the total amount of combined splats.
+  const numSplats = splats.reduce(
+    (acc, splat) => acc + splat.splatData.numSplats,
+    0,
+  );
+  splatEncoder.allocate(numSplats, numSh);
+
+  let newSplatIndex = 0;
+  for (let i = 0; i < splats.length; ++i) {
+    const splatData = splats[i].splatData;
+    if (splatData.numSh !== numSh) {
+      console.error(
+        `SplatUtils: .mergeSplats() failed with splat at index ${i}. All splats must have the same amount of spherical harmonics.`,
+      );
+      return null;
+    }
+
+    if (!("iterateSplats" in splatData)) {
+      console.error(
+        `SplatUtils: .mergeSplats() failed with splat at index ${i}. All splats must have iterable splat data.`,
+      );
+      return null;
+    }
+
+    // Ensure matrix world is up to date
+    splats[i].updateMatrixWorld();
+    const splatScale = splats[i].getWorldScale(new THREE.Vector3());
+    const splatRotation = splats[i].getWorldQuaternion(new THREE.Quaternion());
+
+    const iterableSplatData = splatData as IterableSplatData;
+    iterableSplatData.iterateSplats(
+      (
+        _,
+        x,
+        y,
+        z,
+        scaleX,
+        scaleY,
+        scaleZ,
+        quatX,
+        quatY,
+        quatZ,
+        quatW,
+        opacity,
+        r,
+        g,
+        b,
+        sh,
+      ) => {
+        // Apply splat transform
+        tempCenter.set(x, y, z).applyMatrix4(splats[i].matrixWorld);
+        tempScales.set(scaleX, scaleY, scaleZ).multiplyScalar(splatScale.x); // Assume uniform scaling
+        tempQuat.set(quatX, quatY, quatZ, quatW).premultiply(splatRotation);
+
+        splatEncoder.setSplat(
+          newSplatIndex++,
+          tempCenter.x,
+          tempCenter.y,
+          tempCenter.z,
+          tempScales.x,
+          tempScales.y,
+          tempScales.z,
+          tempQuat.x,
+          tempQuat.y,
+          tempQuat.z,
+          tempQuat.w,
+          opacity,
+          r,
+          g,
+          b,
+        );
+        if (sh) {
+          splatEncoder.setSplatSh(newSplatIndex, sh);
+        }
+      },
+    );
+  }
+
+  return new Splat(splatEncoder.close());
+}

src/encoding/ExtendedSplats.ts

@@ -1,5 +1,5 @@
 import * as THREE from "three";
-import type { SplatData } from "../Splat";
+import type { IterableSplatData, SplatCallback, SplatData } from "../Splat";
 import {
   LN_SCALE_MAX,
   LN_SCALE_MIN,
@@ -8,13 +8,14 @@ import {
 } from "../defines";
 import {
   computeMaxSplats,
+  decodeQuatOctXy88R8,
   encodeQuatOctXy88R8,
   floatBitsToUint,
   floatToUint8,
   getTextureSize,
   uintBitsToFloat,
 } from "../utils";
-import type { ResizableSplatEncoder, SplatEncoder } from "./encoder";
+import type { ResizableSplatEncoder } from "./encoder";
 
 export type ExtendedSplatsOptions = {
   // Reserve space for at least this many splats when constructing the collection
@@ -27,7 +28,7 @@ export type ExtendedSplatsOptions = {
   numSh?: number;
 };
 
-export class ExtendedSplats implements SplatData {
+export class ExtendedSplats implements IterableSplatData {
   maxSplats = 0;
   numSplats = 0;
   numSh = 0;
@@ -150,6 +151,63 @@ export class ExtendedSplats implements SplatData {
     }
   }
 
+  iterateSplats(callback: SplatCallback) {
+    const shCoeffients = SH_DEGREE_TO_NUM_COEFF[this.numSh];
+    const sh = this.numSh > 0 ? new Float32Array(shCoeffients) : undefined;
+
+    for (let i = 0; i < this.numSplats; i++) {
+      const i4 = i * 4;
+      const word0 = this.packedArray1[i4 + 0];
+      const word1 = this.packedArray1[i4 + 1];
+      const word2 = this.packedArray1[i4 + 2];
+      const word3 = this.packedArray1[i4 + 3];
+      const word4 = this.packedArray2[i4 + 0];
+      const word5 = this.packedArray2[i4 + 1];
+      const word6 = this.packedArray2[i4 + 2];
+      const word7 = this.packedArray2[i4 + 3];
+
+      const r = (word5 & 0xff) / 255;
+      const g = ((word5 >>> 8) & 0xff) / 255;
+      const b = ((word5 >>> 16) & 0xff) / 255;
+      const a = ((word5 >>> 24) & 0xff) / 255;
+
+      const lnScaleScale = (LN_SCALE_MAX - LN_SCALE_MIN) / 1023.0;
+      const uScalesX = word3 & 0x3ff;
+      const scaleX = Math.exp(LN_SCALE_MIN + uScalesX * lnScaleScale);
+      const uScalesY = (word3 >>> 10) & 0x3ff;
+      const scaleY = Math.exp(LN_SCALE_MIN + uScalesY * lnScaleScale);
+      const uScalesZ = (word3 >>> 20) & 0x3ff;
+      const scaleZ = Math.exp(LN_SCALE_MIN + uScalesZ * lnScaleScale);
+
+      decodeQuatOctXy88R8(word4, tempQuaternion);
+
+      if (sh && this.packedShArray) {
+        for (let j = 0; j < shCoeffients; j++) {
+          sh[j] = (this.packedShArray[i * shCoeffients + j] - 127) / 127;
+        }
+      }
+
+      callback(
+        i,
+        uintBitsToFloat(word0),
+        uintBitsToFloat(word1),
+        uintBitsToFloat(word2),
+        scaleX,
+        scaleY,
+        scaleZ,
+        tempQuaternion.x,
+        tempQuaternion.y,
+        tempQuaternion.z,
+        tempQuaternion.w,
+        a,
+        r,
+        g,
+        b,
+        sh,
+      );
+    }
+  }
+
   static encodingName = "extended";
 
   static createSplatEncoder(): ResizableSplatEncoder<EncodedExtendedSplats> {