From 25deee0335412ea27753f315b18039ab32bb90b0 Mon Sep 17 00:00:00 2001 From: ooco <80620808+bigbigbig2@users.noreply.github.com> Date: Sat, 19 Jul 2025 14:18:41 +0800 Subject: [PATCH] Add the ability to build an octree and use the octree to accelerate ray detection example --- dist/types/SplatMesh.d.ts | 39 ++ dist/types/SplatTree.d.ts | 53 +++ examples.html | 1 + examples/octree-raycasting/index.html | 444 ++++++++++++++++++++++ index.html | 1 + src/SplatMesh.ts | 145 ++++++- src/SplatTree.ts | 527 ++++++++++++++++++++++++++ 7 files changed, 1209 insertions(+), 1 deletion(-) create mode 100644 dist/types/SplatTree.d.ts create mode 100644 examples/octree-raycasting/index.html create mode 100644 src/SplatTree.ts diff --git a/dist/types/SplatMesh.d.ts b/dist/types/SplatMesh.d.ts index 6503ce7..914edb7 100644 --- a/dist/types/SplatMesh.d.ts +++ b/dist/types/SplatMesh.d.ts @@ -3,6 +3,7 @@ import { RgbaArray } from './RgbaArray'; import { SplatEdit } from './SplatEdit'; import { GsplatModifier, SplatGenerator, SplatTransformer } from './SplatGenerator'; import { SplatFileType } from './SplatLoader'; +import { SplatTree } from './SplatTree'; import { SplatSkinning } from './SplatSkinning'; import { DynoFloat, DynoUsampler2DArray, DynoVal, DynoVec4, Gsplat } from './dyno'; import * as THREE from "three"; @@ -56,6 +57,9 @@ export declare class SplatMesh extends SplatGenerator { private rgbaDisplaceEdits; splatRgba: RgbaArray | null; maxSh: number; + splatTree: SplatTree | null; + baseSplatTree: SplatTree | null; + disposed: boolean; constructor(options?: SplatMeshOptions); asyncInitialize(options: SplatMeshOptions): Promise; static staticInitialized: Promise; @@ -64,6 +68,41 @@ export declare class SplatMesh extends SplatGenerator { static staticInitialize(): Promise; pushSplat(center: THREE.Vector3, scales: THREE.Vector3, quaternion: THREE.Quaternion, opacity: number, color: THREE.Color): void; forEachSplat(callback: (index: number, center: THREE.Vector3, scales: THREE.Vector3, quaternion: THREE.Quaternion, opacity: number, color: THREE.Color) => void): void; + /** + * Retrieves the color and opacity of a specific splat + * @param index - The index of the splat to retrieve color from + * @param out - Output vector to store the color (RGBA format) + * @returns The output vector containing the splat's color and opacity + */ + getSplatColor(index: number, out: THREE.Vector4): THREE.Vector4; + /** + * Retrieves the scale and rotation (quaternion) of a specific splat + * @param index - The index of the splat to retrieve scale and rotation from + * @param out - Output vector to store the scale (x, y, z components) + * @param out2 - Output quaternion to store the rotation + * @returns The output vector containing the splat's scale + */ + getSplatScaleAndRotation(index: number, out: THREE.Vector3, out2: THREE.Quaternion): THREE.Vector3; + /** + * Gets the total number of splats in the mesh + * @returns The total count of splats + */ + getSplatCount(): number; + /** + * Gets the scene index for a specific splat + * Currently returns 0 as a default implementation + * @param index - The index of the splat + * @returns The scene index (always 0 in this implementation) + */ + getSceneIndexForSplat(index: number): number; + /** + * Gets the current scene object + * @returns The current SplatMesh instance as the scene + */ + getScene(): this; + getSplatCenter(index: number, out: THREE.Vector3): THREE.Vector3; + buildSplatTree(minAlphas?: number[], onSplatTreeIndexesUpload?: Function, onSplatTreeConstruction?: Function): Promise; + disposeSplatTree(): void; dispose(): void; constructGenerator(context: SplatMeshContext): void; updateGenerator(): void; diff --git a/dist/types/SplatTree.d.ts b/dist/types/SplatTree.d.ts new file mode 100644 index 0000000..b51caf4 --- /dev/null +++ b/dist/types/SplatTree.d.ts @@ -0,0 +1,53 @@ +import { SplatMesh } from './SplatMesh'; +import * as THREE from 'three'; +declare class SplatTreeNode { + static idGen: number; + min: THREE.Vector3; + max: THREE.Vector3; + boundingBox: THREE.Box3; + center: THREE.Vector3; + depth: number; + children: SplatTreeNode[]; + data: { + indexes: number[]; + }; + id: number; + constructor(min: THREE.Vector3, max: THREE.Vector3, depth: number, id: number); +} +declare class SplatSubTree { + maxDepth: number; + maxCentersPerNode: number; + sceneDimensions: THREE.Vector3; + sceneMin: THREE.Vector3; + sceneMax: THREE.Vector3; + rootNode: SplatTreeNode | null; + nodesWithIndexes: SplatTreeNode[]; + splatMesh: SplatMesh | null; + constructor(maxDepth: number, maxCentersPerNode: number); + static convertWorkerSubTreeNode(workerSubTreeNode: any): SplatTreeNode; + static convertWorkerSubTree(workerSubTree: any, splatMesh: SplatMesh): SplatSubTree; +} +export declare class SplatTree { + maxDepth: number; + maxCentersPerNode: number; + subTrees: SplatSubTree[]; + splatMesh: SplatMesh | null; + disposed: boolean; + splatTreeWorker: Worker | null; + constructor(maxDepth: number, maxCentersPerNode: number); + dispose(): void; + diposeSplatTreeWorker(): void; + /** + * Build SplatTree (octree) from a SplatMesh instance. + * + * @param {SplatMesh} splatMesh SplatMesh instance to build octree from + * @param {function} filterFunc Optional, filter out unwanted splats (points), return true to keep + * @param {function} onIndexesUpload Callback when uploading splat centers to worker (start/end) + * @param {function} onSplatTreeConstruction Callback when worker is building local splat tree (start/end) + * @return {Promise} Promise that resolves when octree building is complete + */ + processSplatMesh(splatMesh: SplatMesh, filterFunc?: (index: number) => boolean, onIndexesUpload?: (isUploading: boolean) => void, onSplatTreeConstruction?: (isBuilding: boolean) => void): Promise; + countLeaves(): number; + visitLeaves(visitFunc: (node: SplatTreeNode) => void): void; +} +export {}; diff --git a/examples.html b/examples.html index d965b02..d839511 100644 --- a/examples.html +++ b/examples.html @@ -430,6 +430,7 @@ Multiple Viewpoints Procedural Splats Raycasting + Octree Raycasting Dynamic Lighting Particle Animation Particle Simulation diff --git a/examples/octree-raycasting/index.html b/examples/octree-raycasting/index.html new file mode 100644 index 0000000..a53d65c --- /dev/null +++ b/examples/octree-raycasting/index.html @@ -0,0 +1,444 @@ + + + + + + + + Spark • Multiple Splats + + + + + + + + + diff --git a/index.html b/index.html index b840051..ef8511b 100644 --- a/index.html +++ b/index.html @@ -139,6 +139,7 @@

Examples

  • Multiple Viewpoints
  • Procedural Splats
  • Raycasting
    • +
  • Octree Raycasting
  • Dynamic Lighting
  • Particle Animation
  • Particle Simulation
    • diff --git a/src/SplatMesh.ts b/src/SplatMesh.ts index e34579f..edac139 100644 --- a/src/SplatMesh.ts +++ b/src/SplatMesh.ts @@ -10,6 +10,7 @@ import { SplatTransformer, } from "./SplatGenerator"; import type { SplatFileType } from "./SplatLoader"; +import { SplatTree } from "./SplatTree"; import type { SplatSkinning } from "./SplatSkinning"; import { DynoFloat, @@ -150,7 +151,10 @@ export class SplatMesh extends SplatGenerator { // Maximum Spherical Harmonics level to use. Call updateGenerator() // after changing. (default: 3) - maxSh = 3; + maxSh = 3; // maximum spherical harmonics level + splatTree: SplatTree | null = null; // splat tree + baseSplatTree: SplatTree | null = null; // base splat tree + disposed: boolean = false; // whether the mesh is disposed constructor(options: SplatMeshOptions = {}) { const transform = new SplatTransformer(); @@ -293,10 +297,149 @@ export class SplatMesh extends SplatGenerator { this.packedSplats.forEachSplat(callback); } + /** + * Retrieves the color and opacity of a specific splat + * @param index - The index of the splat to retrieve color from + * @param out - Output vector to store the color (RGBA format) + * @returns The output vector containing the splat's color and opacity + */ + getSplatColor(index: number, out: THREE.Vector4) { + const splat = this.packedSplats.getSplat(index); + out.set(splat.color.r, splat.color.g, splat.color.b, splat.opacity); + return out; + } + + /** + * Retrieves the scale and rotation (quaternion) of a specific splat + * @param index - The index of the splat to retrieve scale and rotation from + * @param out - Output vector to store the scale (x, y, z components) + * @param out2 - Output quaternion to store the rotation + * @returns The output vector containing the splat's scale + */ + getSplatScaleAndRotation(index: number, out: THREE.Vector3, out2: THREE.Quaternion) { + const splat = this.packedSplats.getSplat(index); + out.set(splat.scales.x, splat.scales.y, splat.scales.z); + out2.set(splat.quaternion.x, splat.quaternion.y, splat.quaternion.z, splat.quaternion.w); + return out; + } + + /** + * Gets the total number of splats in the mesh + * @returns The total count of splats + */ + getSplatCount() { + return this.packedSplats.numSplats; + } + + /** + * Gets the scene index for a specific splat + * Currently returns 0 as a default implementation + * @param index - The index of the splat + * @returns The scene index (always 0 in this implementation) + */ + getSceneIndexForSplat(index: number) { + return 0; + } + + /** + * Gets the current scene object + * @returns The current SplatMesh instance as the scene + */ + getScene() { + return this; + } + + + // get splat center + getSplatCenter(index: number, out: THREE.Vector3) { + const splat = this.packedSplats.getSplat(index); + out.set(splat.center.x, splat.center.y, splat.center.z); + // out.applyMatrix4(this.matrixWorld); // add world transform + return out; + } + + + // build octree for splat mesh + // minAlphas: the minimum alpha value for each scene + // onSplatTreeIndexesUpload: callback function to upload splat tree indexes + // onSplatTreeConstruction: callback function to construct splat tree + // If you do scaling, rotation, or translation after building, you need to rebuild + buildSplatTree( + minAlphas: number[] = [], + onSplatTreeIndexesUpload?: Function, + onSplatTreeConstruction?: Function + ): Promise { + return new Promise((resolve) => { + this.disposeSplatTree(); + this.baseSplatTree = new SplatTree(8, 2000); + const buildStartTime = performance.now(); + const splatColor = new THREE.Vector4(); + this.baseSplatTree.processSplatMesh( + this, + (splatIndex: number) => { + this.getSplatColor(splatIndex, splatColor); + const sceneIndex = this.getSceneIndexForSplat(splatIndex); + const minAlpha = minAlphas[sceneIndex] || 0.1; + return splatColor.w >= minAlpha; + }, + (isUploading: boolean) => { + onSplatTreeIndexesUpload?.(isUploading); + }, + (isBuilding: boolean) => { + onSplatTreeConstruction?.(isBuilding); + } + ).then(() => { + const buildTime = performance.now() - buildStartTime; + console.log('SplatTree build: ' + buildTime + ' ms'); + if (this.disposed) { + resolve(void 0); + } else { + this.splatTree = this.baseSplatTree; + this.baseSplatTree = null; + + let leavesWithVertices = 0; + let avgSplatCount = 0; + let maxSplatCount = 0; + let nodeCount = 0; + + this.splatTree?.visitLeaves((node: any) => { + const nodeSplatCount = node.data.indexes.length; + if (nodeSplatCount > 0) { + avgSplatCount += nodeSplatCount; + maxSplatCount = Math.max(maxSplatCount, nodeSplatCount); + nodeCount++; + leavesWithVertices++; + } + }); + console.log(`SplatTree leaves: ${this.splatTree?.countLeaves()}`); + console.log(`SplatTree leaves with splats:${leavesWithVertices}`); + avgSplatCount = avgSplatCount / nodeCount; + console.log(`Avg splat count per node: ${avgSplatCount}`); + console.log(`Total splat count: ${this.getSplatCount()}`); + resolve(void 0); + } + }); + }); + } + + + disposeSplatTree() { + if (this.splatTree) { + this.splatTree.dispose(); + this.splatTree = null; + } + if (this.baseSplatTree) { + this.baseSplatTree.dispose(); + this.baseSplatTree = null; + } + this.disposed = true; + } + // Call this when you are finished with the SplatMesh and want to free // any buffers it holds (via packedSplats). dispose() { this.packedSplats.dispose(); + this.disposeSplatTree(); } constructGenerator(context: SplatMeshContext) { diff --git a/src/SplatTree.ts b/src/SplatTree.ts new file mode 100644 index 0000000..f1330cc --- /dev/null +++ b/src/SplatTree.ts @@ -0,0 +1,527 @@ +import * as THREE from 'three' +import { SplatMesh } from './SplatMesh'; + + +// delayedExecute function +const delayedExecute = (func: Function, fast?: boolean) => { + return new Promise((resolve) => { + window.setTimeout(() => { + resolve(func ? func() : undefined); + }, fast ? 1 : 50); + }); +}; + + +// SplatTreeNode: Octree node +class SplatTreeNode { + + static idGen = 0; // node id generator + min: THREE.Vector3; // min point of the node + max: THREE.Vector3; // max point of the node + boundingBox: THREE.Box3; // bounding box of the node + center: THREE.Vector3; // center point of the node + depth: number; // depth of the node + children: SplatTreeNode[]; // children nodes + data: { indexes: number[] }; // data of the node + id: number; // id of the node + + constructor(min: THREE.Vector3, max: THREE.Vector3, depth: number, id: number) { + this.min = new THREE.Vector3().copy(min); + this.max = new THREE.Vector3().copy(max); + this.boundingBox = new THREE.Box3(this.min, this.max); + this.center = new THREE.Vector3().copy(this.max).sub(this.min).multiplyScalar(0.5).add(this.min); + this.depth = depth; + this.children = []; + this.data = { indexes: [] }; + this.id = id || SplatTreeNode.idGen++; + } + +} +// SplatSubTree: Octree sub tree, contains root node, parameters, and associated SplatMesh +class SplatSubTree { + maxDepth: number; // max depth of the sub tree + maxCentersPerNode: number; // max centers per node of the sub tree + sceneDimensions: THREE.Vector3; // dimensions of the scene + sceneMin: THREE.Vector3; // min point of the scene + sceneMax: THREE.Vector3; // max point of the scene + rootNode: SplatTreeNode | null; // root node of the sub tree + nodesWithIndexes: SplatTreeNode[]; // nodes with indexes + splatMesh: SplatMesh | null; // associated SplatMesh + + constructor(maxDepth: number, maxCentersPerNode: number) { + this.maxDepth = maxDepth; + this.maxCentersPerNode = maxCentersPerNode; + this.sceneDimensions = new THREE.Vector3(); + this.sceneMin = new THREE.Vector3(); + this.sceneMax = new THREE.Vector3(); + this.rootNode = null; + this.nodesWithIndexes = []; + this.splatMesh = null; + } + + // convert worker sub tree node to main thread node + static convertWorkerSubTreeNode(workerSubTreeNode: any) { + const minVector = new THREE.Vector3().fromArray(workerSubTreeNode.min); + const maxVector = new THREE.Vector3().fromArray(workerSubTreeNode.max); + const convertedNode = new SplatTreeNode(minVector, maxVector, workerSubTreeNode.depth, workerSubTreeNode.id); + if (workerSubTreeNode.data.indexes) { + convertedNode.data = { + 'indexes': [] + }; + for (let index of workerSubTreeNode.data.indexes) { + convertedNode.data.indexes.push(index); + } + } + if (workerSubTreeNode.children && Array.isArray(workerSubTreeNode.children)) { + for (let child of workerSubTreeNode.children) { + // type assertion to avoid type error + (convertedNode.children as SplatTreeNode[]).push(SplatSubTree.convertWorkerSubTreeNode(child)); + } + } + return convertedNode; + } + // convert worker sub tree object to main thread object + static convertWorkerSubTree(workerSubTree: any, splatMesh: SplatMesh) { + const convertedSubTree = new SplatSubTree(workerSubTree.maxDepth, workerSubTree.maxCentersPerNode); + convertedSubTree.sceneMin = new THREE.Vector3().fromArray(workerSubTree.sceneMin); + convertedSubTree.sceneMax = new THREE.Vector3().fromArray(workerSubTree.sceneMax); + + convertedSubTree.splatMesh = splatMesh; + convertedSubTree.rootNode = SplatSubTree.convertWorkerSubTreeNode(workerSubTree.rootNode); + + // collect all leaves (with indexes) + const visitLeavesFromNode = (node: SplatTreeNode, visitFunc: (node: SplatTreeNode) => void) => { + if (node.children.length === 0) visitFunc(node); + for (let child of node.children) { + visitLeavesFromNode(child, visitFunc); + } + }; + + convertedSubTree.nodesWithIndexes = []; + visitLeavesFromNode(convertedSubTree.rootNode, (node) => { + if (node.data && node.data.indexes && node.data.indexes.length > 0) { + convertedSubTree.nodesWithIndexes.push(node); + } + }); + + return convertedSubTree; + } +} +// createSplatTreeWorker: Octree building logic for worker thread (string injection worker) +function createSplatTreeWorker(self: Worker) { + + let WorkerSplatTreeNodeIDGen = 0; + + class WorkerBox3 { + min:number[] + max:number[] + + constructor(min: number[], max: number[]) { + this.min = [min[0], min[1], min[2]]; + this.max = [max[0], max[1], max[2]]; + } + + containsPoint(point: number[]) { + return point[0] >= this.min[0] && point[0] <= this.max[0] && + point[1] >= this.min[1] && point[1] <= this.max[1] && + point[2] >= this.min[2] && point[2] <= this.max[2]; + } + } + + class WorkerSplatSubTree { + maxDepth:number + maxCentersPerNode: number; + sceneDimensions: number[]; + sceneMin: number[]; + sceneMax: number[]; + rootNode!: WorkerSplatTreeNode; + addedIndexes: { [key: number]: boolean }; + nodesWithIndexes: WorkerSplatTreeNode[]; + splatMesh: SplatMesh | null; + disposed: boolean; + + constructor(maxDepth: number, maxCentersPerNode: number) { + this.maxDepth = maxDepth; + this.maxCentersPerNode = maxCentersPerNode; + this.sceneDimensions = []; + this.sceneMin = []; + this.sceneMax = []; + this.addedIndexes = {}; + this.nodesWithIndexes = []; + this.splatMesh = null; + this.disposed = false; + } + + } + + class WorkerSplatTreeNode { + min: number[]; + max: number[]; + center: number[]; + depth: number; + children: WorkerSplatTreeNode[]; + data: { indexes: number[] }; + id: number; + + constructor(min: number[], max: number[], depth: number, id?: number) { + this.min = [min[0], min[1], min[2]]; + this.max = [max[0], max[1], max[2]]; + this.center = [(max[0] - min[0]) * 0.5 + min[0], + (max[1] - min[1]) * 0.5 + min[1], + (max[2] - min[2]) * 0.5 + min[2]]; + this.depth = depth; + this.children = []; + this.data = { indexes: [] }; + this.id = id || WorkerSplatTreeNodeIDGen++; + } + + } + + function processSplatTreeNode(tree: WorkerSplatSubTree, node: WorkerSplatTreeNode, indexToCenter: number[], sceneCenters: Float32Array) { + const splatCount = node.data?.indexes?.length ?? 0; + + if (splatCount < tree.maxCentersPerNode || node.depth > tree.maxDepth) { + const newIndexes = []; + for (let i = 0; i < node.data.indexes.length; i++) { + if (!tree.addedIndexes[node.data.indexes[i]]) { + newIndexes.push(node.data.indexes[i]); + tree.addedIndexes[node.data.indexes[i]] = true; + } + } + node.data.indexes = newIndexes; + node.data.indexes.sort((a, b) => { + if (a > b) return 1; + else return -1; + }); + tree.nodesWithIndexes.push(node); + return; + } + + const nodeDimensions = [node.max[0] - node.min[0], + node.max[1] - node.min[1], + node.max[2] - node.min[2]]; + const halfDimensions = [nodeDimensions[0] * 0.5, + nodeDimensions[1] * 0.5, + nodeDimensions[2] * 0.5]; + const nodeCenter = [node.min[0] + halfDimensions[0], + node.min[1] + halfDimensions[1], + node.min[2] + halfDimensions[2]]; + + const childrenBounds = [ + // top section, clockwise from upper-left (looking from above, +Y) + new WorkerBox3([nodeCenter[0] - halfDimensions[0], nodeCenter[1], nodeCenter[2] - halfDimensions[2]], + [nodeCenter[0], nodeCenter[1] + halfDimensions[1], nodeCenter[2]]), + new WorkerBox3([nodeCenter[0], nodeCenter[1], nodeCenter[2] - halfDimensions[2]], + [nodeCenter[0] + halfDimensions[0], nodeCenter[1] + halfDimensions[1], nodeCenter[2]]), + new WorkerBox3([nodeCenter[0], nodeCenter[1], nodeCenter[2]], + [nodeCenter[0] + halfDimensions[0], nodeCenter[1] + halfDimensions[1], nodeCenter[2] + halfDimensions[2]]), + new WorkerBox3([nodeCenter[0] - halfDimensions[0], nodeCenter[1], nodeCenter[2]], + [nodeCenter[0], nodeCenter[1] + halfDimensions[1], nodeCenter[2] + halfDimensions[2]]), + + // bottom section, clockwise from lower-left (looking from above, +Y) + new WorkerBox3([nodeCenter[0] - halfDimensions[0], nodeCenter[1] - halfDimensions[1], nodeCenter[2] - halfDimensions[2]], + [nodeCenter[0], nodeCenter[1], nodeCenter[2]]), + new WorkerBox3([nodeCenter[0], nodeCenter[1] - halfDimensions[1], nodeCenter[2] - halfDimensions[2]], + [nodeCenter[0] + halfDimensions[0], nodeCenter[1], nodeCenter[2]]), + new WorkerBox3([nodeCenter[0], nodeCenter[1] - halfDimensions[1], nodeCenter[2]], + [nodeCenter[0] + halfDimensions[0], nodeCenter[1], nodeCenter[2] + halfDimensions[2]]), + new WorkerBox3([nodeCenter[0] - halfDimensions[0], nodeCenter[1] - halfDimensions[1], nodeCenter[2]], + [nodeCenter[0], nodeCenter[1], nodeCenter[2] + halfDimensions[2]]), + ]; + + const splatCounts: number[] = []; + const baseIndexes: number[][] = []; + for (let i = 0; i < childrenBounds.length; i++) { + splatCounts[i] = 0; + baseIndexes[i] = []; + } + + const center = [0, 0, 0]; + for (let i = 0; i < splatCount; i++) { + const splatGlobalIndex = node.data.indexes[i]; + const centerBase = indexToCenter[splatGlobalIndex]; + center[0] = sceneCenters[centerBase]; + center[1] = sceneCenters[centerBase + 1]; + center[2] = sceneCenters[centerBase + 2]; + for (let j = 0; j < childrenBounds.length; j++) { + if (childrenBounds[j].containsPoint(center)) { + splatCounts[j]++; + baseIndexes[j].push(splatGlobalIndex); + } + } + } + + + for (let i = 0; i < childrenBounds.length; i++) { + const childNode = new WorkerSplatTreeNode(childrenBounds[i].min, childrenBounds[i].max, node.depth + 1); + childNode.data = { + 'indexes': baseIndexes[i] + }; + node.children.push(childNode); + } + + node.data = { indexes: [] }; + for (let child of node.children) { + processSplatTreeNode(tree, child, indexToCenter, sceneCenters); + } + + // console.log('depth', node.depth, 'splatCount', splatCount, 'children:', splatCounts); + return; + }; + + const buildSubTree = (sceneCenters: Float32Array, maxDepth: number, maxCentersPerNode: number) => { + + const sceneMin = [0, 0, 0]; + const sceneMax = [0, 0, 0]; + const indexes = []; + const centerCount = Math.floor(sceneCenters.length / 4); + for ( let i = 0; i < centerCount; i ++) { + const base = i * 4; + const x = sceneCenters[base]; + const y = sceneCenters[base + 1]; + const z = sceneCenters[base + 2]; + const index = Math.round(sceneCenters[base + 3]); + if (i === 0 || x < sceneMin[0]) sceneMin[0] = x; + if (i === 0 || x > sceneMax[0]) sceneMax[0] = x; + if (i === 0 || y < sceneMin[1]) sceneMin[1] = y; + if (i === 0 || y > sceneMax[1]) sceneMax[1] = y; + if (i === 0 || z < sceneMin[2]) sceneMin[2] = z; + if (i === 0 || z > sceneMax[2]) sceneMax[2] = z; + indexes.push(index); + } + const subTree = new WorkerSplatSubTree(maxDepth, maxCentersPerNode); + subTree.sceneMin = sceneMin; + subTree.sceneMax = sceneMax; + subTree.rootNode = new WorkerSplatTreeNode(subTree.sceneMin, subTree.sceneMax, 0); + subTree.rootNode.data = { + 'indexes': indexes + }; + // console.log('sceneMin', sceneMin, 'sceneMax', sceneMax); + return subTree; + }; + + function createSplatTree(allCenters: Float32Array[], maxDepth: number, maxCentersPerNode: number) { + const indexToCenter = []; + for (let sceneCenters of allCenters) { + const centerCount = Math.floor(sceneCenters.length / 4); + for ( let i = 0; i < centerCount; i ++) { + const base = i * 4; + const index = Math.round(sceneCenters[base + 3]); + indexToCenter[index] = base; + } + } + const subTrees = []; + for (let sceneCenters of allCenters) { + const subTree = buildSubTree(sceneCenters, maxDepth, maxCentersPerNode); + subTrees.push(subTree); + processSplatTreeNode(subTree, subTree.rootNode, indexToCenter, sceneCenters); + } + self.postMessage({ + 'subTrees': subTrees + }); + } + + self.onmessage = (e) => { + if (e.data.process) { + createSplatTree(e.data.process.centers, e.data.process.maxDepth, e.data.process.maxCentersPerNode); + } + }; +} + +function workerProcessCenters(splatTreeWorker: Worker, centers: Float32Array[], transferBuffers: Array, maxDepth: number, maxCentersPerNode: number) { + splatTreeWorker.postMessage({ + 'process': { + 'centers': centers, + 'maxDepth': maxDepth, + 'maxCentersPerNode': maxCentersPerNode + } + }, transferBuffers); +} + +function checkAndCreateWorker() { + const splatTreeWorker = new Worker( + URL.createObjectURL( + new Blob(['(', createSplatTreeWorker.toString(), ')(self)'], { + type: 'application/javascript', + }), + ), + ); + return splatTreeWorker; +} + + +// SplatTree: Octree tailored to splat data from a SplatMesh instance +export class SplatTree { + maxDepth: number; // max depth of the octree + maxCentersPerNode: number; // max centers per node of the octree + subTrees: SplatSubTree[]; // sub trees of the octree + splatMesh: SplatMesh | null; // associated SplatMesh + disposed!: boolean; // whether the octree is disposed + splatTreeWorker!: Worker | null; // worker for octree building + + constructor(maxDepth: number, maxCentersPerNode: number) { + this.maxDepth = maxDepth; + this.maxCentersPerNode = maxCentersPerNode; + this.subTrees = []; + this.splatMesh = null; + } + + + dispose() { + this.diposeSplatTreeWorker(); + this.disposed = true; + } + + diposeSplatTreeWorker() { + if (this.splatTreeWorker) this.splatTreeWorker.terminate(); + this.splatTreeWorker = null; + }; + + /** + * Build SplatTree (octree) from a SplatMesh instance. + * + * @param {SplatMesh} splatMesh SplatMesh instance to build octree from + * @param {function} filterFunc Optional, filter out unwanted splats (points), return true to keep + * @param {function} onIndexesUpload Callback when uploading splat centers to worker (start/end) + * @param {function} onSplatTreeConstruction Callback when worker is building local splat tree (start/end) + * @return {Promise} Promise that resolves when octree building is complete + */ + processSplatMesh( + splatMesh: SplatMesh, + filterFunc: (index: number) => boolean = () => true, + onIndexesUpload?: (isUploading: boolean) => void, + onSplatTreeConstruction?: (isBuilding: boolean) => void + ): Promise { + // if no worker, create one for octree building + if (!this.splatTreeWorker) this.splatTreeWorker = checkAndCreateWorker(); + + this.splatMesh = splatMesh; + this.subTrees = []; + const center = new THREE.Vector3(); + + // tool function: collect all splat centers for a scene + // splatOffset: global starting index, splatCount: number of splats in the scene + const addCentersForScene = (splatOffset: number, splatCount: number) => { + // each splat has 4 floats (x, y, z, index) + const sceneCenters = new Float32Array(splatCount * 4); + let addedCount = 0; + for (let i = 0; i < splatCount; i++) { + // if (i < 100) console.log('center', center.x, center.y, center.z); + const globalSplatIndex = i + splatOffset; + // filter out unwanted splats + if (filterFunc(globalSplatIndex)) { + // get splat center + splatMesh.getSplatCenter(globalSplatIndex, center); + const addBase = addedCount * 4; + sceneCenters[addBase] = center.x; + sceneCenters[addBase + 1] = center.y; + sceneCenters[addBase + 2] = center.z; + sceneCenters[addBase + 3] = globalSplatIndex; + addedCount++; + } + } + // console.log('sceneCenters',sceneCenters) + return sceneCenters; + }; + + return new Promise((resolve) => { + + const checkForEarlyExit = () => { + if (this.disposed) { + this.diposeSplatTreeWorker(); + resolve(); + return true; + } + return false; + }; + + // notify external "upload indexes" start + if (onIndexesUpload) onIndexesUpload(false); + + delayedExecute(() => { + + if (checkForEarlyExit()) return; + + const allCenters: Float32Array[] = []; + // add centers for single scene + const sceneCenters = addCentersForScene(0, splatMesh.getSplatCount()); + allCenters.push(sceneCenters); + + // worker process completed callback + if (this.splatTreeWorker) { + this.splatTreeWorker.onmessage = (e) => { + + if (checkForEarlyExit()) return; + + if (e.data && e.data.subTrees) { + + // notify external "build octree" start + if (onSplatTreeConstruction) onSplatTreeConstruction(false); + + delayedExecute(() => { + + if (checkForEarlyExit()) return; + + // convert worker returned sub tree structure to main thread object + for (let workerSubTree of e.data.subTrees) { + const convertedSubTree = SplatSubTree.convertWorkerSubTree(workerSubTree, splatMesh); + this.subTrees.push(convertedSubTree); + } + // release worker after building + this.diposeSplatTreeWorker(); + + // notify external "build octree" end + if (onSplatTreeConstruction) onSplatTreeConstruction(true); + + // finally resolve + delayedExecute(() => { + resolve(); + }); + + }); + } + }; + } + + // really start uploading data to worker + delayedExecute(() => { + if (checkForEarlyExit()) return; + if (onIndexesUpload) onIndexesUpload(true); + // pass all scene centers data to worker + const transferBuffers = allCenters.map((array) => array.buffer); + workerProcessCenters(this.splatTreeWorker as Worker, allCenters, transferBuffers, this.maxDepth, this.maxCentersPerNode); + }); + + }); + + }); + + }; + + // count leaves + countLeaves() { + + let leafCount = 0; + this.visitLeaves(() => { + leafCount++; + }); + + return leafCount; + } + + // visit leaves + visitLeaves(visitFunc: (node: SplatTreeNode) => void) { + + const visitLeavesFromNode = (node: SplatTreeNode, visitFunc: (node: SplatTreeNode) => void) => { + if (node.children.length === 0) visitFunc(node); + for (let child of node.children) { + visitLeavesFromNode(child, visitFunc); + } + }; + + for (let subTree of this.subTrees) { + visitLeavesFromNode(subTree.rootNode as unknown as SplatTreeNode, visitFunc); + } + } + +}