src/splat.wgsl
与本次数学页面一起发布的源码快照。SHA-256: 0a43511f10c402c69096374f5ca6a7647695ec8ab77bcaacaf1f84a74f10e7b5
1struct Camera {
2 view_proj: mat4x4<f32>,
3 viewport: vec4<f32>,
4 scene: vec4<f32>,
5};
6
7// Exact 32-float row layout of the public STG-Lite PLY model.
8struct SourceSplat {
9 r0: vec4<f32>, // xyz, temporal center
10 r1: vec4<f32>, // temporal log-scale, unused normal xyz
11 r2: vec4<f32>, // linear motion xyz, quadratic x
12 r3: vec4<f32>, // quadratic yz, cubic xy
13 r4: vec4<f32>, // cubic z, RGB feature
14 r5: vec4<f32>, // opacity logit, log scale xyz
15 r6: vec4<f32>, // quaternion wxyz
16 r7: vec4<f32>, // quaternion velocity wxyz
17};
18
19@group(0) @binding(0) var<uniform> camera: Camera;
20@group(0) @binding(1) var<storage, read> splats: array<SourceSplat>;
21
22struct VertexInput {
23 @builtin(vertex_index) vertex_index: u32,
24 @location(0) source_index: u32,
25};
26
27struct VertexOutput {
28 @builtin(position) position: vec4<f32>,
29 @location(0) local: vec2<f32>,
30 @location(1) color_opacity: vec4<f32>,
31};
32
33fn 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}
40
41fn 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-space
43 // 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.w
48 ) / w2;
49 return ndc_derivative * camera.viewport.xy * 0.5;
50}
51
52@vertex
53fn 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.xyz
61 + motion_linear * dt
62 + motion_quadratic * dt2
63 + motion_cubic * dt2 * dt;
64
65 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;
74
75 // 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 );
92
93 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);
97
98 // 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));
105
106 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.x
118 + minor_axis * sigma_minor * corner.y
119 );
120 let ndc_offset = pixel_offset * vec2<f32>(2.0 * camera.viewport.z, 2.0 * camera.viewport.w);
121
122 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}
128
129@fragment
130fn 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