bevy_gaussian_splatting 2.4.2

bevy gaussian splatting render pipeline plugin
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

use std::io::BufRead;

use ply_rs::{
    ply::{
        Property,
        PropertyAccess,
    },
    parser::Parser,
};

use crate::{
    material::spherical_harmonics::{
        SH_CHANNELS,
        SH_COEFF_COUNT,
        SH_COEFF_COUNT_PER_CHANNEL,
    },
    gaussian::packed::Gaussian,
};


pub const MAX_SIZE_VARIANCE: f32 = 5.0;

impl PropertyAccess for Gaussian {
    fn new() -> Self {
        Gaussian::default()
    }

    fn set_property(&mut self, key: String, property: Property) {
        match (key.as_ref(), property) {
            ("x", Property::Float(v))           => self.position_visibility.position[0] = v,
            ("y", Property::Float(v))           => self.position_visibility.position[1] = v,
            ("z", Property::Float(v))           => self.position_visibility.position[2] = v,
            ("f_dc_0", Property::Float(v))      => self.spherical_harmonic.set(0, v),
            ("f_dc_1", Property::Float(v))      => self.spherical_harmonic.set(1, v),
            ("f_dc_2", Property::Float(v))      => self.spherical_harmonic.set(2, v),
            ("scale_0", Property::Float(v))     => self.scale_opacity.scale[0] = v,
            ("scale_1", Property::Float(v))     => self.scale_opacity.scale[1] = v,
            ("scale_2", Property::Float(v))     => self.scale_opacity.scale[2] = v,
            ("opacity", Property::Float(v))     => self.scale_opacity.opacity = 1.0 / (1.0 + (-v).exp()),
            ("rot_0", Property::Float(v))       => self.rotation.rotation[0] = v,
            ("rot_1", Property::Float(v))       => self.rotation.rotation[1] = v,
            ("rot_2", Property::Float(v))       => self.rotation.rotation[2] = v,
            ("rot_3", Property::Float(v))       => self.rotation.rotation[3] = v,
            (_, Property::Float(v)) if key.starts_with("f_rest_") => {
                let i = key[7..].parse::<usize>().unwrap();

                // interleaved
                // if (i + 3) < SH_COEFF_COUNT {
                //     self.spherical_harmonic.coefficients[i + 3] = v;
                // }

                // planar
                let channel = i / SH_COEFF_COUNT_PER_CHANNEL;
                let coefficient = if SH_COEFF_COUNT_PER_CHANNEL == 1 {
                    1
                } else {
                    (i % (SH_COEFF_COUNT_PER_CHANNEL - 1)) + 1
                };

                let interleaved_idx = coefficient * SH_CHANNELS + channel;

                if interleaved_idx < SH_COEFF_COUNT {
                    self.spherical_harmonic.set(interleaved_idx, v);
                } else {
                    // TODO: convert higher degree SH to lower degree SH
                }
            }
            (_, _) => {},
        }
    }
}

pub fn parse_ply(mut reader: &mut dyn BufRead) -> Result<Vec<Gaussian>, std::io::Error> {
    let gaussian_parser = Parser::<Gaussian>::new();
    let header = gaussian_parser.read_header(&mut reader)?;

    let mut cloud = Vec::new();

    for (_ignore_key, element) in &header.elements {
        if element.name == "vertex" {
            cloud = gaussian_parser.read_payload_for_element(&mut reader, element, &header)?;
        }
    }

    for gaussian in &mut cloud {
        gaussian.position_visibility.visibility = 1.0;

        let mean_scale = (gaussian.scale_opacity.scale[0] + gaussian.scale_opacity.scale[1] + gaussian.scale_opacity.scale[2]) / 3.0;
        for i in 0..3 {
            gaussian.scale_opacity.scale[i] = gaussian.scale_opacity.scale[i]
                .max(mean_scale - MAX_SIZE_VARIANCE)
                .min(mean_scale + MAX_SIZE_VARIANCE)
                .exp();
        }

        let norm = (0..4).map(|i| gaussian.rotation.rotation[i].powf(2.0)).sum::<f32>().sqrt();
        for i in 0..4 {
            gaussian.rotation.rotation[i] /= norm;
        }
    }

    Ok(cloud)
}