Introduce SPZ transcoding for viewer example

Author: mrxz

examples/viewer/index.html

@@ -185,7 +185,7 @@
   </script>
   <script type="module">
     import * as THREE from "three";
-    import { SplatLoader } from "@sparkjsdev/spark";
+    import { SplatLoader, transcodeSpz } from "@sparkjsdev/spark";
     import { SparkControls } from "/examples/js/controls.js";
     import { OrbitControls } from "three/addons/controls/OrbitControls.js";
 
@@ -272,7 +272,8 @@
     async function loadSplatFile(splatFile) {
       fileBytes = await splatFile.arrayBuffer();
       fileName = splatFile.name;
-      setSplatFile(fileBytes, fileName);
+      // Create copy of fileBytes to transfer to the worker
+      setSplatFile(fileBytes.slice(), fileName);
     }
 
     var loadedSplat;
@@ -317,7 +318,7 @@
 
     document.querySelector('.spz-button').addEventListener('click', async function () {
       const fileInfo =  {
-        fileBytes: fileBytes,
+        fileBytes: new Uint8Array(fileBytes),
         pathOrUrl: fileName,
       };
 

src/Splat.ts

@@ -327,7 +327,6 @@ export class Splat extends THREE.Mesh<SplatGeometry, THREE.ShaderMaterial> {
         lastOriginToCamera: new THREE.Matrix4(),
         sortJob: null,
         ordering: new Uint32Array(this.splatData.maxSplats),
-        // FIXME: Use FreeList instead of double buffering
         pendingOrdering: new Uint32Array(this.splatData.maxSplats),
         activeSplats: 0,
         orderingId: globalOrderingId++,

src/formats/spz.ts

@@ -1,3 +1,4 @@
+import * as THREE from "three";
 import { SH_C0, SH_DEGREE_TO_NUM_COEFF } from "../defines";
 import type { SplatEncoder, UnpackResult } from "../encoding/encoder";
 import { GunzipReader, fromHalf } from "../utils";
@@ -232,3 +233,201 @@ export class SpzReader {
     }
   }
 }
+
+// SPZ file format writer
+
+export const SPZ_MAGIC = 0x5053474e; // NGSP = Niantic gaussian splat
+export const SPZ_VERSION = 3;
+export const FLAG_ANTIALIASED = 0x1;
+
+export class SpzWriter {
+  private buffer: ArrayBuffer;
+  private view: DataView;
+  private numSplats: number;
+  readonly shDegree: number;
+  private fractionalBits: number;
+  private fraction: number;
+  private flagAntiAlias: boolean;
+  clippedCount = 0;
+
+  constructor({
+    numSplats,
+    shDegree,
+    fractionalBits = 12,
+    flagAntiAlias = true,
+  }: {
+    numSplats: number;
+    shDegree: number;
+    fractionalBits?: number;
+    flagAntiAlias?: boolean;
+  }) {
+    const splatSize =
+      9 + // Position
+      1 + // Opacity
+      3 + // Scale
+      3 + // DC-rgb
+      4 + // Rotation
+      SH_DEGREE_TO_NUM_COEFF[shDegree];
+    const bufferSize = 16 + numSplats * splatSize;
+    this.buffer = new ArrayBuffer(bufferSize);
+    this.view = new DataView(this.buffer);
+
+    this.view.setUint32(0, SPZ_MAGIC, true); // NGSP
+    this.view.setUint32(4, SPZ_VERSION, true);
+    this.view.setUint32(8, numSplats, true);
+    this.view.setUint8(12, shDegree);
+    this.view.setUint8(13, fractionalBits);
+    this.view.setUint8(14, flagAntiAlias ? FLAG_ANTIALIASED : 0);
+    this.view.setUint8(15, 0); // Reserved
+
+    this.numSplats = numSplats;
+    this.shDegree = shDegree;
+    this.fractionalBits = fractionalBits;
+    this.fraction = 1 << fractionalBits;
+    this.flagAntiAlias = flagAntiAlias;
+  }
+
+  setCenter(index: number, x: number, y: number, z: number) {
+    // Divide by this.fraction and round to nearest integer,
+    // then write as 3-bytes per x then y then z.
+    const xRounded = Math.round(x * this.fraction);
+    const xInt = Math.max(-0x7fffff, Math.min(0x7fffff, xRounded));
+    const yRounded = Math.round(y * this.fraction);
+    const yInt = Math.max(-0x7fffff, Math.min(0x7fffff, yRounded));
+    const zRounded = Math.round(z * this.fraction);
+    const zInt = Math.max(-0x7fffff, Math.min(0x7fffff, zRounded));
+    const clipped = xRounded !== xInt || yRounded !== yInt || zRounded !== zInt;
+    if (clipped) {
+      this.clippedCount += 1;
+    }
+    const i9 = index * 9;
+    const base = 16 + i9;
+    this.view.setUint8(base, xInt & 0xff);
+    this.view.setUint8(base + 1, (xInt >> 8) & 0xff);
+    this.view.setUint8(base + 2, (xInt >> 16) & 0xff);
+    this.view.setUint8(base + 3, yInt & 0xff);
+    this.view.setUint8(base + 4, (yInt >> 8) & 0xff);
+    this.view.setUint8(base + 5, (yInt >> 16) & 0xff);
+    this.view.setUint8(base + 6, zInt & 0xff);
+    this.view.setUint8(base + 7, (zInt >> 8) & 0xff);
+    this.view.setUint8(base + 8, (zInt >> 16) & 0xff);
+  }
+
+  setAlpha(index: number, alpha: number) {
+    const base = 16 + this.numSplats * 9 + index;
+    this.view.setUint8(
+      base,
+      Math.max(0, Math.min(255, Math.round(alpha * 255))),
+    );
+  }
+
+  static scaleRgb(r: number) {
+    const v = ((r - 0.5) / (SH_C0 / 0.15) + 0.5) * 255;
+    return Math.max(0, Math.min(255, Math.round(v)));
+  }
+
+  setRgb(index: number, r: number, g: number, b: number) {
+    const base = 16 + this.numSplats * 10 + index * 3;
+    this.view.setUint8(base, SpzWriter.scaleRgb(r));
+    this.view.setUint8(base + 1, SpzWriter.scaleRgb(g));
+    this.view.setUint8(base + 2, SpzWriter.scaleRgb(b));
+  }
+
+  setScale(index: number, scaleX: number, scaleY: number, scaleZ: number) {
+    const base = 16 + this.numSplats * 13 + index * 3;
+    this.view.setUint8(
+      base,
+      Math.max(0, Math.min(255, Math.round((Math.log(scaleX) + 10) * 16))),
+    );
+    this.view.setUint8(
+      base + 1,
+      Math.max(0, Math.min(255, Math.round((Math.log(scaleY) + 10) * 16))),
+    );
+    this.view.setUint8(
+      base + 2,
+      Math.max(0, Math.min(255, Math.round((Math.log(scaleZ) + 10) * 16))),
+    );
+  }
+
+  setQuat(
+    index: number,
+    ...q: [number, number, number, number] // x, y, z, w
+  ) {
+    const base = 16 + this.numSplats * 16 + index * 4;
+
+    const quat = normalize(q);
+
+    // Find largest component
+    let iLargest = 0;
+    for (let i = 1; i < 4; ++i) {
+      if (Math.abs(quat[i]) > Math.abs(quat[iLargest])) {
+        iLargest = i;
+      }
+    }
+
+    // Since -quat represents the same rotation as quat, transform the quaternion so the largest element
+    // is positive. This avoids having to send its sign bit.
+    const negate = quat[iLargest] < 0 ? 1 : 0;
+
+    // Do compression using sign bit and 9-bit precision per element.
+    let comp = iLargest;
+    for (let i = 0; i < 4; ++i) {
+      if (i !== iLargest) {
+        const negbit = (quat[i] < 0 ? 1 : 0) ^ negate;
+        const mag = Math.floor(
+          ((1 << 9) - 1) * (Math.abs(quat[i]) / Math.SQRT1_2) + 0.5,
+        );
+        comp = (comp << 10) | (negbit << 9) | mag;
+      }
+    }
+
+    this.view.setUint8(base, comp & 0xff);
+    this.view.setUint8(base + 1, (comp >> 8) & 0xff);
+    this.view.setUint8(base + 2, (comp >> 16) & 0xff);
+    this.view.setUint8(base + 3, (comp >>> 24) & 0xff);
+  }
+
+  static quantizeSh(sh: number, bits: number) {
+    const value = Math.round(sh * 128) + 128;
+    const bucketSize = 1 << (8 - bits);
+    const quantized =
+      Math.floor((value + bucketSize / 2) / bucketSize) * bucketSize;
+    return Math.max(0, Math.min(255, quantized));
+  }
+
+  setSh(index: number, sh: ArrayLike<number>) {
+    const base =
+      16 + this.numSplats * 20 + index * SH_DEGREE_TO_NUM_COEFF[this.shDegree];
+    for (let i = 0; i < SH_DEGREE_TO_NUM_COEFF[this.shDegree]; ++i) {
+      this.view.setUint8(base + i, SpzWriter.quantizeSh(sh[i], i >= 9 ? 4 : 5));
+    }
+  }
+
+  async finalize(): Promise<Uint8Array> {
+    const input = new Uint8Array(this.buffer);
+    const stream = new ReadableStream({
+      async start(controller) {
+        controller.enqueue(input);
+        controller.close();
+      },
+    });
+    const compressed = stream.pipeThrough(new CompressionStream("gzip"));
+    const response = new Response(compressed);
+    const buffer = await response.arrayBuffer();
+    console.log(
+      "Compressed",
+      input.length,
+      "bytes to",
+      buffer.byteLength,
+      "bytes",
+    );
+    return new Uint8Array(buffer);
+  }
+}
+
+const tempQuat = new THREE.Quaternion();
+function normalize(
+  quat: [number, number, number, number],
+): [number, number, number, number] {
+  return tempQuat.fromArray(quat).normalize().toArray(quat);
+}

src/index.ts

@@ -11,3 +11,5 @@ export * from "./raycast";
 export type { SplatEncoder, ResizableSplatEncoder } from "./encoding/encoder";
 export { PackedSplats } from "./encoding/PackedSplats";
 export { ExtendedSplats } from "./encoding/ExtendedSplats";
+
+export { transcodeSpz } from "./transcode";