storytold/spark
2
0
Fork 0
mirror of https://github.com/storytold/spark.git synced 2026-10-13 23:52:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use DecompressionStream for decompressing .spz files (#181)

Browse Source
This commit is contained in:
Noeri Huisman
2025-09-25 17:25:57 +02:00
committed by GitHub
parent 2defaf3f64
commit c6bb49d74d
3 changed files with 48 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/spz.ts
+30 -18
View File
@@ -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 });
}
const header = new DataView(this.reader.read(16).buffer);
async parseHeader() {
if (this.headerParsed) {
throw new Error("SPZ file header already parsed");
}
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,6 +82,9 @@ 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");
}
@@ -81,7 +92,7 @@ export class SpzReader {
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
+12 -18
View File
@@ -1344,10 +1344,9 @@ export class GunzipReader {
fileBytes: Uint8Array;
chunkBytes: number;
offset: number;
chunks: Uint8Array[];
totalBytes: number;
gunzip: Gunzip;
reader: ReadableStreamDefaultReader;
constructor({
fileBytes,
@@ -1355,28 +1354,23 @@ export class GunzipReader {
}: { 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
+6 -5
View File
@@ -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);
},