Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 30 additions & 18 deletions src/spz.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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");
}
Expand All @@ -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,
Expand All @@ -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;
Expand All @@ -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 =
Expand Down Expand Up @@ -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;
Expand All @@ -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);
Expand All @@ -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];
Expand Down Expand Up @@ -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;
Expand All @@ -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,
);

Expand Down Expand Up @@ -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);
Expand Down
30 changes: 12 additions & 18 deletions src/utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down
11 changes: 6 additions & 5 deletions src/worker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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);
},
Expand Down
Loading