Skip to content

Commit d53bea9

Browse files
mrxzmrxz
authored
Pre-compute lookup tables when parsing SOGS files (#159)
1 parent 935bc2e commit d53bea9

1 file changed

Lines changed: 84 additions & 49 deletions

File tree

‎src/pcsogs.ts‎

Lines changed: 84 additions & 49 deletions
Original file line numberDiff line numberDiff line change
@@ -54,26 +54,39 @@ export async function unpackPcSogs(
5454

5555
const scalesPromise = decodeImageRgba(extraFiles[json.scales.files[0]]).then(
5656
(scales) => {
57+
const xLookup = new Array(256)
58+
.fill(0)
59+
.map(
60+
(_, i) =>
61+
json.scales.mins[0] +
62+
(json.scales.maxs[0] - json.scales.mins[0]) * (i / 255),
63+
)
64+
.map((x) => Math.exp(x));
65+
const yLookup = new Array(256)
66+
.fill(0)
67+
.map(
68+
(_, i) =>
69+
json.scales.mins[1] +
70+
(json.scales.maxs[1] - json.scales.mins[1]) * (i / 255),
71+
)
72+
.map((x) => Math.exp(x));
73+
const zLookup = new Array(256)
74+
.fill(0)
75+
.map(
76+
(_, i) =>
77+
json.scales.mins[2] +
78+
(json.scales.maxs[2] - json.scales.mins[2]) * (i / 255),
79+
)
80+
.map((x) => Math.exp(x));
81+
5782
for (let i = 0; i < numSplats; ++i) {
5883
const i4 = i * 4;
59-
const fx = scales[i4 + 0] / 255;
60-
const fy = scales[i4 + 1] / 255;
61-
const fz = scales[i4 + 2] / 255;
62-
const x =
63-
json.scales.mins[0] +
64-
(json.scales.maxs[0] - json.scales.mins[0]) * fx;
65-
const y =
66-
json.scales.mins[1] +
67-
(json.scales.maxs[1] - json.scales.mins[1]) * fy;
68-
const z =
69-
json.scales.mins[2] +
70-
(json.scales.maxs[2] - json.scales.mins[2]) * fz;
7184
setPackedSplatScales(
7285
packedArray,
7386
i,
74-
Math.exp(x),
75-
Math.exp(y),
76-
Math.exp(z),
87+
xLookup[scales[i4 + 0]],
88+
yLookup[scales[i4 + 1]],
89+
zLookup[scales[i4 + 2]],
7790
splatEncoding,
7891
);
7992
}
@@ -83,11 +96,15 @@ export async function unpackPcSogs(
8396
const quatsPromise = decodeImageRgba(extraFiles[json.quats.files[0]]).then(
8497
(quats) => {
8598
const SQRT2 = Math.sqrt(2);
99+
const lookup = new Array(256)
100+
.fill(0)
101+
.map((_, i) => (i / 255 - 0.5) * SQRT2);
102+
86103
for (let i = 0; i < numSplats; ++i) {
87104
const i4 = i * 4;
88-
const r0 = (quats[i4 + 0] / 255 - 0.5) * SQRT2;
89-
const r1 = (quats[i4 + 1] / 255 - 0.5) * SQRT2;
90-
const r2 = (quats[i4 + 2] / 255 - 0.5) * SQRT2;
105+
const r0 = lookup[quats[i4 + 0]];
106+
const r1 = lookup[quats[i4 + 1]];
107+
const r2 = lookup[quats[i4 + 2]];
91108
const rr = Math.sqrt(Math.max(0, 1.0 - r0 * r0 - r1 * r1 - r2 * r2));
92109
const rOrder = quats[i4 + 3] - 252;
93110
const quatX = rOrder === 0 ? r0 : rOrder === 1 ? rr : r1;
@@ -101,25 +118,50 @@ export async function unpackPcSogs(
101118
const sh0Promise = decodeImageRgba(extraFiles[json.sh0.files[0]]).then(
102119
(sh0) => {
103120
const SH_C0 = 0.28209479177387814;
121+
const rLookup = new Array(256)
122+
.fill(0)
123+
.map(
124+
(_, i) =>
125+
json.sh0.mins[0] +
126+
(json.sh0.maxs[0] - json.sh0.mins[0]) * (i / 255),
127+
)
128+
.map((x) => SH_C0 * x + 0.5);
129+
const gLookup = new Array(256)
130+
.fill(0)
131+
.map(
132+
(_, i) =>
133+
json.sh0.mins[1] +
134+
(json.sh0.maxs[1] - json.sh0.mins[1]) * (i / 255),
135+
)
136+
.map((x) => SH_C0 * x + 0.5);
137+
const bLookup = new Array(256)
138+
.fill(0)
139+
.map(
140+
(_, i) =>
141+
json.sh0.mins[2] +
142+
(json.sh0.maxs[2] - json.sh0.mins[2]) * (i / 255),
143+
)
144+
.map((x) => SH_C0 * x + 0.5);
145+
const aLookup = new Array(256)
146+
.fill(0)
147+
.map(
148+
(_, i) =>
149+
json.sh0.mins[3] +
150+
(json.sh0.maxs[3] - json.sh0.mins[3]) * (i / 255),
151+
)
152+
.map((x) => 1.0 / (1.0 + Math.exp(-x)));
153+
104154
for (let i = 0; i < numSplats; ++i) {
105155
const i4 = i * 4;
106-
const f0 = sh0[i4 + 0] / 255;
107-
const f1 = sh0[i4 + 1] / 255;
108-
const f2 = sh0[i4 + 2] / 255;
109-
const f3 = sh0[i4 + 3] / 255;
110-
const dc0 =
111-
json.sh0.mins[0] + (json.sh0.maxs[0] - json.sh0.mins[0]) * f0;
112-
const dc1 =
113-
json.sh0.mins[1] + (json.sh0.maxs[1] - json.sh0.mins[1]) * f1;
114-
const dc2 =
115-
json.sh0.mins[2] + (json.sh0.maxs[2] - json.sh0.mins[2]) * f2;
116-
const opa =
117-
json.sh0.mins[3] + (json.sh0.maxs[3] - json.sh0.mins[3]) * f3;
118-
const r = SH_C0 * dc0 + 0.5;
119-
const g = SH_C0 * dc1 + 0.5;
120-
const b = SH_C0 * dc2 + 0.5;
121-
const a = 1.0 / (1.0 + Math.exp(-opa));
122-
setPackedSplatRgba(packedArray, i, r, g, b, a, splatEncoding);
156+
setPackedSplatRgba(
157+
packedArray,
158+
i,
159+
rLookup[sh0[i4 + 0]],
160+
gLookup[sh0[i4 + 1]],
161+
bLookup[sh0[i4 + 2]],
162+
aLookup[sh0[i4 + 3]],
163+
splatEncoding,
164+
);
123165
}
124166
},
125167
);
@@ -143,6 +185,10 @@ export async function unpackPcSogs(
143185
decodeImage(extraFiles[json.shN.files[0]]),
144186
decodeImage(extraFiles[json.shN.files[1]]),
145187
]).then(([centroids, labels]) => {
188+
const lookup = new Array(256)
189+
.fill(0)
190+
.map((_, i) => shN.mins + (shN.maxs - shN.mins) * (i / 255));
191+
146192
for (let i = 0; i < numSplats; ++i) {
147193
const i4 = i * 4;
148194
const label = labels.rgba[i4 + 0] + (labels.rgba[i4 + 1] << 8);
@@ -153,30 +199,19 @@ export async function unpackPcSogs(
153199
for (let d = 0; d < 3; ++d) {
154200
if (useSH1) {
155201
for (let k = 0; k < 3; ++k) {
156-
sh1[k * 3 + d] =
157-
shN.mins +
158-
((shN.maxs - shN.mins) * centroids.rgba[(offset + k) * 4 + d]) /
159-
255;
202+
sh1[k * 3 + d] = lookup[centroids.rgba[(offset + k) * 4 + d]];
160203
}
161204
}
162205

163206
if (useSH2) {
164207
for (let k = 0; k < 5; ++k) {
165-
sh2[k * 3 + d] =
166-
shN.mins +
167-
((shN.maxs - shN.mins) *
168-
centroids.rgba[(offset + 3 + k) * 4 + d]) /
169-
255;
208+
sh2[k * 3 + d] = lookup[centroids.rgba[(offset + 3 + k) * 4 + d]];
170209
}
171210
}
172211

173212
if (useSH3) {
174213
for (let k = 0; k < 7; ++k) {
175-
sh3[k * 3 + d] =
176-
shN.mins +
177-
((shN.maxs - shN.mins) *
178-
centroids.rgba[(offset + 8 + k) * 4 + d]) /
179-
255;
214+
sh3[k * 3 + d] = lookup[centroids.rgba[(offset + 8 + k) * 4 + d]];
180215
}
181216
}
182217
}

0 commit comments

Comments
 (0)
⚡