Accept extended splats when writing SPZ

rust/spark-rs/src/lib.rs

@@ -529,6 +529,7 @@ pub fn extsplats_to_gsplatarray(num_splats: u32, ext1: Uint32Array, ext2: Uint32
     };
     Ok(GsplatArray::new(splats))
 }
+stub_fn!(feature = "gsplat", extsplats_to_gsplatarray);
 
 #[wasm_bindgen]
 #[cfg(all(feature = "csplat", feature = "tiny_lod"))]

src/SplatLoader.ts

@@ -5,7 +5,6 @@ import { PackedSplats, type PackedSplatsOptions } from "./PackedSplats";
 import { SplatMesh } from "./SplatMesh";
 import { workerPool } from "./SplatWorker";
 import { type ExtResult, type PackedResult, SplatFileType } from "./defines";
-import type { SpzWriteVersion } from "./spz";
 import { decompressPartialGzip, getTextureSize } from "./utils";
 
 // SplatLoader implements the THREE.Loader interface and supports loading a variety
@@ -625,19 +624,3 @@ export class SplatData {
     }
   }
 }
-
-export type FileInput = {
-  fileBytes: Uint8Array;
-  fileType?: SplatFileType;
-  pathOrUrl?: string;
-  transform?: { translate?: number[]; quaternion?: number[]; scale?: number };
-};
-
-export type TranscodeSpzInput = {
-  inputs: FileInput[];
-  maxSh?: number;
-  clipXyz?: { min: number[]; max: number[] };
-  fractionalBits?: number;
-  opacityThreshold?: number;
-  version?: SpzWriteVersion;
-};

src/index.ts

@@ -17,6 +17,8 @@ export {
   writeSpz,
   type SpzWriteVersion,
   type WriteSpzOptions,
+  type TranscodeSpzFileInput,
+  type TranscodeSpzInput,
 } from "./spz";
 
 export { PackedSplats, type PackedSplatsOptions } from "./PackedSplats";

src/spz.ts

@@ -1,15 +1,33 @@
-import type { PackedSplats } from "./PackedSplats";
-import {
-  type TranscodeSpzInput,
-  getSplatFileType,
-  getSplatFileTypeFromPath,
-} from "./SplatLoader";
+import { ExtSplats } from "./ExtSplats";
+import { PackedSplats } from "./PackedSplats";
+import { getSplatFileType, getSplatFileTypeFromPath } from "./SplatLoader";
+import type { SplatFileType } from "./defines";
 
-import { decode_to_gsplatarray, packedsplats_to_gsplatarray } from "spark-rs";
+import {
+  decode_to_gsplatarray,
+  extsplats_to_gsplatarray,
+  packedsplats_to_gsplatarray,
+} from "spark-rs";
 import * as wasm from "./wasm";
 
 export type SpzWriteVersion = 2 | 3;
 
+export type TranscodeSpzFileInput = {
+  fileBytes: Uint8Array;
+  fileType?: SplatFileType;
+  pathOrUrl?: string;
+  transform?: { translate?: number[]; quaternion?: number[]; scale?: number };
+};
+
+export type TranscodeSpzInput = {
+  inputs: TranscodeSpzFileInput[];
+  maxSh?: number;
+  clipXyz?: { min: number[]; max: number[] };
+  fractionalBits?: number;
+  opacityThreshold?: number;
+  version?: SpzWriteVersion;
+};
+
 export async function transcodeSpz(input: TranscodeSpzInput) {
   await wasm.initialization;
 
@@ -78,36 +96,52 @@ export type WriteSpzOptions = {
 };
 
 export function writeSpz(
-  packedSplats: PackedSplats,
+  splats: PackedSplats | ExtSplats,
   maxSh?: number,
   fractionalBits?: number,
 ): { fileBytes: Uint8Array };
 export function writeSpz(
-  packedSplats: PackedSplats,
+  splats: PackedSplats | ExtSplats,
   options?: WriteSpzOptions,
 ): { fileBytes: Uint8Array };
 export function writeSpz(
-  packedSplats: PackedSplats,
+  splats: PackedSplats | ExtSplats,
   maxShOrOptions?: number | WriteSpzOptions,
   fractionalBits?: number,
 ) {
-  if (!packedSplats.packedArray) {
-    throw new Error("");
-  }
   const options: WriteSpzOptions =
     typeof maxShOrOptions === "number"
       ? { maxSh: maxShOrOptions, fractionalBits }
       : (maxShOrOptions ?? {});
-  const gsplats = packedsplats_to_gsplatarray(
-    packedSplats.numSplats,
-    packedSplats.packedArray,
-    packedSplats.extra,
-    packedSplats.splatEncoding,
-  );
-  const spzBytes = gsplats.encode_to_spz(
-    options.maxSh ?? 3,
-    options.fractionalBits ?? 12,
-    options.version,
-  );
-  return { fileBytes: spzBytes };
+  const shDegree = options.maxSh ?? 3;
+  const bits = options.fractionalBits ?? 12;
+
+  if (splats instanceof ExtSplats) {
+    const gsplats = extsplats_to_gsplatarray(
+      splats.numSplats,
+      splats.extArrays[0],
+      splats.extArrays[1],
+      splats.extra,
+    );
+    return {
+      fileBytes: gsplats.encode_to_spz(shDegree, bits, options.version),
+    };
+  }
+
+  if (splats instanceof PackedSplats) {
+    if (!splats.packedArray) {
+      throw new Error("PackedSplats has no splat data");
+    }
+    const gsplats = packedsplats_to_gsplatarray(
+      splats.numSplats,
+      splats.packedArray,
+      splats.extra,
+      splats.splatEncoding,
+    );
+    return {
+      fileBytes: gsplats.encode_to_spz(shDegree, bits, options.version),
+    };
+  }
+
+  throw new Error("writeSpz requires PackedSplats or ExtSplats");
 }