Introduce writeSpz function to replace SpzWriter for splat-painter example

Author: mrxz

examples/editor/index.html

@@ -69,7 +69,7 @@
     import * as THREE from "three";
     import { OrbitControls } from "three/addons/controls/OrbitControls.js";
     import { GUI } from "lil-gui";
-    import { constructGrid, SparkControls, SparkRenderer, SplatMesh, textSplats, dyno, transcodeSpz, isMobile, isPcSogs, LN_SCALE_MIN, LN_SCALE_MAX } from "@sparkjsdev/spark";
+    import { constructGrid, SparkControls, SparkRenderer, SplatMesh, textSplats, dyno, transcodeSpz, isMobile, LN_SCALE_MIN, LN_SCALE_MAX } from "@sparkjsdev/spark";
     import { getAssetFileURL } from "../js/get-asset-url.js";
 
     const scene = new THREE.Scene();

examples/splat-painter/index.html

@@ -62,7 +62,7 @@
       RgbaArray,
       readRgbaArray,
       SparkControls,
-      SpzWriter,
+      writeSpz,
       unpackSplat,
       PackedSplats,
     } from "@sparkjsdev/spark";
@@ -405,7 +405,7 @@
           alert("No splat mesh loaded for export.");
           return;
         }
-        
+
         try {
           console.log("Starting SPZ export with painted changes...");
 
@@ -416,7 +416,7 @@
             });
             currentSplatMesh.updateGenerator();
           }
-          
+
           const rgba = new RgbaArray();
           rgba.render({
             renderer,
@@ -470,7 +470,7 @@
             unpacked.color.g = rgbaBytes[rgbaOffset + 1] / 255;
             unpacked.color.b = rgbaBytes[rgbaOffset + 2] / 255;
             unpacked.opacity = opacity;
-            
+
             // Push to new PackedSplats
             newPackedSplats.pushSplat(
               unpacked.center,
@@ -482,39 +482,13 @@
 
             processedCount++;
           }
-          
+
           console.log(`Processed ${processedCount} splats`);
 
           // Now export the PackedSplats to SPZ
-          console.log("Creating SPZ writer...");
-          const maxSh = ioOptions.maxSh;
-          const spzWriter = new SpzWriter({
-            numSplats: nonZeroCount,
-            shDegree: maxSh,
-            fractionalBits: ioOptions.fractionalBits,
-            flagAntiAlias: true,
-          });
-          
           console.log("Writing splats to SPZ...");
-          // Iterate through the new packed array
-          for (let i = 0; i < nonZeroCount; i++) {
-            const unpacked = unpackSplat(
-              newPackedSplats.packedArray,
-              i,
-              newPackedSplats.splatEncoding
-            );
-            
-            spzWriter.setCenter(i, unpacked.center.x, unpacked.center.y, unpacked.center.z);
-            spzWriter.setScale(i, unpacked.scales.x, unpacked.scales.y, unpacked.scales.z);
-            spzWriter.setQuat(i, unpacked.quaternion.x, unpacked.quaternion.y, unpacked.quaternion.z, unpacked.quaternion.w);
-            spzWriter.setAlpha(i, unpacked.opacity);
-            spzWriter.setRgb(i, unpacked.color.r, unpacked.color.g, unpacked.color.b);
-          }
-          const spzBytes = await spzWriter.finalize();
-          if (spzWriter.clippedCount > 0) {
-            console.log(`Clipped ${spzWriter.clippedCount} splats. Consider decreasing fractional-bits from ${ioOptions.fractionalBits} to reduce clipping.`);
-          }
-          
+          const { fileBytes: spzBytes } = writeSpz(newPackedSplats, ioOptions.maxSh, ioOptions.fractionalBits);
+
           console.log("Creating download...");
           const blob = new Blob([spzBytes], { type: "application/octet-stream" });
           const url = URL.createObjectURL(blob);
@@ -531,7 +505,7 @@
         }
       },
     };
-    
+
     const ioFolder = gui.addFolder("I/O");
     ioFolder.add(ioOptions, "loadFile").name("Load Splats (SPZ/PLY)");
     ioFolder.add(ioOptions, "saveToSpz").name("Save Splats (SPZ)");

src/index.ts

@@ -12,7 +12,7 @@ export {
   SplatLoader,
   getSplatFileType,
 } from "./SplatLoader";
-export { transcodeSpz } from "./spz";
+export { transcodeSpz, writeSpz } from "./spz";
 
 export { PackedSplats, type PackedSplatsOptions } from "./PackedSplats";
 export { ExtSplats, type ExtSplatsOptions } from "./ExtSplats";

src/spz.ts

@@ -1,10 +1,11 @@
+import type { PackedSplats } from "./PackedSplats";
 import {
   type TranscodeSpzInput,
   getSplatFileType,
   getSplatFileTypeFromPath,
 } from "./SplatLoader";
 
-import { decode_to_gsplatarray } from "spark-rs";
+import { decode_to_gsplatarray, packedsplats_to_gsplatarray } from "spark-rs";
 
 export async function transcodeSpz(input: TranscodeSpzInput) {
   const splatArrays = [];
@@ -59,3 +60,21 @@ export async function transcodeSpz(input: TranscodeSpzInput) {
 
   return { fileBytes: spzBytes, clippedCount: 0 };
 }
+
+export function writeSpz(
+  packedSplats: PackedSplats,
+  maxSh?: number,
+  fractionalBits?: number,
+) {
+  if (!packedSplats.packedArray) {
+    throw new Error("");
+  }
+  const gsplats = packedsplats_to_gsplatarray(
+    packedSplats.numSplats,
+    packedSplats.packedArray,
+    packedSplats.extra,
+    packedSplats.splatEncoding,
+  );
+  const spzBytes = gsplats.encode_to_spz(maxSh ?? 3, fractionalBits ?? 12);
+  return { fileBytes: spzBytes };
+}