src/splat.wgsl
Repository source snapshot · Plain text · GitHub
1struct Camera {2 view_proj: mat4x4<f32>,3 viewport: vec4<f32>,4 scene: vec4<f32>,5};67// Exact 32-float row layout of the public STG-Lite PLY model.8struct SourceSplat {9 r0: vec4<f32>, // xyz, temporal center10 r1: vec4<f32>, // temporal log-scale, unused normal xyz11 r2: vec4<f32>, // linear motion xyz, quadratic x12 r3: vec4<f32>, // quadratic yz, cubic xy13 r4: vec4<f32>, // cubic z, RGB feature14 r5: vec4<f32>, // opacity logit, log scale xyz15 r6: vec4<f32>, // quaternion wxyz16 r7: vec4<f32>, // quaternion velocity wxyz17};1819@group(0) @binding(0) var<uniform> camera: Camera;20@group(0) @binding(1) var<storage, read> splats: array<SourceSplat>;2122struct VertexInput {23 @builtin(vertex_index) vertex_index: u32,24 @location(0) source_index: u32,25};2627struct VertexOutput {28 @builtin(position) position: vec4<f32>,29 @location(0) local: vec2<f32>,30 @location(1) color_opacity: vec4<f32>,31};3233fn quad_corner(index: u32) -> vec2<f32> {34 var corners = array<vec2<f32>, 6>(35 vec2<f32>(-1.0, -1.0), vec2<f32>( 1.0, -1.0), vec2<f32>(-1.0, 1.0),36 vec2<f32>(-1.0, 1.0), vec2<f32>( 1.0, -1.0), vec2<f32>( 1.0, 1.0)37 );38 return corners[index];39}4041fn projected_axis(center_clip: vec4<f32>, axis: vec3<f32>) -> vec2<f32> {42 // Exact first-order derivative of ndc = clip.xy / clip.w along a world-space43 // covariance axis. This is the homogeneous form of the perspective Jacobian.44 let axis_clip = camera.view_proj * vec4<f32>(axis, 0.0);45 let w2 = max(center_clip.w * center_clip.w, 0.00000001);46 let ndc_derivative = (47 axis_clip.xy * center_clip.w - center_clip.xy * axis_clip.w48 ) / w2;49 return ndc_derivative * camera.viewport.xy * 0.5;50}5152@vertex53fn vs_main(input: VertexInput) -> VertexOutput {54 let splat = splats[input.source_index];55 let dt = camera.scene.x - splat.r0.w;56 let dt2 = dt * dt;57 let motion_linear = splat.r2.xyz;58 let motion_quadratic = vec3<f32>(splat.r2.w, splat.r3.x, splat.r3.y);59 let motion_cubic = vec3<f32>(splat.r3.z, splat.r3.w, splat.r4.x);60 let center = splat.r0.xyz61 + motion_linear * dt62 + motion_quadratic * dt263 + motion_cubic * dt2 * dt;6465 let temporal_scale = max(exp(splat.r1.x), 0.000001);66 let temporal_opacity = exp(-pow(dt / temporal_scale, 2.0));67 let opacity = (1.0 / (1.0 + exp(-splat.r5.x))) * temporal_opacity;68 let scales = exp(splat.r5.yzw);69 let quaternion = normalize(splat.r6 + dt * splat.r7);70 let w = quaternion.x;71 let x = quaternion.y;72 let y = quaternion.z;73 let z = quaternion.w;7475 // Rows of GLM's column-major 3DGS quaternion matrix, scaled into covariance axes.76 // This yields Sigma = R^T S^2 R, matching the reference covariance construction.77 let axis0 = scales.x * vec3<f32>(78 1.0 - 2.0 * (y * y + z * z),79 2.0 * (x * y + w * z),80 2.0 * (x * z - w * y)81 );82 let axis1 = scales.y * vec3<f32>(83 2.0 * (x * y - w * z),84 1.0 - 2.0 * (x * x + z * z),85 2.0 * (y * z + w * x)86 );87 let axis2 = scales.z * vec3<f32>(88 2.0 * (x * z + w * y),89 2.0 * (y * z - w * x),90 1.0 - 2.0 * (x * x + y * y)91 );9293 let center_clip = camera.view_proj * vec4<f32>(center, 1.0);94 let d0 = projected_axis(center_clip, axis0);95 let d1 = projected_axis(center_clip, axis1);96 let d2 = projected_axis(center_clip, axis2);9798 // Project J W Sigma W^T J^T by summing the projected covariance axes.99 let cov_a = dot(vec3<f32>(d0.x, d1.x, d2.x), vec3<f32>(d0.x, d1.x, d2.x)) + 0.3;100 let cov_b = d0.x * d0.y + d1.x * d1.y + d2.x * d2.y;101 let cov_c = dot(vec3<f32>(d0.y, d1.y, d2.y), vec3<f32>(d0.y, d1.y, d2.y)) + 0.3;102 let discriminant = sqrt(max(0.0, (cov_a - cov_c) * (cov_a - cov_c) + 4.0 * cov_b * cov_b));103 let lambda_major = max(0.1, 0.5 * (cov_a + cov_c + discriminant));104 let lambda_minor = max(0.1, 0.5 * (cov_a + cov_c - discriminant));105106 var major_axis = vec2<f32>(1.0, 0.0);107 if (abs(cov_b) > 0.00001) {108 major_axis = normalize(vec2<f32>(cov_b, lambda_major - cov_a));109 } else if (cov_c > cov_a) {110 major_axis = vec2<f32>(0.0, 1.0);111 }112 let minor_axis = vec2<f32>(-major_axis.y, major_axis.x);113 let corner = quad_corner(input.vertex_index);114 let sigma_major = min(sqrt(lambda_major), 192.0);115 let sigma_minor = min(sqrt(lambda_minor), 192.0);116 let pixel_offset = 3.0 * (117 major_axis * sigma_major * corner.x118 + minor_axis * sigma_minor * corner.y119 );120 let ndc_offset = pixel_offset * vec2<f32>(2.0 * camera.viewport.z, 2.0 * camera.viewport.w);121122 var output: VertexOutput;123 output.position = center_clip + vec4<f32>(ndc_offset * center_clip.w, 0.0, 0.0);124 output.local = corner * 3.0;125 output.color_opacity = vec4<f32>(clamp(splat.r4.yzw, vec3<f32>(0.0), vec3<f32>(1.0)), opacity);126 return output;127}128129@fragment130fn fs_main(input: VertexOutput) -> @location(0) vec4<f32> {131 let power = -0.5 * dot(input.local, input.local);132 let alpha = min(0.99, input.color_opacity.a * exp(power));133 if (alpha < 0.0039215686) {134 discard;135 }136 return vec4<f32>(input.color_opacity.rgb * alpha, alpha);137}138