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

Allow image decoding to happen in parallel when parsing SOGS (#122)

Browse Source
This commit is contained in:
Noeri Huisman
2025-07-16 18:43:09 +02:00
committed by GitHub
parent 1e697b397c
commit c803096259
1 changed files with 133 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/pcsogs.ts
+133 -103
View File
@@ -28,74 +28,100 @@ export async function unpackPcSogs(
const packedArray = new Uint32Array(maxSplats * 4);
const extra: Record<string, unknown> = {};
const means = await Promise.all([
const meansPromise = Promise.all([
decodeImageRgba(extraFiles[json.means.files[0]]),
decodeImageRgba(extraFiles[json.means.files[1]]),
]);
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const fx = (means[0][i4 + 0] + (means[1][i4 + 0] << 8)) / 65535;
const fy = (means[0][i4 + 1] + (means[1][i4 + 1] << 8)) / 65535;
const fz = (means[0][i4 + 2] + (means[1][i4 + 2] << 8)) / 65535;
let x = json.means.mins[0] + (json.means.maxs[0] - json.means.mins[0]) * fx;
let y = json.means.mins[1] + (json.means.maxs[1] - json.means.mins[1]) * fy;
let z = json.means.mins[2] + (json.means.maxs[2] - json.means.mins[2]) * fz;
x = Math.sign(x) * (Math.exp(Math.abs(x)) - 1);
y = Math.sign(y) * (Math.exp(Math.abs(y)) - 1);
z = Math.sign(z) * (Math.exp(Math.abs(z)) - 1);
setPackedSplatCenter(packedArray, i, x, y, z);
}
]).then((means) => {
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const fx = (means[0][i4 + 0] + (means[1][i4 + 0] << 8)) / 65535;
const fy = (means[0][i4 + 1] + (means[1][i4 + 1] << 8)) / 65535;
const fz = (means[0][i4 + 2] + (means[1][i4 + 2] << 8)) / 65535;
let x =
json.means.mins[0] + (json.means.maxs[0] - json.means.mins[0]) * fx;
let y =
json.means.mins[1] + (json.means.maxs[1] - json.means.mins[1]) * fy;
let z =
json.means.mins[2] + (json.means.maxs[2] - json.means.mins[2]) * fz;
x = Math.sign(x) * (Math.exp(Math.abs(x)) - 1);
y = Math.sign(y) * (Math.exp(Math.abs(y)) - 1);
z = Math.sign(z) * (Math.exp(Math.abs(z)) - 1);
setPackedSplatCenter(packedArray, i, x, y, z);
}
});
const scales = await decodeImageRgba(extraFiles[json.scales.files[0]]);
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const fx = scales[i4 + 0] / 255;
const fy = scales[i4 + 1] / 255;
const fz = scales[i4 + 2] / 255;
const x =
json.scales.mins[0] + (json.scales.maxs[0] - json.scales.mins[0]) * fx;
const y =
json.scales.mins[1] + (json.scales.maxs[1] - json.scales.mins[1]) * fy;
const z =
json.scales.mins[2] + (json.scales.maxs[2] - json.scales.mins[2]) * fz;
setPackedSplatScales(packedArray, i, Math.exp(x), Math.exp(y), Math.exp(z));
}
const scalesPromise = decodeImageRgba(extraFiles[json.scales.files[0]]).then(
(scales) => {
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const fx = scales[i4 + 0] / 255;
const fy = scales[i4 + 1] / 255;
const fz = scales[i4 + 2] / 255;
const x =
json.scales.mins[0] +
(json.scales.maxs[0] - json.scales.mins[0]) * fx;
const y =
json.scales.mins[1] +
(json.scales.maxs[1] - json.scales.mins[1]) * fy;
const z =
json.scales.mins[2] +
(json.scales.maxs[2] - json.scales.mins[2]) * fz;
setPackedSplatScales(
packedArray,
i,
Math.exp(x),
Math.exp(y),
Math.exp(z),
);
}
},
);
const quats = await decodeImageRgba(extraFiles[json.quats.files[0]]);
const SQRT2 = Math.sqrt(2);
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const r0 = (quats[i4 + 0] / 255 - 0.5) * SQRT2;
const r1 = (quats[i4 + 1] / 255 - 0.5) * SQRT2;
const r2 = (quats[i4 + 2] / 255 - 0.5) * SQRT2;
const rr = Math.sqrt(Math.max(0, 1.0 - r0 * r0 - r1 * r1 - r2 * r2));
const rOrder = quats[i4 + 3] - 252;
const quatX = rOrder === 0 ? r0 : rOrder === 1 ? rr : r1;
const quatY = rOrder <= 1 ? r1 : rOrder === 2 ? rr : r2;
const quatZ = rOrder <= 2 ? r2 : rr;
const quatW = rOrder === 0 ? rr : r0;
setPackedSplatQuat(packedArray, i, quatX, quatY, quatZ, quatW);
}
const sh0 = await decodeImageRgba(extraFiles[json.sh0.files[0]]);
const SH_C0 = 0.28209479177387814;
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const f0 = sh0[i4 + 0] / 255;
const f1 = sh0[i4 + 1] / 255;
const f2 = sh0[i4 + 2] / 255;
const f3 = sh0[i4 + 3] / 255;
const dc0 = json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
const dc1 = json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
const dc2 = json.sh0.mins[2] + (json.sh0.maxs[2] - json.sh0.mins[2]) * f2;
const opa = json.sh0.mins[3] + (json.sh0.maxs[3] - json.sh0.mins[3]) * f3;
const r = SH_C0 * dc0 + 0.5;
const g = SH_C0 * dc1 + 0.5;
const b = SH_C0 * dc2 + 0.5;
const a = 1.0 / (1.0 + Math.exp(-opa));
setPackedSplatRgba(packedArray, i, r, g, b, a);
}
const quatsPromise = decodeImageRgba(extraFiles[json.quats.files[0]]).then(
(quats) => {
const SQRT2 = Math.sqrt(2);
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const r0 = (quats[i4 + 0] / 255 - 0.5) * SQRT2;
const r1 = (quats[i4 + 1] / 255 - 0.5) * SQRT2;
const r2 = (quats[i4 + 2] / 255 - 0.5) * SQRT2;
const rr = Math.sqrt(Math.max(0, 1.0 - r0 * r0 - r1 * r1 - r2 * r2));
const rOrder = quats[i4 + 3] - 252;
const quatX = rOrder === 0 ? r0 : rOrder === 1 ? rr : r1;
const quatY = rOrder <= 1 ? r1 : rOrder === 2 ? rr : r2;
const quatZ = rOrder <= 2 ? r2 : rr;
const quatW = rOrder === 0 ? rr : r0;
setPackedSplatQuat(packedArray, i, quatX, quatY, quatZ, quatW);
}
},
);
const sh0Promise = decodeImageRgba(extraFiles[json.sh0.files[0]]).then(
(sh0) => {
const SH_C0 = 0.28209479177387814;
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const f0 = sh0[i4 + 0] / 255;
const f1 = sh0[i4 + 1] / 255;
const f2 = sh0[i4 + 2] / 255;
const f3 = sh0[i4 + 3] / 255;
const dc0 =
json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
const dc1 =
json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
const dc2 =
json.sh0.mins[2] + (json.sh0.maxs[2] - json.sh0.mins[2]) * f2;
const opa =
json.sh0.mins[3] + (json.sh0.maxs[3] - json.sh0.mins[3]) * f3;
const r = SH_C0 * dc0 + 0.5;
const g = SH_C0 * dc1 + 0.5;
const b = SH_C0 * dc2 + 0.5;
const a = 1.0 / (1.0 + Math.exp(-opa));
setPackedSplatRgba(packedArray, i, r, g, b, a);
}
},
);
const promises = [meansPromise, scalesPromise, quatsPromise, sh0Promise];
if (json.shN) {
const useSH3 = json.shN.shape[1] >= 48 - 3;
const useSH2 = json.shN.shape[1] >= 27 - 3;
@@ -109,55 +135,59 @@ export async function unpackPcSogs(
const sh2 = new Float32Array(15);
const sh3 = new Float32Array(21);
const [centroids, labels] = await Promise.all([
const shN = json.shN;
const shNPromise = Promise.all([
decodeImage(extraFiles[json.shN.files[0]]),
decodeImage(extraFiles[json.shN.files[1]]),
]);
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const label = labels.rgba[i4 + 0] + (labels.rgba[i4 + 1] << 8);
const col = (label & 63) * 15;
const row = label >>> 6;
const offset = row * centroids.width + col;
]).then(([centroids, labels]) => {
for (let i = 0; i < numSplats; ++i) {
const i4 = i * 4;
const label = labels.rgba[i4 + 0] + (labels.rgba[i4 + 1] << 8);
const col = (label & 63) * 15;
const row = label >>> 6;
const offset = row * centroids.width + col;
for (let d = 0; d < 3; ++d) {
if (useSH1) {
for (let k = 0; k < 3; ++k) {
sh1[k * 3 + d] =
json.shN.mins +
((json.shN.maxs - json.shN.mins) *
centroids.rgba[(offset + k) * 4 + d]) /
255;
for (let d = 0; d < 3; ++d) {
if (useSH1) {
for (let k = 0; k < 3; ++k) {
sh1[k * 3 + d] =
shN.mins +
((shN.maxs - shN.mins) * centroids.rgba[(offset + k) * 4 + d]) /
255;
}
}
if (useSH2) {
for (let k = 0; k < 5; ++k) {
sh2[k * 3 + d] =
shN.mins +
((shN.maxs - shN.mins) *
centroids.rgba[(offset + 3 + k) * 4 + d]) /
255;
}
}
if (useSH3) {
for (let k = 0; k < 7; ++k) {
sh3[k * 3 + d] =
shN.mins +
((shN.maxs - shN.mins) *
centroids.rgba[(offset + 8 + k) * 4 + d]) /
255;
}
}
}
if (useSH2) {
for (let k = 0; k < 5; ++k) {
sh2[k * 3 + d] =
json.shN.mins +
((json.shN.maxs - json.shN.mins) *
centroids.rgba[(offset + 3 + k) * 4 + d]) /
255;
}
}
if (useSH3) {
for (let k = 0; k < 7; ++k) {
sh3[k * 3 + d] =
json.shN.mins +
((json.shN.maxs - json.shN.mins) *
centroids.rgba[(offset + 8 + k) * 4 + d]) /
255;
}
}
if (useSH1) encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1);
if (useSH2) encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
if (useSH3) encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
}
if (useSH1) encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1);
if (useSH2) encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
if (useSH3) encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
}
});
promises.push(shNPromise);
}
await Promise.all(promises);
return { packedArray, numSplats, extra };
}