From dbcd8df94f6d60c95695d8f689545e37c4bfcf0b Mon Sep 17 00:00:00 2001 From: mosure Date: Sun, 28 Jul 2024 19:59:40 -0500 Subject: [PATCH 01/21] feat: render-mode cli arg --- src/primitive.rs | 2 +- src/render/depth.wgsl | 2 +- src/render/mod.rs | 38 ++++++++++++++++++++++++++++++++------ src/scene/room.rs | 43 ++++++++++++++----------------------------- tools/viewer.rs | 7 +++++++ 5 files changed, 55 insertions(+), 37 deletions(-) diff --git a/src/primitive.rs b/src/primitive.rs index f4f6e7c..0df28c1 100644 --- a/src/primitive.rs +++ b/src/primitive.rs @@ -438,7 +438,7 @@ fn build_primitive( } -fn process_primitives( +pub fn process_primitives( mut commands: Commands, mut meshes: ResMut>, mut standard_materials: ResMut>, diff --git a/src/render/depth.wgsl b/src/render/depth.wgsl index fdab01e..15c9272 100644 --- a/src/render/depth.wgsl +++ b/src/render/depth.wgsl @@ -13,5 +13,5 @@ fn depth_to_rgb(depth: f32) -> vec3 { @fragment fn fragment(in: VertexOutput) -> @location(0) vec4 { - return vec4(depth_to_rgb(in.position.w), 1.0); + return vec4(vec3(in.position.w), 1.0); } diff --git a/src/render/mod.rs b/src/render/mod.rs index 7a417c3..4e60571 100644 --- a/src/render/mod.rs +++ b/src/render/mod.rs @@ -2,12 +2,29 @@ use bevy::{ prelude::*, render::render_resource::Face, }; +use bevy_args::{ + Deserialize, + Serialize, + ValueEnum, +}; + +use crate::primitive::process_primitives; pub mod depth; pub mod normal; -#[derive(Default, Debug, Resource, Reflect)] +#[derive( + Debug, + Default, + Clone, + PartialEq, + Serialize, + Deserialize, + Reflect, + Resource, + ValueEnum, +)] #[reflect(Resource)] pub enum RenderMode { #[default] @@ -29,11 +46,20 @@ impl Plugin for RenderPlugin { app.add_plugins(normal::NormalPlugin); // TODO: add wireframe depth, pbr disable, normals - app.add_systems(PostUpdate, auto_disable_pbr_material::); - app.add_systems(PostUpdate, auto_disable_pbr_material::); - - app.add_systems(PostUpdate, apply_render_modes); - app.add_systems(PostUpdate, enable_pbr_material); + app.add_systems( + Update, + apply_render_modes.after(process_primitives), + ); + app.add_systems( + Update, + ( + auto_disable_pbr_material::, + auto_disable_pbr_material::, + enable_pbr_material, + ) + .after(apply_render_modes) + .after(process_primitives) + ); } } diff --git a/src/scene/room.rs b/src/scene/room.rs index 3ec2f6c..2f78887 100644 --- a/src/scene/room.rs +++ b/src/scene/room.rs @@ -52,6 +52,7 @@ impl Plugin for ZeroverseRoomPlugin { #[derive(Clone, Debug, Reflect, Resource)] #[reflect(Resource)] pub struct ZeroverseRoomSettings { + pub camera_floor_padding: f32, pub camera_wall_padding: f32, pub center_primitive_count: CountSampler, pub center_primitive_scale_sampler: ScaleSampler, @@ -69,6 +70,7 @@ impl Default for ZeroverseRoomSettings { // }; Self { + camera_floor_padding: 3.0, camera_wall_padding: 1.0, center_primitive_count: CountSampler::Bounded(4, 10), center_primitive_scale_sampler: ScaleSampler::Bounded( @@ -78,7 +80,7 @@ impl Default for ZeroverseRoomSettings { center_primitive_settings: ZeroversePrimitiveSettings::default(), looking_at_sampler: LookingAtSampler::Sphere { geometry: Sphere::new(4.0), - transform: Transform::from_translation(Vec3::new(0.0, 1.0, 0.0)), + transform: Transform::from_translation(Vec3::new(0.0, 3.0, 0.0)), }, room_size: ScaleSampler::Bounded( Vec3::new(12.0, 8.0, 12.0), @@ -105,27 +107,6 @@ fn setup_scene( .insert(SpatialBundle::default()) .with_children(|commands| { {// outer walls - // // bounding box - // let outer_walls_settings = ZeroversePrimitiveSettings { - // invert_normals: true, - // available_types: vec![ZeroversePrimitives::Cuboid], // TODO: change to plane, support multi-material hull - // components: 1, - // wireframe_probability: 1.0, - // noise_probability: 0.0, - // cast_shadows: false, - // position_sampler: PositionSampler::Exact { - // position: Vec3::new(0.0, scale.y / 2.0, 0.0), - // }, - // rotation_sampler: RotationSampler::Identity, - // scale_sampler: ScaleSampler::Exact(scale), - // ..default() - // }; - - // commands.spawn(PrimitiveBundle { - // settings: outer_walls_settings, - // ..default() - // }); - let base_plane_settings = ZeroversePrimitiveSettings { cull_mode: Some(Face::Front), available_types: vec![ZeroversePrimitives::Plane], @@ -259,9 +240,9 @@ fn setup_scene( let center_object_height = 8.0; let center_object_sampler = PositionSampler::Cube { extents: Vec3::new( - scale.x / 2.0, + scale.x / 1.5, center_object_height / 2.0, - scale.z / 2.0, + scale.z / 1.5, ), }; @@ -293,10 +274,10 @@ fn setup_scene( }); for _ in 0..scene_settings.num_cameras { - let size = Vec3::new( - scale.x - room_settings.camera_wall_padding, - scale.y / 2.0, - scale.z - room_settings.camera_wall_padding, + let size: Vec3 = Vec3::new( + scale.x - room_settings.camera_wall_padding * 2.0, + scale.y - room_settings.camera_floor_padding, + scale.z - room_settings.camera_wall_padding * 2.0, ); commands.spawn(ZeroverseCamera { @@ -304,7 +285,11 @@ fn setup_scene( sampler_type: CameraPositionSamplerType::Band { size, rotation: Quat::IDENTITY, - translate: Vec3::new(0.0, scale.y / 2.0 + size.y / 2.0, 0.0), + translate: Vec3::new( + 0.0, + room_settings.camera_floor_padding + size.y / 2.0, + 0.0, + ), }, looking_at: room_settings.looking_at_sampler.clone(), }, diff --git a/tools/viewer.rs b/tools/viewer.rs index da6a701..afd28d6 100644 --- a/tools/viewer.rs +++ b/tools/viewer.rs @@ -31,6 +31,7 @@ use bevy_zeroverse::{ }, plucker::ZeroversePluckerSettings, primitive::ScaleSampler, + render::RenderMode, scene::{ room::ZeroverseRoomSettings, RegenerateSceneEvent, @@ -94,6 +95,9 @@ struct BevyZeroverseViewer { #[arg(long, default_value = "0.0")] yaw_speed: f32, + #[arg(long, value_enum, default_value_t = RenderMode::Color)] + render_mode: RenderMode, + #[arg(long, value_enum, default_value_t = ZeroverseSceneType::Object)] scene_type: ZeroverseSceneType, } @@ -112,6 +116,7 @@ impl Default for BevyZeroverseViewer { name: "bevy_zeroverse".to_string(), regenerate_ms: 0, yaw_speed: 0.0, + render_mode: Default::default(), scene_type: Default::default(), } } @@ -181,6 +186,8 @@ fn viewer_app() { app.add_plugins(BevyZeroversePlugin); + app.insert_resource(args.render_mode); + app.insert_resource(DefaultZeroverseCamera { resolution: UVec2::new(args.width as u32, args.height as u32).into(), }); From e2f6dcb79cd1a9e8fbc01a0697b1607f2afe93b0 Mon Sep 17 00:00:00 2001 From: mosure Date: Sun, 28 Jul 2024 20:22:29 -0500 Subject: [PATCH 02/21] docs: python bindings --- python/README.md | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) create mode 100644 python/README.md diff --git a/python/README.md b/python/README.md new file mode 100644 index 0000000..d775823 --- /dev/null +++ b/python/README.md @@ -0,0 +1,27 @@ +# bevy_zeroverse python bindings + +install to your python environment with `pip install .` + + +## dataloader + +torch dataloader API for online bevy_zeroverse batch generation. + +```python +from bevy_zeroverse.dataloader import ZeroverseDataloader + +# create a torch compatible dataloader +dataloader = ZeroverseDataloader( + width=256, + height=144, + num_cameras=4, + render_modes=['color', 'depth', 'normal'], + seed=0, + scene_type='room', +) + +for batch in dataloader: + print(batch) + break + +``` From 86e3c82f573c97c2ae9da7b1c6134875d8acde3f Mon Sep 17 00:00:00 2001 From: mosure Date: Sun, 28 Jul 2024 22:13:55 -0500 Subject: [PATCH 03/21] feat: ffi scaffold --- .gitignore | 2 - Cargo.toml | 2 +- ffi/bevy_zeroverse/Cargo.toml | 47 ++++++++++++++++++++ {python => ffi/bevy_zeroverse}/README.md | 2 +- ffi/bevy_zeroverse/pyproject.toml | 3 ++ ffi/bevy_zeroverse/python/dataloader.py | 18 ++++++++ ffi/bevy_zeroverse/setup.py | 13 ++++++ ffi/bevy_zeroverse/src/lib.rs | 50 +++++++++++++++++++++ tools/viewer.rs => src/app.rs | 56 ++++++++++++++---------- src/lib.rs | 1 + src/viewer.rs | 7 +++ 11 files changed, 173 insertions(+), 28 deletions(-) create mode 100644 ffi/bevy_zeroverse/Cargo.toml rename {python => ffi/bevy_zeroverse}/README.md (82%) create mode 100644 ffi/bevy_zeroverse/pyproject.toml create mode 100644 ffi/bevy_zeroverse/python/dataloader.py create mode 100644 ffi/bevy_zeroverse/setup.py create mode 100644 ffi/bevy_zeroverse/src/lib.rs rename tools/viewer.rs => src/app.rs (93%) create mode 100644 src/viewer.rs diff --git a/.gitignore b/.gitignore index 3984489..b479621 100644 --- a/.gitignore +++ b/.gitignore @@ -6,8 +6,6 @@ Cargo.lock **/*.rs.bk *.pdb -*.py - *.ply *.gcloud diff --git a/Cargo.toml b/Cargo.toml index 9385ff4..c3943eb 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -143,5 +143,5 @@ path = "src/lib.rs" [[bin]] name = "viewer" -path = "tools/viewer.rs" +path = "src/viewer.rs" required-features = ["viewer"] diff --git a/ffi/bevy_zeroverse/Cargo.toml b/ffi/bevy_zeroverse/Cargo.toml new file mode 100644 index 0000000..e16d675 --- /dev/null +++ b/ffi/bevy_zeroverse/Cargo.toml @@ -0,0 +1,47 @@ +[package] +name = "bevy_zeroverse_ffi" +version = "0.1.0" +edition = "2021" + +[toolchain] +channel = "nightly" + + +[dependencies] +pyo3 = { version = "0.22" } +bevy_zeroverse = { path = "../../" } + + +[dependencies.bevy] +version = "0.14" +default-features = false +features = [ + "bevy_asset", + "bevy_render", + "bevy_ui", + "bevy_winit", +] + + +[profile.dev.package."*"] +opt-level = 3 + +[profile.dev] +opt-level = 1 + +[profile.release] +lto = "thin" +codegen-units = 1 +opt-level = 3 + +[profile.wasm-release] +inherits = "release" +opt-level = "z" +lto = "fat" +codegen-units = 1 + + + +[lib] +name = "bevy_zeroverse" +path = "src/lib.rs" diff --git a/python/README.md b/ffi/bevy_zeroverse/README.md similarity index 82% rename from python/README.md rename to ffi/bevy_zeroverse/README.md index d775823..0d59559 100644 --- a/python/README.md +++ b/ffi/bevy_zeroverse/README.md @@ -1,6 +1,6 @@ # bevy_zeroverse python bindings -install to your python environment with `pip install .` +install to your python environment with `pip install ./ffi/bevy_zeroverse` (from repository root) ## dataloader diff --git a/ffi/bevy_zeroverse/pyproject.toml b/ffi/bevy_zeroverse/pyproject.toml new file mode 100644 index 0000000..8619d87 --- /dev/null +++ b/ffi/bevy_zeroverse/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools", "wheel", "setuptools-rust"] +build-backend = "setuptools.build_meta" diff --git a/ffi/bevy_zeroverse/python/dataloader.py b/ffi/bevy_zeroverse/python/dataloader.py new file mode 100644 index 0000000..251c8c4 --- /dev/null +++ b/ffi/bevy_zeroverse/python/dataloader.py @@ -0,0 +1,18 @@ +import bevy_zeroverse + + +# bevy_zeroverse.initialize() + + +dataloader = bevy_zeroverse.ZeroverseDataloader( + width=256, + height=144, + num_cameras=4, + render_modes=['color', 'depth', 'normal'], + seed=0, + scene_type='room', +) + +for batch in dataloader: + print(batch) + break diff --git a/ffi/bevy_zeroverse/setup.py b/ffi/bevy_zeroverse/setup.py new file mode 100644 index 0000000..5304edc --- /dev/null +++ b/ffi/bevy_zeroverse/setup.py @@ -0,0 +1,13 @@ +from setuptools import find_packages, setup +from setuptools_rust import RustExtension + + +setup( + name="bevy_zeroverse", + version="0.1", + packages=find_packages(where="python"), + package_dir={"": "python"}, + rust_extensions=[ + RustExtension("bevy_zeroverse") + ], +) diff --git a/ffi/bevy_zeroverse/src/lib.rs b/ffi/bevy_zeroverse/src/lib.rs new file mode 100644 index 0000000..7d9e60a --- /dev/null +++ b/ffi/bevy_zeroverse/src/lib.rs @@ -0,0 +1,50 @@ +use bevy::prelude::*; +use pyo3::prelude::*; + + +// TODO: create Dataloader torch class (or a render `n` frames and return capture fn, used within a python wrapper dataloader, wrapper requires setup.py to include the python module) + + +pub fn setup_and_run_app( + new_thread: bool, +) { + let ready = Arc::new(AtomicBool::new(false)); + + let mut startup = { + let ready = Arc::clone(&ready); + + move || { + let mut app = viewer_app(None); + + ready.store(true, Ordering::Release); + + app.run(); + } + }; + + if new_thread { + thread::spawn(startup); + + while !ready.load(Ordering::Acquire) { + thread::yield_now(); + } + + return; + } + + startup(); +} + + +// TODO: add BevyZeroverseViewer struct parameter +#[pyfunction] +fn main(new_thread: bool) { + setup_and_run_app(new_thread); +} + + +#[pymodule] +fn bevy_zeroverse(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_function(wrap_pyfunction!(main, m)?)?; + Ok(()) +} diff --git a/tools/viewer.rs b/src/app.rs similarity index 93% rename from tools/viewer.rs rename to src/app.rs index afd28d6..9695f24 100644 --- a/tools/viewer.rs +++ b/src/app.rs @@ -17,7 +17,7 @@ use bevy_panorbit_camera::{ PanOrbitCameraPlugin, }; -use bevy_zeroverse::{ +use crate::{ BevyZeroversePlugin, camera::{ DefaultZeroverseCamera, @@ -43,6 +43,7 @@ use bevy_zeroverse::{ }; +// TODO: register as a python class #[derive( Debug, Resource, @@ -53,59 +54,64 @@ use bevy_zeroverse::{ )] #[command(about = "bevy_zeroverse viewer", version, long_about = None)] #[reflect(Resource)] -struct BevyZeroverseViewer { +pub struct BevyZeroverseViewer { /// enable the bevy inspector #[arg(long, default_value = "true")] - editor: bool, + pub editor: bool, + + /// no window will be shown + #[arg(long, default_value = "true")] + pub headless: bool, /// view available material basecolor textures in a grid #[arg(long, default_value = "false")] - material_grid: bool, + pub material_grid: bool, /// view plücker embeddings #[arg(long, default_value = "false")] - plucker_visualization: bool, + pub plucker_visualization: bool, /// enable closing the window with the escape key (doesn't work in web) #[arg(long, default_value = "true")] - press_esc_close: bool, + pub press_esc_close: bool, #[arg(long, default_value = "1920.0")] - width: f32, + pub width: f32, #[arg(long, default_value = "1080.0")] - height: f32, + pub height: f32, #[arg(long, default_value = "0")] - num_cameras: usize, + pub num_cameras: usize, /// display a grid of Zeroverse cameras #[arg(long, default_value = "false")] - camera_grid: bool, + pub camera_grid: bool, /// window title #[arg(long, default_value = "bevy_zeroverse")] - name: String, + pub name: String, /// move to the next scene after `regenerate_ms` milliseconds #[arg(long, default_value = "0")] - regenerate_ms: u32, + pub regenerate_ms: u32, /// automatically rotate the root scene object in the y axis #[arg(long, default_value = "0.0")] - yaw_speed: f32, + pub yaw_speed: f32, #[arg(long, value_enum, default_value_t = RenderMode::Color)] - render_mode: RenderMode, + pub render_mode: RenderMode, #[arg(long, value_enum, default_value_t = ZeroverseSceneType::Object)] - scene_type: ZeroverseSceneType, + pub scene_type: ZeroverseSceneType, } impl Default for BevyZeroverseViewer { fn default() -> BevyZeroverseViewer { BevyZeroverseViewer { editor: true, + headless: false, material_grid: false, plucker_visualization: false, press_esc_close: true, @@ -123,9 +129,16 @@ impl Default for BevyZeroverseViewer { } -fn viewer_app() { - let args = parse_args::(); - info!("{:?}", args); +pub fn viewer_app( + override_args: Option, +) -> App { + let args = if override_args.is_some() { + override_args.unwrap() + } else { + parse_args::() + }; + + info!("args: {:?}", args); let mut app = App::new(); @@ -211,7 +224,7 @@ fn viewer_app() { setup_camera_grid, )); - app.run(); + app } @@ -480,8 +493,3 @@ fn press_esc_close( exit.send(AppExit::Success); } } - - -pub fn main() { - viewer_app(); -} diff --git a/src/lib.rs b/src/lib.rs index f59ef56..842cec4 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,5 +1,6 @@ use bevy::prelude::*; +pub mod app; pub mod camera; pub mod manifold; pub mod material; diff --git a/src/viewer.rs b/src/viewer.rs new file mode 100644 index 0000000..1bc832f --- /dev/null +++ b/src/viewer.rs @@ -0,0 +1,7 @@ +use bevy_zeroverse::app::viewer_app; + + +pub fn main() { + viewer_app(None) + .run(); +} From 2438b125f0a2faac8cead773d5beed4235434230 Mon Sep 17 00:00:00 2001 From: mosure Date: Mon, 29 Jul 2024 17:08:10 -0500 Subject: [PATCH 04/21] fix: lint --- src/app.rs | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/app.rs b/src/app.rs index 9695f24..485b45b 100644 --- a/src/app.rs +++ b/src/app.rs @@ -132,10 +132,9 @@ impl Default for BevyZeroverseViewer { pub fn viewer_app( override_args: Option, ) -> App { - let args = if override_args.is_some() { - override_args.unwrap() - } else { - parse_args::() + let args = match override_args { + Some(args) => args, + None => parse_args::(), }; info!("args: {:?}", args); From c19b5738a4ef52d760d7134f645094e0eed2c433 Mon Sep 17 00:00:00 2001 From: mosure Date: Mon, 29 Jul 2024 22:44:30 -0500 Subject: [PATCH 05/21] fix: logging and threading --- .gitignore | 3 +++ Cargo.toml | 7 ++++++ ffi/bevy_zeroverse/Cargo.toml | 25 ++----------------- ffi/bevy_zeroverse/README.md | 12 +++++++++- ffi/bevy_zeroverse/pyproject.toml | 2 +- ffi/bevy_zeroverse/python/dataloader.py | 26 ++++++++++---------- ffi/bevy_zeroverse/setup.py | 8 +++++++ ffi/bevy_zeroverse/src/lib.rs | 32 ++++++++++++++++++------- src/app.rs | 7 +++--- 9 files changed, 73 insertions(+), 49 deletions(-) diff --git a/.gitignore b/.gitignore index b479621..a156040 100644 --- a/.gitignore +++ b/.gitignore @@ -6,6 +6,9 @@ Cargo.lock **/*.rs.bk *.pdb +build/ +*.egg-info/ + *.ply *.gcloud diff --git a/Cargo.toml b/Cargo.toml index c3943eb..111bac0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -145,3 +145,10 @@ path = "src/lib.rs" name = "viewer" path = "src/viewer.rs" required-features = ["viewer"] + + +[workspace] +members = [ + "ffi/bevy_zeroverse", + ".", +] diff --git a/ffi/bevy_zeroverse/Cargo.toml b/ffi/bevy_zeroverse/Cargo.toml index e16d675..0c6dda5 100644 --- a/ffi/bevy_zeroverse/Cargo.toml +++ b/ffi/bevy_zeroverse/Cargo.toml @@ -3,13 +3,11 @@ name = "bevy_zeroverse_ffi" version = "0.1.0" edition = "2021" -[toolchain] -channel = "nightly" - [dependencies] -pyo3 = { version = "0.22" } bevy_zeroverse = { path = "../../" } +pyo3 = { version = "0.22", features = ["extension-module"] } +pyo3-log = "0.11" [dependencies.bevy] @@ -23,25 +21,6 @@ features = [ ] -[profile.dev.package."*"] -opt-level = 3 - -[profile.dev] -opt-level = 1 - -[profile.release] -lto = "thin" -codegen-units = 1 -opt-level = 3 - -[profile.wasm-release] -inherits = "release" -opt-level = "z" -lto = "fat" -codegen-units = 1 - - - [lib] name = "bevy_zeroverse" path = "src/lib.rs" diff --git a/ffi/bevy_zeroverse/README.md b/ffi/bevy_zeroverse/README.md index 0d59559..e4f2be4 100644 --- a/ffi/bevy_zeroverse/README.md +++ b/ffi/bevy_zeroverse/README.md @@ -1,12 +1,14 @@ # bevy_zeroverse python bindings -install to your python environment with `pip install ./ffi/bevy_zeroverse` (from repository root) +install to your python environment with `pip install ffi/bevy_zeroverse` (from repository root) ## dataloader torch dataloader API for online bevy_zeroverse batch generation. +run the test script with `python ffi/bevy_zeroverse/python/dataloader.py` to see the dataloader in action. + ```python from bevy_zeroverse.dataloader import ZeroverseDataloader @@ -25,3 +27,11 @@ for batch in dataloader: break ``` + + +### macos setup + +```bash +LIBTORCH_PATH=$(python3 -c "import site; print(site.getsitepackages()[0] + '/torch/lib')") +export DYLD_LIBRARY_PATH=$LIBTORCH_PATH:$DYLD_LIBRARY_PATH +``` diff --git a/ffi/bevy_zeroverse/pyproject.toml b/ffi/bevy_zeroverse/pyproject.toml index 8619d87..1a5585c 100644 --- a/ffi/bevy_zeroverse/pyproject.toml +++ b/ffi/bevy_zeroverse/pyproject.toml @@ -1,3 +1,3 @@ [build-system] -requires = ["setuptools", "wheel", "setuptools-rust"] +requires = ["setuptools", "setuptools-rust", "torch==2.3.0", "wheel"] build-backend = "setuptools.build_meta" diff --git a/ffi/bevy_zeroverse/python/dataloader.py b/ffi/bevy_zeroverse/python/dataloader.py index 251c8c4..dd3b0ef 100644 --- a/ffi/bevy_zeroverse/python/dataloader.py +++ b/ffi/bevy_zeroverse/python/dataloader.py @@ -1,18 +1,20 @@ import bevy_zeroverse -# bevy_zeroverse.initialize() +print('initializing bevy zeroverse...') +bevy_zeroverse.initialize() +print('bevy zeroverse initialized!') -dataloader = bevy_zeroverse.ZeroverseDataloader( - width=256, - height=144, - num_cameras=4, - render_modes=['color', 'depth', 'normal'], - seed=0, - scene_type='room', -) +# dataloader = bevy_zeroverse.ZeroverseDataloader( +# width=256, +# height=144, +# num_cameras=4, +# render_modes=['color', 'depth', 'normal'], +# seed=0, +# scene_type='room', +# ) -for batch in dataloader: - print(batch) - break +# for batch in dataloader: +# print(batch) +# break diff --git a/ffi/bevy_zeroverse/setup.py b/ffi/bevy_zeroverse/setup.py index 5304edc..b2f967a 100644 --- a/ffi/bevy_zeroverse/setup.py +++ b/ffi/bevy_zeroverse/setup.py @@ -1,7 +1,15 @@ +import os from setuptools import find_packages, setup from setuptools_rust import RustExtension +# TODO: remove tch-rs dependency and output raw bytes (do tensor conversion in python, this has the overhead of GPU -> CPU -> GPU transfer) +os.environ["LIBTORCH_USE_PYTORCH"] = "1" + +libtorch_path = os.popen("python -c 'import site; print(site.getsitepackages()[0] + \"/torch/lib\")'").read().strip() +os.environ["LD_LIBRARY_PATH"] = libtorch_path + ":" + os.environ.get("LD_LIBRARY_PATH", "") + + setup( name="bevy_zeroverse", version="0.1", diff --git a/ffi/bevy_zeroverse/src/lib.rs b/ffi/bevy_zeroverse/src/lib.rs index 7d9e60a..8d79c2f 100644 --- a/ffi/bevy_zeroverse/src/lib.rs +++ b/ffi/bevy_zeroverse/src/lib.rs @@ -1,6 +1,16 @@ -use bevy::prelude::*; +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Arc, + }, + thread, +}; + +// use bevy::prelude::*; use pyo3::prelude::*; +use ::bevy_zeroverse::app::viewer_app; + // TODO: create Dataloader torch class (or a render `n` frames and return capture fn, used within a python wrapper dataloader, wrapper requires setup.py to include the python module) @@ -10,7 +20,7 @@ pub fn setup_and_run_app( ) { let ready = Arc::new(AtomicBool::new(false)); - let mut startup = { + let startup = { let ready = Arc::clone(&ready); move || { @@ -28,23 +38,27 @@ pub fn setup_and_run_app( while !ready.load(Ordering::Acquire) { thread::yield_now(); } - - return; + } else { + startup(); } - - startup(); } // TODO: add BevyZeroverseViewer struct parameter #[pyfunction] -fn main(new_thread: bool) { - setup_and_run_app(new_thread); +fn initialize( + py: Python<'_>, +) { + py.allow_threads(|| { + setup_and_run_app(true); + }); } #[pymodule] fn bevy_zeroverse(m: &Bound<'_, PyModule>) -> PyResult<()> { - m.add_function(wrap_pyfunction!(main, m)?)?; + pyo3_log::init(); + + m.add_function(wrap_pyfunction!(initialize, m)?)?; Ok(()) } diff --git a/src/app.rs b/src/app.rs index 485b45b..6b86237 100644 --- a/src/app.rs +++ b/src/app.rs @@ -6,7 +6,6 @@ use bevy::{ }; use bevy_args::{ parse_args, - BevyArgsPlugin, Deserialize, Parser, Serialize, @@ -45,6 +44,7 @@ use crate::{ // TODO: register as a python class #[derive( + Clone, Debug, Resource, Serialize, @@ -137,9 +137,11 @@ pub fn viewer_app( None => parse_args::(), }; + let mut app = App::new(); + info!("args: {:?}", args); + app.insert_resource(args.clone()); - let mut app = App::new(); #[cfg(target_arch = "wasm32")] let primary_window = Some(Window { @@ -182,7 +184,6 @@ pub fn viewer_app( app.add_plugins(default_plugins); - app.add_plugins(BevyArgsPlugin::::default()); app.add_plugins(PanOrbitCameraPlugin); app.insert_resource(Msaa::Sample8); From 6f5a488937469b850507799d52ace4c4986b0a7c Mon Sep 17 00:00:00 2001 From: mosure Date: Tue, 30 Jul 2024 18:27:09 -0500 Subject: [PATCH 06/21] feat: scaffold for ffi io --- .github/workflows/clippy.yml | 2 +- Cargo.toml | 18 +- ffi/bevy_zeroverse/Cargo.toml | 2 + ffi/bevy_zeroverse/python/dataloader.py | 21 +- ffi/bevy_zeroverse/src/lib.rs | 127 ++++++++- src/app.rs | 74 +++-- src/io.rs | 342 ++++++++++++++++++++++++ src/lib.rs | 1 + src/render/mod.rs | 2 + src/scene/mod.rs | 2 + 10 files changed, 545 insertions(+), 46 deletions(-) create mode 100644 src/io.rs diff --git a/.github/workflows/clippy.yml b/.github/workflows/clippy.yml index 0c07027..7c18160 100644 --- a/.github/workflows/clippy.yml +++ b/.github/workflows/clippy.yml @@ -34,4 +34,4 @@ jobs: enable-sccache: "true" - name: lint - run: cargo clippy -- -Dwarnings + run: cargo clippy --all -- -Dwarnings diff --git a/Cargo.toml b/Cargo.toml index 111bac0..7f1e8f8 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,6 +31,13 @@ exclude = [ default-run = "viewer" +[workspace] +members = [ + ".", + "ffi/bevy_zeroverse", +] + + [features] default = [ "asset_pipeline", @@ -71,14 +78,18 @@ bevy_args = "1.6" bevy-inspector-egui = { version = "0.25", optional = true } bevy_panorbit_camera = { version = "0.19", optional = true, features = ["bevy_egui"] } clap = { version = "4.4", features = ["derive"] } +futures-intrusive = "0.5" glob = "0.3" itertools = "0.13" noise = { version = "0.9" } +pollster = "0.3" +pyo3 = { version = "0.22", features = ["extension-module"] } rand = "0.8" rayon = { version = "1.10", optional = true } serde = "1.0" strum = "0.26" strum_macros = "0.26" +wgpu = "0.20.0" [target.'cfg(target_arch = "wasm32")'.dependencies] @@ -145,10 +156,3 @@ path = "src/lib.rs" name = "viewer" path = "src/viewer.rs" required-features = ["viewer"] - - -[workspace] -members = [ - "ffi/bevy_zeroverse", - ".", -] diff --git a/ffi/bevy_zeroverse/Cargo.toml b/ffi/bevy_zeroverse/Cargo.toml index 0c6dda5..afb7422 100644 --- a/ffi/bevy_zeroverse/Cargo.toml +++ b/ffi/bevy_zeroverse/Cargo.toml @@ -6,6 +6,8 @@ edition = "2021" [dependencies] bevy_zeroverse = { path = "../../" } +ndarray = { version = "0.15", features = ["blas"] } +once_cell = "1.19" pyo3 = { version = "0.22", features = ["extension-module"] } pyo3-log = "0.11" diff --git a/ffi/bevy_zeroverse/python/dataloader.py b/ffi/bevy_zeroverse/python/dataloader.py index dd3b0ef..9389d80 100644 --- a/ffi/bevy_zeroverse/python/dataloader.py +++ b/ffi/bevy_zeroverse/python/dataloader.py @@ -1,20 +1,13 @@ import bevy_zeroverse -print('initializing bevy zeroverse...') -bevy_zeroverse.initialize() -print('bevy zeroverse initialized!') +config = bevy_zeroverse.BevyZeroverseConfig() +config.headless = True +config.num_cameras = 4 -# dataloader = bevy_zeroverse.ZeroverseDataloader( -# width=256, -# height=144, -# num_cameras=4, -# render_modes=['color', 'depth', 'normal'], -# seed=0, -# scene_type='room', -# ) +bevy_zeroverse.initialize(config) -# for batch in dataloader: -# print(batch) -# break + +sample = bevy_zeroverse.next() +print(sample) diff --git a/ffi/bevy_zeroverse/src/lib.rs b/ffi/bevy_zeroverse/src/lib.rs index 8d79c2f..9f704c5 100644 --- a/ffi/bevy_zeroverse/src/lib.rs +++ b/ffi/bevy_zeroverse/src/lib.rs @@ -1,15 +1,93 @@ use std::{ sync::{ - atomic::{AtomicBool, Ordering}, Arc, + Mutex, + atomic::{ + AtomicBool, + Ordering, + }, + mpsc::{ + Sender, + Receiver, + RecvTimeoutError, + }, }, thread, + time::Duration, }; -// use bevy::prelude::*; -use pyo3::prelude::*; +use bevy::prelude::*; +use ndarray::{Array2, Array3}; +use once_cell::sync::OnceCell; +use pyo3::{ + prelude::*, + exceptions::PyTimeoutError, + types::PyList, +}; + +use ::bevy_zeroverse::app::{ + viewer_app, + BevyZeroverseConfig, +}; + + +type ColorImage = Array3; +type DepthImage = Array2; +type NormalImage = Array2<[f32; 3]>; + + +#[derive(Clone)] +#[pyclass] +struct View { + color: ColorImage, + depth: DepthImage, + normal: NormalImage, + + #[pyo3(get, set)] + view_from_world: [[f32; 4]; 4], + + #[pyo3(get, set)] + fovy: f32, +} + +#[pymethods] +impl View { + // TODO: upgrade rust-numpy to latest once pyo3 0.22 support is available + #[getter] + fn color<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> { + let color_list: Vec<_> = self.color.iter().map(|&v| v.into_py(py)).collect(); + PyList::new_bound(py, color_list) + } + + #[getter] + fn depth<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> { + let depth_list: Vec<_> = self.depth.iter().map(|&v| v.into_py(py)).collect(); + PyList::new_bound(py, depth_list) + } + + #[getter] + fn normal<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> { + let normal_list: Vec<_> = self.normal.iter().map(|&v| v.into_py(py)).collect(); + PyList::new_bound(py, normal_list) + } +} -use ::bevy_zeroverse::app::viewer_app; +#[pyclass] +struct Sample { + views: Vec, +} + +#[pymethods] +impl Sample { + #[getter] + fn views<'py>(&self, py: Python<'py>) -> Bound<'py, PyList> { + let views_list: Vec<_> = self.views.iter().map(|v| v.clone().into_py(py)).collect(); + PyList::new_bound(py, views_list) + } +} + +static SAMPLE_RECEIVER: OnceCell>>> = OnceCell::new(); +// static SAMPLE_SENDER: OnceCell> = OnceCell::new(); // TODO: create Dataloader torch class (or a render `n` frames and return capture fn, used within a python wrapper dataloader, wrapper requires setup.py to include the python module) @@ -17,6 +95,7 @@ use ::bevy_zeroverse::app::viewer_app; pub fn setup_and_run_app( new_thread: bool, + override_args: Option, ) { let ready = Arc::new(AtomicBool::new(false)); @@ -24,7 +103,7 @@ pub fn setup_and_run_app( let ready = Arc::clone(&ready); move || { - let mut app = viewer_app(None); + let mut app = viewer_app(override_args); ready.store(true, Ordering::Release); @@ -33,6 +112,8 @@ pub fn setup_and_run_app( }; if new_thread { + info!("starting bevy_zeroverse in a new thread"); + thread::spawn(startup); while !ready.load(Ordering::Acquire) { @@ -44,21 +125,53 @@ pub fn setup_and_run_app( } -// TODO: add BevyZeroverseViewer struct parameter #[pyfunction] +#[pyo3(signature = (override_args=None))] fn initialize( py: Python<'_>, + override_args: Option, ) { py.allow_threads(|| { - setup_and_run_app(true); + setup_and_run_app(true, override_args); }); } +// TODO: support batch dimension (e.g. single array allocation for multiple samples) +// TODO: add systems for pushing and receiving camera outputs + metadata to python +// TODO: add options to bevy_zeroverse.next (e.g. render_modes, scene parameters, etc.) +#[pyfunction] +fn next() -> PyResult { + // TODO: advance 'n' frames - requires app reference + // TODO: capture frame - requires app system registration to write to textures and readback, triggered by app event after 'n' frames + // TODO: send frame + + let receiver = SAMPLE_RECEIVER.get().unwrap(); + let receiver = receiver.lock().unwrap(); + + let timeout = Duration::from_secs(5); + + match receiver.recv_timeout(timeout) { + Ok(sample) => Ok(sample), + Err(RecvTimeoutError::Timeout) => { + Err(PyTimeoutError::new_err("receive operation timed out")) + } + Err(RecvTimeoutError::Disconnected) => { + Err(PyTimeoutError::new_err("channel disconnected")) + } + } +} + + #[pymodule] fn bevy_zeroverse(m: &Bound<'_, PyModule>) -> PyResult<()> { pyo3_log::init(); + m.add_class::()?; + m.add_class::()?; + m.add_class::()?; + m.add_function(wrap_pyfunction!(initialize, m)?)?; + m.add_function(wrap_pyfunction!(next, m)?)?; Ok(()) } diff --git a/src/app.rs b/src/app.rs index 6b86237..fa8e286 100644 --- a/src/app.rs +++ b/src/app.rs @@ -3,6 +3,7 @@ use bevy::{ app::AppExit, time::Stopwatch, render::camera::RenderTarget, + winit::WinitPlugin, }; use bevy_args::{ parse_args, @@ -15,6 +16,7 @@ use bevy_panorbit_camera::{ PanOrbitCamera, PanOrbitCameraPlugin, }; +use pyo3::prelude::*; use crate::{ BevyZeroversePlugin, @@ -54,62 +56,93 @@ use crate::{ )] #[command(about = "bevy_zeroverse viewer", version, long_about = None)] #[reflect(Resource)] -pub struct BevyZeroverseViewer { +#[pyclass] +pub struct BevyZeroverseConfig { /// enable the bevy inspector + #[pyo3(get, set)] #[arg(long, default_value = "true")] pub editor: bool, /// no window will be shown - #[arg(long, default_value = "true")] + #[pyo3(get, set)] + #[arg(long, default_value = "false")] pub headless: bool, /// view available material basecolor textures in a grid + #[pyo3(get, set)] #[arg(long, default_value = "false")] pub material_grid: bool, /// view plücker embeddings + #[pyo3(get, set)] #[arg(long, default_value = "false")] pub plucker_visualization: bool, /// enable closing the window with the escape key (doesn't work in web) + #[pyo3(get, set)] #[arg(long, default_value = "true")] pub press_esc_close: bool, + #[pyo3(get, set)] #[arg(long, default_value = "1920.0")] pub width: f32, + #[pyo3(get, set)] #[arg(long, default_value = "1080.0")] pub height: f32, + #[pyo3(get, set)] #[arg(long, default_value = "0")] pub num_cameras: usize, /// display a grid of Zeroverse cameras + #[pyo3(get, set)] #[arg(long, default_value = "false")] pub camera_grid: bool, /// window title + #[pyo3(get, set)] #[arg(long, default_value = "bevy_zeroverse")] pub name: String, /// move to the next scene after `regenerate_ms` milliseconds + #[pyo3(get, set)] #[arg(long, default_value = "0")] pub regenerate_ms: u32, /// automatically rotate the root scene object in the y axis + #[pyo3(get, set)] #[arg(long, default_value = "0.0")] pub yaw_speed: f32, + #[pyo3(get, set)] #[arg(long, value_enum, default_value_t = RenderMode::Color)] pub render_mode: RenderMode, + #[pyo3(get, set)] #[arg(long, value_enum, default_value_t = ZeroverseSceneType::Object)] pub scene_type: ZeroverseSceneType, } -impl Default for BevyZeroverseViewer { - fn default() -> BevyZeroverseViewer { - BevyZeroverseViewer { +#[pymethods] +impl BevyZeroverseConfig { + #[new] + pub fn new() -> Self { + Default::default() + } + + fn __str__(&self) -> PyResult { + Ok(format!("{:?}", self)) + } + + fn __repr__(&self) -> PyResult { + Ok(format!("{:?}", self)) + } +} + +impl Default for BevyZeroverseConfig { + fn default() -> BevyZeroverseConfig { + BevyZeroverseConfig { editor: true, headless: false, material_grid: false, @@ -130,11 +163,11 @@ impl Default for BevyZeroverseViewer { pub fn viewer_app( - override_args: Option, + override_args: Option, ) -> App { let args = match override_args { Some(args) => args, - None => parse_args::(), + None => parse_args::(), }; let mut app = App::new(); @@ -175,12 +208,19 @@ pub fn viewer_app( }); app.insert_resource(ClearColor(Color::srgba(0.0, 0.0, 0.0, 0.0))); + let default_plugins = DefaultPlugins - .set(ImagePlugin::default_nearest()) - .set(WindowPlugin { + .set(ImagePlugin::default_nearest()); + + let default_plugins = if args.headless { + default_plugins + .disable::() + } else { + default_plugins.set(WindowPlugin { primary_window, ..default() - }); + }) + }; app.add_plugins(default_plugins); @@ -189,7 +229,7 @@ pub fn viewer_app( app.insert_resource(Msaa::Sample8); if args.editor { - app.register_type::(); + app.register_type::(); app.add_plugins(WorldInspectorPlugin::new()); } @@ -232,7 +272,7 @@ pub fn viewer_app( struct MaterialGridCameraMarker; fn setup_camera( - args: Res, + args: Res, mut commands: Commands, material_grid_cameras: Query>, editor_cameras: Query>, @@ -290,7 +330,7 @@ struct CameraGrid; fn setup_camera_grid( mut commands: Commands, - args: Res, + args: Res, camera_grids: Query>, zeroverse_cameras: Query< (Entity, &Camera), @@ -360,7 +400,7 @@ struct MaterialGrid; fn setup_material_grid( mut commands: Commands, - args: Res, + args: Res, standard_materials: Res>, zeroverse_materials: Res, material_grids: Query>, @@ -429,7 +469,7 @@ fn setup_scene( #[allow(clippy::too_many_arguments)] fn regenerate_scene_system( - args: Res, + args: Res, keys: Res>, time: Res