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.ply spherical harmonics #7127

Closed
wants to merge 6 commits into from
+494 −117
Closed

Support compressed.ply spherical harmonics #7127

Show file tree
Hide file tree
Changes from 5 commits
Commits
Show all changes
6 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
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
108 changes: 54 additions & 54 deletions src/framework/parsers/ply.js
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -212,11 +212,30 @@ const isCompressedPly = (elements) => {
'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 band1Properties = ['0', '1', '2'].map(x => `coeff_${x}`);
const band2Properties = ['3', '4', '5', '6', '7'].map(x => `coeff_${x}`);
const band3Properties = ['8', '9', '10', '11', '12', '13', '14'].map(x => `coeff_${x}`);
const indicesProperties = ['0', '1', '2', '3'].map(x => `packed_sh_${x}`);

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_band_1' &&
elements[2].properties.every((p, i) => p.name === band1Properties[i] && p.type === 'ushort') &&
elements[3].name === 'sh_band_2' &&
elements[3].properties.every((p, i) => p.name === band2Properties[i] && p.type === 'ushort') &&
elements[4].name === 'sh_band_3' &&
elements[4].properties.every((p, i) => p.name === band3Properties[i] && p.type === 'ushort') &&
elements[5].name === 'vertex_sh' &&
elements[5].properties.every((p, i) => p.name === indicesProperties[i] && p.type === 'uint');
};

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

const isFloatPly = (elements) => {
Expand All @@ -230,10 +249,7 @@ const readCompressedPly = async (streamBuf, elements, littleEndian) => {
const result = new GSplatCompressedData();

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

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 +260,48 @@ const readCompressedPly = async (streamBuf, elements, littleEndian) => {
};

// allocate result
result.numSplats = elements[1].count;
result.numSplats = numVertices;
result.chunkData = new Float32Array(evalStorageSize(numChunks) * 12);
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 length bytes of data into buffer
const read = async (buffer, length) => {
const target = new Uint8Array(buffer);
let cursor = 0;

// read the next chunk of data
const toRead = Math.min(numChunks - chunks, Math.floor(streamBuf.remaining / chunkSize));
while (cursor < length) {
while (streamBuf.remaining === 0) {
/* eslint-disable no-await-in-loop */
await streamBuf.read();
}

const dstOffset = chunks * 12;
const srcOffset = streamBuf.head / 4;
for (let i = 0; i < toRead * 12; ++i) {
uint32ChunkData[dstOffset + i] = uint32StreamData[srcOffset + i];
const toCopy = Math.min(length - cursor, streamBuf.remaining);
const src = streamBuf.data;
for (let i = 0; i < toCopy; ++i) {
target[cursor++] = src[streamBuf.head++];
}
}
};

streamBuf.head += toRead * chunkSize;
chunks += toRead;
}
// read chunk data
await read(result.chunkData.buffer, numChunks * 12 * 4);

// read vertices
let vertices = 0;
while (vertices < numVertices) {
while (streamBuf.remaining < vertexSize) {
/* eslint-disable no-await-in-loop */
await streamBuf.read();
}
// read packed vertices
await read(result.vertexData.buffer, numVertices * 4 * 4);

// 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 sh data
if (elements.length === 6) {
result.band1Data = new Uint16Array(elements[2].count * 3);
await read(result.band1Data.buffer, result.band1Data.byteLength);

// read the next chunk of data
const toRead = Math.min(numVertices - vertices, Math.floor(streamBuf.remaining / vertexSize));
result.band2Data = new Uint16Array(elements[3].count * 5);
await read(result.band2Data.buffer, result.band2Data.byteLength);

const dstOffset = vertices * 4;
const srcOffset = streamBuf.head / 4;
for (let i = 0; i < toRead * 4; ++i) {
uint32VertexData[dstOffset + i] = uint32StreamData[srcOffset + i];
}
result.band3Data = new Uint16Array(elements[4].count * 7);
await read(result.band3Data.buffer, result.band3Data.byteLength);

streamBuf.head += toRead * vertexSize;
vertices += toRead;
result.packedSHData = new Uint32Array(evalStorageSize(numVertices) * 4);
await read(result.packedSHData.buffer, numVertices * 4 * 4);
}

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]
slimbuck marked this conversation as resolved.
Show resolved Hide resolved
* - 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: 129 additions & 6 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 @@ -9,9 +9,34 @@ import { GSplatData } from './gsplat-data.js';

const SH_C0 = 0.28209479177387814;

// https://fgiesen.wordpress.com/2012/03/28/half-to-float-done-quic/
slimbuck marked this conversation as resolved.
Show resolved Hide resolved
const halfToFloat = (() => {
const f32 = new Float32Array(1);
const magicF = new Float32Array(1);
const wasInfNanF = new Float32Array(1);

const u32 = new Uint32Array(f32.buffer);
const magicU = new Uint32Array(magicF.buffer);
const wasInfNanU = new Uint32Array(wasInfNanF.buffer);

magicU[0] = (254 - 15) << 23;
wasInfNanU[0] = (127 + 16) << 23;

return (h) => {
u32[0] = (h & 0x7fff) << 13;
f32[0] *= magicF[0];

if (f32[0] >= wasInfNanF[0]) {
u32[0] |= 255 << 23;
}
u32[0] |= (h & 0x8000) << 16; // sign bit
return f32[0];
};
})();

// 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,8 +73,7 @@ class SplatCompressedIterator {

const lerp = (a, b, t) => a * (1 - t) + b * t;

const chunkData = gsplatData.chunkData;
const vertexData = gsplatData.vertexData;
const { chunkData, vertexData, band1Data, band2Data, band3Data, packedSHData } = gsplatData;

this.read = (i) => {
const ci = Math.floor(i / 256) * 12;
Expand All @@ -75,6 +99,54 @@ class SplatCompressedIterator {
if (c) {
unpack8888(c, vertexData[i * 4 + 3]);
}

if (sh) {
const bits0 = packedSHData[i * 4 + 0];
const bits1 = packedSHData[i * 4 + 1];
const bits2 = packedSHData[i * 4 + 2];
const bits3 = packedSHData[i * 4 + 3];

const b1 = [
bits0 >>> 21,
(bits0 >>> 10) & 0x7ff,
bits0 & 0x3ff
];

const b2 = [
bits1 >>> 17,
(bits1 >>> 2) & 0x7fff,
((bits1 << 13) | (bits2 >>> 19)) & 0x7fff
];

const b3 = [
(bits2 >>> 2) & 0x1ffff,
((bits2 << 15) | (bits3 >>> 17)) & 0x1ffff,
(bits3 & 0x1ffff)
];

for (let i = 0; i < 3; ++i) {
const i1 = b1[i] * 3;
sh[i * 15 + 0] = halfToFloat(band1Data[i1 + 0]);
sh[i * 15 + 1] = halfToFloat(band1Data[i1 + 1]);
sh[i * 15 + 2] = halfToFloat(band1Data[i1 + 2]);

const i2 = b2[i] * 5;
sh[i * 15 + 3] = halfToFloat(band2Data[i2 + 0]);
sh[i * 15 + 4] = halfToFloat(band2Data[i2 + 1]);
sh[i * 15 + 5] = halfToFloat(band2Data[i2 + 2]);
sh[i * 15 + 6] = halfToFloat(band2Data[i2 + 3]);
sh[i * 15 + 7] = halfToFloat(band2Data[i2 + 4]);

const i3 = b3[i] * 7;
sh[i * 15 + 8] = halfToFloat(band3Data[i3 + 0]);
sh[i * 15 + 9] = halfToFloat(band3Data[i3 + 1]);
sh[i * 15 + 10] = halfToFloat(band3Data[i3 + 2]);
sh[i * 15 + 11] = halfToFloat(band3Data[i3 + 3]);
sh[i * 15 + 12] = halfToFloat(band3Data[i3 + 4]);
sh[i * 15 + 13] = halfToFloat(band3Data[i3 + 5]);
sh[i * 15 + 14] = halfToFloat(band3Data[i3 + 6]);
}
}
};
}
}
Expand Down Expand Up @@ -102,17 +174,48 @@ class GSplatCompressedData {
*/
vertexData;

// optional spherical harmonics data

/**
* Contains 3 half values per band 1 palette entry
* @type {Uint16Array}
*/
band1Data;

/**
* Contains 5 half values per band 2 palette entry
* @type {Uint16Array}
*/
band2Data;

/**
* Contains 7 half values per band 3 palette entry
* @type {Uint16Array}
*/
band3Data;

/**
* Contains 4 uint32 per vertex with packed bits referencing the SH palette:
* packed_sh_0
* packed_sh_1
* packed_sh_2
* packed_sh_3
* @type {Uint32Array}
*/
packedSHData;

/**
* 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 Down Expand Up @@ -209,6 +312,10 @@ class GSplatCompressedData {
return true;
}

get hasSHData() {
return !!this.packedSHData;
}

// decompress into GSplatData
decompress() {
const members = [
Expand All @@ -218,6 +325,15 @@ class GSplatCompressedData {
'rot_0', 'rot_1', 'rot_2', 'rot_3'
];

// allocate spherical harmonics data
if (this.hasSHData) {
const shMembers = [];
for (let i = 0; i < 45; ++i) {
shMembers.push(`f_rest_${i}`);
}
members.splice(9, 0, ...shMembers);
}

// allocate uncompressed data
const data = {};
members.forEach((name) => {
Expand All @@ -228,8 +344,9 @@ class GSplatCompressedData {
const r = new Quat();
const s = new Vec3();
const c = new Vec4();
const sh = this.hasSHData ? 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 +369,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 (this.hasSHData) {
for (let c = 0; c < 45; ++c) {
data[`f_rest_${c}`][i] = sh[c];
}
}
}

return new GSplatData([{
Expand Down
Loading