wgpu-3dgs-viewer 0.7.0

A 3D Gaussian splatting viewing library written in Rust using wgpu.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

import package::camera::{ Camera, world_to_camera };
import wgpu_3dgs_core::{
    gaussian::{
        Gaussian,
        gaussian_unpack_color,
        gaussian_unpack_cov3d,
        gaussian_unpack_sh,
    },
    model_transform::{
        ModelTransform,
        model_to_world,
        model_transform_mat,
        model_scale_rot_mat,
    },
};

// Determine if a Gaussian should be culled.
fn cull(ndc_pos: vec3<f32>) -> bool {
    return !(all(ndc_pos >= vec3<f32>(-1.0, -1.0, 0.0)) && all(ndc_pos <= vec3<f32>(1.0)));
}

// Project a 3D covariance matrix into a 2D covariance matrix.
//
// The returned 2D covariance matrix is represented as (cov.xx, cov.xy, cov.yy).
fn cov2d(gaussian: Gaussian, model_transform: ModelTransform, camera: Camera) -> vec3<f32> {
    let cov3d = gaussian_unpack_cov3d(gaussian);
    let sr = model_scale_rot_mat(model_transform);

    let vrk = mat3x3<f32>(
        cov3d[0], cov3d[1], cov3d[2],
        cov3d[1], cov3d[3], cov3d[4],
        cov3d[2], cov3d[4], cov3d[5],
    );

    let focal = vec2<f32>(camera.proj[0][0], camera.proj[1][1]) * camera.size * 0.5;

    let t = camera.view * model_transform_mat(model_transform) * vec4<f32>(gaussian.pos, 1.0);
    let j = transpose(mat3x3<f32>(
        focal.x / t.z, 0.0, -(focal.x * t.x) / (t.z * t.z),
        0.0, focal.y / t.z, -(focal.y * t.y) / (t.z * t.z),
        0.0, 0.0, 0.0,
    ));
    let w = mat3x3<f32>(camera.view[0].xyz, camera.view[1].xyz, camera.view[2].xyz);

    let cov2d = (j * w * sr) * vrk * transpose(j * w * sr);

    return vec3<f32>(cov2d[0][0], cov2d[0][1], cov2d[1][1]);
}

// Calculate the diagonal axes of a 2D covariance matrix.
//
// The first two components is the major axis, last two components is the minor axis.
// If the returned vector is (0, 0, 0, 0), it means the Gaussian is not visible.
fn cov2d_axes(
    gaussian: Gaussian,
    model_transform: ModelTransform,
    camera: Camera,
    std_dev: f32,
) -> vec4<f32> {
    let cov2d = cov2d(gaussian, model_transform, camera);
    let mid = 0.5 * (cov2d.x + cov2d.z);
    let radius = length(vec2<f32>(0.5 * (cov2d.x - cov2d.z), cov2d.y));
    let major_lambda = mid + radius;
    let minor_lambda = mid - radius;

    if minor_lambda < 0.0 {
        return vec4<f32>(0.0);
    }

    let diag = vec2<f32>(cov2d.y, major_lambda - cov2d.x);
    let diag_dir = select(normalize(diag), vec2<f32>(0.0, 1.0), all(diag == vec2<f32>(0.0)));
    let ortho_diag_dir = vec2<f32>(diag_dir.y, -diag_dir.x);
    let major_len = min(std_dev * sqrt(major_lambda), 1024.0);
    let minor_len = min(std_dev * sqrt(minor_lambda), 1024.0);
    let major_axis = major_len * diag_dir;
    let minor_axis = minor_len * ortho_diag_dir;

    return vec4<f32>(major_axis, minor_axis);
}

// Calculate the color of a Gaussian based on its SH coefficients and view direction.
fn view_color(gaussian: Gaussian, dir: vec3<f32>, sh_deg: u32, no_sh0: bool) -> vec4<f32> {
    const sh_c1 = 0.4886025;
    const sh_c2 = array<f32, 5>(1.0925484, -1.0925484, 0.3153916, -1.0925484, 0.5462742);
    const sh_c3 = array<f32, 7>(
        -0.5900436, 2.8906114, -0.4570458, 0.3731763, -0.4570458, 1.4453057, -0.5900436
    );

    let x = dir.x;
    let y = dir.y;
    let z = dir.z;

    let c = gaussian_unpack_color(gaussian);
    var result = c.rgb; // 0.5 + SH_C0 * sh[0] already precomputed

    if no_sh0 {
        result = vec3<f32>(0.5);
    }

    if sh_deg >= 1u {
        result += sh_c1 * (
            -gaussian_unpack_sh(gaussian, 0u) * y +
            gaussian_unpack_sh(gaussian, 1u) * z -
            gaussian_unpack_sh(gaussian, 2u) * x
        );

        if sh_deg >= 2u {
            let xx = x * x;
            let yy = y * y;
            let zz = z * z;
            let xy = x * y;
            let yz = y * z;
            let xz = x * z;

            result += 
                sh_c2[0] * xy * gaussian_unpack_sh(gaussian, 3u) +
                sh_c2[1] * yz * gaussian_unpack_sh(gaussian, 4u) +
                sh_c2[2] * (2.0 * zz - xx - yy) * gaussian_unpack_sh(gaussian, 5u) +
                sh_c2[3] * xz * gaussian_unpack_sh(gaussian, 6u) +
                sh_c2[4] * (xx - yy) * gaussian_unpack_sh(gaussian, 7u);

            if sh_deg >= 3u {
                result += 
                    sh_c3[0] * y * (3.0 * xx - yy) * gaussian_unpack_sh(gaussian, 8u) +
                    sh_c3[1] * xy * z * gaussian_unpack_sh(gaussian, 9u) +
                    sh_c3[2] * y * (4.0 * zz - xx - yy) * gaussian_unpack_sh(gaussian, 10u) +
                    sh_c3[3] * z * (2.0 * zz - 3.0 * xx - 3.0 * yy) * gaussian_unpack_sh(gaussian, 11u) +
                    sh_c3[4] * x * (4.0 * zz - xx - yy) * gaussian_unpack_sh(gaussian, 12u) +
                    sh_c3[5] * z * (xx - yy) * gaussian_unpack_sh(gaussian, 13u) +
                    sh_c3[6] * x * (xx - 3.0 * yy) * gaussian_unpack_sh(gaussian, 14u);
            }
        }
    }

    return vec4<f32>(max(result, vec3<f32>(0.0)), c.a);
}