diff --git a/.github/workflows/bench.yml b/.github/workflows/bench.yml new file mode 100644 index 00000000..fa4ab25b --- /dev/null +++ b/.github/workflows/bench.yml @@ -0,0 +1,41 @@ +name: bench + +on: + pull_request: + types: [ labeled, synchronize ] + branches: [ "main" ] + +env: + CARGO_TERM_COLOR: always + RUST_BACKTRACE: 1 + +jobs: + bench: + if: contains(github.event.pull_request.labels.*.name, 'bench') + + strategy: + fail-fast: false + matrix: + os: [windows-latest, macos-latest] + + runs-on: ${{ matrix.os }} + timeout-minutes: 120 + + steps: + - uses: actions/checkout@v3 + - uses: actions/cache@v3 + with: + path: | + ~/.cargo/bin/ + ~/.cargo/registry/index/ + ~/.cargo/registry/cache/ + ~/.cargo/git/db/ + target/ + key: ${{ runner.os }}-cargo-build-stable-${{ hashFiles('**/Cargo.toml') }} + + - name: io benchmark + uses: boa-dev/criterion-compare-action@v3.2.4 + with: + benchName: "io" + branchName: ${{ github.base_ref }} + token: ${{ secrets.GITHUB_TOKEN }} diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 7deb85cd..058b9203 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -33,9 +33,15 @@ jobs: target/ key: ${{ runner.os }}-cargo-build-stable-${{ hashFiles('**/Cargo.toml') }} + - name: check + run: cargo check + - name: build run: cargo build + - name: build_tools + run: cargo build --bin ply_to_gcloud + - name: lint run: cargo clippy diff --git a/Cargo.toml b/Cargo.toml index 6812476c..ee21d62d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,22 +14,31 @@ exclude = [".devcontainer", ".github", "docs", "dist", "build", "assets", "credi default-run = "viewer" -# TODO: separate dev-dependencies +[features] +default = ["io_flexbuffers", "io_ply"] + +io_bincode2 = ["bincode2", "flate2"] +io_flexbuffers = ["flexbuffers"] +io_ply = ["ply-rs"] + + [dependencies] bevy-inspector-egui = "0.21" -bevy_panorbit_camera = "0.9.0" -bincode2 = "2.0.1" -bytemuck = "1.14.0" -flate2 = "1.0.28" -ply-rs = "0.1.3" -rand = "0.8.5" -serde = "1.0.189" +bevy_panorbit_camera = "0.9" +bincode2 = { version = "2.0", optional = true } +byte-unit = "4.0" +bytemuck = "1.14" +flate2 = { version = "1.0", optional = true } +flexbuffers = { version = "2.0", optional = true } +ply-rs = { version = "0.1", optional = true } +rand = "0.8" +serde = "1.0" wgpu = "0.17.1" [target.'cfg(target_arch = "wasm32")'.dependencies] -console_error_panic_hook = "0.1.7" -wasm-bindgen = "0.2.87" +console_error_panic_hook = "0.1" +wasm-bindgen = "0.2" # TODO: use minimal bevy features @@ -49,6 +58,10 @@ features = [ ] +[dev-dependencies] +criterion = { version = "0.5", features = ["html_reports"] } + + [profile.dev.package."*"] opt-level = 3 @@ -77,3 +90,9 @@ path = "viewer/viewer.rs" [[bin]] name = "ply_to_gcloud" path = "tools/ply_to_gcloud.rs" +required-features = ["io_ply"] + + +[[bench]] +name = "io" +harness = false diff --git a/README.md b/README.md index 38b9c9bc..87bece76 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ bevy gaussian splatting render pipeline plugin - [ ] skeletons - [ ] volume masks - [ ] level of detail +- [ ] lighting and shadows - [ ] gaussian cloud particle effects (accelerated spatial queries) - [ ] bevy_openxr support - [ ] bevy 3D camera to gaussian cloud pipeline @@ -94,6 +95,7 @@ to build wasm run: - [making gaussian splats smaller](https://aras-p.info/blog/2023/09/13/Making-Gaussian-Splats-smaller/) - [masked-spacetime-hashing](https://github.com/masked-spacetime-hashing/msth) - [onesweep](https://arxiv.org/ftp/arxiv/papers/2206/2206.01784.pdf) +- [pasture](https://github.com/Mortano/pasture) - [point-visualizer](https://github.com/mosure/point-visualizer) - [rusty-automata](https://github.com/mosure/rusty-automata) - [splat](https://github.com/antimatter15/splat) diff --git a/assets/scenes/icecream.gcloud b/assets/scenes/icecream.gcloud index 884a0c86..b3cb469f 100644 Binary files a/assets/scenes/icecream.gcloud and b/assets/scenes/icecream.gcloud differ diff --git a/benches/io.rs b/benches/io.rs new file mode 100644 index 00000000..4e060d93 --- /dev/null +++ b/benches/io.rs @@ -0,0 +1,47 @@ +use criterion::{ + BenchmarkId, + criterion_group, + criterion_main, + Criterion, + Throughput, +}; + +use bevy_gaussian_splatting::{ + Gaussian, + GaussianCloud, + io::codec::GaussianCloudCodec, + random_gaussians, +}; + + +const GAUSSIAN_COUNTS: [usize; 4] = [ + 1000, + 10000, + 84_348, + 1_244_819, + // 6_131_954, +]; + +fn gaussian_cloud_decode_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("encode gaussian clouds"); + for count in GAUSSIAN_COUNTS.iter() { + group.throughput(Throughput::Bytes(*count as u64 * std::mem::size_of::() as u64)); + group.bench_with_input( + BenchmarkId::new("decode", count), + &count, + |b, &count| { + let gaussians = random_gaussians(*count); + let bytes = gaussians.encode(); + + b.iter(|| GaussianCloud::decode(bytes.as_slice())); + }, + ); + } +} + +criterion_group!{ + name = io_benches; + config = Criterion::default().sample_size(10); + targets = gaussian_cloud_decode_benchmark +} +criterion_main!(io_benches); diff --git a/src/gaussian.rs b/src/gaussian.rs index cae4458a..92682154 100644 --- a/src/gaussian.rs +++ b/src/gaussian.rs @@ -1,31 +1,19 @@ -use rand::{seq::SliceRandom, prelude::Distribution, Rng}; -use std::{ - io::{ - BufReader, - Cursor, - ErrorKind, - }, - marker::Copy, +use rand::{ + seq::SliceRandom, + prelude::Distribution, + Rng, }; +use std::marker::Copy; use bevy::{ prelude::*, - asset::{ - AssetLoader, - AsyncReadExt, - LoadContext, - io::Reader, - }, reflect::TypeUuid, render::render_resource::ShaderType, - utils::BoxedFuture, }; -use bincode2::deserialize_from; use bytemuck::{ Pod, Zeroable, }; -use flate2::read::GzDecoder; use serde::{ Deserialize, Serialize, @@ -33,8 +21,6 @@ use serde::{ ser::SerializeTuple, }; -use crate::ply::parse_ply; - const fn num_sh_coefficients(degree: usize) -> usize { if degree == 0 { @@ -51,6 +37,7 @@ pub const MAX_SH_COEFF_COUNT: usize = MAX_SH_COEFF_COUNT_PER_CHANNEL * SH_CHANNE Clone, Copy, Debug, + PartialEq, Reflect, ShaderType, Pod, @@ -121,6 +108,7 @@ pub const MAX_SIZE_VARIANCE: f32 = 5.0; Debug, Default, Copy, + PartialEq, Reflect, ShaderType, Pod, @@ -141,6 +129,7 @@ pub struct Gaussian { Asset, Clone, Debug, + PartialEq, Reflect, TypeUuid, Serialize, @@ -231,53 +220,6 @@ impl Default for GaussianCloudSettings { } } - -#[derive(Default)] -pub struct GaussianCloudLoader; - -impl AssetLoader for GaussianCloudLoader { - type Asset = GaussianCloud; - type Settings = (); - type Error = std::io::Error; - - fn load<'a>( - &'a self, - reader: &'a mut Reader, - _settings: &'a Self::Settings, - load_context: &'a mut LoadContext, - ) -> BoxedFuture<'a, Result> { - - Box::pin(async move { - let mut bytes = Vec::new(); - reader.read_to_end(&mut bytes).await?; - - 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); - - Ok(cloud) - }, - Some(ext) if ext == "gcloud" => { - let decompressed = GzDecoder::new(bytes.as_slice()); - let cloud: GaussianCloud = deserialize_from(decompressed).expect("failed to decode cloud"); - - Ok(cloud) - }, - _ => Err(std::io::Error::new(ErrorKind::Other, "only .ply and .gcloud supported")), - } - }) - } - - fn extensions(&self) -> &[&str] { - &["ply", "gcloud"] - } -} - - impl Distribution for rand::distributions::Standard { fn sample(&self, rng: &mut R) -> Gaussian { Gaussian { diff --git a/src/io/codec.rs b/src/io/codec.rs new file mode 100644 index 00000000..01c720e5 --- /dev/null +++ b/src/io/codec.rs @@ -0,0 +1,6 @@ + +// TODO: support streamed codecs +pub trait GaussianCloudCodec { + fn encode(&self) -> Vec; + fn decode(data: &[u8]) -> Self; +} diff --git a/src/io/gcloud/bincode2.rs b/src/io/gcloud/bincode2.rs new file mode 100644 index 00000000..7e9b0874 --- /dev/null +++ b/src/io/gcloud/bincode2.rs @@ -0,0 +1,35 @@ +use bincode2::{ + deserialize_from, + serialize_into, +}; +use flate2::{ + Compression, + read::GzDecoder, + write::GzEncoder, +}; + +use crate::{ + gaussian::GaussianCloud, + io::codec::GaussianCloudCodec, +}; + + +impl GaussianCloudCodec for GaussianCloud { + fn encode(&self) -> Vec { + let mut output = Vec::new(); + + { + let mut gz_encoder = GzEncoder::new(&mut output, Compression::default()); + serialize_into(&mut gz_encoder, &self).expect("failed to encode cloud"); + } + + output + } + + fn decode(data: &[u8]) -> Self { + let decompressed = GzDecoder::new(data); + let cloud: GaussianCloud = deserialize_from(decompressed).expect("failed to decode cloud"); + + cloud + } +} diff --git a/src/io/gcloud/flexbuffers.rs b/src/io/gcloud/flexbuffers.rs new file mode 100644 index 00000000..a2490d80 --- /dev/null +++ b/src/io/gcloud/flexbuffers.rs @@ -0,0 +1,30 @@ +use flexbuffers::{ + FlexbufferSerializer, + Reader, +}; +use serde::{ + Deserialize, + Serialize, +}; + +use crate::{ + GaussianCloud, + io::codec::GaussianCloudCodec, +}; + + +impl GaussianCloudCodec for GaussianCloud { + fn encode(&self) -> Vec { + let mut serializer = FlexbufferSerializer::new(); + self.serialize(&mut serializer).expect("failed to serialize cloud"); + + serializer.view().to_vec() + } + + fn decode(data: &[u8]) -> Self { + let reader = Reader::get_root(data).expect("failed to read flexbuffer"); + let cloud = GaussianCloud::deserialize(reader).expect("deserialization failed"); + + cloud + } +} diff --git a/src/io/gcloud/mod.rs b/src/io/gcloud/mod.rs new file mode 100644 index 00000000..50c5e235 --- /dev/null +++ b/src/io/gcloud/mod.rs @@ -0,0 +1,5 @@ +#[cfg(feature = "io_bincode2")] +pub mod bincode2; + +#[cfg(feature = "io_flexbuffers")] +pub mod flexbuffers; diff --git a/src/io/loader.rs b/src/io/loader.rs new file mode 100644 index 00000000..1fedea26 --- /dev/null +++ b/src/io/loader.rs @@ -0,0 +1,73 @@ +use std::io::{ + BufReader, + Cursor, + ErrorKind, +}; + +use bevy::{ + asset::{ + AssetLoader, + AsyncReadExt, + LoadContext, + io::Reader, + }, + utils::BoxedFuture, +}; + +use crate::{ + GaussianCloud, + io::codec::GaussianCloudCodec, +}; + + +#[derive(Default)] +pub struct GaussianCloudLoader; + +impl AssetLoader for GaussianCloudLoader { + type Asset = GaussianCloud; + type Settings = (); + type Error = std::io::Error; + + fn load<'a>( + &'a self, + reader: &'a mut Reader, + _settings: &'a Self::Settings, + load_context: &'a mut LoadContext, + ) -> BoxedFuture<'a, Result> { + + Box::pin(async move { + let mut bytes = Vec::new(); + reader.read_to_end(&mut bytes).await?; + + match load_context.path().extension() { + Some(ext) if ext == "ply" => { + #[cfg(feature = "io_ply")] + { + let cursor = Cursor::new(bytes); + let mut f = BufReader::new(cursor); + + let ply_cloud = crate::io::ply::parse_ply(&mut f)?; + let cloud = GaussianCloud(ply_cloud); + + Ok(cloud) + } + + #[cfg(not(feature = "io_ply"))] + { + Err(std::io::Error::new(ErrorKind::Other, "ply support not enabled, enable with io_ply feature")) + } + }, + Some(ext) if ext == "gcloud" => { + let cloud = GaussianCloud::decode(bytes.as_slice()); + + Ok(cloud) + }, + _ => Err(std::io::Error::new(ErrorKind::Other, "only .ply and .gcloud supported")), + } + }) + } + + fn extensions(&self) -> &[&str] { + &["ply", "gcloud"] + } +} diff --git a/src/io/mod.rs b/src/io/mod.rs new file mode 100644 index 00000000..cfe1eb24 --- /dev/null +++ b/src/io/mod.rs @@ -0,0 +1,6 @@ +pub mod codec; +pub mod gcloud; +pub mod loader; + +#[cfg(feature = "io_ply")] +pub mod ply; diff --git a/src/ply.rs b/src/io/ply.rs similarity index 100% rename from src/ply.rs rename to src/io/ply.rs diff --git a/src/lib.rs b/src/lib.rs index f5d9f3f4..56ad130b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -3,16 +3,17 @@ use bevy::prelude::*; pub use gaussian::{ Gaussian, GaussianCloud, - GaussianCloudLoader, GaussianCloudSettings, SphericalHarmonicCoefficients, random_gaussians, }; +use io::loader::GaussianCloudLoader; + use render::RenderPipelinePlugin; pub mod gaussian; -pub mod ply; +pub mod io; pub mod render; pub mod utils; diff --git a/tests/io.rs b/tests/io.rs new file mode 100644 index 00000000..5cb2a9fd --- /dev/null +++ b/tests/io.rs @@ -0,0 +1,17 @@ +use bevy_gaussian_splatting::{ + GaussianCloud, + io::codec::GaussianCloudCodec, + random_gaussians, +}; + + +#[test] +fn test_codec() { + let count = 100; + + let gaussians = random_gaussians(count); + let encoded = gaussians.encode(); + let decoded = GaussianCloud::decode(encoded.as_slice()); + + assert_eq!(gaussians, decoded); +} diff --git a/tools/ply_to_gcloud.rs b/tools/ply_to_gcloud.rs index be65bfe1..482cd7ea 100644 --- a/tools/ply_to_gcloud.rs +++ b/tools/ply_to_gcloud.rs @@ -1,12 +1,13 @@ -use bincode2::serialize_into; -use flate2::{ - Compression, - write::GzEncoder, -}; +use std::io::Write; + +use byte_unit::Byte; use bevy_gaussian_splatting::{ GaussianCloud, - ply::parse_ply, + io::{ + codec::GaussianCloudCodec, + ply::parse_ply, + }, }; @@ -15,23 +16,20 @@ fn main() { println!("converting {}", filename); - // filepath to BufRead let file = std::fs::File::open(&filename).expect("failed to open file"); let mut reader = std::io::BufReader::new(file); let cloud = GaussianCloud(parse_ply(&mut reader).expect("failed to parse ply file")); - // write cloud to .gcloud file (remove .ply) let base_filename = filename.split('.').next().expect("no extension").to_string(); let gcloud_filename = base_filename + ".gcloud"; - // let gcloud_file = std::fs::File::create(&gcloud_filename).expect("failed to create file"); - // let mut writer = std::io::BufWriter::new(gcloud_file); - // serialize_into(&mut writer, &cloud).expect("failed to encode cloud"); + let gcloud_file = std::fs::File::create(&gcloud_filename).expect("failed to create file"); + let mut gcloud_writer = std::io::BufWriter::new(gcloud_file); + + let data = cloud.encode(); + gcloud_writer.write_all(data.as_slice()).expect("failed to write to gcloud file"); - // write gloud.gz - let gz_file = std::fs::File::create(&gcloud_filename).expect("failed to create file"); - let mut gz_writer = std::io::BufWriter::new(gz_file); - let mut gz_encoder = GzEncoder::new(&mut gz_writer, Compression::default()); - serialize_into(&mut gz_encoder, &cloud).expect("failed to encode cloud"); + let post_encode_bytes = Byte::from_bytes(std::fs::metadata(&gcloud_filename).expect("failed to get metadata").len() as u128); + println!("output file size: {}", post_encode_bytes.get_appropriate_unit(true).to_string()); }