Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Support compressed spherical harmonics #7143

Merged
merged 10 commits into from
Nov 27, 2024
Merged

Support compressed spherical harmonics #7143

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
a60052b
first
slimbuck Nov 20, 2024
9c2dba3
support sh decompress
slimbuck Nov 20, 2024
5409f3c
small
slimbuck Nov 20, 2024
bd314ae
lint
slimbuck Nov 20, 2024
a8baa1d
small type fix
slimbuck Nov 20, 2024
d92c95a
changes based on feedback
slimbuck Nov 22, 2024
1358370
reworking
slimbuck Nov 26, 2024
346421d
Merge branch 'main' into sh-dev-2
slimbuck Nov 26, 2024
acb9624
small
slimbuck Nov 26, 2024
77f6514
small
slimbuck Nov 26, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
104 changes: 45 additions & 59 deletions src/framework/parsers/ply.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -205,18 +205,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 @@ -230,10 +243,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 @@ -244,64 +255,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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,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
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,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 @@ -48,11 +48,11 @@ class SplatCompressedIterator {

const lerp = (a, b, t) => 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 @@ -74,6 +74,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 @@ -83,11 +96,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 @@ -102,17 +117,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 @@ -123,29 +145,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 @@ -158,20 +176,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 @@ -190,17 +206,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 @@ -209,6 +225,23 @@ 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 = [
Expand All @@ -218,6 +251,17 @@ class GSplatCompressedData {
'rot_0', 'rot_1', 'rot_2', 'rot_3'
];

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 @@ -228,8 +272,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 @@ -252,6 +297,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