Use DecompressionStream for decompressing .spz files (#181)

Author: mrxz

src/spz.ts

@@ -18,21 +18,28 @@ export class SpzReader {
   fileBytes: Uint8Array;
   reader: GunzipReader;
 
-  version: number;
-  numSplats: number;
-  shDegree: number;
-  fractionalBits: number;
-  flags: number;
-  flagAntiAlias: boolean;
-  reserved: number;
-  parsed: boolean;
+  version = -1;
+  numSplats = 0;
+  shDegree = 0;
+  fractionalBits = 0;
+  flags = 0;
+  flagAntiAlias = false;
+  reserved = 0;
+  headerParsed = false;
+  parsed = false;
 
   constructor({ fileBytes }: { fileBytes: Uint8Array | ArrayBuffer }) {
     this.fileBytes =
       fileBytes instanceof ArrayBuffer ? new Uint8Array(fileBytes) : fileBytes;
     this.reader = new GunzipReader({ fileBytes: this.fileBytes });
+  }
+
+  async parseHeader() {
+    if (this.headerParsed) {
+      throw new Error("SPZ file header already parsed");
+    }
 
-    const header = new DataView(this.reader.read(16).buffer);
+    const header = new DataView((await this.reader.read(16)).buffer);
     if (header.getUint32(0, true) !== 0x5053474e) {
       throw new Error("Invalid SPZ file");
     }
@@ -47,10 +54,11 @@ export class SpzReader {
     this.flags = header.getUint8(14);
     this.flagAntiAlias = (this.flags & 0x01) !== 0;
     this.reserved = header.getUint8(15);
+    this.headerParsed = true;
     this.parsed = false;
   }
 
-  parseSplats(
+  async parseSplats(
     centerCallback?: (index: number, x: number, y: number, z: number) => void,
     alphaCallback?: (index: number, alpha: number) => void,
     rgbCallback?: (index: number, r: number, g: number, b: number) => void,
@@ -74,14 +82,17 @@ export class SpzReader {
       sh3?: Float32Array,
     ) => void,
   ) {
+    if (!this.headerParsed) {
+      throw new Error("SPZ file header must be parsed first");
+    }
     if (this.parsed) {
       throw new Error("SPZ file already parsed");
     }
     this.parsed = true;
 
     if (this.version === 1) {
       // float16 centers
-      const centerBytes = this.reader.read(this.numSplats * 3 * 2);
+      const centerBytes = await this.reader.read(this.numSplats * 3 * 2);
       const centerUint16 = new Uint16Array(centerBytes.buffer);
       for (let i = 0; i < this.numSplats; i++) {
         const i3 = i * 3;
@@ -93,7 +104,7 @@ export class SpzReader {
     } else if (this.version === 2 || this.version === 3) {
       // 24-bit fixed-point centers
       const fixed = 1 << this.fractionalBits;
-      const centerBytes = this.reader.read(this.numSplats * 3 * 3);
+      const centerBytes = await this.reader.read(this.numSplats * 3 * 3);
       for (let i = 0; i < this.numSplats; i++) {
         const i9 = i * 9;
         const x =
@@ -121,13 +132,13 @@ export class SpzReader {
     }
 
     {
-      const bytes = this.reader.read(this.numSplats);
+      const bytes = await this.reader.read(this.numSplats);
       for (let i = 0; i < this.numSplats; i++) {
         alphaCallback?.(i, bytes[i] / 255);
       }
     }
     {
-      const rgbBytes = this.reader.read(this.numSplats * 3);
+      const rgbBytes = await this.reader.read(this.numSplats * 3);
       const scale = SH_C0 / 0.15;
       for (let i = 0; i < this.numSplats; i++) {
         const i3 = i * 3;
@@ -138,7 +149,7 @@ export class SpzReader {
       }
     }
     {
-      const scalesBytes = this.reader.read(this.numSplats * 3);
+      const scalesBytes = await this.reader.read(this.numSplats * 3);
       for (let i = 0; i < this.numSplats; i++) {
         const i3 = i * 3;
         const scaleX = Math.exp(scalesBytes[i3] / 16 - 10);
@@ -160,7 +171,7 @@ export class SpzReader {
       // v^2 + v^2 = 1
       // v = 1 / sqrt(2);
       const maxValue = 1 / Math.sqrt(2); // 0.7071
-      const quatBytes = this.reader.read(this.numSplats * 4);
+      const quatBytes = await this.reader.read(this.numSplats * 4);
       for (let i = 0; i < this.numSplats; i++) {
         const i3 = i * 4;
         const quaternion = [0, 0, 0, 0];
@@ -211,7 +222,7 @@ export class SpzReader {
         );
       }
     } else {
-      const quatBytes = this.reader.read(this.numSplats * 3);
+      const quatBytes = await this.reader.read(this.numSplats * 3);
       for (let i = 0; i < this.numSplats; i++) {
         const i3 = i * 3;
         const quatX = quatBytes[i3] / 127.5 - 1;
@@ -228,7 +239,7 @@ export class SpzReader {
       const sh1 = new Float32Array(3 * 3);
       const sh2 = this.shDegree >= 2 ? new Float32Array(5 * 3) : undefined;
       const sh3 = this.shDegree >= 3 ? new Float32Array(7 * 3) : undefined;
-      const shBytes = this.reader.read(
+      const shBytes = await this.reader.read(
         this.numSplats * SH_DEGREE_TO_VECS[this.shDegree] * 3,
       );
 
@@ -592,6 +603,7 @@ export async function transcodeSpz(input: TranscodeSpzInput) {
       }
       case SplatFileType.SPZ: {
         const spz = new SpzReader({ fileBytes: input.fileBytes });
+        await spz.parseHeader();
         const mapping = new Int32Array(spz.numSplats);
         mapping.fill(-1);
         const centers = new Float32Array(spz.numSplats * 3);

src/utils.ts

@@ -1344,39 +1344,33 @@ export class GunzipReader {
   fileBytes: Uint8Array;
   chunkBytes: number;
 
-  offset: number;
   chunks: Uint8Array[];
   totalBytes: number;
-  gunzip: Gunzip;
+  reader: ReadableStreamDefaultReader;
 
   constructor({
     fileBytes,
     chunkBytes = 64 * 1024,
   }: { fileBytes: Uint8Array; chunkBytes?: number }) {
     this.fileBytes = fileBytes;
     this.chunkBytes = chunkBytes;
-    this.offset = 0;
     this.chunks = [];
     this.totalBytes = 0;
 
-    this.gunzip = new Gunzip((chunk, _final) => {
-      this.chunks.push(chunk);
-      this.totalBytes += chunk.length;
-    });
+    const ds = new DecompressionStream("gzip");
+    const decompressionStream = new Blob([fileBytes]).stream().pipeThrough(ds);
+    this.reader = decompressionStream.getReader();
   }
 
-  read(numBytes: number): Uint8Array {
-    while (this.totalBytes < numBytes && this.offset < this.fileBytes.length) {
-      const end = Math.min(
-        this.offset + this.chunkBytes,
-        this.fileBytes.length,
-      );
-      this.gunzip.push(this.fileBytes.subarray(this.offset, end), false);
-      this.offset = end;
-    }
+  async read(numBytes: number): Promise<Uint8Array> {
+    while (this.totalBytes < numBytes) {
+      const { value: chunk, done: readerDone } = await this.reader.read();
+      if (readerDone) {
+        break;
+      }
 
-    if (this.totalBytes < numBytes && this.offset >= this.fileBytes.length) {
-      this.gunzip.push(new Uint8Array(0), true);
+      this.chunks.push(chunk);
+      this.totalBytes += chunk.length;
     }
 
     if (this.totalBytes < numBytes) {

src/worker.ts

@@ -61,7 +61,7 @@ async function onMessage(event: MessageEvent) {
           fileBytes: Uint8Array;
           splatEncoding: SplatEncoding;
         };
-        const decoded = unpackSpz(fileBytes, splatEncoding);
+        const decoded = await unpackSpz(fileBytes, splatEncoding);
         result = {
           id,
           numSplats: decoded.numSplats,
@@ -379,21 +379,22 @@ async function unpackPly({
   return { packedArray, numSplats, extra };
 }
 
-function unpackSpz(
+async function unpackSpz(
   fileBytes: Uint8Array,
   splatEncoding: SplatEncoding,
-): {
+): Promise<{
   packedArray: Uint32Array;
   numSplats: number;
   extra: Record<string, unknown>;
-} {
+}> {
   const spz = new SpzReader({ fileBytes });
+  await spz.parseHeader();
   const numSplats = spz.numSplats;
   const maxSplats = computeMaxSplats(numSplats);
   const packedArray = new Uint32Array(maxSplats * 4);
   const extra: Record<string, unknown> = {};
 
-  spz.parseSplats(
+  await spz.parseSplats(
     (index, x, y, z) => {
       setPackedSplatCenter(packedArray, index, x, y, z);
     },