Skip to content

Commit c803096

Browse files
mrxzmrxz
authored
Allow image decoding to happen in parallel when parsing SOGS (#122)
1 parent 1e697b3 commit c803096

1 file changed

Lines changed: 130 additions & 100 deletions

File tree

‎src/pcsogs.ts‎

Lines changed: 130 additions & 100 deletions
Original file line numberDiff line numberDiff line change
@@ -28,74 +28,100 @@ export async function unpackPcSogs(
2828
const packedArray = new Uint32Array(maxSplats * 4);
2929
const extra: Record<string, unknown> = {};
3030

31-
const means = await Promise.all([
31+
const meansPromise = Promise.all([
3232
decodeImageRgba(extraFiles[json.means.files[0]]),
3333
decodeImageRgba(extraFiles[json.means.files[1]]),
34-
]);
35-
for (let i = 0; i < numSplats; ++i) {
36-
const i4 = i * 4;
37-
const fx = (means[0][i4 + 0] + (means[1][i4 + 0] << 8)) / 65535;
38-
const fy = (means[0][i4 + 1] + (means[1][i4 + 1] << 8)) / 65535;
39-
const fz = (means[0][i4 + 2] + (means[1][i4 + 2] << 8)) / 65535;
40-
let x = json.means.mins[0] + (json.means.maxs[0] - json.means.mins[0]) * fx;
41-
let y = json.means.mins[1] + (json.means.maxs[1] - json.means.mins[1]) * fy;
42-
let z = json.means.mins[2] + (json.means.maxs[2] - json.means.mins[2]) * fz;
43-
x = Math.sign(x) * (Math.exp(Math.abs(x)) - 1);
44-
y = Math.sign(y) * (Math.exp(Math.abs(y)) - 1);
45-
z = Math.sign(z) * (Math.exp(Math.abs(z)) - 1);
46-
setPackedSplatCenter(packedArray, i, x, y, z);
47-
}
48-
49-
const scales = await decodeImageRgba(extraFiles[json.scales.files[0]]);
50-
for (let i = 0; i < numSplats; ++i) {
51-
const i4 = i * 4;
52-
const fx = scales[i4 + 0] / 255;
53-
const fy = scales[i4 + 1] / 255;
54-
const fz = scales[i4 + 2] / 255;
55-
const x =
56-
json.scales.mins[0] + (json.scales.maxs[0] - json.scales.mins[0]) * fx;
57-
const y =
58-
json.scales.mins[1] + (json.scales.maxs[1] - json.scales.mins[1]) * fy;
59-
const z =
60-
json.scales.mins[2] + (json.scales.maxs[2] - json.scales.mins[2]) * fz;
61-
setPackedSplatScales(packedArray, i, Math.exp(x), Math.exp(y), Math.exp(z));
62-
}
34+
]).then((means) => {
35+
for (let i = 0; i < numSplats; ++i) {
36+
const i4 = i * 4;
37+
const fx = (means[0][i4 + 0] + (means[1][i4 + 0] << 8)) / 65535;
38+
const fy = (means[0][i4 + 1] + (means[1][i4 + 1] << 8)) / 65535;
39+
const fz = (means[0][i4 + 2] + (means[1][i4 + 2] << 8)) / 65535;
40+
let x =
41+
json.means.mins[0] + (json.means.maxs[0] - json.means.mins[0]) * fx;
42+
let y =
43+
json.means.mins[1] + (json.means.maxs[1] - json.means.mins[1]) * fy;
44+
let z =
45+
json.means.mins[2] + (json.means.maxs[2] - json.means.mins[2]) * fz;
46+
x = Math.sign(x) * (Math.exp(Math.abs(x)) - 1);
47+
y = Math.sign(y) * (Math.exp(Math.abs(y)) - 1);
48+
z = Math.sign(z) * (Math.exp(Math.abs(z)) - 1);
49+
setPackedSplatCenter(packedArray, i, x, y, z);
50+
}
51+
});
6352

64-
const quats = await decodeImageRgba(extraFiles[json.quats.files[0]]);
65-
const SQRT2 = Math.sqrt(2);
66-
for (let i = 0; i < numSplats; ++i) {
67-
const i4 = i * 4;
68-
const r0 = (quats[i4 + 0] / 255 - 0.5) * SQRT2;
69-
const r1 = (quats[i4 + 1] / 255 - 0.5) * SQRT2;
70-
const r2 = (quats[i4 + 2] / 255 - 0.5) * SQRT2;
71-
const rr = Math.sqrt(Math.max(0, 1.0 - r0 * r0 - r1 * r1 - r2 * r2));
72-
const rOrder = quats[i4 + 3] - 252;
73-
const quatX = rOrder === 0 ? r0 : rOrder === 1 ? rr : r1;
74-
const quatY = rOrder <= 1 ? r1 : rOrder === 2 ? rr : r2;
75-
const quatZ = rOrder <= 2 ? r2 : rr;
76-
const quatW = rOrder === 0 ? rr : r0;
77-
setPackedSplatQuat(packedArray, i, quatX, quatY, quatZ, quatW);
78-
}
53+
const scalesPromise = decodeImageRgba(extraFiles[json.scales.files[0]]).then(
54+
(scales) => {
55+
for (let i = 0; i < numSplats; ++i) {
56+
const i4 = i * 4;
57+
const fx = scales[i4 + 0] / 255;
58+
const fy = scales[i4 + 1] / 255;
59+
const fz = scales[i4 + 2] / 255;
60+
const x =
61+
json.scales.mins[0] +
62+
(json.scales.maxs[0] - json.scales.mins[0]) * fx;
63+
const y =
64+
json.scales.mins[1] +
65+
(json.scales.maxs[1] - json.scales.mins[1]) * fy;
66+
const z =
67+
json.scales.mins[2] +
68+
(json.scales.maxs[2] - json.scales.mins[2]) * fz;
69+
setPackedSplatScales(
70+
packedArray,
71+
i,
72+
Math.exp(x),
73+
Math.exp(y),
74+
Math.exp(z),
75+
);
76+
}
77+
},
78+
);
7979

80-
const sh0 = await decodeImageRgba(extraFiles[json.sh0.files[0]]);
81-
const SH_C0 = 0.28209479177387814;
82-
for (let i = 0; i < numSplats; ++i) {
83-
const i4 = i * 4;
84-
const f0 = sh0[i4 + 0] / 255;
85-
const f1 = sh0[i4 + 1] / 255;
86-
const f2 = sh0[i4 + 2] / 255;
87-
const f3 = sh0[i4 + 3] / 255;
88-
const dc0 = json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
89-
const dc1 = json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
90-
const dc2 = json.sh0.mins[2] + (json.sh0.maxs[2] - json.sh0.mins[2]) * f2;
91-
const opa = json.sh0.mins[3] + (json.sh0.maxs[3] - json.sh0.mins[3]) * f3;
92-
const r = SH_C0 * dc0 + 0.5;
93-
const g = SH_C0 * dc1 + 0.5;
94-
const b = SH_C0 * dc2 + 0.5;
95-
const a = 1.0 / (1.0 + Math.exp(-opa));
96-
setPackedSplatRgba(packedArray, i, r, g, b, a);
97-
}
80+
const quatsPromise = decodeImageRgba(extraFiles[json.quats.files[0]]).then(
81+
(quats) => {
82+
const SQRT2 = Math.sqrt(2);
83+
for (let i = 0; i < numSplats; ++i) {
84+
const i4 = i * 4;
85+
const r0 = (quats[i4 + 0] / 255 - 0.5) * SQRT2;
86+
const r1 = (quats[i4 + 1] / 255 - 0.5) * SQRT2;
87+
const r2 = (quats[i4 + 2] / 255 - 0.5) * SQRT2;
88+
const rr = Math.sqrt(Math.max(0, 1.0 - r0 * r0 - r1 * r1 - r2 * r2));
89+
const rOrder = quats[i4 + 3] - 252;
90+
const quatX = rOrder === 0 ? r0 : rOrder === 1 ? rr : r1;
91+
const quatY = rOrder <= 1 ? r1 : rOrder === 2 ? rr : r2;
92+
const quatZ = rOrder <= 2 ? r2 : rr;
93+
const quatW = rOrder === 0 ? rr : r0;
94+
setPackedSplatQuat(packedArray, i, quatX, quatY, quatZ, quatW);
95+
}
96+
},
97+
);
98+
const sh0Promise = decodeImageRgba(extraFiles[json.sh0.files[0]]).then(
99+
(sh0) => {
100+
const SH_C0 = 0.28209479177387814;
101+
for (let i = 0; i < numSplats; ++i) {
102+
const i4 = i * 4;
103+
const f0 = sh0[i4 + 0] / 255;
104+
const f1 = sh0[i4 + 1] / 255;
105+
const f2 = sh0[i4 + 2] / 255;
106+
const f3 = sh0[i4 + 3] / 255;
107+
const dc0 =
108+
json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
109+
const dc1 =
110+
json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
111+
const dc2 =
112+
json.sh0.mins[2] + (json.sh0.maxs[2] - json.sh0.mins[2]) * f2;
113+
const opa =
114+
json.sh0.mins[3] + (json.sh0.maxs[3] - json.sh0.mins[3]) * f3;
115+
const r = SH_C0 * dc0 + 0.5;
116+
const g = SH_C0 * dc1 + 0.5;
117+
const b = SH_C0 * dc2 + 0.5;
118+
const a = 1.0 / (1.0 + Math.exp(-opa));
119+
setPackedSplatRgba(packedArray, i, r, g, b, a);
120+
}
121+
},
122+
);
98123

124+
const promises = [meansPromise, scalesPromise, quatsPromise, sh0Promise];
99125
if (json.shN) {
100126
const useSH3 = json.shN.shape[1] >= 48 - 3;
101127
const useSH2 = json.shN.shape[1] >= 27 - 3;
@@ -109,55 +135,59 @@ export async function unpackPcSogs(
109135
const sh2 = new Float32Array(15);
110136
const sh3 = new Float32Array(21);
111137

112-
const [centroids, labels] = await Promise.all([
138+
const shN = json.shN;
139+
const shNPromise = Promise.all([
113140
decodeImage(extraFiles[json.shN.files[0]]),
114141
decodeImage(extraFiles[json.shN.files[1]]),
115-
]);
116-
for (let i = 0; i < numSplats; ++i) {
117-
const i4 = i * 4;
118-
const label = labels.rgba[i4 + 0] + (labels.rgba[i4 + 1] << 8);
119-
const col = (label & 63) * 15;
120-
const row = label >>> 6;
121-
const offset = row * centroids.width + col;
142+
]).then(([centroids, labels]) => {
143+
for (let i = 0; i < numSplats; ++i) {
144+
const i4 = i * 4;
145+
const label = labels.rgba[i4 + 0] + (labels.rgba[i4 + 1] << 8);
146+
const col = (label & 63) * 15;
147+
const row = label >>> 6;
148+
const offset = row * centroids.width + col;
122149

123-
for (let d = 0; d < 3; ++d) {
124-
if (useSH1) {
125-
for (let k = 0; k < 3; ++k) {
126-
sh1[k * 3 + d] =
127-
json.shN.mins +
128-
((json.shN.maxs - json.shN.mins) *
129-
centroids.rgba[(offset + k) * 4 + d]) /
130-
255;
150+
for (let d = 0; d < 3; ++d) {
151+
if (useSH1) {
152+
for (let k = 0; k < 3; ++k) {
153+
sh1[k * 3 + d] =
154+
shN.mins +
155+
((shN.maxs - shN.mins) * centroids.rgba[(offset + k) * 4 + d]) /
156+
255;
157+
}
131158
}
132-
}
133159

134-
if (useSH2) {
135-
for (let k = 0; k < 5; ++k) {
136-
sh2[k * 3 + d] =
137-
json.shN.mins +
138-
((json.shN.maxs - json.shN.mins) *
139-
centroids.rgba[(offset + 3 + k) * 4 + d]) /
140-
255;
160+
if (useSH2) {
161+
for (let k = 0; k < 5; ++k) {
162+
sh2[k * 3 + d] =
163+
shN.mins +
164+
((shN.maxs - shN.mins) *
165+
centroids.rgba[(offset + 3 + k) * 4 + d]) /
166+
255;
167+
}
141168
}
142-
}
143169

144-
if (useSH3) {
145-
for (let k = 0; k < 7; ++k) {
146-
sh3[k * 3 + d] =
147-
json.shN.mins +
148-
((json.shN.maxs - json.shN.mins) *
149-
centroids.rgba[(offset + 8 + k) * 4 + d]) /
150-
255;
170+
if (useSH3) {
171+
for (let k = 0; k < 7; ++k) {
172+
sh3[k * 3 + d] =
173+
shN.mins +
174+
((shN.maxs - shN.mins) *
175+
centroids.rgba[(offset + 8 + k) * 4 + d]) /
176+
255;
177+
}
151178
}
152179
}
153-
}
154180

155-
if (useSH1) encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1);
156-
if (useSH2) encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
157-
if (useSH3) encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
158-
}
181+
if (useSH1) encodeSh1Rgb(extra.sh1 as Uint32Array, i, sh1);
182+
if (useSH2) encodeSh2Rgb(extra.sh2 as Uint32Array, i, sh2);
183+
if (useSH3) encodeSh3Rgb(extra.sh3 as Uint32Array, i, sh3);
184+
}
185+
});
186+
promises.push(shNPromise);
159187
}
160188

189+
await Promise.all(promises);
190+
161191
return { packedArray, numSplats, extra };
162192
}
163193

0 commit comments

Comments
 (0)