Skip to content

Commit

Permalink
feat: oriented bounding box and editor UI (#3)
Browse files Browse the repository at this point in the history
* feat: aligned rects (unstable eigenvectors)

* refactor: OBB and AABB switching

* feat: editor and example scaffolding

* chore: fix todo tracker permissions

* feat: inspector UI & global transform implementation

* feat: pipeline specialization through UI

* feat: ply -> gcloud conversion and loading (#4)
  • Loading branch information
mosure authored Oct 15, 2023
1 parent 0847923 commit 5068524
Show file tree
Hide file tree
Showing 12 changed files with 407 additions and 282 deletions.
3 changes: 3 additions & 0 deletions .github/workflows/todo_tracker.yml
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,9 @@ on:

jobs:
build:
permissions:
issues: write

name: todo_tracker
runs-on: [ubuntu-latest]

Expand Down
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -6,5 +6,6 @@ Cargo.lock


*.ply
*.gcloud

.DS_Store
20 changes: 19 additions & 1 deletion Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,14 +10,20 @@ categories = ["computer-vision", "graphics", "rendering", "rendering::data-forma
homepage = "https://github.com/mosure/bevy_gaussian_splatting"
repository = "https://github.com/mosure/bevy_gaussian_splatting"
readme = "README.md"
include = ["tools"]
exclude = [".devcontainer", ".github", "docs", "dist", "build", "assets", "credits"]
default-run = "bevy_gaussian_splatting"

# TODO: use minimal bevy features
[dependencies]
bevy = "0.11.2"
bevy = "0.11.3"
bevy-inspector-egui = "0.20.0"
bevy_panorbit_camera = "0.8.0"
bincode2 = "2.0.1"
bytemuck = "1.14.0"
flate2 = "1.0.28"
ply-rs = "0.1.3"
serde = "1.0.189"


[target.'cfg(target_arch = "wasm32")'.dependencies]
Expand Down Expand Up @@ -51,3 +57,15 @@ inherits = "release"
opt-level = "z"
lto = "fat"
codegen-units = 1


[lib]
path = "src/lib.rs"

[[bin]]
name = "bevy_gaussian_splatting"
path = "src/main.rs"

[[bin]]
name = "ply_to_gcloud"
path = "tools/ply_to_gcloud.rs"
8 changes: 7 additions & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -12,9 +12,13 @@

bevy gaussian splatting render pipeline plugin

`cargo run -- {path to ply or gcloud.gz file}`

## capabilities

- [ ] bevy gaussian cloud render pipeline
- [X] ply to gcloud converter
- [X] gcloud and ply asset loaders
- [X] bevy gaussian cloud render pipeline
- [ ] 4D gaussian clouds via morph targets
- [ ] bevy 3D camera to gaussian cloud pipeline

Expand Down Expand Up @@ -55,11 +59,13 @@ fn setup_gaussian_cloud(

# credits

- [4d gaussians](https://github.com/hustvl/4DGaussians)
- [bevy](https://github.com/bevyengine/bevy)
- [bevy-hanabi](https://github.com/djeedai/bevy_hanabi)
- [diff-gaussian-rasterization](https://github.com/graphdeco-inria/diff-gaussian-rasterization)
- [dreamgaussian](https://github.com/dreamgaussian/dreamgaussian)
- [dynamic-3d-gaussians](https://github.com/JonathonLuiten/Dynamic3DGaussians)
- [ewa splatting](https://www.cs.umd.edu/~zwicker/publications/EWASplatting-TVCG02.pdf)
- [gaussian-splatting](https://github.com/graphdeco-inria/gaussian-splatting)
- [gaussian-splatting-web](https://github.com/cvlab-epfl/gaussian-splatting-web)
- [making gaussian splats smaller](https://aras-p.info/blog/2023/09/13/Making-Gaussian-Splats-smaller/)
Expand Down
Binary file added assets/scenes/icecream.gcloud
Binary file not shown.
131 changes: 110 additions & 21 deletions src/gaussian.rs
Original file line number Diff line number Diff line change
Expand Up @@ -13,17 +13,22 @@ use bevy::{
LoadContext,
LoadedAsset,
},
reflect::{
TypePath,
TypeUuid,
},
reflect::TypeUuid,
render::render_resource::ShaderType,
utils::BoxedFuture,
};
use bincode2::deserialize_from;
use bytemuck::{
Pod,
Zeroable,
};
use flate2::read::GzDecoder;
use serde::{
Deserialize,
Serialize,
Serializer,
ser::SerializeTuple,
};

use crate::ply::parse_ply;

Expand All @@ -37,9 +42,19 @@ const fn num_sh_coefficients(degree: usize) -> usize {
}
const SH_DEGREE: usize = 3;
pub const MAX_SH_COEFF_COUNT: usize = num_sh_coefficients(SH_DEGREE) * 3;
#[derive(Clone, Copy, ShaderType, Pod, Zeroable)]
#[derive(
Clone,
Copy,
Reflect,
ShaderType,
Pod,
Zeroable,
Serialize,
Deserialize,
)]
#[repr(C)]
pub struct SphericalHarmonicCoefficients {
#[serde(serialize_with = "coefficients_serializer", deserialize_with = "coefficients_deserializer")]
pub coefficients: [f32; MAX_SH_COEFF_COUNT],
}
impl Default for SphericalHarmonicCoefficients {
Expand All @@ -49,12 +64,64 @@ impl Default for SphericalHarmonicCoefficients {
}
}
}
fn coefficients_serializer<S>(n: &[f32; MAX_SH_COEFF_COUNT], s: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
let mut tup = s.serialize_tuple(MAX_SH_COEFF_COUNT)?;
for &x in n.iter() {
tup.serialize_element(&x)?;
}

tup.end()
}

fn coefficients_deserializer<'de, D>(d: D) -> Result<[f32; MAX_SH_COEFF_COUNT], D::Error>
where
D: serde::Deserializer<'de>,
{
struct CoefficientsVisitor;

impl<'de> serde::de::Visitor<'de> for CoefficientsVisitor {
type Value = [f32; MAX_SH_COEFF_COUNT];

fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
formatter.write_str("an array of floats")
}

fn visit_seq<A>(self, mut seq: A) -> Result<[f32; MAX_SH_COEFF_COUNT], A::Error>
where
A: serde::de::SeqAccess<'de>,
{
let mut coefficients = [0.0; MAX_SH_COEFF_COUNT];
for i in 0..MAX_SH_COEFF_COUNT {
coefficients[i] = seq
.next_element()?
.ok_or_else(|| serde::de::Error::invalid_length(i, &self))?;
}
Ok(coefficients)
}
}

d.deserialize_tuple(MAX_SH_COEFF_COUNT, CoefficientsVisitor)
}


#[derive(Clone, Default, Copy, ShaderType, Pod, Zeroable)]
pub const MAX_SIZE_VARIANCE: f32 = 5.0;

#[derive(
Clone,
Default,
Copy,
Reflect,
ShaderType,
Pod,
Zeroable,
Serialize,
Deserialize,
)]
#[repr(C)]
pub struct Gaussian {
//pub anisotropic_covariance: AnisotropicCovariance,
//pub normal: Vec3,
pub rotation: [f32; 4],
pub position: Vec3,
pub scale: Vec3,
Expand All @@ -63,7 +130,13 @@ pub struct Gaussian {
padding: f32,
}

#[derive(Clone, TypeUuid, TypePath)]
#[derive(
Clone,
Reflect,
TypeUuid,
Serialize,
Deserialize,
)]
#[uuid = "ac2f08eb-bc32-aabb-ff21-51571ea332d5"]
pub struct GaussianCloud(pub Vec<Gaussian>);

Expand Down Expand Up @@ -103,9 +176,9 @@ impl GaussianCloud {
};
let mut cloud = GaussianCloud(Vec::new());

for &x in [-1.0, 1.0].iter() {
for &y in [-1.0, 1.0].iter() {
for &z in [-1.0, 1.0].iter() {
for &x in [-0.5, 0.5].iter() {
for &y in [-0.5, 0.5].iter() {
for &z in [-0.5, 0.5].iter() {
let mut g = origin.clone();
g.position = Vec3::new(x, y, z);
cloud.0.push(g);
Expand All @@ -120,15 +193,19 @@ impl GaussianCloud {

#[derive(Component, Reflect, Clone)]
pub struct GaussianCloudSettings {
pub aabb: bool,
pub global_scale: f32,
pub global_transform: GlobalTransform,
pub visualize_bounding_box: bool,
}

impl Default for GaussianCloudSettings {
fn default() -> Self {
Self {
aabb: false,
global_scale: 1.0,
global_transform: Transform::IDENTITY.into(),
visualize_bounding_box: false,
}
}
}
Expand All @@ -144,18 +221,30 @@ impl AssetLoader for GaussianCloudLoader {
load_context: &'a mut LoadContext,
) -> BoxedFuture<'a, Result<(), bevy::asset::Error>> {
Box::pin(async move {
let cursor = Cursor::new(bytes);
let mut f = BufReader::new(cursor);

let ply_cloud = parse_ply(&mut f)?;
let cloud = GaussianCloud(ply_cloud);

load_context.set_default_asset(LoadedAsset::new(cloud));
Ok(())
match load_context.path().extension() {
Some(ext) if ext == "ply" => {
let cursor = Cursor::new(bytes);
let mut f = BufReader::new(cursor);

let ply_cloud = parse_ply(&mut f)?;
let cloud = GaussianCloud(ply_cloud);

load_context.set_default_asset(LoadedAsset::new(cloud));
return Ok(());
},
Some(ext) if ext == "gcloud" => {
let decompressed = GzDecoder::new(bytes);
let cloud: GaussianCloud = deserialize_from(decompressed).expect("failed to decode cloud");

load_context.set_default_asset(LoadedAsset::new(cloud));
return Ok(());
},
_ => Ok(()),
}
})
}

fn extensions(&self) -> &[&str] {
&["ply"]
&["ply", "gcloud"]
}
}
5 changes: 4 additions & 1 deletion src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ pub mod utils;

#[derive(Bundle, Default, Reflect)]
pub struct GaussianSplattingBundle {
pub settings: GaussianCloudSettings, // TODO: implement global transform
pub settings: GaussianCloudSettings,
pub cloud: Handle<GaussianCloud>,
}

Expand All @@ -31,9 +31,12 @@ pub struct GaussianSplattingPlugin;

impl Plugin for GaussianSplattingPlugin {
fn build(&self, app: &mut App) {
// TODO: allow hot reloading of GaussianCloud handle through inspector UI
app.add_asset::<GaussianCloud>();
app.init_asset_loader::<GaussianCloudLoader>();

app.register_asset_reflect::<GaussianCloud>();
app.register_type::<GaussianCloudSettings>();
app.register_type::<GaussianSplattingBundle>();

app.add_plugins((
Expand Down
31 changes: 26 additions & 5 deletions src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -6,20 +6,23 @@ use bevy::{
FrameTimeDiagnosticsPlugin,
},
};
use bevy_inspector_egui::quick::WorldInspectorPlugin;
use bevy_panorbit_camera::{
PanOrbitCamera,
PanOrbitCameraPlugin,
};

use bevy_gaussian_splatting::{
GaussianCloud,
GaussianCloudSettings,
GaussianSplattingBundle,
GaussianSplattingPlugin,
utils::setup_hooks,
};


pub struct GaussianSplattingViewer {
pub editor: bool,
pub esc_close: bool,
pub show_fps: bool,
pub width: f32,
Expand All @@ -30,6 +33,7 @@ pub struct GaussianSplattingViewer {
impl Default for GaussianSplattingViewer {
fn default() -> GaussianSplattingViewer {
GaussianSplattingViewer {
editor: true,
esc_close: true,
show_fps: true,
width: 1920.0,
Expand All @@ -42,14 +46,27 @@ impl Default for GaussianSplattingViewer {

fn setup_gaussian_cloud(
mut commands: Commands,
_asset_server: Res<AssetServer>,
asset_server: Res<AssetServer>,
mut gaussian_assets: ResMut<Assets<GaussianCloud>>,
) {
let cloud = gaussian_assets.add(GaussianCloud::test_model());
let cloud: Handle<GaussianCloud>;
let settings = GaussianCloudSettings {
aabb: true,
visualize_bounding_box: false,
..default()
};

let filename = std::env::args().nth(1);
if let Some(filename) = filename {
println!("loading {}", filename);
cloud = asset_server.load(filename.as_str());
} else {
cloud = gaussian_assets.add(GaussianCloud::test_model());
}

commands.spawn(GaussianSplattingBundle {
cloud,
// cloud: _asset_server.load("scenes/icecream.ply"),
..Default::default()
settings,
});

commands.spawn((
Expand Down Expand Up @@ -82,12 +99,16 @@ fn example_app() {
..default()
}),
..default()
})
}),
);
app.add_plugins((
PanOrbitCameraPlugin,
));

if config.editor {
app.add_plugins(WorldInspectorPlugin::new());
}

if config.esc_close {
app.add_systems(Update, esc_close);
}
Expand Down
Loading

0 comments on commit 5068524

Please sign in to comment.