Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,10 @@ class GaussianSplattingMaterialDefines extends MaterialDefines {
public IS_COMPOUND = false;
/** Defines the maximum number of parts (computed from engine caps at runtime) */
public MAX_PART_COUNT = GaussianSplattingMaxPartCount;
/** Defines whether SOG raw-texture in-shader dequantization is enabled */
public USE_SOG = false;
/** Defines whether SOG v2 (codebook) dequantization is enabled */
public USE_SOG_V2 = false;

/**
* Constructor of the defines.
Expand Down Expand Up @@ -184,6 +188,10 @@ export class GaussianSplattingMaterial extends PushMaterial {
"shTexture3",
"shTexture4",
"partIndicesTexture",
"sogQuatsTexture",
"sogShNCentroidsTexture",
"sogShNLabelsTexture",
"sogCodebookTexture",
];
protected static _UniformBuffers = ["Scene", "Mesh"];
protected static _VoxelUniforms = [
Expand Down Expand Up @@ -216,6 +224,15 @@ export class GaussianSplattingMaterial extends PushMaterial {
"depthValues",
"partWorld",
"partVisibility",
"sogMeansMin",
"sogMeansMax",
"sogScalesMin",
"sogScalesMax",
"sogSh0Min",
"sogSh0Max",
"sogShnMin",
"sogShnMax",
"sogShCoeffCount",
];
private _sourceMesh: GaussianSplattingMesh | null = null;
/**
Expand Down Expand Up @@ -304,6 +321,8 @@ export class GaussianSplattingMaterial extends PushMaterial {

defines["IS_COMPOUND"] = gsMesh.isCompound;
defines["MAX_PART_COUNT"] = GetGaussianSplattingMaxPartCount(engine);
defines["USE_SOG"] = gsMesh.useSog;
defines["USE_SOG_V2"] = gsMesh.useSog && gsMesh.sogParams?.version === 2;

// Compensation
const splatMaterial = gsMesh.material as GaussianSplattingMaterial;
Expand Down Expand Up @@ -462,7 +481,9 @@ export class GaussianSplattingMaterial extends PushMaterial {
effect.setTexture("centersTexture", gsMesh.centersTexture);
effect.setTexture("colorsTexture", gsMesh.colorsTexture);

if (gsMesh.shTextures) {
if (gsMesh.useSog) {
GaussianSplattingMaterial._BindSogUniforms(gsMesh, effect);
} else if (gsMesh.shTextures) {
for (let i = 0; i < gsMesh.shTextures.length; i++) {
effect.setTexture(`shTexture${i}`, gsMesh.shTextures[i]);
}
Expand All @@ -472,6 +493,38 @@ export class GaussianSplattingMaterial extends PushMaterial {
gsMesh.bindExtraEffectUniforms(effect);
}
}

/**
* Bind SOG dequantization uniforms + raw textures.
* @internal
*/
protected static _BindSogUniforms(gsMesh: GaussianSplattingMesh, effect: Effect): void {
const p = gsMesh.sogParams;
if (!p) {
return;
}
effect.setTexture("sogQuatsTexture", gsMesh.rotationsATexture);
if (gsMesh.shTextures && gsMesh.shTextures.length >= 2) {
effect.setTexture("sogShNCentroidsTexture", gsMesh.shTextures[0]);
effect.setTexture("sogShNLabelsTexture", gsMesh.shTextures[1]);
}
if (p.codebookTexture) {
effect.setTexture("sogCodebookTexture", p.codebookTexture);
}
effect.setFloat3("sogMeansMin", p.meansMin[0], p.meansMin[1], p.meansMin[2]);
effect.setFloat3("sogMeansMax", p.meansMax[0], p.meansMax[1], p.meansMax[2]);
if (p.scalesMin && p.scalesMax) {
effect.setFloat3("sogScalesMin", p.scalesMin[0], p.scalesMin[1], p.scalesMin[2]);
effect.setFloat3("sogScalesMax", p.scalesMax[0], p.scalesMax[1], p.scalesMax[2]);
}
if (p.sh0Min && p.sh0Max) {
effect.setFloat4("sogSh0Min", p.sh0Min[0], p.sh0Min[1], p.sh0Min[2], p.sh0Min[3]);
effect.setFloat4("sogSh0Max", p.sh0Max[0], p.sh0Max[1], p.sh0Max[2], p.sh0Max[3]);
}
effect.setFloat("sogShnMin", p.shnMin ?? 0);
effect.setFloat("sogShnMax", p.shnMax ?? 0);
effect.setFloat("sogShCoeffCount", p.shCoeffCount ?? 0);
}
/**
* Binds the submesh to this material by preparing the effect and shader to draw
* @param world defines the world transformation matrix
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,31 @@ import { type INative } from "core/Engines/Native/nativeInterfaces";
// eslint-disable-next-line @typescript-eslint/naming-convention
declare const _native: INative;

/** Internal mirror of ISogTexturePack (defined in loaders) — avoids a circular import. */
interface ISogPackInternal {
version: 1 | 2;
splatCount: number;
shDegree: number;
meansTextureL: BaseTexture;
meansTextureU: BaseTexture;
scalesTexture: BaseTexture;
quatsTexture: BaseTexture;
sh0Texture: BaseTexture;
shCentroidsTexture?: BaseTexture;
shLabelsTexture?: BaseTexture;
codebookTexture?: BaseTexture;
meansMin: [number, number, number];
meansMax: [number, number, number];
scalesMin?: [number, number, number];
scalesMax?: [number, number, number];
sh0Min?: [number, number, number, number];
sh0Max?: [number, number, number, number];
shnMin?: number;
shnMax?: number;
shCoeffCount: number;
positions: Float32Array;
}

const IsNative = typeof _native !== "undefined";
const Native = IsNative ? _native : null;
interface IDelayedTextureUpdate {
Expand Down Expand Up @@ -394,6 +419,8 @@ export class GaussianSplattingMeshBase extends Mesh {

private _delayedTextureUpdate: Nullable<IDelayedTextureUpdate> = null;
protected _useRGBACovariants = false;
protected _useSog = false;
protected _sogParams: Nullable<ISogPackInternal> = null;
private _material: Nullable<Material> = null;

private _tmpCovariances = [0, 0, 0, 0, 0, 0];
Expand Down Expand Up @@ -632,6 +659,107 @@ export class GaussianSplattingMeshBase extends Mesh {
return this._shTextures;
}

/**
* True when this mesh holds raw SOG webp textures (dequantized in-shader) rather than the
* pre-decoded covariance/center/color textures produced by the standard splat loader.
*/
public get useSog(): boolean {
return this._useSog;
}

/**
* SOG dequantization parameters paired with the raw textures.
* Set by the splat loader when `useSogTextures: true`. Null otherwise.
*/
public get sogParams(): Nullable<ISogPackInternal> {
return this._sogParams;
}

/**
* Install a set of raw SOG webp textures and bind the mesh to the in-shader dequantization path.
* @param pack SOG texture pack produced by ParseSogMetaAsTextures.
* @internal
*/
public setSogTextureData(pack: ISogPackInternal): void {
this._useSog = true;
this._sogParams?.codebookTexture?.dispose();
this._sogParams = pack;
Comment thread
CedricGuillemet marked this conversation as resolved.
this._vertexCount = pack.splatCount;
this._shDegree = pack.shDegree ?? 0;
Comment thread
CedricGuillemet marked this conversation as resolved.
this._maxShDegree = this._shDegree;

// Stride-4 (xyz + 1) — required by the depth-sort worker and the centers texture path.
this._splatPositions = pack.positions;

// Reuse existing texture slots for SOG textures (the shader, under USE_SOG, samples them as RGBA8).
this._covariancesATexture?.dispose();
this._covariancesBTexture?.dispose();
this._centersTexture?.dispose();
this._colorsTexture?.dispose();
this._rotationsATexture?.dispose();
if (this._shTextures) {
for (const t of this._shTextures) {
t.dispose();
}
}

this._centersTexture = pack.meansTextureL;
this._covariancesATexture = pack.meansTextureU;
this._covariancesBTexture = pack.scalesTexture;
this._rotationsATexture = pack.quatsTexture;
this._colorsTexture = pack.sh0Texture;

const shTextures: BaseTexture[] = [];
if (pack.shCentroidsTexture) {
shTextures.push(pack.shCentroidsTexture);
}
if (pack.shLabelsTexture) {
shTextures.push(pack.shLabelsTexture);
}
this._shTextures = shTextures.length ? shTextures : null;

// Force pipeline rebuild so the USE_SOG define and extra samplers are picked up.
this._material?.resetDrawCache();

const size = pack.meansTextureL.getSize();
this._textureSize.x = size.width;
this._textureSize.y = size.height;

this._updateSplatIndexBuffer(this._vertexCount);
this._instantiateWorker();

// Compute bounds from the CPU-decoded positions (stride-4) so the mesh is not frustum-culled.
const positions = pack.positions as Float32Array;
const minimum = new Vector3(Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY, Number.POSITIVE_INFINITY);
const maximum = new Vector3(Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY, Number.NEGATIVE_INFINITY);
for (let i = 0; i < this._vertexCount; i++) {
const x = positions[i * 4 + 0];
const y = positions[i * 4 + 1];
const z = positions[i * 4 + 2];
if (x < minimum.x) {
minimum.x = x;
}
if (y < minimum.y) {
minimum.y = y;
}
if (z < minimum.z) {
minimum.z = z;
}
if (x > maximum.x) {
maximum.x = x;
}
if (y > maximum.y) {
maximum.y = y;
}
if (z > maximum.z) {
maximum.z = z;
}
}
this.getBoundingInfo().reConstruct(minimum, maximum, this.getWorldMatrix());
this.setEnabled(true);
this._sortIsDirty = true;
}

/**
* Gets the kernel size
* Documentation and mathematical explanations here:
Expand Down Expand Up @@ -1855,6 +1983,7 @@ export class GaussianSplattingMeshBase extends Mesh {
this._rotationsATexture?.dispose();
this._rotationsBTexture?.dispose();
this._rotationScaleTexture?.dispose();
this._sogParams?.codebookTexture?.dispose();
this._rotationsATexture = null;
this._rotationsBTexture = null;
this._rotationScaleTexture = null;
Expand Down
Loading