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

src/pcsogs.ts

@@ -28,40 +28,57 @@ 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]]),
-  ]);
+  ]).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 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 =
+        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]]);
+  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;
+          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;
+          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;
+          json.scales.mins[2] +
-    setPackedSplatScales(packedArray, i, Math.exp(x), Math.exp(y), Math.exp(z));
+          (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(
+    (quats) => {
       const SQRT2 = Math.sqrt(2);
       for (let i = 0; i < numSplats; ++i) {
         const i4 = i * 4;
@@ -76,8 +93,10 @@ export async function unpackPcSogs(
         const quatW = rOrder === 0 ? rr : r0;
         setPackedSplatQuat(packedArray, i, quatX, quatY, quatZ, quatW);
       }
-
+    },
-  const sh0 = await decodeImageRgba(extraFiles[json.sh0.files[0]]);
+  );
+  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;
@@ -85,17 +104,24 @@ export async function unpackPcSogs(
         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 dc0 =
-    const dc1 = json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
+          json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
-    const dc2 = json.sh0.mins[2] + (json.sh0.maxs[2] - json.sh0.mins[2]) * f2;
+        const dc1 =
-    const opa = json.sh0.mins[3] + (json.sh0.maxs[3] - json.sh0.mins[3]) * f3;
+          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,10 +135,11 @@ 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]]),
-    ]);
+    ]).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);
@@ -124,9 +151,8 @@ export async function unpackPcSogs(
           if (useSH1) {
             for (let k = 0; k < 3; ++k) {
               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;
             }
           }
@@ -134,8 +160,8 @@ export async function unpackPcSogs(
           if (useSH2) {
             for (let k = 0; k < 5; ++k) {
               sh2[k * 3 + d] =
-              json.shN.mins +
+                shN.mins +
-              ((json.shN.maxs - json.shN.mins) *
+                ((shN.maxs - shN.mins) *
                   centroids.rgba[(offset + 3 + k) * 4 + d]) /
                   255;
             }
@@ -144,8 +170,8 @@ export async function unpackPcSogs(
           if (useSH3) {
             for (let k = 0; k < 7; ++k) {
               sh3[k * 3 + d] =
-              json.shN.mins +
+                shN.mins +
-              ((json.shN.maxs - json.shN.mins) *
+                ((shN.maxs - shN.mins) *
                   centroids.rgba[(offset + 8 + k) * 4 + d]) /
                   255;
             }
@@ -156,8 +182,12 @@ export async function unpackPcSogs(
         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 };
 }