Skip to content

Commit

Permalink
Support compressed spherical harmonics (playcanvas#7143)
Browse files Browse the repository at this point in the history
  • Loading branch information
@slimbuck
slimbuck committed Nov 27, 2024
1 parent bb268f9 commit 66cbb9f
Show file tree
Hide file tree
Showing 7 changed files with 383 additions and 118 deletions.
104 changes: 45 additions & 59 deletions src/framework/parsers/ply.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -198,18 +198,31 @@ const isCompressedPly = (elements) => {
'min_x', 'min_y', 'min_z',
'max_x', 'max_y', 'max_z',
'min_scale_x', 'min_scale_y', 'min_scale_z',
'max_scale_x', 'max_scale_y', 'max_scale_z'
'max_scale_x', 'max_scale_y', 'max_scale_z',
'min_r', 'min_g', 'min_b',
'max_r', 'max_g', 'max_b'
];

const vertexProperties = [
'packed_position', 'packed_rotation', 'packed_scale', 'packed_color'
];

return elements.length === 2 &&
elements[0].name === 'chunk' &&
elements[0].properties.every((p, i) => p.name === chunkProperties[i] && p.type === 'float') &&
elements[1].name === 'vertex' &&
elements[1].properties.every((p, i) => p.name === vertexProperties[i] && p.type === 'uint');
const shProperties = new Array(45).fill('').map((_, i) => `f_rest_${i}`);

const hasBaseElements = () => {
return elements[0].name === 'chunk' &&
elements[0].properties.every((p, i) => p.name === chunkProperties[i] && p.type === 'float') &&
elements[1].name === 'vertex' &&
elements[1].properties.every((p, i) => p.name === vertexProperties[i] && p.type === 'uint');
};

const hasSHElements = () => {
return elements[2].name === 'sh' &&
[9, 24, 45].indexOf(elements[2].properties.length) !== -1 &&
elements[2].properties.every((p, i) => p.name === shProperties[i] && p.type === 'uchar');
};

return (elements.length === 2 && hasBaseElements()) || (elements.length === 3 && hasBaseElements() && hasSHElements());
};

const isFloatPly = (elements) => {
Expand All @@ -223,10 +236,8 @@ const readCompressedPly = async (streamBuf, elements, littleEndian) => {
const result = new GSplatCompressedData();

const numChunks = elements[0].count;
const chunkSize = 12 * 4;

const numChunkProperties = elements[0].properties.length;
const numVertices = elements[1].count;
const vertexSize = 4 * 4;

// evaluate the storage size for the given count (this must match the
// texture size calculation in GSplatCompressed).
Expand All @@ -237,64 +248,39 @@ const readCompressedPly = async (streamBuf, elements, littleEndian) => {
};

// allocate result
result.numSplats = elements[1].count;
result.chunkData = new Float32Array(evalStorageSize(numChunks) * 12);
result.numSplats = numVertices;
result.chunkData = new Float32Array(numChunks * numChunkProperties);
result.vertexData = new Uint32Array(evalStorageSize(numVertices) * 4);

let uint32StreamData;
const uint32ChunkData = new Uint32Array(result.chunkData.buffer);
const uint32VertexData = result.vertexData;

// read chunks
let chunks = 0;
while (chunks < numChunks) {
while (streamBuf.remaining < chunkSize) {
/* eslint-disable no-await-in-loop */
await streamBuf.read();
}

// ensure the uint32 view is still valid
if (uint32StreamData?.buffer !== streamBuf.data.buffer) {
uint32StreamData = new Uint32Array(streamBuf.data.buffer, 0, Math.floor(streamBuf.data.buffer.byteLength / 4));
}

// read the next chunk of data
const toRead = Math.min(numChunks - chunks, Math.floor(streamBuf.remaining / chunkSize));
// read length bytes of data into buffer
const read = async (buffer, length) => {
const target = new Uint8Array(buffer);
let cursor = 0;

const dstOffset = chunks * 12;
const srcOffset = streamBuf.head / 4;
for (let i = 0; i < toRead * 12; ++i) {
uint32ChunkData[dstOffset + i] = uint32StreamData[srcOffset + i];
}

streamBuf.head += toRead * chunkSize;
chunks += toRead;
}

// read vertices
let vertices = 0;
while (vertices < numVertices) {
while (streamBuf.remaining < vertexSize) {
/* eslint-disable no-await-in-loop */
await streamBuf.read();
}
while (cursor < length) {
while (streamBuf.remaining === 0) {
/* eslint-disable no-await-in-loop */
await streamBuf.read();
}

// ensure the uint32 view is still valid
if (uint32StreamData?.buffer !== streamBuf.data.buffer) {
uint32StreamData = new Uint32Array(streamBuf.data.buffer, 0, Math.floor(streamBuf.data.buffer.byteLength / 4));
const toCopy = Math.min(length - cursor, streamBuf.remaining);
const src = streamBuf.data;
for (let i = 0; i < toCopy; ++i) {
target[cursor++] = src[streamBuf.head++];
}
}
};

// read the next chunk of data
const toRead = Math.min(numVertices - vertices, Math.floor(streamBuf.remaining / vertexSize));
// read chunk data
await read(result.chunkData.buffer, numChunks * numChunkProperties * 4);

const dstOffset = vertices * 4;
const srcOffset = streamBuf.head / 4;
for (let i = 0; i < toRead * 4; ++i) {
uint32VertexData[dstOffset + i] = uint32StreamData[srcOffset + i];
}
// read packed vertices
await read(result.vertexData.buffer, numVertices * 4 * 4);

streamBuf.head += toRead * vertexSize;
vertices += toRead;
// read sh data
if (elements.length === 3) {
result.shData = new Uint8Array(elements[2].count * elements[2].properties.length);
await read(result.shData.buffer, result.shData.byteLength);
}

return result;
Expand Down
2 changes: 1 addition & 1 deletion src/platform/graphics/texture.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -163,7 +163,7 @@ class Texture {
* - {@link FUNC_NOTEQUAL}
*
* Defaults to {@link FUNC_LESS}.
* @param {Uint8Array[]|HTMLCanvasElement[]|HTMLImageElement[]|HTMLVideoElement[]|Uint8Array[][]} [options.levels]
* @param {Uint8Array[]|Uint16Array[]|Uint32Array[]|Float32Array[]|HTMLCanvasElement[]|HTMLImageElement[]|HTMLVideoElement[]|Uint8Array[][]} [options.levels]
* - Array of Uint8Array or other supported browser interface; or a two-dimensional array
* of Uint8Array if options.arrayLength is defined and greater than zero.
* @param {boolean} [options.storage] - Defines if texture can be used as a storage texture by
Expand Down
135 changes: 93 additions & 42 deletions src/scene/gsplat/gsplat-compressed-data.js
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ const SH_C0 = 0.28209479177387814;

// iterator for accessing compressed splat data
class SplatCompressedIterator {
constructor(gsplatData, p, r, s, c) {
constructor(gsplatData, p, r, s, c, sh) {
const unpackUnorm = (value, bits) => {
const t = (1 << bits) - 1;
return (value & t) / t;
Expand Down Expand Up @@ -47,11 +47,11 @@ class SplatCompressedIterator {
return (a === b) ? a : a * (1 - t) + b * t;
};

const chunkData = gsplatData.chunkData;
const vertexData = gsplatData.vertexData;
const { chunkData, chunkSize, vertexData, shData, shBands } = gsplatData;
const shCoeffs = [3, 8, 15][shBands - 1];

this.read = (i) => {
const ci = Math.floor(i / 256) * 12;
const ci = Math.floor(i / 256) * chunkSize;

if (p) {
unpack111011(p, vertexData[i * 4 + 0]);
Expand All @@ -73,6 +73,19 @@ class SplatCompressedIterator {

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) {
for (let j = 0; j < 3; ++j) {
for (let k = 0; k < 15; ++k) {
sh[j * 15 + k] = (k < shCoeffs) ? (shData[(i * 3 + j) * shCoeffs + k] * (8 / 255) - 4) : 0;
}
}
}
};
}
Expand All @@ -82,11 +95,13 @@ class GSplatCompressedData {
numSplats;

/**
* Contains 12 floats per chunk:
* Contains either 12 or 18 floats per chunk:
* min_x, min_y, min_z,
* max_x, max_y, max_z,
* min_scale_x, min_scale_y, min_scale_z,
* max_scale_x, max_scale_y, max_scale_z
* min_r, min_g, min_b,
* max_r, max_g, max_b
* @type {Float32Array}
*/
chunkData;
Expand All @@ -101,17 +116,24 @@ class GSplatCompressedData {
*/
vertexData;

/**
* Contains optional quantized spherical harmonic data for up to 3 bands.
* @type {Uint8Array}
*/
shData;

/**
* Create an iterator for accessing splat data
*
* @param {Vec3|null} [p] - the vector to receive splat position
* @param {Quat|null} [r] - the quaternion to receive splat rotation
* @param {Vec3|null} [s] - the vector to receive splat scale
* @param {Vec4|null} [c] - the vector to receive splat color
* @param {Float32Array|null} [sh] - the array to receive spherical harmonics data
* @returns {SplatCompressedIterator} - The iterator
*/
createIter(p, r, s, c) {
return new SplatCompressedIterator(this, p, r, s, c);
createIter(p, r, s, c, sh) {
return new SplatCompressedIterator(this, p, r, s, c, sh);
}

/**
Expand All @@ -122,29 +144,25 @@ class GSplatCompressedData {
* @returns {boolean} - Whether the calculation was successful.
*/
calcAabb(result) {
let mx, my, mz, Mx, My, Mz;

// fast bounds calc using chunk data
const numChunks = Math.ceil(this.numSplats / 256);

const chunkData = this.chunkData;
const { chunkData, numChunks, chunkSize } = this;

let s = Math.exp(Math.max(chunkData[9], chunkData[10], chunkData[11]));
mx = chunkData[0] - s;
my = chunkData[1] - s;
mz = chunkData[2] - s;
Mx = chunkData[3] + s;
My = chunkData[4] + s;
Mz = chunkData[5] + s;
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) {
s = Math.exp(Math.max(chunkData[i * 12 + 9], chunkData[i * 12 + 10], chunkData[i * 12 + 11]));
mx = Math.min(mx, chunkData[i * 12 + 0] - s);
my = Math.min(my, chunkData[i * 12 + 1] - s);
mz = Math.min(mz, chunkData[i * 12 + 2] - s);
Mx = Math.max(Mx, chunkData[i * 12 + 3] + s);
My = Math.max(My, chunkData[i * 12 + 4] + s);
Mz = Math.max(Mz, chunkData[i * 12 + 5] + s);
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);
Expand All @@ -157,20 +175,18 @@ class GSplatCompressedData {
* @param {Float32Array} result - Array containing the centers.
*/
getCenters(result) {
const chunkData = this.chunkData;
const vertexData = this.vertexData;

const numChunks = Math.ceil(this.numSplats / 256);
const { vertexData, chunkData, numChunks, chunkSize } = this;

let mx, my, mz, Mx, My, Mz;

for (let c = 0; c < numChunks; ++c) {
mx = chunkData[c * 12 + 0];
my = chunkData[c * 12 + 1];
mz = chunkData[c * 12 + 2];
Mx = chunkData[c * 12 + 3];
My = chunkData[c * 12 + 4];
Mz = chunkData[c * 12 + 5];
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) {
Expand All @@ -189,17 +205,17 @@ class GSplatCompressedData {
* @param {Vec3} result - The result.
*/
calcFocalPoint(result) {
const chunkData = this.chunkData;
const numChunks = Math.ceil(this.numSplats / 256);
const { chunkData, numChunks, chunkSize } = this;

result.x = 0;
result.y = 0;
result.z = 0;

for (let i = 0; i < numChunks; ++i) {
result.x += chunkData[i * 12 + 0] + chunkData[i * 12 + 3];
result.y += chunkData[i * 12 + 1] + chunkData[i * 12 + 4];
result.z += chunkData[i * 12 + 2] + chunkData[i * 12 + 5];
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);
}
Expand All @@ -208,10 +224,38 @@ class GSplatCompressedData {
return true;
}

get numChunks() {
return Math.ceil(this.numSplats / 256);
}

get chunkSize() {
return this.chunkData.length / this.numChunks;
}

get shBands() {
const sizes = {
3: 1,
8: 2,
15: 3
};
return sizes[this.shData?.length / this.numSplats / 3] ?? 0;
}

// decompress into GSplatData
decompress() {
const members = ['x', 'y', 'z', 'f_dc_0', 'f_dc_1', 'f_dc_2', 'opacity', 'rot_0', 'rot_1', 'rot_2', 'rot_3', 'scale_0', 'scale_1', 'scale_2'];

const { shBands } = this;

// allocate spherical harmonics data
if (shBands > 0) {
const shMembers = [];
for (let i = 0; i < 45; ++i) {
shMembers.push(`f_rest_${i}`);
}
members.splice(members.indexOf('f_dc_0') + 1, 0, ...shMembers);
}

// allocate uncompressed data
const data = {};
members.forEach((name) => {
Expand All @@ -222,8 +266,9 @@ class GSplatCompressedData {
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);
const iter = this.createIter(p, r, s, c, sh);

for (let i = 0; i < this.numSplats; ++i) {
iter.read(i);
Expand All @@ -246,6 +291,12 @@ class GSplatCompressedData {
data.f_dc_2[i] = (c.z - 0.5) / SH_C0;
// convert opacity to log sigmoid taking into account infinities at 0 and 1
data.opacity[i] = (c.w <= 0) ? -40 : (c.w >= 1) ? 40 : -Math.log(1 / c.w - 1);

if (sh) {
for (let c = 0; c < 45; ++c) {
data[`f_rest_${c}`][i] = sh[c];
}
}
}

return new GSplatData([{
Expand Down
Loading

0 comments on commit 66cbb9f

Please sign in to comment.