From d19d15c393f8064a248d5a53638b8d723523af0f Mon Sep 17 00:00:00 2001 From: mosure Date: Thu, 2 Jan 2025 20:29:07 -0600 Subject: [PATCH] refactor: better format separation --- src/gaussian/f16.rs | 6 +- src/gaussian/f32.rs | 5 +- src/gaussian/formats/planar_3d.rs | 101 +++++++++++++- src/gaussian/formats/planar_4d.rs | 117 +++++++++++++++- src/gaussian/formats/planar_4d_hierarchy.rs | 2 +- src/gaussian/mod.rs | 2 - src/gaussian/packed.rs | 87 ------------ src/gaussian/rand.rs | 145 -------------------- src/io/gcloud/flexbuffers.rs | 6 +- src/io/loader.rs | 2 +- src/io/ply.rs | 5 +- src/lib.rs | 22 +-- src/query/select.rs | 5 +- src/sort/mod.rs | 8 +- src/sort/rayon.rs | 6 +- src/sort/std.rs | 6 +- 16 files changed, 262 insertions(+), 263 deletions(-) delete mode 100644 src/gaussian/packed.rs delete mode 100644 src/gaussian/rand.rs diff --git a/src/gaussian/f16.rs b/src/gaussian/f16.rs index 7da1ab32..e1b04288 100644 --- a/src/gaussian/f16.rs +++ b/src/gaussian/f16.rs @@ -21,9 +21,9 @@ use crate::gaussian::{ Rotation, ScaleOpacity, }, - packed::{ - Gaussian3d, - Gaussian4d, + formats::{ + planar_3d::Gaussian3d, + planar_4d::Gaussian4d, }, }; diff --git a/src/gaussian/f32.rs b/src/gaussian/f32.rs index b9709e4b..1cf0e9fc 100644 --- a/src/gaussian/f32.rs +++ b/src/gaussian/f32.rs @@ -15,7 +15,10 @@ use serde::{ use crate::gaussian::{ covariance::compute_covariance_3d, - packed::{Gaussian3d, Gaussian4d}, + formats::{ + planar_3d::Gaussian3d, + planar_4d::Gaussian4d, + }, }; diff --git a/src/gaussian/formats/planar_3d.rs b/src/gaussian/formats/planar_3d.rs index 35feffa1..420fd702 100644 --- a/src/gaussian/formats/planar_3d.rs +++ b/src/gaussian/formats/planar_3d.rs @@ -1,9 +1,20 @@ +use std::marker::Copy; use rand::{ + prelude::Distribution, seq::SliceRandom, Rng, }; +use bevy::prelude::*; use bevy_interleave::prelude::*; +use bytemuck::{ + Pod, + Zeroable, +}; +use serde::{ + Deserialize, + Serialize, +}; #[allow(unused_imports)] use crate::{ @@ -19,7 +30,6 @@ use crate::{ TestCloud, }, iter::PositionIter, - packed::{Gaussian3d, PlanarGaussian3d}, settings::CloudSettings, }, material::spherical_harmonics::{ @@ -29,6 +39,35 @@ use crate::{ }, }; + + +#[derive( + Clone, + Debug, + Default, + Copy, + PartialEq, + Planar, + ReflectInterleaved, + StorageBindings, + Reflect, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct Gaussian3d { + pub position_visibility: PositionVisibility, + pub spherical_harmonic: SphericalHarmonicCoefficients, + pub rotation: Rotation, + pub scale_opacity: ScaleOpacity, +} + +pub type Gaussian2d = Gaussian3d; // GaussianMode::Gaussian2d /w Gaussian3d structure + + + // #[allow(unused_imports)] // #[cfg(feature = "f16")] // use crate::gaussian::f16::{ @@ -94,6 +133,66 @@ impl From> for PlanarGaussian3d { } +impl Distribution for rand::distributions::Standard { + fn sample(&self, rng: &mut R) -> Gaussian3d { + Gaussian3d { + rotation: [ + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + ].into(), + position_visibility: [ + rng.gen_range(-20.0..20.0), + rng.gen_range(-20.0..20.0), + rng.gen_range(-20.0..20.0), + 1.0, + ].into(), + scale_opacity: [ + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..0.8), + ].into(), + spherical_harmonic: SphericalHarmonicCoefficients { + coefficients: { + // #[cfg(feature = "f16")] + // { + // let mut coefficients: [u32; HALF_SH_COEFF_COUNT] = [0; HALF_SH_COEFF_COUNT]; + // for coefficient in coefficients.iter_mut() { + // let upper = rng.gen_range(-1.0..1.0); + // let lower = rng.gen_range(-1.0..1.0); + + // *coefficient = pack_f32s_to_u32(upper, lower); + // } + // coefficients + // } + + { + let mut coefficients = [0.0; SH_COEFF_COUNT]; + for coefficient in coefficients.iter_mut() { + *coefficient = rng.gen_range(-1.0..1.0); + } + coefficients + } + }, + }, + } + } +} + +pub fn random_gaussians_3d(n: usize) -> PlanarGaussian3d { + let mut rng = rand::thread_rng(); + let mut gaussians: Vec = Vec::with_capacity(n); + + for _ in 0..n { + gaussians.push(rng.gen()); + } + + PlanarGaussian3d::from_interleaved(gaussians) +} + + impl TestCloud for PlanarGaussian3d { fn test_model() -> Self { let mut rng = rand::thread_rng(); diff --git a/src/gaussian/formats/planar_4d.rs b/src/gaussian/formats/planar_4d.rs index ad2293b5..df986301 100644 --- a/src/gaussian/formats/planar_4d.rs +++ b/src/gaussian/formats/planar_4d.rs @@ -1,17 +1,76 @@ +use std::marker::Copy; + +use bevy::prelude::*; use bevy_interleave::prelude::*; +use bytemuck::{ + Pod, + Zeroable, +}; +use rand::{ + prelude::Distribution, + Rng, +}; +use serde::{ + Deserialize, + Serialize, +}; use crate::{ gaussian::{ + f32::{ + IsotropicRotations, + PositionVisibility, + ScaleOpacity, + TimestampTimescale, + }, interface::{ CommonCloud, TestCloud, }, iter::PositionIter, - packed::{Gaussian4d, PlanarGaussian4d}, }, - random_gaussians_4d, + material::spherindrical_harmonics::{ + SH_4D_COEFF_COUNT, + SpherindricalHarmonicCoefficients, + }, }; + + +#[derive( + Clone, + Debug, + Default, + Copy, + PartialEq, + Planar, + ReflectInterleaved, + StorageBindings, + Reflect, + Pod, + Zeroable, + Serialize, + Deserialize, +)] +#[repr(C)] +pub struct Gaussian4d { + pub position_visibility: PositionVisibility, + pub spherindrical_harmonic: SpherindricalHarmonicCoefficients, + pub isotropic_rotations: IsotropicRotations, + pub scale_opacity: ScaleOpacity, + pub timestamp_timescale: TimestampTimescale, +} + +// // TODO: GaussianSpacetime, determine temporal position/rotation structure +// pub struct GaussianSpacetime { +// pub position_visibility: PositionVisibility, +// pub color_mlp: ColorMlp, +// pub isotropic_rotations: IsotropicRotations, +// pub scale_opacity: ScaleOpacity, +// pub timestamp_timescale: TimestampTimescale, +// } + + // TODO: quantize 4d representation // #[derive( // Debug, @@ -197,8 +256,62 @@ impl From> for PlanarGaussian4d { } +impl Distribution for rand::distributions::Standard { + fn sample(&self, rng: &mut R) -> Gaussian4d { + let mut coefficients = [0.0; SH_4D_COEFF_COUNT]; + for coefficient in coefficients.iter_mut() { + *coefficient = rng.gen_range(-1.0..1.0); + } + + Gaussian4d { + isotropic_rotations: [ + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + rng.gen_range(-1.0..1.0), + ].into(), + position_visibility: [ + rng.gen_range(-20.0..20.0), + rng.gen_range(-20.0..20.0), + rng.gen_range(-20.0..20.0), + 1.0, + ].into(), + scale_opacity: [ + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..1.0), + rng.gen_range(0.0..0.8), + ].into(), + spherindrical_harmonic: coefficients.into(), + timestamp_timescale: [ + rng.gen_range(0.0..1.0), + rng.gen_range(-1.0..1.0), + 0.0, + 0.0, + ].into(), + } + } +} + +pub fn random_gaussians_4d(n: usize) -> PlanarGaussian4d { + let mut rng = rand::thread_rng(); + let mut gaussians: Vec = Vec::with_capacity(n); + + for _ in 0..n { + gaussians.push(rng.gen()); + } + + PlanarGaussian4d::from_interleaved(gaussians) +} + + impl TestCloud for PlanarGaussian4d { fn test_model() -> Self { random_gaussians_4d(512) } } + diff --git a/src/gaussian/formats/planar_4d_hierarchy.rs b/src/gaussian/formats/planar_4d_hierarchy.rs index 96671f85..6bf813ff 100644 --- a/src/gaussian/formats/planar_4d_hierarchy.rs +++ b/src/gaussian/formats/planar_4d_hierarchy.rs @@ -1,5 +1,5 @@ // TODO: gaussian cloud 4d with temporal hierarchy -use crate::gaussian::packed::PlanarGaussian4dHandle; +use crate::gaussian::formats::planar_4d::PlanarGaussian4dHandle; pub struct TemporalGaussianLevel { diff --git a/src/gaussian/mod.rs b/src/gaussian/mod.rs index fefcabc9..15b0083d 100644 --- a/src/gaussian/mod.rs +++ b/src/gaussian/mod.rs @@ -7,8 +7,6 @@ pub mod f32; pub mod formats; pub mod interface; pub mod iter; -pub mod packed; -pub mod rand; pub mod settings; diff --git a/src/gaussian/packed.rs b/src/gaussian/packed.rs deleted file mode 100644 index a079013d..00000000 --- a/src/gaussian/packed.rs +++ /dev/null @@ -1,87 +0,0 @@ -use std::marker::Copy; - -use bevy::prelude::*; -use bevy_interleave::prelude::*; -use bytemuck::{ - Pod, - Zeroable, -}; -use serde::{ - Deserialize, - Serialize, -}; - -use crate::{ - gaussian::f32::{ - IsotropicRotations, - PositionVisibility, - Rotation, - ScaleOpacity, - TimestampTimescale, - }, - material::{ - spherical_harmonics::SphericalHarmonicCoefficients, - spherindrical_harmonics::SpherindricalHarmonicCoefficients, - }, -}; - -#[derive( - Clone, - Debug, - Default, - Copy, - PartialEq, - Planar, - ReflectInterleaved, - StorageBindings, - Reflect, - Pod, - Zeroable, - Serialize, - Deserialize, -)] -#[repr(C)] -pub struct Gaussian3d { - pub position_visibility: PositionVisibility, - pub spherical_harmonic: SphericalHarmonicCoefficients, - pub rotation: Rotation, - pub scale_opacity: ScaleOpacity, -} - -// GaussianMode::Gaussian2d /w Gaussian3d structure -pub type Gaussian2d = Gaussian3d; - - -#[derive( - Clone, - Debug, - Default, - Copy, - PartialEq, - Planar, - ReflectInterleaved, - StorageBindings, - Reflect, - Pod, - Zeroable, - Serialize, - Deserialize, -)] -#[repr(C)] -pub struct Gaussian4d { - pub position_visibility: PositionVisibility, - pub spherindrical_harmonic: SpherindricalHarmonicCoefficients, - pub isotropic_rotations: IsotropicRotations, - pub scale_opacity: ScaleOpacity, - pub timestamp_timescale: TimestampTimescale, -} - - -// // TODO: GaussianSpacetime, determine temporal position/rotation structure -// pub struct GaussianSpacetime { -// pub position_visibility: PositionVisibility, -// pub color_mlp: ColorMlp, -// pub isotropic_rotations: IsotropicRotations, -// pub scale_opacity: ScaleOpacity, -// pub timestamp_timescale: TimestampTimescale, -// } diff --git a/src/gaussian/rand.rs b/src/gaussian/rand.rs deleted file mode 100644 index 08e88ba3..00000000 --- a/src/gaussian/rand.rs +++ /dev/null @@ -1,145 +0,0 @@ -use bevy_interleave::prelude::Planar; -use rand::{ - prelude::Distribution, - Rng, -}; - -// #[cfg(feature = "f16")] -// use crate::gaussian::f16::pack_f32s_to_u32; - -#[allow(unused_imports)] -use crate::{ - gaussian::packed::{ - Gaussian3d, - Gaussian4d, - PlanarGaussian3d, - PlanarGaussian4d, - }, - material::{ - spherical_harmonics::{ - HALF_SH_COEFF_COUNT, - SH_COEFF_COUNT, - SphericalHarmonicCoefficients, - }, - spherindrical_harmonics::{ - HALF_SH_4D_COEFF_COUNT, - SH_4D_COEFF_COUNT, - SpherindricalHarmonicCoefficients, - }, - }, -}; - - -impl Distribution for rand::distributions::Standard { - fn sample(&self, rng: &mut R) -> Gaussian3d { - Gaussian3d { - rotation: [ - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - ].into(), - position_visibility: [ - rng.gen_range(-20.0..20.0), - rng.gen_range(-20.0..20.0), - rng.gen_range(-20.0..20.0), - 1.0, - ].into(), - scale_opacity: [ - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..0.8), - ].into(), - spherical_harmonic: SphericalHarmonicCoefficients { - coefficients: { - // #[cfg(feature = "f16")] - // { - // let mut coefficients: [u32; HALF_SH_COEFF_COUNT] = [0; HALF_SH_COEFF_COUNT]; - // for coefficient in coefficients.iter_mut() { - // let upper = rng.gen_range(-1.0..1.0); - // let lower = rng.gen_range(-1.0..1.0); - - // *coefficient = pack_f32s_to_u32(upper, lower); - // } - // coefficients - // } - - { - let mut coefficients = [0.0; SH_COEFF_COUNT]; - for coefficient in coefficients.iter_mut() { - *coefficient = rng.gen_range(-1.0..1.0); - } - coefficients - } - }, - }, - } - } -} - - -impl Distribution for rand::distributions::Standard { - fn sample(&self, rng: &mut R) -> Gaussian4d { - let mut coefficients = [0.0; SH_4D_COEFF_COUNT]; - for coefficient in coefficients.iter_mut() { - *coefficient = rng.gen_range(-1.0..1.0); - } - - Gaussian4d { - isotropic_rotations: [ - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - rng.gen_range(-1.0..1.0), - ].into(), - position_visibility: [ - rng.gen_range(-20.0..20.0), - rng.gen_range(-20.0..20.0), - rng.gen_range(-20.0..20.0), - 1.0, - ].into(), - scale_opacity: [ - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..1.0), - rng.gen_range(0.0..0.8), - ].into(), - spherindrical_harmonic: coefficients.into(), - timestamp_timescale: [ - rng.gen_range(0.0..1.0), - rng.gen_range(-1.0..1.0), - 0.0, - 0.0, - ].into(), - } - } -} - - - -pub fn random_gaussians_3d(n: usize) -> PlanarGaussian3d { - let mut rng = rand::thread_rng(); - let mut gaussians: Vec = Vec::with_capacity(n); - - for _ in 0..n { - gaussians.push(rng.gen()); - } - - PlanarGaussian3d::from_interleaved(gaussians) -} - -pub fn random_gaussians_4d(n: usize) -> PlanarGaussian4d { - let mut rng = rand::thread_rng(); - let mut gaussians: Vec = Vec::with_capacity(n); - - for _ in 0..n { - gaussians.push(rng.gen()); - } - - PlanarGaussian4d::from_interleaved(gaussians) -} diff --git a/src/io/gcloud/flexbuffers.rs b/src/io/gcloud/flexbuffers.rs index 21c97167..bdd1faf0 100644 --- a/src/io/gcloud/flexbuffers.rs +++ b/src/io/gcloud/flexbuffers.rs @@ -9,9 +9,9 @@ use serde::{ use crate::{ io::codec::CloudCodec, - gaussian::packed::{ - PlanarGaussian3d, - PlanarGaussian4d, + gaussian::formats::{ + planar_3d::PlanarGaussian3d, + planar_4d::PlanarGaussian4d, }, }; diff --git a/src/io/loader.rs b/src/io/loader.rs index 7737f0be..7de238bd 100644 --- a/src/io/loader.rs +++ b/src/io/loader.rs @@ -13,7 +13,7 @@ use bevy::asset::{ use crate::{ io::codec::CloudCodec, - gaussian::packed::PlanarGaussian3d, + gaussian::formats::planar_3d::PlanarGaussian3d, }; diff --git a/src/io/ply.rs b/src/io/ply.rs index b8a97249..f4a8b6c6 100644 --- a/src/io/ply.rs +++ b/src/io/ply.rs @@ -10,7 +10,10 @@ use ply_rs::{ }; use crate::{ - gaussian::packed::{Gaussian3d, PlanarGaussian3d}, + gaussian::formats::planar_3d::{ + Gaussian3d, + PlanarGaussian3d, + }, material::spherical_harmonics::{ SH_CHANNELS, SH_COEFF_COUNT, diff --git a/src/lib.rs b/src/lib.rs index 83a09050..3e392391 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -7,16 +7,20 @@ use bevy_interleave::prelude::*; pub use camera::GaussianCamera; pub use gaussian::{ - packed::{ - Gaussian3d, - Gaussian4d, - PlanarGaussian3d, - PlanarGaussian4d, - PlanarGaussian3dHandle, - PlanarGaussian4dHandle, + formats::{ + planar_3d::{ + Gaussian3d, + PlanarGaussian3d, + PlanarGaussian3dHandle, + random_gaussians_3d, + }, + planar_4d::{ + Gaussian4d, + PlanarGaussian4d, + PlanarGaussian4dHandle, + random_gaussians_4d, + }, }, - rand::random_gaussians_3d, - rand::random_gaussians_4d, settings::{ RasterizeMode, CloudSettings, diff --git a/src/query/select.rs b/src/query/select.rs index 4e947117..227e279d 100644 --- a/src/query/select.rs +++ b/src/query/select.rs @@ -4,7 +4,10 @@ use bevy_interleave::prelude::*; use crate::{ gaussian::{ interface::CommonCloud, - packed::{Gaussian3d, Gaussian4d}, + formats::{ + planar_3d::Gaussian3d, + planar_4d::Gaussian4d, + }, }, io::codec::CloudCodec, }; diff --git a/src/sort/mod.rs b/src/sort/mod.rs index c308cd4f..1edfde9d 100644 --- a/src/sort/mod.rs +++ b/src/sort/mod.rs @@ -35,11 +35,11 @@ use crate::{ camera::GaussianCamera, CloudSettings, gaussian::{ - interface::CommonCloud, - packed::{ - Gaussian3d, - Gaussian4d, + formats::{ + planar_3d::Gaussian3d, + planar_4d::Gaussian4d, }, + interface::CommonCloud, }, }; diff --git a/src/sort/rayon.rs b/src/sort/rayon.rs index df8e272d..d483afbc 100644 --- a/src/sort/rayon.rs +++ b/src/sort/rayon.rs @@ -10,8 +10,11 @@ use crate::{ camera::GaussianCamera, CloudSettings, gaussian::{ + formats::{ + planar_3d::Gaussian3d, + planar_4d::Gaussian4d, + }, interface::CommonCloud, - packed::Gaussian3d, }, sort::{ SortConfig, @@ -29,6 +32,7 @@ pub struct RayonSortPlugin; impl Plugin for RayonSortPlugin { fn build(&self, app: &mut App) { app.add_systems(Update, rayon_sort::); + app.add_systems(Update, rayon_sort::); } } diff --git a/src/sort/std.rs b/src/sort/std.rs index ddc00a0c..54942b83 100644 --- a/src/sort/std.rs +++ b/src/sort/std.rs @@ -9,8 +9,11 @@ use crate::{ camera::GaussianCamera, CloudSettings, gaussian::{ + formats::{ + planar_3d::Gaussian3d, + planar_4d::Gaussian4d, + }, interface::CommonCloud, - packed::Gaussian3d, }, sort::{ SortConfig, @@ -28,6 +31,7 @@ pub struct StdSortPlugin; impl Plugin for StdSortPlugin { fn build(&self, app: &mut App) { app.add_systems(Update, std_sort::); + app.add_systems(Update, std_sort::); } }