@threepipe/plugin-gaussian-splatting 0.1.0

package/src/three-gaussian-splat/loaders/SplatLoader.ts

@@ -0,0 +1,51 @@
+import {FileLoader, ILoader, IMaterial, Loader, LoadingManager} from 'threepipe'
+import {SortWorkerManager, SPLAT_ROW_LENGTH} from '../cpp-sorter/SortWorkerManager'
+import {GaussianSplatGeometry} from '../geometry/GaussianSplatGeometry'
+import {GaussianSplatMesh} from '../mesh/GaussianSplatMesh'
+import {GaussianSplatMaterialUnlit} from '../materials/GaussianSplatMaterialUnlit'
+
+export class SplatLoader extends Loader implements ILoader {
+    sortWorkerManager: SortWorkerManager
+
+    materialConstructor = (_: GaussianSplatGeometry): IMaterial|undefined => new GaussianSplatMaterialUnlit()
+
+    constructor(manager?: LoadingManager) {
+        super(manager)
+    }
+
+    onGeometryLoad = (_: GaussianSplatGeometry)=>{
+        return
+    }
+
+    public load(url: string, onLoad?: (data: GaussianSplatMesh) => void, onProgress?: (event: ProgressEvent) => void, onError?: (event: ErrorEvent) => void): void {
+        // const path = ( scope.path === '' ) ? LoaderUtils.extractUrlBase( url ) : scope.path;
+        const loader = new FileLoader(this.manager)
+        loader.setPath(this.path)
+        loader.setResponseType('arraybuffer')
+        loader.setRequestHeader(this.requestHeader)
+        loader.setWithCredentials(this.withCredentials)
+        loader.load(url, async(buffer) => {
+            try {
+                const data = new Uint8Array(buffer as ArrayBuffer)
+                const maxSplats = 1000000
+                const vertexCount = Math.floor(data.length / SPLAT_ROW_LENGTH)
+                const worker = await this.sortWorkerManager.createWorker(data, maxSplats)
+                const geometry = new GaussianSplatGeometry(worker, vertexCount, maxSplats, this.onGeometryLoad)
+                const mesh = new GaussianSplatMesh(geometry, this.materialConstructor(geometry) as any)
+                // const mesh = new GaussianSplatMesh(geometry, new UnlitMaterial() as any)
+                mesh.rotation.x = Math.PI
+                if (!url.startsWith('blob:') && !url.startsWith('data:'))
+                    mesh.name = url.split('/').pop()!.split('?')[0]
+                onLoad && onLoad(mesh)
+            } catch (e) {
+                if (onError) onError(e)
+                else console.error(e)
+                this.manager.itemError(url)
+            }
+        }, onProgress, onError)
+    }
+
+    public async loadAsync(url: string, onProgress?: (event: ProgressEvent) => void): Promise<GaussianSplatMesh> {
+        return new Promise((resolve, reject) => this.load(url, resolve, onProgress, reject))
+    }
+}

package/src/three-gaussian-splat/materials/GaussianSplatMaterialExtension.ts

@@ -0,0 +1,105 @@
+import {
+    Camera,
+    IMaterial,
+    MaterialExtension,
+    PerspectiveCamera,
+    Shader,
+    shaderReplaceString,
+    Vector2,
+    WebGLRenderer,
+} from 'threepipe'
+import {computeFocalLengths} from './util'
+import {gaussianSplatShaders} from '../shaders'
+
+export class GaussianSplatMaterialExtension implements MaterialExtension {
+    readonly isGaussianSplatMaterialExtension = true
+    extraUniforms = {
+        viewport: {value: new Vector2()},
+        focal: {value: new Vector2()},
+        minAlpha: {value: 0.02},
+    }
+    parsFragmentSnippet = gaussianSplatShaders.pars_frag
+    parsVertexSnippet = gaussianSplatShaders.pars_vert
+    shaderExtender = (shader: Shader, material: IMaterial) => {
+        shader.vertexShader = shaderReplaceString(shader.vertexShader,
+            '#include <begin_vertex>',
+            'vec3 transformed = vec3( center );')
+        if (shader.vertexShader.includes('#include <beginnormal_vertex>')) {
+            // todo: get correct normal by rendering to depth and then sampling the normal
+            shader.vertexShader = shaderReplaceString(shader.vertexShader,
+                '#include <beginnormal_vertex>',
+                '\nobjectNormal = vec3(1,0,0);\n', {append: true})
+        }
+        shader.vertexShader = shaderReplaceString(shader.vertexShader,
+            '#include <project_vertex>',
+            gaussianSplatShaders.main_vert, {append: true})
+
+        if (!material.isGBufferMaterial && !material.userData.isGBufferMaterial) {
+            shader.fragmentShader = shaderReplaceString(shader.fragmentShader,
+                '#include <map_fragment>',
+                shaderReplaceString(
+                    gaussianSplatShaders.main_frag,
+                    /gl_FragColor\s*=/, // because of minification
+                    'diffuseColor*='
+                ), {prepend: true})
+        }
+        // eslint-disable-next-line no-constant-condition
+        if (0)
+            shader.fragmentShader = shaderReplaceString(shader.fragmentShader,
+                '#include <output_fragment>',
+                '\ngl_FragColor=diffuseColor;',
+                // '\ngl_FragColor=vec4(reflectedLight.directDiffuse, 1.);',
+                // '\ngl_FragColor=vec4(geometryNormal, 1.);',
+                {append: true})
+        return shader
+    }
+    isCompatible = (material: IMaterial) =>
+        !!material.isPhysicalMaterial ||
+        !!material.isUnlitMaterial ||
+        !!material.isGBufferMaterial ||
+        !!material.userData.isGBufferMaterial
+
+    setDirty?: () => void
+
+    private _currentCamera?: PerspectiveCamera | Camera
+    private _renderer?: WebGLRenderer
+
+    public set minAlpha(value: number) {
+        this.extraUniforms.minAlpha.value = value
+        this.setDirty && this.setDirty()
+    }
+
+
+    constructor() {
+        window.addEventListener('resize', this._refresh)
+    }
+
+    dispose() {
+        // todo: add again on added to mesh?
+        window.removeEventListener('resize', this._refresh)
+    }
+
+    update(camera: PerspectiveCamera | Camera, renderer: WebGLRenderer): void {
+        if (this._currentCamera === camera && this._renderer === renderer) return
+        this._renderer = renderer
+        this._currentCamera = camera
+        this._refresh()
+    }
+
+    private _refresh = (): void => {
+        if (!this._currentCamera) return
+        const size = new Vector2()
+        this._renderer?.getSize(size)
+        const dpr = this._renderer?.getPixelRatio() || 1
+        let fov = 75
+        let aspect = size.x / size.y
+
+        if (this._currentCamera instanceof PerspectiveCamera) {
+            fov = this._currentCamera.fov
+            aspect = this._currentCamera.aspect
+        }
+
+        this.extraUniforms.focal.value = computeFocalLengths(size.x, size.y, fov, aspect, dpr)
+        this.extraUniforms.viewport.value = new Vector2(size.x * dpr, size.y * dpr)
+    }
+}

package/src/three-gaussian-splat/materials/GaussianSplatMaterialPhysical.ts

@@ -0,0 +1,25 @@
+import {PhysicalMaterial} from 'threepipe'
+import {GaussianSplatMaterialExtension} from './GaussianSplatMaterialExtension'
+
+export class GaussianSplatMaterialPhysical extends PhysicalMaterial {
+    readonly isGaussianSplatMaterialPhysical = true
+
+    gsplatExtension = new GaussianSplatMaterialExtension()
+
+    constructor() {
+        super({
+            depthTest: true,
+            depthWrite: false,
+            transparent: true,
+            vertexColors: false,
+        })
+        // this.userData.renderToGBuffer = true
+        // this.userData.renderToDepth = true
+        this.registerMaterialExtensions([this.gsplatExtension])
+    }
+
+    dispose() {
+        this.gsplatExtension.dispose()
+        return super.dispose()
+    }
+}

package/src/three-gaussian-splat/materials/GaussianSplatMaterialRaw.ts

@@ -0,0 +1,73 @@
+import {Camera, PerspectiveCamera, ShaderMaterial2, Vector2, WebGLRenderer} from 'threepipe'
+import {computeFocalLengths} from './util'
+import {gaussianSplatShaders} from '../shaders'
+
+export class GaussianSplatMaterialRaw extends ShaderMaterial2 {
+    private _currentCamera?: PerspectiveCamera | Camera
+    private _renderer?: WebGLRenderer
+
+    readonly isGaussianSplatMaterialRaw = true
+
+    public set minAlpha(value: number) {
+        this.uniforms.minAlpha.value = value
+        this.needsUpdate = true
+    }
+
+    constructor() {
+        super({
+            uniforms: {
+                viewport: {value: new Vector2()},
+                focal: {value: new Vector2()},
+                minAlpha: {value: 0.02},
+            },
+            fragmentShader: `${gaussianSplatShaders.pars_frag}
+void main () {
+${gaussianSplatShaders.main_frag}
+}`,
+            vertexShader: `${gaussianSplatShaders.pars_vert}
+void main () {
+    gl_Position = projectionMatrix * modelViewMatrix * vec4(center, 1);
+    ${gaussianSplatShaders.main_vert}
+}`,
+            depthTest: true,
+            depthWrite: false,
+            transparent: true,
+        }, true)
+
+        window.addEventListener('resize', this._refresh)
+    }
+
+    private _refresh = (): void => {
+        if (!this._currentCamera) return
+
+        const size = new Vector2()
+        this._renderer?.getSize(size)
+        const width = size.x
+        const height = size.y
+
+        const dpr = this._renderer?.getPixelRatio() || 1
+
+        let fov = 75
+        let aspect = width / height
+
+        if (this._currentCamera instanceof PerspectiveCamera) {
+            fov = this._currentCamera.fov
+            aspect = this._currentCamera.aspect
+        }
+
+        this.uniforms.focal.value = computeFocalLengths(width, height, fov, aspect, dpr)
+        this.uniforms.viewport.value = new Vector2(width * dpr, height * dpr)
+    }
+
+    dispose() {
+        // todo: add again on added to mesh?
+        window.removeEventListener('resize', this._refresh)
+        return super.dispose()
+    }
+
+    update(camera: PerspectiveCamera | Camera, renderer: WebGLRenderer): void {
+        this._renderer = renderer
+        this._currentCamera = camera
+        this._refresh()
+    }
+}

package/src/three-gaussian-splat/materials/GaussianSplatMaterialUnlit.ts

@@ -0,0 +1,26 @@
+import {UnlitMaterial} from 'threepipe'
+
+import {GaussianSplatMaterialExtension} from './GaussianSplatMaterialExtension'
+
+export class GaussianSplatMaterialUnlit extends UnlitMaterial {
+    readonly isGaussianSplatMaterialUnlit = true
+    gsplatExtension = new GaussianSplatMaterialExtension()
+
+    constructor() {
+        super({
+            depthTest: true,
+            depthWrite: false,
+            transparent: true,
+            vertexColors: false,
+        })
+        // this.userData.renderToGBuffer = true
+        // this.userData.renderToDepth = true
+        this.registerMaterialExtensions([this.gsplatExtension])
+    }
+
+    dispose() {
+        this.gsplatExtension.dispose()
+        return super.dispose()
+    }
+}
+

package/src/three-gaussian-splat/materials/util.ts

@@ -0,0 +1,9 @@
+import {MathUtils, Vector2} from 'threepipe'
+
+export const computeFocalLengths = (width: number, height: number, fov: number, aspect: number, dpr: number) => {
+    const fovRad = MathUtils.degToRad(fov)
+    const fovXRad = 2 * Math.atan(Math.tan(fovRad / 2) * aspect)
+    const fy = dpr * height / (2 * Math.tan(fovRad / 2))
+    const fx = dpr * width / (2 * Math.tan(fovXRad / 2))
+    return new Vector2(fx, fy)
+}

package/src/three-gaussian-splat/mesh/GaussianSplatMesh.ts

@@ -0,0 +1,96 @@
+import {GaussianSplatGeometry} from '../geometry/GaussianSplatGeometry'
+import {Camera, IMaterial, Mesh2, PerspectiveCamera, WebGLRenderer} from 'threepipe'
+import type {GaussianSplatMaterialRaw} from '../materials/GaussianSplatMaterialRaw'
+import {GaussianSplatMaterialUnlit} from '../materials/GaussianSplatMaterialUnlit'
+import {Matrix4, Ray, Raycaster, Sphere, Vector3} from 'three'
+import {GaussianSplatMaterialExtension} from '../materials/GaussianSplatMaterialExtension'
+
+export class GaussianSplatMesh extends Mesh2<GaussianSplatGeometry, IMaterial> {
+    readonly isGaussianSplatMesh = true
+
+    constructor(geometry: GaussianSplatGeometry, material: IMaterial) {
+        super(geometry, material)
+        this.frustumCulled = false
+    }
+
+    public async update(camera: PerspectiveCamera | Camera, renderer: WebGLRenderer) {
+        if ((this.material as any as GaussianSplatMaterialRaw)?.isGaussianSplatMaterialRaw) {
+            (this.material as any as GaussianSplatMaterialRaw).update(camera, renderer)
+        } else if (this.material) {
+            const ext =
+                (this.material as GaussianSplatMaterialUnlit).gsplatExtension ??
+                this.material.materialExtensions.find(e=>
+                    (e as GaussianSplatMaterialExtension).isGaussianSplatMaterialExtension) as GaussianSplatMaterialExtension
+            ext && ext.update(camera, renderer)
+        }
+        this.updateMatrixWorld(true)
+        return this.geometry.update(camera, this.matrixWorld)
+    }
+
+    raycast(raycaster: Raycaster, intersects: any[]) {
+
+        const geometry = this.geometry
+        const matrixWorld = this.matrixWorld
+        // const threshold = raycaster.params.Points?.threshold ?? .01
+        const threshold = .02
+
+        // Checking boundingSphere distance to ray
+
+        if (geometry.boundingSphere === null) geometry.computeBoundingSphere()
+
+        const sphere = new Sphere()
+        sphere.copy(geometry.boundingSphere!)
+        sphere.applyMatrix4(matrixWorld)
+        sphere.radius += threshold
+
+        if (!raycaster.ray.intersectsSphere(sphere)) return
+
+        //
+
+        const inverseMatrix = new Matrix4()
+        const ray = new Ray()
+
+        inverseMatrix.copy(matrixWorld).invert()
+        ray.copy(raycaster.ray).applyMatrix4(inverseMatrix)
+
+        const localThreshold = threshold / ((this.scale.x + this.scale.y + this.scale.z) / 3)
+        const localThresholdSq = localThreshold * localThreshold
+
+        const attributes = geometry.attributes
+        const positionAttribute = attributes.center
+
+        const position = new Vector3()
+        for (let i = 0, l = positionAttribute.count; i < l; i++) {
+
+            position.fromBufferAttribute(positionAttribute, i)
+
+            const rayPointDistanceSq = ray.distanceSqToPoint(position)
+
+            if (rayPointDistanceSq < localThresholdSq) {
+
+                const intersectPoint = new Vector3()
+
+                ray.closestPointToPoint(position, intersectPoint)
+                intersectPoint.applyMatrix4(matrixWorld)
+
+                const distance = raycaster.ray.origin.distanceTo(intersectPoint)
+
+                if (distance < raycaster.near || distance > raycaster.far) return
+
+                intersects.push({
+
+                    distance: distance,
+                    distanceToRay: Math.sqrt(rayPointDistanceSq),
+                    point: intersectPoint,
+                    face: null,
+                    object: this,
+
+                })
+
+
+            }
+        }
+
+    }
+
+}

package/src/three-gaussian-splat/shaders/gsplat.main.frag.glsl

@@ -0,0 +1,18 @@
+vec2 d = (vCenter - 2.0 * (gl_FragCoord.xy/viewport - vec2(0.5, 0.5))) * viewport * 0.5;
+
+float power = -0.5 * (vConic.x * d.x * d.x + vConic.z * d.y * d.y) + vConic.y * d.x * d.y;
+
+if (power > 0.0) discard;
+float alpha = min(0.99, vColor.a * exp(power));
+if(alpha < minAlpha) discard;
+
+gl_FragColor = vec4(vColor.rgb, alpha);
+//vec2 d = (vCenter - 2.0 * (gl_FragCoord.xy/viewport - vec2(0.5, 0.5))) * viewport * 0.5;
+//
+//float power = -0.5 * (vConic.x * d.x * d.x + vConic.z * d.y * d.y) + vConic.y * d.x * d.y;
+//
+//if (power > 0.0) discard;
+//float alpha = min(0.99, vColor.a * exp(power));
+//if(alpha < minAlpha) discard;
+//
+//diffuseColor *= vec4(vColor.rgb, alpha);

package/src/three-gaussian-splat/shaders/gsplat.main.vert.glsl

@@ -0,0 +1,15 @@
+vec3 cov2d = compute_cov2d(center, scale, quat);
+float det = cov2d.x * cov2d.z - cov2d.y * cov2d.y;
+vec3 conic = vec3(cov2d.z, cov2d.y, cov2d.x) / det;
+float mid = 0.5 * (cov2d.x + cov2d.z);
+float lambda1 = mid + sqrt(max(0.1, mid * mid - det));
+float lambda2 = mid - sqrt(max(0.1, mid * mid - det));
+vec2 v1 = 7.0 * sqrt(lambda1) * normalize(vec2(cov2d.y, lambda1 - cov2d.x));
+vec2 v2 = 7.0 * sqrt(lambda2) * normalize(vec2(-(lambda1 - cov2d.x),cov2d.y));
+
+vColor = color;
+vConic = conic;
+vCenter = vec2(gl_Position) / gl_Position.w;
+
+vPosition = vec2(vCenter + position.x * (position.y < 0.0 ? v1 : v2) / viewport);
+gl_Position = vec4(vPosition, gl_Position.z / gl_Position.w, 1);

package/src/three-gaussian-splat/shaders/gsplat.pars.frag.glsl

@@ -0,0 +1,10 @@
+// https://github.com/vincent-lecrubier-skydio/react-three-fiber-gaussian-splat
+precision mediump float;
+
+varying vec4 vColor;
+varying vec3 vConic;
+varying vec2 vCenter;
+
+uniform vec2 viewport;
+uniform vec2 focal;
+uniform float minAlpha;

package/src/three-gaussian-splat/shaders/gsplat.pars.vert.glsl

@@ -0,0 +1,58 @@
+// https://github.com/vincent-lecrubier-skydio/react-three-fiber-gaussian-splat
+precision mediump float;
+
+#ifndef SHADER_NAME // isRawShaderMaterial
+attribute vec3 position;
+uniform mat4 modelViewMatrix;
+uniform mat4 projectionMatrix;
+mat3 transpose(mat3 m) { return mat3(m[0][0], m[1][0], m[2][0], m[0][1], m[1][1], m[2][1], m[0][2], m[1][2], m[2][2]); }
+#endif
+
+attribute vec4 color;
+attribute vec4 quat;
+attribute vec3 scale;
+attribute vec3 center;
+
+uniform vec2 focal;
+uniform vec2 viewport;
+//uniform vec3 sphereCenter;
+//uniform vec3 planeNormal;
+//uniform float planeDistance;
+
+varying vec4 vColor;
+varying vec3 vConic;
+varying vec2 vCenter;
+varying vec2 vPosition;
+//varying float vDistance;
+//varying float vPlaneSide;
+
+mat3 compute_cov3d(vec3 scale, vec4 rot) {
+    mat3 S = mat3(
+    scale.x, 0.0, 0.0,
+    0.0, scale.y, 0.0,
+    0.0, 0.0, scale.z
+    );
+    mat3 R = mat3(
+    1.0 - 2.0 * (rot.z * rot.z + rot.w * rot.w), 2.0 * (rot.y * rot.z - rot.x * rot.w), 2.0 * (rot.y * rot.w + rot.x * rot.z),
+    2.0 * (rot.y * rot.z + rot.x * rot.w), 1.0 - 2.0 * (rot.y * rot.y + rot.w * rot.w), 2.0 * (rot.z * rot.w - rot.x * rot.y),
+    2.0 * (rot.y * rot.w - rot.x * rot.z), 2.0 * (rot.z * rot.w + rot.x * rot.y), 1.0 - 2.0 * (rot.y * rot.y + rot.z * rot.z)
+    );
+    mat3 M = S * R;
+    return transpose(M) * M;
+}
+
+vec3 compute_cov2d(vec3 center, vec3 scale, vec4 rot){
+    mat3 Vrk = compute_cov3d(scale, rot);
+    vec4 t = modelViewMatrix * vec4(center, 1.0);
+    vec2 lims = 1.3 * 0.5 * viewport / focal;
+    t.xy = min(lims, max(-lims, t.xy / t.z)) * t.z;
+    mat3 J = mat3(
+    focal.x / t.z, 0., -(focal.x * t.x) / (t.z * t.z),
+    0., focal.y / t.z, -(focal.y * t.y) / (t.z * t.z),
+    0., 0., 0.
+    );
+    mat3 W = transpose(mat3(modelViewMatrix));
+    mat3 T = W * J;
+    mat3 cov = transpose(T) * transpose(Vrk) * T;
+    return vec3(cov[0][0] + 0.3, cov[0][1], cov[1][1] + 0.3);
+}

package/src/three-gaussian-splat/shaders/index.ts

@@ -0,0 +1,11 @@
+import fragmentParsShaderSource from '../shaders/gsplat.pars.frag.glsl'
+import vertexParsShaderSource from '../shaders/gsplat.pars.vert.glsl'
+import vertexShaderSource from '../shaders/gsplat.main.vert.glsl'
+import fragmentShaderSource from '../shaders/gsplat.main.frag.glsl'
+
+export const gaussianSplatShaders = {
+    main_frag: '\n' + fragmentShaderSource + '\n',
+    main_vert: '\n' + vertexShaderSource + '\n',
+    pars_frag: fragmentParsShaderSource,
+    pars_vert: vertexParsShaderSource,
+}