playcanvas
Version:
Open-source WebGL/WebGPU 3D engine for the web
260 lines (259 loc) • 7.44 kB
JavaScript
import { Quat } from "../../core/math/quat.js";
import { Vec3 } from "../../core/math/vec3.js";
import { Vec4 } from "../../core/math/vec4.js";
import { GSplatData } from "./gsplat-data.js";
const SH_C0 = 0.28209479177387814;
class SplatCompressedIterator {
constructor(gsplatData, p, r, s, c, sh) {
const unpackUnorm = (value, bits) => {
const t = (1 << bits) - 1;
return (value & t) / t;
};
const unpack111011 = (result, value) => {
result.x = unpackUnorm(value >>> 21, 11);
result.y = unpackUnorm(value >>> 11, 10);
result.z = unpackUnorm(value, 11);
};
const unpack8888 = (result, value) => {
result.x = unpackUnorm(value >>> 24, 8);
result.y = unpackUnorm(value >>> 16, 8);
result.z = unpackUnorm(value >>> 8, 8);
result.w = unpackUnorm(value, 8);
};
const unpackRot = (result, value) => {
const norm = Math.SQRT2;
const a = (unpackUnorm(value >>> 20, 10) - 0.5) * norm;
const b = (unpackUnorm(value >>> 10, 10) - 0.5) * norm;
const c2 = (unpackUnorm(value, 10) - 0.5) * norm;
const m = Math.sqrt(1 - (a * a + b * b + c2 * c2));
switch (value >>> 30) {
case 0:
result.set(a, b, c2, m);
break;
case 1:
result.set(m, b, c2, a);
break;
case 2:
result.set(b, m, c2, a);
break;
case 3:
result.set(b, c2, m, a);
break;
}
};
const lerp = (a, b, t) => a * (1 - t) + b * t;
const { chunkData, chunkSize, vertexData, shData0, shData1, shData2, shBands } = gsplatData;
const shCoeffs = [3, 8, 15][shBands - 1];
this.read = (i) => {
const ci = Math.floor(i / 256) * chunkSize;
if (p) {
unpack111011(p, vertexData[i * 4 + 0]);
p.x = lerp(chunkData[ci + 0], chunkData[ci + 3], p.x);
p.y = lerp(chunkData[ci + 1], chunkData[ci + 4], p.y);
p.z = lerp(chunkData[ci + 2], chunkData[ci + 5], p.z);
}
if (r) {
unpackRot(r, vertexData[i * 4 + 1]);
}
if (s) {
unpack111011(s, vertexData[i * 4 + 2]);
s.x = lerp(chunkData[ci + 6], chunkData[ci + 9], s.x);
s.y = lerp(chunkData[ci + 7], chunkData[ci + 10], s.y);
s.z = lerp(chunkData[ci + 8], chunkData[ci + 11], s.z);
}
if (c) {
unpack8888(c, vertexData[i * 4 + 3]);
if (chunkSize > 12) {
c.x = lerp(chunkData[ci + 12], chunkData[ci + 15], c.x);
c.y = lerp(chunkData[ci + 13], chunkData[ci + 16], c.y);
c.z = lerp(chunkData[ci + 14], chunkData[ci + 17], c.z);
}
}
if (sh && shBands > 0) {
const shData = [shData0, shData1, shData2];
for (let j = 0; j < 3; ++j) {
for (let k = 0; k < 15; ++k) {
sh[j * 15 + k] = k < shCoeffs ? shData[j][i * 16 + k] * (8 / 255) - 4 : 0;
}
}
}
};
}
}
class GSplatCompressedData {
numSplats;
comments;
chunkData;
vertexData;
shData0;
shData1;
shData2;
shBands;
createIter(p, r, s, c, sh) {
return new SplatCompressedIterator(this, p, r, s, c, sh);
}
calcAabb(result) {
const { chunkData, numChunks, chunkSize } = this;
let s = Math.exp(Math.max(chunkData[9], chunkData[10], chunkData[11]));
let mx = chunkData[0] - s;
let my = chunkData[1] - s;
let mz = chunkData[2] - s;
let Mx = chunkData[3] + s;
let My = chunkData[4] + s;
let Mz = chunkData[5] + s;
for (let i = 1; i < numChunks; ++i) {
const off = i * chunkSize;
s = Math.exp(Math.max(chunkData[off + 9], chunkData[off + 10], chunkData[off + 11]));
mx = Math.min(mx, chunkData[off + 0] - s);
my = Math.min(my, chunkData[off + 1] - s);
mz = Math.min(mz, chunkData[off + 2] - s);
Mx = Math.max(Mx, chunkData[off + 3] + s);
My = Math.max(My, chunkData[off + 4] + s);
Mz = Math.max(Mz, chunkData[off + 5] + s);
}
result.center.set((mx + Mx) * 0.5, (my + My) * 0.5, (mz + Mz) * 0.5);
result.halfExtents.set((Mx - mx) * 0.5, (My - my) * 0.5, (Mz - mz) * 0.5);
return true;
}
getCenters() {
const { vertexData, chunkData, numChunks, chunkSize } = this;
const result = new Float32Array(this.numSplats * 3);
let mx, my, mz, Mx, My, Mz;
for (let c = 0; c < numChunks; ++c) {
const off = c * chunkSize;
mx = chunkData[off + 0];
my = chunkData[off + 1];
mz = chunkData[off + 2];
Mx = chunkData[off + 3];
My = chunkData[off + 4];
Mz = chunkData[off + 5];
const end = Math.min(this.numSplats, (c + 1) * 256);
for (let i = c * 256; i < end; ++i) {
const p = vertexData[i * 4];
const px = (p >>> 21) / 2047;
const py = (p >>> 11 & 1023) / 1023;
const pz = (p & 2047) / 2047;
result[i * 3 + 0] = (1 - px) * mx + px * Mx;
result[i * 3 + 1] = (1 - py) * my + py * My;
result[i * 3 + 2] = (1 - pz) * mz + pz * Mz;
}
}
return result;
}
getChunks(result) {
const { chunkData, numChunks, chunkSize } = this;
let mx, my, mz, Mx, My, Mz;
for (let c = 0; c < numChunks; ++c) {
const off = c * chunkSize;
mx = chunkData[off + 0];
my = chunkData[off + 1];
mz = chunkData[off + 2];
Mx = chunkData[off + 3];
My = chunkData[off + 4];
Mz = chunkData[off + 5];
result[c * 6 + 0] = mx;
result[c * 6 + 1] = my;
result[c * 6 + 2] = mz;
result[c * 6 + 3] = Mx;
result[c * 6 + 4] = My;
result[c * 6 + 5] = Mz;
}
}
calcFocalPoint(result) {
const { chunkData, numChunks, chunkSize } = this;
result.x = 0;
result.y = 0;
result.z = 0;
for (let i = 0; i < numChunks; ++i) {
const off = i * chunkSize;
result.x += chunkData[off + 0] + chunkData[off + 3];
result.y += chunkData[off + 1] + chunkData[off + 4];
result.z += chunkData[off + 2] + chunkData[off + 5];
}
result.mulScalar(0.5 / numChunks);
}
get isCompressed() {
return true;
}
get numChunks() {
return Math.ceil(this.numSplats / 256);
}
get chunkSize() {
return this.chunkData.length / this.numChunks;
}
// decompress into GSplatData
decompress() {
const members = [
"x",
"y",
"z",
"f_dc_0",
"f_dc_1",
"f_dc_2",
"opacity",
"scale_0",
"scale_1",
"scale_2",
"rot_0",
"rot_1",
"rot_2",
"rot_3"
];
const { shBands } = this;
if (shBands > 0) {
const shMembers = [];
for (let i = 0; i < 45; ++i) {
shMembers.push(`f_rest_${i}`);
}
const location = Math.max(...["f_dc_0", "f_dc_1", "f_dc_2"].map((name) => members.indexOf(name)));
members.splice(location + 1, 0, ...shMembers);
}
const data = {};
members.forEach((name) => {
data[name] = new Float32Array(this.numSplats);
});
const p = new Vec3();
const r = new Quat();
const s = new Vec3();
const c = new Vec4();
const sh = shBands > 0 ? new Float32Array(45) : null;
const iter = this.createIter(p, r, s, c, sh);
for (let i = 0; i < this.numSplats; ++i) {
iter.read(i);
data.x[i] = p.x;
data.y[i] = p.y;
data.z[i] = p.z;
data.rot_1[i] = r.x;
data.rot_2[i] = r.y;
data.rot_3[i] = r.z;
data.rot_0[i] = r.w;
data.scale_0[i] = s.x;
data.scale_1[i] = s.y;
data.scale_2[i] = s.z;
data.f_dc_0[i] = (c.x - 0.5) / SH_C0;
data.f_dc_1[i] = (c.y - 0.5) / SH_C0;
data.f_dc_2[i] = (c.z - 0.5) / SH_C0;
data.opacity[i] = c.w <= 0 ? -40 : c.w >= 1 ? 40 : -Math.log(1 / c.w - 1);
if (sh) {
for (let c2 = 0; c2 < 45; ++c2) {
data[`f_rest_${c2}`][i] = sh[c2];
}
}
}
return new GSplatData([{
name: "vertex",
count: this.numSplats,
properties: members.map((name) => {
return {
name,
type: "float",
byteSize: 4,
storage: data[name]
};
})
}], this.comments);
}
}
export {
GSplatCompressedData
};