10000 feat: improve gcloud load times by mosure · Pull Request #28 · mosure/bevy_gaussian_splatting · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

feat: improve gcloud load times #28

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Nov 15, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
41 changes: 41 additions & 0 deletions .github/workflows/bench.yml
Original file line number Diff line number Diff line change
@@ -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 }}
6 changes: 6 additions & 0 deletions .github/workflows/test.yml
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
39 changes: 29 additions & 10 deletions Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -49,6 +58,10 @@ features = [
]


[dev-dependencies]
criterion = { version = "0.5", features = ["html_reports"] }


[profile.dev.package."*"]
opt-level = 3

Expand Down Expand Up @@ -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
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
Binary file modified assets/scenes/icecream.gcloud
Binary file not shown.
47 changes: 47 additions & 0 deletions benches/io.rs
Original file line number Diff line number Diff line change
@@ -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::<Gaussian>() 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);
74 changes: 8 additions & 66 deletions src/gaussian.rs
Original file line number Diff line number Diff line change
@@ -1,40 +1,26 @@
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,
Serializer,
ser::SerializeTuple,
};

use crate::ply::parse_ply;


const fn num_sh_coefficients(degree: usize) -> usize {
if degree == 0 {
Expand All @@ -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,
Expand Down Expand Up @@ -121,6 +108,7 @@ pub const MAX_SIZE_VARIANCE: f32 = 5.0;
Debug,
Default,
Copy,
PartialEq,
Reflect,
ShaderType,
Pod,
Expand All @@ -141,6 +129,7 @@ pub struct Gaussian {
Asset,
Clone,
Debug,
PartialEq,
Reflect,
TypeUuid,
Serialize,
Expand Down Expand Up @@ -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<Self::Asset, Self::Error>> {

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<Gaussian> for rand::distributions::Standard {
fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Gaussian {
Gaussian {
Expand Down
6 changes: 6 additions & 0 deletions src/io/codec.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,6 @@

// TODO: support streamed codecs
pub trait GaussianCloudCodec {
fn encode(&self) -> Vec<u8>;
fn decode(data: &[u8]) -> Self;
}
35 changes: 35 additions & 0 deletions src/io/gcloud/bincode2.rs
Original file line number Diff line number Diff line change
@@ -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<u8> {
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
}
}
30 changes: 30 additions & 0 deletions src/io/gcloud/flexbuffers.rs
Original file line number Diff line number Diff line change
@@ -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<u8> {
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
}
}
5 changes: 5 additions & 0 deletions src/io/gcloud/mod.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
#[cfg(feature = "io_bincode2")]
pub mod bincode2;

#[cfg(feature = "io_flexbuffers")]
pub mod flexbuffers;
Loading
0