UNPKG

@animech-public/playcanvas

Version:
311 lines (308 loc) 8.26 kB
import { Color } from '../../core/math/color.js'; import { Mat4 } from '../../core/math/mat4.js'; import { Quat } from '../../core/math/quat.js'; import { Vec3 } from '../../core/math/vec3.js'; import { BoundingBox } from '../../core/shape/bounding-box.js'; const mat4 = new Mat4(); const quat = new Quat(); const aabb = new BoundingBox(); const aabb2 = new BoundingBox(); const debugColor = new Color(1, 1, 0, 0.4); const SH_C0 = 0.28209479177387814; class SplatIterator { constructor(gsplatData, p, r, s, c) { const x = gsplatData.getProp('x'); const y = gsplatData.getProp('y'); const z = gsplatData.getProp('z'); const rx = gsplatData.getProp('rot_1'); const ry = gsplatData.getProp('rot_2'); const rz = gsplatData.getProp('rot_3'); const rw = gsplatData.getProp('rot_0'); const sx = gsplatData.getProp('scale_0'); const sy = gsplatData.getProp('scale_1'); const sz = gsplatData.getProp('scale_2'); const cr = gsplatData.getProp('f_dc_0'); const cg = gsplatData.getProp('f_dc_1'); const cb = gsplatData.getProp('f_dc_2'); const ca = gsplatData.getProp('opacity'); const sigmoid = v => { if (v > 0) { return 1 / (1 + Math.exp(-v)); } const t = Math.exp(v); return t / (1 + t); }; this.read = i => { if (p) { p.x = x[i]; p.y = y[i]; p.z = z[i]; } if (r) { r.set(rx[i], ry[i], rz[i], rw[i]); } if (s) { s.set(Math.exp(sx[i]), Math.exp(sy[i]), Math.exp(sz[i])); } if (c) { c.set(0.5 + cr[i] * SH_C0, 0.5 + cg[i] * SH_C0, 0.5 + cb[i] * SH_C0, sigmoid(ca[i])); } }; } } const calcSplatMat = (result, p, r) => { quat.set(r.x, r.y, r.z, r.w).normalize(); result.setTRS(p, quat, Vec3.ONE); }; class GSplatData { constructor(elements) { this.elements = void 0; this.numSplats = void 0; this.elements = elements; this.numSplats = this.getElement('vertex').count; } static calcSplatAabb(result, p, r, s) { calcSplatMat(mat4, p, r); aabb.center.set(0, 0, 0); aabb.halfExtents.set(s.x * 2, s.y * 2, s.z * 2); result.setFromTransformedAabb(aabb, mat4); } getProp(name, elementName = 'vertex') { var _this$getElement; return (_this$getElement = this.getElement(elementName)) == null || (_this$getElement = _this$getElement.properties.find(p => p.name === name)) == null ? void 0 : _this$getElement.storage; } getElement(name) { return this.elements.find(e => e.name === name); } addProp(name, storage) { this.getElement('vertex').properties.push({ type: 'float', name, storage, byteSize: 4 }); } createIter(p, r, s, c) { return new SplatIterator(this, p, r, s, c); } calcAabb(result, pred) { let mx, my, mz, Mx, My, Mz; let first = true; const x = this.getProp('x'); const y = this.getProp('y'); const z = this.getProp('z'); const sx = this.getProp('scale_0'); const sy = this.getProp('scale_1'); const sz = this.getProp('scale_2'); for (let i = 0; i < this.numSplats; ++i) { if (pred && !pred(i)) { continue; } const scaleVal = 2.0 * Math.exp(Math.max(sx[i], sy[i], sz[i])); const px = x[i]; const py = y[i]; const pz = z[i]; if (first) { first = false; mx = px - scaleVal; my = py - scaleVal; mz = pz - scaleVal; Mx = px + scaleVal; My = py + scaleVal; Mz = pz + scaleVal; } else { mx = Math.min(mx, px - scaleVal); my = Math.min(my, py - scaleVal); mz = Math.min(mz, pz - scaleVal); Mx = Math.max(Mx, px + scaleVal); My = Math.max(My, py + scaleVal); Mz = Math.max(Mz, pz + scaleVal); } } if (!first) { result.center.set((mx + Mx) * 0.5, (my + My) * 0.5, (mz + Mz) * 0.5); result.halfExtents.set((Mx - mx) * 0.5, (My - my) * 0.5, (Mz - mz) * 0.5); } return !first; } calcAabbExact(result, pred) { const p = new Vec3(); const r = new Quat(); const s = new Vec3(); const iter = this.createIter(p, r, s); let first = true; for (let i = 0; i < this.numSplats; ++i) { if (pred && !pred(i)) { continue; } iter.read(i); if (first) { first = false; GSplatData.calcSplatAabb(result, p, r, s); } else { GSplatData.calcSplatAabb(aabb2, p, r, s); result.add(aabb2); } } return !first; } getCenters(result) { const x = this.getProp('x'); const y = this.getProp('y'); const z = this.getProp('z'); for (let i = 0; i < this.numSplats; ++i) { result[i * 3 + 0] = x[i]; result[i * 3 + 1] = y[i]; result[i * 3 + 2] = z[i]; } } calcFocalPoint(result, pred) { const x = this.getProp('x'); const y = this.getProp('y'); const z = this.getProp('z'); const sx = this.getProp('scale_0'); const sy = this.getProp('scale_1'); const sz = this.getProp('scale_2'); result.x = 0; result.y = 0; result.z = 0; let sum = 0; for (let i = 0; i < this.numSplats; ++i) { if (pred && !pred(i)) { continue; } const weight = 1.0 / (1.0 + Math.exp(Math.max(sx[i], sy[i], sz[i]))); result.x += x[i] * weight; result.y += y[i] * weight; result.z += z[i] * weight; sum += weight; } result.mulScalar(1 / sum); } renderWireframeBounds(scene, worldMat) { const p = new Vec3(); const r = new Quat(); const s = new Vec3(); const min = new Vec3(); const max = new Vec3(); const iter = this.createIter(p, r, s); for (let i = 0; i < this.numSplats; ++i) { iter.read(i); calcSplatMat(mat4, p, r); mat4.mul2(worldMat, mat4); min.set(s.x * -2.0, s.y * -2.0, s.z * -2.0); max.set(s.x * 2.0, s.y * 2.0, s.z * 2.0); scene.immediate.drawWireAlignedBox(min, max, debugColor, true, scene.defaultDrawLayer, mat4); } } get isCompressed() { return false; } get hasSHData() { for (let i = 0; i < 45; ++i) { if (!this.getProp(`f_rest_${i}`)) { return false; } } return true; } calcMortonOrder() { const calcMinMax = arr => { let min = arr[0]; let max = arr[0]; for (let i = 1; i < arr.length; i++) { if (arr[i] < min) min = arr[i]; if (arr[i] > max) max = arr[i]; } return { min, max }; }; const encodeMorton3 = (x, y, z) => { const Part1By2 = x => { x &= 0x000003ff; x = (x ^ x << 16) & 0xff0000ff; x = (x ^ x << 8) & 0x0300f00f; x = (x ^ x << 4) & 0x030c30c3; x = (x ^ x << 2) & 0x09249249; return x; }; return (Part1By2(z) << 2) + (Part1By2(y) << 1) + Part1By2(x); }; const x = this.getProp('x'); const y = this.getProp('y'); const z = this.getProp('z'); const { min: minX, max: maxX } = calcMinMax(x); const { min: minY, max: maxY } = calcMinMax(y); const { min: minZ, max: maxZ } = calcMinMax(z); const sizeX = minX === maxX ? 0 : 1024 / (maxX - minX); const sizeY = minY === maxY ? 0 : 1024 / (maxY - minY); const sizeZ = minZ === maxZ ? 0 : 1024 / (maxZ - minZ); const codes = new Map(); for (let i = 0; i < this.numSplats; i++) { const ix = Math.floor((x[i] - minX) * sizeX); const iy = Math.floor((y[i] - minY) * sizeY); const iz = Math.floor((z[i] - minZ) * sizeZ); const code = encodeMorton3(ix, iy, iz); const val = codes.get(code); if (val) { val.push(i); } else { codes.set(code, [i]); } } const keys = Array.from(codes.keys()).sort((a, b) => a - b); const indices = new Uint32Array(this.numSplats); let idx = 0; for (let i = 0; i < keys.length; ++i) { const val = codes.get(keys[i]); for (let j = 0; j < val.length; ++j) { indices[idx++] = val[j]; } } return indices; } reorder(order) { const cache = new Map(); const getStorage = size => { if (cache.has(size)) { const buffer = cache.get(size); cache.delete(size); return buffer; } return new ArrayBuffer(size); }; const returnStorage = buffer => { cache.set(buffer.byteLength, buffer); }; const reorder = data => { const result = new data.constructor(getStorage(data.byteLength)); for (let i = 0; i < order.length; i++) { result[i] = data[order[i]]; } returnStorage(data.buffer); return result; }; this.elements.forEach(element => { element.properties.forEach(property => { if (property.storage) { property.storage = reorder(property.storage); } }); }); } reorderData() { this.reorder(this.calcMortonOrder()); } } export { GSplatData };