mirror of
https://github.com/storytold/spark.git
synced 2026-10-19 07:54:46 +00:00
Allow image decoding to happen in parallel when parsing SOGS (#122)
This commit is contained in:
Noeri Huisman
+133
-103
@@ -28,74 +28,100 @@ export async function unpackPcSogs(
|
|||||||
const packedArray = new Uint32Array(maxSplats * 4);
|
const packedArray = new Uint32Array(maxSplats * 4);
|
||||||
const extra: Record<string, unknown> = {};
|
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[0]]),
|
||||||
decodeImageRgba(extraFiles[json.means.files[1]]),
|
decodeImageRgba(extraFiles[json.means.files[1]]),
|
||||||
]);
|
]).then((means) => {
|
||||||
for (let i = 0; i < numSplats; ++i) {
|
for (let i = 0; i < numSplats; ++i) {
|
||||||
const i4 = i * 4;
|
const i4 = i * 4;
|
||||||
const fx = (means[0][i4 + 0] + (means[1][i4 + 0] << 8)) / 65535;
|
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 fy = (means[0][i4 + 1] + (means[1][i4 + 1] << 8)) / 65535;
|
||||||
const fz = (means[0][i4 + 2] + (means[1][i4 + 2] << 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 x =
|
||||||
let y = json.means.mins[1] + (json.means.maxs[1] - json.means.mins[1]) * fy;
|
json.means.mins[0] + (json.means.maxs[0] - json.means.mins[0]) * fx;
|
||||||
let z = json.means.mins[2] + (json.means.maxs[2] - json.means.mins[2]) * fz;
|
let y =
|
||||||
x = Math.sign(x) * (Math.exp(Math.abs(x)) - 1);
|
json.means.mins[1] + (json.means.maxs[1] - json.means.mins[1]) * fy;
|
||||||
y = Math.sign(y) * (Math.exp(Math.abs(y)) - 1);
|
let z =
|
||||||
z = Math.sign(z) * (Math.exp(Math.abs(z)) - 1);
|
json.means.mins[2] + (json.means.maxs[2] - json.means.mins[2]) * fz;
|
||||||
setPackedSplatCenter(packedArray, i, x, y, z);
|
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]]);
|
const scalesPromise = decodeImageRgba(extraFiles[json.scales.files[0]]).then(
|
||||||
for (let i = 0; i < numSplats; ++i) {
|
(scales) => {
|
||||||
const i4 = i * 4;
|
for (let i = 0; i < numSplats; ++i) {
|
||||||
const fx = scales[i4 + 0] / 255;
|
const i4 = i * 4;
|
||||||
const fy = scales[i4 + 1] / 255;
|
const fx = scales[i4 + 0] / 255;
|
||||||
const fz = scales[i4 + 2] / 255;
|
const fy = scales[i4 + 1] / 255;
|
||||||
const x =
|
const fz = scales[i4 + 2] / 255;
|
||||||
json.scales.mins[0] + (json.scales.maxs[0] - json.scales.mins[0]) * fx;
|
const x =
|
||||||
const y =
|
json.scales.mins[0] +
|
||||||
json.scales.mins[1] + (json.scales.maxs[1] - json.scales.mins[1]) * fy;
|
(json.scales.maxs[0] - json.scales.mins[0]) * fx;
|
||||||
const z =
|
const y =
|
||||||
json.scales.mins[2] + (json.scales.maxs[2] - json.scales.mins[2]) * fz;
|
json.scales.mins[1] +
|
||||||
setPackedSplatScales(packedArray, i, Math.exp(x), Math.exp(y), Math.exp(z));
|
(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 quatsPromise = decodeImageRgba(extraFiles[json.quats.files[0]]).then(
|
||||||
const SQRT2 = Math.sqrt(2);
|
(quats) => {
|
||||||
for (let i = 0; i < numSplats; ++i) {
|
const SQRT2 = Math.sqrt(2);
|
||||||
const i4 = i * 4;
|
for (let i = 0; i < numSplats; ++i) {
|
||||||
const r0 = (quats[i4 + 0] / 255 - 0.5) * SQRT2;
|
const i4 = i * 4;
|
||||||
const r1 = (quats[i4 + 1] / 255 - 0.5) * SQRT2;
|
const r0 = (quats[i4 + 0] / 255 - 0.5) * SQRT2;
|
||||||
const r2 = (quats[i4 + 2] / 255 - 0.5) * SQRT2;
|
const r1 = (quats[i4 + 1] / 255 - 0.5) * SQRT2;
|
||||||
const rr = Math.sqrt(Math.max(0, 1.0 - r0 * r0 - r1 * r1 - r2 * r2));
|
const r2 = (quats[i4 + 2] / 255 - 0.5) * SQRT2;
|
||||||
const rOrder = quats[i4 + 3] - 252;
|
const rr = Math.sqrt(Math.max(0, 1.0 - r0 * r0 - r1 * r1 - r2 * r2));
|
||||||
const quatX = rOrder === 0 ? r0 : rOrder === 1 ? rr : r1;
|
const rOrder = quats[i4 + 3] - 252;
|
||||||
const quatY = rOrder <= 1 ? r1 : rOrder === 2 ? rr : r2;
|
const quatX = rOrder === 0 ? r0 : rOrder === 1 ? rr : r1;
|
||||||
const quatZ = rOrder <= 2 ? r2 : rr;
|
const quatY = rOrder <= 1 ? r1 : rOrder === 2 ? rr : r2;
|
||||||
const quatW = rOrder === 0 ? rr : r0;
|
const quatZ = rOrder <= 2 ? r2 : rr;
|
||||||
setPackedSplatQuat(packedArray, i, quatX, quatY, quatZ, quatW);
|
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 sh0Promise = decodeImageRgba(extraFiles[json.sh0.files[0]]).then(
|
||||||
const i4 = i * 4;
|
(sh0) => {
|
||||||
const f0 = sh0[i4 + 0] / 255;
|
const SH_C0 = 0.28209479177387814;
|
||||||
const f1 = sh0[i4 + 1] / 255;
|
for (let i = 0; i < numSplats; ++i) {
|
||||||
const f2 = sh0[i4 + 2] / 255;
|
const i4 = i * 4;
|
||||||
const f3 = sh0[i4 + 3] / 255;
|
const f0 = sh0[i4 + 0] / 255;
|
||||||
const dc0 = json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
|
const f1 = sh0[i4 + 1] / 255;
|
||||||
const dc1 = json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
|
const f2 = sh0[i4 + 2] / 255;
|
||||||
const dc2 = json.sh0.mins[2] + (json.sh0.maxs[2] - json.sh0.mins[2]) * f2;
|
const f3 = sh0[i4 + 3] / 255;
|
||||||
const opa = json.sh0.mins[3] + (json.sh0.maxs[3] - json.sh0.mins[3]) * f3;
|
const dc0 =
|
||||||
const r = SH_C0 * dc0 + 0.5;
|
json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
|
||||||
const g = SH_C0 * dc1 + 0.5;
|
const dc1 =
|
||||||
const b = SH_C0 * dc2 + 0.5;
|
json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
|
||||||
const a = 1.0 / (1.0 + Math.exp(-opa));
|
const dc2 =
|
||||||
setPackedSplatRgba(packedArray, i, r, g, b, a);
|
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) {
|
if (json.shN) {
|
||||||
const useSH3 = json.shN.shape[1] >= 48 - 3;
|
const useSH3 = json.shN.shape[1] >= 48 - 3;
|
||||||
const useSH2 = json.shN.shape[1] >= 27 - 3;
|
const useSH2 = json.shN.shape[1] >= 27 - 3;
|
||||||
@@ -109,55 +135,59 @@ export async function unpackPcSogs(
|
|||||||
const sh2 = new Float32Array(15);
|
const sh2 = new Float32Array(15);
|
||||||
const sh3 = new Float32Array(21);
|
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[0]]),
|
||||||
decodeImage(extraFiles[json.shN.files[1]]),
|
decodeImage(extraFiles[json.shN.files[1]]),
|
||||||
]);
|
]).then(([centroids, labels]) => {
|
||||||
for (let i = 0; i < numSplats; ++i) {
|
for (let i = 0; i < numSplats; ++i) {
|
||||||
const i4 = i * 4;
|
const i4 = i * 4;
|
||||||
const label = labels.rgba[i4 + 0] + (labels.rgba[i4 + 1] << 8);
|
const label = labels.rgba[i4 + 0] + (labels.rgba[i4 + 1] << 8);
|
||||||
const col = (label & 63) * 15;
|
const col = (label & 63) * 15;
|
||||||
const row = label >>> 6;
|
const row = label >>> 6;
|
||||||
const offset = row * centroids.width + col;
|
const offset = row * centroids.width + col;
|
||||||
|
|
||||||
for (let d = 0; d < 3; ++d) {
|
for (let d = 0; d < 3; ++d) {
|
||||||
if (useSH1) {
|
if (useSH1) {
|
||||||
for (let k = 0; k < 3; ++k) {
|
for (let k = 0; k < 3; ++k) {
|
||||||
sh1[k * 3 + d] =
|
sh1[k * 3 + d] =
|
||||||
json.shN.mins +
|
shN.mins +
|
||||||
((json.shN.maxs - json.shN.mins) *
|
((shN.maxs - shN.mins) * centroids.rgba[(offset + k) * 4 + d]) /
|
||||||
centroids.rgba[(offset + k) * 4 + d]) /
|
255;
|
||||||
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) {
|
if (useSH1) encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1);
|
||||||
for (let k = 0; k < 5; ++k) {
|
if (useSH2) encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
|
||||||
sh2[k * 3 + d] =
|
if (useSH3) encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
|
||||||
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);
|
promises.push(shNPromise);
|
||||||
if (useSH2) encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
|
|
||||||
if (useSH3) encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
await Promise.all(promises);
|
||||||
|
|
||||||
return { packedArray, numSplats, extra };
|
return { packedArray, numSplats, extra };
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user